"""Frozen Teacher projectors that rebuild regular-RoPE and ProPE Context K/V."""

from __future__ import annotations

import copy

import torch
from einops import rearrange
from torch import nn

from hyvideo.models.transformers.modules.posemb_layers import (
    apply_rotary_emb,
    get_nd_rotary_pos_embed,
)
from hyvideo.prope.camera_rope import _prepare_apply_fns_all_dim


class FrozenTeacherContextProjector(nn.Module):
    """Rebuild the exact cache layout ``[regular batches, ProPE batches]``."""

    def __init__(self, teacher_transformer: nn.Module, block_ids: tuple[int, ...]) -> None:
        super().__init__()
        self.block_ids = tuple(int(value) for value in block_ids)
        self.heads = int(teacher_transformer.heads_num)
        self.head_dim = int(teacher_transformer.hidden_size // self.heads)
        self.rope_dim_list = tuple(teacher_transformer.rope_dim_list)
        self.rope_theta = float(teacher_transformer.rope_theta)
        self.patch_size = tuple(teacher_transformer.patch_size)
        self.latent_height = int(getattr(teacher_transformer, "latent_height", 30))
        self.latent_width = int(getattr(teacher_transformer, "latent_width", 52))
        self.k_proj = nn.ModuleDict()
        self.v_proj = nn.ModuleDict()
        self.k_norm = nn.ModuleDict()
        for block_id in self.block_ids:
            block = teacher_transformer.double_blocks[block_id]
            key = str(block_id)
            self.k_proj[key] = copy.deepcopy(block.img_attn_k)
            self.v_proj[key] = copy.deepcopy(block.img_attn_v)
            self.k_norm[key] = copy.deepcopy(block.img_attn_k_norm)
        self.eval().requires_grad_(False)

    def _rope(self, frames: int, height: int, width: int, device: torch.device):
        cos, sin = get_nd_rotary_pos_embed(
            list(self.rope_dim_list),
            (frames // self.patch_size[0], height, width),
            theta=self.rope_theta,
            use_real=True,
            theta_rescale_factor=1,
        )
        return cos.to(device), sin.to(device)

    def forward_block(
        self,
        block_id: int,
        img_modulated: torch.Tensor,
        context_viewmats: torch.Tensor,
        context_Ks: torch.Tensor,
        context_valid_mask: torch.Tensor,
        *,
        latent_height: int | None = None,
        latent_width: int | None = None,
    ) -> dict[str, torch.Tensor]:
        key = str(int(block_id))
        if key not in self.k_proj:
            raise ValueError(f"Projector has no block {block_id}")
        latent_height = self.latent_height if latent_height is None else int(latent_height)
        latent_width = self.latent_width if latent_width is None else int(latent_width)
        batch, tokens, _ = img_modulated.shape
        tokens_per_frame = latent_height * latent_width
        if tokens % tokens_per_frame:
            raise ValueError("Context features are not frame-aligned")
        frames = tokens // tokens_per_frame
        if tuple(context_valid_mask.shape) != (batch, tokens):
            raise ValueError("Context valid mask shape mismatch")
        if tuple(context_viewmats.shape) != (batch, frames, 4, 4):
            raise ValueError("Context viewmats shape mismatch")
        if tuple(context_Ks.shape) != (batch, frames, 3, 3):
            raise ValueError("Context Ks shape mismatch")
        if tokens == 0:
            empty = img_modulated.new_empty((2 * batch, self.heads, 0, self.head_dim))
            return {"k_vision": empty, "v_vision": empty.clone()}

        k = rearrange(
            self.k_proj[key](img_modulated),
            "b s (h d) -> b s h d",
            h=self.heads,
        )
        v = rearrange(
            self.v_proj[key](img_modulated),
            "b s (h d) -> b s h d",
            h=self.heads,
        )
        k = self.k_norm[key](k).to(v)
        cos, sin = self._rope(frames, latent_height, latent_width, img_modulated.device)
        _, regular_k = apply_rotary_emb(k, k, (cos, sin), head_first=False)
        regular_k = regular_k.permute(0, 2, 1, 3)
        regular_v = v.permute(0, 2, 1, 3)

        token_viewmats = context_viewmats.repeat_interleave(tokens_per_frame, dim=1)
        token_Ks = context_Ks.repeat_interleave(tokens_per_frame, dim=1)
        _, apply_kv, _ = _prepare_apply_fns_all_dim(
            head_dim=self.head_dim,
            viewmats=token_viewmats,
            Ks=token_Ks,
            patches_x=None,
            patches_y=None,
            image_width=None,
            image_height=None,
        )
        prope_k = apply_kv(k.permute(0, 2, 1, 3))
        prope_v = apply_kv(v.permute(0, 2, 1, 3))
        k_cache = torch.cat([regular_k, prope_k], dim=0)
        v_cache = torch.cat([regular_v, prope_v], dim=0)
        branch_mask = torch.cat([context_valid_mask, context_valid_mask], dim=0)
        branch_mask = branch_mask[:, None, :, None]
        k_cache = k_cache.masked_fill(~branch_mask, 0)
        v_cache = v_cache.masked_fill(~branch_mask, 0)
        return {"k_vision": k_cache, "v_vision": v_cache}

    def forward(self, context_by_block: dict[int, dict[str, torch.Tensor]]):
        return {
            int(block_id): self.forward_block(
                int(block_id),
                values["img_modulated"],
                values["context_viewmats"],
                values["context_Ks"],
                values["context_valid_mask"],
            )
            for block_id, values in context_by_block.items()
        }