import unittest
from unittest.mock import patch

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

from predictor_training.atc_fusion import ATCFusion, MultiInputMLP, StateTransport
from predictor_training.single_block import SingleBlockPredictor
from wan.modules.model import rope_params


def small_freqs(head_dim: int = 8) -> torch.Tensor:
    return torch.cat(
        [
            rope_params(32, head_dim - 4 * (head_dim // 6)),
            rope_params(32, 2 * (head_dim // 6)),
            rope_params(32, 2 * (head_dim // 6)),
        ],
        dim=1,
    )


def reference_attention(
    q: torch.Tensor, k: torch.Tensor, v: torch.Tensor
) -> torch.Tensor:
    scores = torch.einsum("blhd,bshd->bhls", q, k) / q.shape[-1] ** 0.5
    probability = scores.float().softmax(dim=-1).to(v.dtype)
    return torch.einsum("bhls,bshd->blhd", probability, v)


class DummyBlock(nn.Module):
    num_heads = 1

    def forward(self, value: torch.Tensor, **_kwargs) -> torch.Tensor:
        return value


class ATCFusionTest(unittest.TestCase):
    def setUp(self) -> None:
        torch.manual_seed(0)
        self.grid = torch.tensor([[2, 2, 2]], dtype=torch.long)
        self.freqs = small_freqs()
        self.current = torch.randn(1, 8, 8)
        self.previous = torch.randn(1, 8, 8)

    def test_multi_input_mlp_matches_explicit_concat(self) -> None:
        module = MultiInputMLP(3, dim=8, hidden_dim=12)
        values = [torch.randn(2, 5, 8) for _ in range(3)]
        combined_weight = torch.cat(
            [layer.weight for layer in module.inputs], dim=1
        )
        combined_bias = sum(layer.bias for layer in module.inputs)
        reference_hidden = F.linear(
            torch.cat(values, dim=-1), combined_weight, combined_bias
        )
        reference = module.out(module.activation(reference_hidden))
        torch.testing.assert_close(module(*values), reference)

    @patch(
        "predictor_training.atc_fusion.attention",
        side_effect=reference_attention,
    )
    def test_last_frame_scope_cannot_read_earlier_frames(self, _mock) -> None:
        transport = StateTransport(
            dim=8,
            num_heads=1,
            previous_scope="last_frame",
            gradient_checkpointing=False,
            diagnostic_queries=4,
        ).eval()
        first, _ = transport(
            self.current, self.previous, self.grid, self.freqs, current_start=8
        )
        perturbed_early = self.previous.clone()
        perturbed_early[:, :4] += torch.randn_like(perturbed_early[:, :4])
        second, _ = transport(
            self.current,
            perturbed_early,
            self.grid,
            self.freqs,
            current_start=8,
        )
        torch.testing.assert_close(first, second, rtol=0, atol=0)

        perturbed_last = self.previous.clone()
        perturbed_last[:, 4:] += torch.randn_like(perturbed_last[:, 4:])
        third, _ = transport(
            self.current,
            perturbed_last,
            self.grid,
            self.freqs,
            current_start=8,
        )
        self.assertFalse(torch.equal(first, third))

    @patch(
        "predictor_training.atc_fusion.attention",
        side_effect=reference_attention,
    )
    def test_zero_initialized_correction_bypass(self, _mock) -> None:
        fusion = ATCFusion(
            dim=8,
            num_heads=1,
            freq_dim=8,
            mlp_hidden_dim=16,
            gate_hidden_dim=4,
            previous_scope="chunk",
            gradient_checkpointing=False,
        ).eval()
        condition = torch.randn_like(self.current)
        backbone, gated_correction = fusion(
            self.current,
            torch.randn_like(self.current),
            self.previous,
            condition,
            torch.tensor([62.5]),
            self.grid,
            self.freqs,
            current_start=8,
        )
        self.assertEqual(backbone.shape, self.current.shape)
        torch.testing.assert_close(
            gated_correction,
            torch.zeros_like(gated_correction),
            rtol=0,
            atol=0,
        )
        self.assertAlmostEqual(float(fusion.last_gate.mean()), 0.3, delta=0.01)

    def test_atc_token_gate_is_a_stage1_parameter(self) -> None:
        predictor = SingleBlockPredictor(
            DummyBlock(),
            dim=8,
            gradient_checkpointing=False,
            input_variant="atc",
            atc_freq_dim=8,
            atc_mlp_hidden_dim=16,
            atc_gate_hidden_dim=4,
        )
        token_gate_ids = {
            id(parameter)
            for parameter in predictor.fusion.token_gate.parameters()
        }
        stage1_ids = {
            id(parameter) for parameter in predictor.stage1_input_parameters()
        }
        self.assertTrue(token_gate_ids)
        self.assertTrue(token_gate_ids.issubset(stage1_ids))
        self.assertEqual(predictor.gate_parameters(), [])

    @patch(
        "predictor_training.atc_fusion.attention",
        side_effect=reference_attention,
    )
    def test_zero_initialized_predictor_returns_anchor(self, _mock) -> None:
        predictor = SingleBlockPredictor(
            DummyBlock(),
            dim=8,
            gradient_checkpointing=False,
            input_variant="atc",
            atc_freq_dim=8,
            atc_mlp_hidden_dim=16,
            atc_gate_hidden_dim=4,
        ).eval()
        anchor = torch.randn_like(self.current)
        empty_history = torch.empty(1, 0, 1, 8)
        prediction = predictor(
            current_tokens=self.current,
            anchor_hidden=anchor,
            previous_hidden=self.previous,
            timestep_modulation=torch.randn(1, 2, 6, 8),
            grid_sizes=self.grid,
            freqs=self.freqs,
            history_k=empty_history,
            history_v=empty_history,
            cross_k=torch.empty(1, 0, 1, 8),
            cross_v=torch.empty(1, 0, 1, 8),
            current_start=8,
            condition_tokens=torch.randn_like(self.current),
            anchor_distance=torch.tensor([62.5]),
        )
        torch.testing.assert_close(prediction, anchor, rtol=0, atol=0)


if __name__ == "__main__":
    unittest.main()