"""SCAIL-2 in-context mask conditioning (mechanism 2) for LTX-2.

Encodes ``K+1`` semantic pixel-space masks (1 environment switch + ``K`` character
binding slots) into per-token conditioning channels and writes them onto the
sequence via :attr:`LatentState.cond_channels`. Following the paper, mask signals
are carried by the *driving* (and reference) tokens while the noisy target keeps
an all-zero mask -- so by default the channels are written onto the last ``N``
tokens of the sequence (the appended driving group), which means **this item must
be applied after the driving conditioning**.

Channel expansion (LTX-2 faithful port of the paper's ``4(K+1)``): each semantic
mask is spatially downsampled to the latent grid and temporally stacked along the
channel dimension by the VAE temporal factor ``t`` (8 for LTX-2, vs 4 for the
paper's Wan-2.1 backbone), giving ``t*(K+1)`` channels -- ``8*7 = 56`` for the
default ``K=6``. The model's ``patchify_proj`` must be built with a matching
``mask_conditioning_channels`` (see :class:`ltx_core.model.transformer.model.LTXModel`).
"""

from __future__ import annotations

from dataclasses import replace

import torch
import torch.nn.functional as F
from einops import rearrange

from ltx_core.conditioning.item import ConditioningItem
from ltx_core.tools import VideoLatentTools
from ltx_core.types import LatentState


def encode_mask_channels(
    masks: torch.Tensor,
    temporal_factor: int,
    height_lat: int,
    width_lat: int,
    frames_lat: int,
) -> torch.Tensor:
    """Encode ``[B, K+1, F_pix, H_pix, W_pix]`` masks into ``[B, t*(K+1), F_lat, H_lat, W_lat]``.

    Each semantic mask is area-downsampled to the latent spatial grid, then the pixel frames that
    map to one latent frame are stacked along the channel dimension (temporal factor ``t``). The
    causal first latent frame corresponds to a single pixel frame, which is replicated across its
    ``t`` stacked channels so every latent frame yields a uniform ``t`` channels per semantic class.
    """
    b, s, f_pix, _h_pix, _w_pix = masks.shape
    t = temporal_factor
    expected_f_pix = (frames_lat - 1) * t + 1
    if f_pix != expected_f_pix:
        raise ValueError(
            f"mask pixel frames ({f_pix}) incompatible with latent frames ({frames_lat}) at temporal "
            f"factor {t}: expected (F_lat - 1) * {t} + 1 = {expected_f_pix}."
        )

    # Spatial downsample every (semantic, frame) mask to the latent grid.
    flat = rearrange(masks.to(dtype=torch.float32), "b s f h w -> (b s f) 1 h w")
    down = F.interpolate(flat, size=(height_lat, width_lat), mode="area")
    down = rearrange(down, "(b s f) 1 h w -> b s f h w", b=b, s=s)

    # Temporal stacking. Latent frame 0 = pixel frame 0 (causal), replicated across t channels;
    # latent frames 1.. group t consecutive pixel frames.
    first = down[:, :, :1].repeat(1, 1, t, 1, 1).unsqueeze(2)  # [B, S, 1, t, H, W]
    rest = rearrange(down[:, :, 1:], "b s (fl t) h w -> b s fl t h w", t=t)  # [B, S, F_lat-1, t, H, W]
    stacked = torch.cat([first, rest], dim=2)  # [B, S, F_lat, t, H, W]

    # Fold (semantic, temporal) into a single channel axis, semantic-major: channel = s * t + ti.
    return rearrange(stacked, "b s fl t h w -> b (s t) fl h w")  # [B, t*(K+1), F_lat, H_lat, W_lat]


class VideoConditionByMaskChannels(ConditioningItem):
    """Write SCAIL-2 in-context mask channels onto the driving/reference tokens.

    Args:
        masks: Pixel-space semantic masks ``[B, K+1, F_pix, H_pix, W_pix]``. Channel 0 is the
            environment switch (whether the environment comes from the reference vs the driving
            video); channels ``1..K`` are the character binding slots (regions sharing a slot share
            motion). ``F_pix`` must equal ``(F_lat - 1) * temporal_factor + 1``.
        applies_to_last_n: Number of trailing tokens to write the mask onto. ``None`` (default) uses
            the target token count, i.e. the appended driving group when this item runs right after
            the driving conditioning. The noisy target tokens keep an all-zero mask (paper-faithful).
    """

    def __init__(self, masks: torch.Tensor, applies_to_last_n: int | None = None):
        self.masks = masks
        self.applies_to_last_n = applies_to_last_n

    def apply_to(self, latent_state: LatentState, latent_tools: VideoLatentTools) -> LatentState:
        shape = latent_tools.target_shape
        cond = encode_mask_channels(
            masks=self.masks,
            temporal_factor=latent_tools.scale_factors.time,
            height_lat=shape.height,
            width_lat=shape.width,
            frames_lat=shape.frames,
        )
        tokens = latent_tools.patchifier.patchify(cond)  # [B, T_region, C_mask]
        b, n_region, c_mask = tokens.shape

        total_tokens = latent_state.latent.shape[1]
        n = self.applies_to_last_n if self.applies_to_last_n is not None else n_region
        if n != n_region:
            raise ValueError(
                f"applies_to_last_n ({n}) must equal the encoded mask token count ({n_region})."
            )
        if n > total_tokens:
            raise ValueError(
                f"cannot write {n} mask tokens onto a sequence of only {total_tokens} tokens."
            )

        cond_channels = latent_state.cond_channels
        if cond_channels is None:
            cond_channels = tokens.new_zeros(b, total_tokens, c_mask)
        else:
            if cond_channels.shape[2] != c_mask:
                raise ValueError(
                    f"existing cond_channels width ({cond_channels.shape[2]}) != mask channels ({c_mask})."
                )
            cond_channels = cond_channels.clone()

        start = total_tokens - n
        cond_channels[:, start : start + n] = tokens.to(dtype=cond_channels.dtype)
        return replace(latent_state, cond_channels=cond_channels)
