"""Compositional Latent Predictor φ.

Takes the view-variant stream of one clip and the view-invariant stream of
another, fuses them per time step, and runs a small causal transformer
(paper §3.1–3.3). One token per frame::

    token_t = concat(z_vv_t, z_vi_t) ∈ R^{2·d_z}

with a standard causal mask (frame ``t`` attends to ``≤ t``) and a key-padding
mask blocking frames where either stream is padded.

Three linear heads on the final hidden state ``x ∈ R^{B×T×2·d_z}``:

  - ``cls_head`` (2·d_z → d_t): at the last valid pair-frame → the compositional
    semantic latent ``s`` aligned to language (``L_decomp``).
  - ``vi_head`` (2·d_z → d_z): per t → ``ẑ_vi`` predicting the view-invariant
    input stream's next frame (``L_temp``).
  - ``vv_head`` (2·d_z → d_z): per t → ``ẑ_vv`` predicting the view-variant
    input stream's next frame (``L_temp``).
"""

from __future__ import annotations

import torch
from torch import nn

from .layers import TemporalBlock, build_sin_pos_embed, causal_mask


class CompositionalPredictor(nn.Module):
    def __init__(
        self,
        d_z: int = 512,
        d_t: int = 1024,
        max_frames: int = 128,
        depth: int = 4,
        num_heads: int = 8,
        mlp_ratio: float = 4.0,
    ):
        super().__init__()
        self.d_z = d_z
        self.d_t = d_t
        self.max_frames = max_frames
        self.d_in = 2 * d_z  # channel-concat [z_vv ‖ z_vi]

        self.register_buffer(
            "pos_embed", build_sin_pos_embed(max_frames, self.d_in), persistent=False
        )
        self.input_norm = nn.LayerNorm(self.d_in)
        self.blocks = nn.ModuleList(
            [TemporalBlock(self.d_in, num_heads, mlp_ratio) for _ in range(depth)]
        )
        self.final_norm = nn.LayerNorm(self.d_in)

        self.cls_head = nn.Linear(self.d_in, d_t)   # → s   (compositional latent)
        self.vi_head = nn.Linear(self.d_in, d_z)    # → ẑ_vi (next-frame, V-I)
        self.vv_head = nn.Linear(self.d_in, d_z)    # → ẑ_vv (next-frame, V-V)

    def forward(
        self,
        z_vv: torch.Tensor,
        z_vi: torch.Tensor,
        valid_vv: torch.Tensor,
        valid_vi: torch.Tensor,
    ) -> dict:
        """
        z_vv, z_vi: ``(B, T, d_z)`` — view-variant / view-invariant streams,
            generally sourced from two different clips.
        valid_vv, valid_vi: ``(B, T)`` bool — True where the frame is real.

        Returns dict:
            s:          ``(B, d_t)``    cls head at the last valid pair-frame.
            z_vi_pred:  ``(B, T, d_z)`` vi head per t (predicts z_vi_{t+1}).
            z_vv_pred:  ``(B, T, d_z)`` vv head per t (predicts z_vv_{t+1}).
            pair_valid: ``(B, T)``      valid_vv & valid_vi.
        """
        B, T, _ = z_vv.shape
        device = z_vv.device
        assert T <= self.max_frames, f"T={T} > max_frames={self.max_frames}"
        assert z_vi.shape == z_vv.shape, "z_vv and z_vi must match shape"

        x = torch.cat([z_vv, z_vi], dim=-1)                  # (B, T, 2·d_z)
        x = x + self.pos_embed[:, :T, :].to(dtype=x.dtype)
        x = self.input_norm(x)

        attn_mask = causal_mask(T, device=device)            # (T, T) True = blocked
        pair_valid = valid_vv & valid_vi                     # (B, T)
        key_padding_mask = ~pair_valid                       # True = blocked
        for blk in self.blocks:
            x = blk(x, attn_mask=attn_mask, key_padding_mask=key_padding_mask)
        x = self.final_norm(x)                               # (B, T, 2·d_z)

        # cls head at the last valid pair-frame per sample. (Samples with no
        # valid pair-frame are filtered out of the loss downstream.)
        arange_T = torch.arange(T, device=device).unsqueeze(0).expand(B, -1)
        scored = arange_T.where(pair_valid, torch.full_like(arange_T, -1))
        last_valid = scored.max(dim=1).values.clamp(min=0)   # (B,)
        last_tokens = x[torch.arange(B, device=device), last_valid, :]
        s = self.cls_head(last_tokens)                       # (B, d_t)

        return {
            "s": s,
            "z_vi_pred": self.vi_head(x),                    # (B, T, d_z)
            "z_vv_pred": self.vv_head(x),                    # (B, T, d_z)
            "pair_valid": pair_valid,
        }