"""Driving-video conditioning for SCAIL-2-style end-to-end character animation.

Ports mechanism 1 + 3 of SCAIL-2 (arXiv:2606.10804) to LTX-2: the *driving*
video latent is concatenated directly into the DiT token sequence (no skeleton /
pose intermediate), and carries a fixed spatial offset ``width_offset`` (the
paper's ΔW) on the RoPE width axis so its coordinates stay detached from the
main video tokens. This is the inference-only PoC path -- it reuses the existing
frozen-reference-token machinery and does not touch the transformer or RoPE.

This mirrors :class:`ltx_core.conditioning.types.reference_video_cond.VideoConditionByReferenceLatent`
(same patchify -> positions -> append -> attention-mask flow); the only new
behaviour is the mode-aware coordinate assignment.

Scope note (Phase 1): SCAIL-2's in-context mask channels (mechanism 2) and the
reference-latent height shift ΔH_ref of Replacement Mode require model-weight
surgery / a separate reference token group and are intentionally out of scope
here. Because LTX handles the reference image as a frame-0 in-place replacement
(not a separate token group), the driving-token placement is identical for both
:class:`DrivingMode` values in Phase 1 -- ``mode`` is stored for forward
compatibility and to document intent, but does not yet alter the driving
coordinates. See the plan for the deferred Phase 2 work.
"""

from __future__ import annotations

from enum import Enum

import torch

from ltx_core.components.patchifiers import get_pixel_coords
from ltx_core.conditioning.cond_channels import extend_cond_channels
from ltx_core.conditioning.item import ConditioningItem
from ltx_core.conditioning.mask_utils import update_attention_mask
from ltx_core.tools import VideoLatentTools
from ltx_core.types import LatentState, VideoLatentShape

# Default normalization ceiling for the RoPE width axis. Must stay in sync with
# the model's ``positional_embedding_max_pos[2]`` (see rope.py:precompute_freqs_cis
# default ``max_pos=[20, 2048, 2048]`` and model.py `_init_` default). Driving
# width coordinates that reach or exceed this value would wrap under RoPE.
DEFAULT_MAX_WIDTH_POSITION = 2048


class DrivingMode(Enum):
    """SCAIL-2 conditioning mode. See module docstring for the Phase 1 caveat."""

    ANIMATION = "animation"
    REPLACEMENT = "replacement"


class VideoConditionByDrivingLatent(ConditioningItem):
    """Append driving-video tokens with a width-axis RoPE offset (SCAIL-2 ΔW).

    The driving tokens are appended after the target sequence as clean latents
    (placeholder zeros in the noisy latent), kept frozen (``denoise_mask =
    1 - strength``), temporally aligned to the target, and shifted along the
    width axis by ``width_offset`` so they occupy ``[ΔW, ΔW + Wv)`` while the
    target stays at ``[0, Wv)``.

    Args:
        latent: Driving video latents ``[B, C, F, H, W]``. Must match the target
            shape (same F/H/W) so tokens align frame-for-frame with the target.
        mode: SCAIL-2 mode (reserved for Phase 2; see module docstring).
        width_offset: ΔW in RoPE pixel-space width units. ``None`` (default) uses
            the target's pixel width (``target_shape.width * scale_factors.width``),
            placing the driving tokens immediately to the right of the target.
        strength: 1.0 keeps the driving latent fully clean (frozen); 0.0 would
            denoise it. Default 1.0.
        max_width_position: RoPE width normalization ceiling; validation raises if
            the shifted driving coordinates would reach it. Keep in sync with the
            model's ``positional_embedding_max_pos[2]``.
    """

    def __init__(
        self,
        latent: torch.Tensor,
        mode: DrivingMode = DrivingMode.ANIMATION,
        width_offset: float | None = None,
        strength: float = 1.0,
        max_width_position: int = DEFAULT_MAX_WIDTH_POSITION,
    ):
        self.latent = latent
        self.mode = mode
        self.width_offset = width_offset
        self.strength = strength
        self.max_width_position = max_width_position

    def apply_to(
        self,
        latent_state: LatentState,
        latent_tools: VideoLatentTools,
    ) -> LatentState:
        """Append driving tokens with target-aligned time and a ΔW width shift."""
        tokens = latent_tools.patchifier.patchify(self.latent)

        num_target_tokens = latent_tools.patchifier.get_token_count(latent_tools.target_shape)
        if tokens.shape[1] != num_target_tokens:
            raise ValueError(
                "VideoConditionByDrivingLatent expects the driving latent to match the target shape "
                f"(same F/H/W): got {tokens.shape[1]} driving tokens vs {num_target_tokens} target tokens. "
                "Resize/resample the driving video to the target resolution and frame count."
            )

        # Compute the driving tokens' own pixel-space coordinates (same flow as
        # the base reference conditioning and create_initial_state).
        latent_coords = latent_tools.patchifier.get_patch_grid_bounds(
            output_shape=VideoLatentShape.from_torch_shape(self.latent.shape),
            device=self.latent.device,
        )
        positions = get_pixel_coords(
            latent_coords=latent_coords,
            scale_factors=latent_tools.scale_factors,
            causal_fix=latent_tools.causal_fix,
        ).to(dtype=torch.float32)

        # Temporal alignment: copy the target's time coordinates so the driving
        # tokens sit on exactly the same time grid as z_t (robust to causal_fix /
        # fps nuances). Token order is a flattened (f h w) grid identical to the
        # target's, so a token-wise copy is frame-aligned.
        positions[:, 0:1, :] = latent_state.positions[:, 0:1, :num_target_tokens].to(dtype=torch.float32)

        # ΔW: shift the driving tokens along the width axis so they stay spatially
        # detached from the target tokens. Default offset = target pixel width,
        # giving target=[0, Wv), driving=[Wv, 2*Wv).
        width_offset = self.width_offset
        if width_offset is None:
            width_offset = float(latent_tools.target_shape.width * latent_tools.scale_factors.width)
        positions[:, 2, ...] = positions[:, 2, ...] + width_offset

        max_width = positions[:, 2, ...].max().item()
        if max_width >= self.max_width_position:
            raise ValueError(
                f"Driving width coordinate {max_width:.1f} reaches the RoPE ceiling "
                f"{self.max_width_position} and would wrap. Reduce width_offset "
                f"(currently {width_offset:.1f}) or lower the output width."
            )

        denoise_mask = torch.full(
            size=(*tokens.shape[:2], 1),
            fill_value=1.0 - self.strength,
            device=self.latent.device,
            dtype=self.latent.dtype,
        )

        new_attention_mask = update_attention_mask(
            latent_state=latent_state,
            attention_mask=None,
            num_noisy_tokens=num_target_tokens,
            num_new_tokens=tokens.shape[1],
            batch_size=tokens.shape[0],
            device=self.latent.device,
            dtype=self.latent.dtype,
        )

        return LatentState(
            latent=torch.cat([latent_state.latent, torch.zeros_like(tokens)], dim=1),
            denoise_mask=torch.cat([latent_state.denoise_mask, denoise_mask], dim=1),
            positions=torch.cat([latent_state.positions, positions], dim=2),
            clean_latent=torch.cat([latent_state.clean_latent, tokens], dim=1),
            attention_mask=new_attention_mask,
            cond_channels=extend_cond_channels(latent_state, tokens.shape[1]),
        )
