"""Rebuild clean Wan K/V from saved prefeatures and compare with live cache."""

from __future__ import annotations

import torch
import torch.nn.functional as F

from .schema import (
    CHUNK_FRAMES,
    LATENT_HEIGHT,
    LATENT_WIDTH,
    TOKENS_PER_FRAME,
)


@torch.no_grad()
def rebuild_clean_kv(
    model: torch.nn.Module,
    block_id: int,
    self_attn_input: torch.Tensor,
    *,
    start_frame: int,
    num_frames: int = CHUNK_FRAMES,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Apply the frozen Teacher's K/V projections and causal RoPE."""
    from wan.modules.causal_model import causal_rope_apply

    if num_frames != CHUNK_FRAMES:
        raise ValueError(f"expected {CHUNK_FRAMES} clean frames, got {num_frames}")
    block = model.blocks[int(block_id)]
    attention = block.self_attn
    device = attention.k.weight.device
    dtype = attention.k.weight.dtype
    value = self_attn_input.to(device=device, dtype=dtype)
    batch, tokens, _ = value.shape
    if tokens != num_frames * TOKENS_PER_FRAME:
        raise ValueError("clean prefeature token count is not frame-aligned")
    grid_sizes = torch.tensor(
        [[num_frames, LATENT_HEIGHT // 2, LATENT_WIDTH // 2]],
        device=device,
        dtype=torch.long,
    ).expand(batch, -1)
    freqs = model.freqs.to(device=device)
    key = attention.norm_k(attention.k(value)).view(
        batch, tokens, attention.num_heads, attention.head_dim
    )
    rebuilt_k = causal_rope_apply(
        key,
        grid_sizes,
        freqs,
        start_frame=int(start_frame),
    ).type_as(value)
    rebuilt_v = attention.v(value).view(
        batch, tokens, attention.num_heads, attention.head_dim
    )
    return rebuilt_k, rebuilt_v


def comparison_metrics(actual: torch.Tensor, expected: torch.Tensor) -> dict[str, float]:
    actual_flat = actual.float().flatten()
    expected_flat = expected.float().flatten()
    difference = actual_flat - expected_flat
    return {
        "relative_l2": float(
            difference.norm() / expected_flat.norm().clamp_min(1e-12)
        ),
        "cosine": float(F.cosine_similarity(actual_flat, expected_flat, dim=0)),
        "max_abs": float(difference.abs().max()),
    }


@torch.no_grad()
def validate_against_live_cache(
    model: torch.nn.Module,
    kv_cache: list[dict[str, torch.Tensor]],
    clean_features: dict[int, torch.Tensor],
    *,
    start_frame: int,
    num_frames: int = CHUNK_FRAMES,
    relative_l2_limit: float = 5e-3,
    cosine_limit: float = 0.9999,
) -> dict[str, float]:
    """Validate reconstructed K/V against the clean pass cache slice."""
    start_token = int(start_frame) * TOKENS_PER_FRAME
    end_token = start_token + int(num_frames) * TOKENS_PER_FRAME
    metrics: dict[str, float] = {}
    for block_id, feature in clean_features.items():
        rebuilt_k, rebuilt_v = rebuild_clean_kv(
            model,
            block_id,
            feature,
            start_frame=start_frame,
            num_frames=num_frames,
        )
        for field, rebuilt in (("k", rebuilt_k), ("v", rebuilt_v)):
            expected = kv_cache[block_id][field][:, start_token:end_token]
            values = comparison_metrics(rebuilt, expected)
            for metric_name, value in values.items():
                metrics[f"block_{block_id:02d}_{field}_{metric_name}"] = value
            if (
                values["relative_l2"] > relative_l2_limit
                or values["cosine"] < cosine_limit
            ):
                raise RuntimeError(
                    "clean KV rebuild failed "
                    f"block={block_id} field={field}: "
                    f"relative_l2={values['relative_l2']:.6g}, "
                    f"cosine={values['cosine']:.8f}"
                )
    return metrics