"""Predictor-DisCa: current/anchor fusion without AR Context Vision K/V."""

from __future__ import annotations

from typing import Mapping

import torch

from .predictor_v2 import HYWorldPlayPredictorV2, PredictorV2Config


class HYWorldPlayPredictorDisCa(HYWorldPlayPredictorV2):
    """Two-block Predictor using only current/anchor image features and text K/V.

    This is the project-local ``predictor-disca`` variant.  Unlike Predictor-v4,
    it neither fuses previous-chunk hidden states nor injects Context Vision K/V.
    The cached text K/V remains available because HY-WorldPlay AR vision blocks
    require text conditioning.
    """

    uses_context_kv = True
    uses_context_vision_kv = False
    uses_context_prefeature = False
    requires_previous_chunk_hidden = False

    def _text_only_cache(
        self,
        values: Mapping[int, Mapping[str, torch.Tensor]],
    ) -> dict[int, dict[str, torch.Tensor | None]]:
        result: dict[int, dict[str, torch.Tensor | None]] = {}
        for block_id in self.predictor_config.source_block_ids:
            if block_id not in values:
                raise ValueError(f"Missing Predictor-DisCa text block {block_id}")
            block = values[block_id]
            missing = {"k_txt", "v_txt"}.difference(block)
            if missing:
                raise ValueError(
                    f"Predictor-DisCa block {block_id} lacks {sorted(missing)}"
                )
            result[block_id] = {
                "k_txt": block["k_txt"],
                "v_txt": block["v_txt"],
                "text_valid_mask": block.get("text_valid_mask"),
                # Deliberately discard online Context Vision K/V and construct
                # zero-length vision caches in HYWorldPlayPredictorV2._make_cache.
                "k_vision": None,
                "v_vision": None,
            }
        return result

    def forward(
        self,
        *,
        target_model_input: torch.Tensor,
        anchor_hidden: torch.Tensor,
        target_frame_condition: torch.Tensor,
        target_viewmats: torch.Tensor,
        target_Ks: torch.Tensor,
        rope_temporal_size: torch.Tensor | int,
        start_rope_start_idx: torch.Tensor | int,
        text_kv: Mapping[int, Mapping[str, torch.Tensor]] | None = None,
        context_kv: Mapping[int, Mapping[str, torch.Tensor]] | None = None,
        previous_chunk_hidden: torch.Tensor | None = None,
    ) -> dict[str, torch.Tensor]:
        if previous_chunk_hidden is not None:
            raise ValueError("Predictor-DisCa does not accept previous_chunk_hidden")
        source = text_kv if text_kv is not None else context_kv
        if source is None:
            raise ValueError("Predictor-DisCa requires text_kv or online context_kv")
        return super().forward(
            target_model_input=target_model_input,
            anchor_hidden=anchor_hidden,
            target_frame_condition=target_frame_condition,
            target_viewmats=target_viewmats,
            target_Ks=target_Ks,
            rope_temporal_size=rope_temporal_size,
            start_rope_start_idx=start_rope_start_idx,
            context_kv=self._text_only_cache(source),
            previous_chunk_hidden=None,
        )