#!/usr/bin/env python3
"""Validate the active three-encoder MMD composition and gradients on CPU."""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

import torch

sys.dont_write_bytecode = True


def assert_same(left: torch.Tensor, right: torch.Tensor, label: str) -> None:
    torch.testing.assert_close(left, right, rtol=0, atol=0, msg=label)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--release", choices=("1.3B", "14B"), required=True)
    args = parser.parse_args()

    root = Path(__file__).resolve().parents[1]
    sys.path.insert(0, str(root / args.release / "code"))

    from losses.multi_encoder_mmd import (
        MultiEncoderMMD,
        normalize_mmd_with_detached_scale,
    )
    from losses.video_nystrom_mmd import (
        VideoNystromMMD,
        build_video_nystrom_reference,
    )

    shapes = {
        "dinov3": (1, 2),
        "vjepa2": (2, 3),
        "videomae_v2": (1, 4),
    }

    def objective() -> MultiEncoderMMD:
        branches = {}
        for index, (name, shape) in enumerate(shapes.items()):
            generator = torch.Generator().manual_seed(9000 + index)
            reference = torch.randn(12, *shape, generator=generator)
            pack = build_video_nystrom_reference(
                reference,
                num_landmarks=5,
                metadata={
                    "feature": {"name": name, "matrix_shape": list(shape)},
                    "reference": {"dataset": "active-mmd-check"},
                },
                bandwidth=1.7,
                kmeans_iterations=3,
                ridge=1.0e-5,
                chunk_size=16,
            )
            branches[name] = VideoNystromMMD(
                pack,
                kernel_chunk_size=3,
                attraction_weight=2.0,
                current_teacher_attraction_enabled=True,
                current_teacher_attraction_weight=1.0,
            )
        return MultiEncoderMMD(
            branches,
            weights={"dinov3": 10.0, "vjepa2": 1.0, "videomae_v2": 1.0},
            normalization_eps=1.0e-4,
            aggregation_divisor=3.0,
        )

    generator = torch.Generator().manual_seed(1862643)
    for mode in ("full", "partial"):
        combined = objective()
        generated = {
            name: torch.randn(4, *shape, generator=generator)
            for name, shape in shapes.items()
        }
        teacher = {
            name: torch.randn(
                4 if mode == "full" else 2, *shape, generator=generator)
            for name, shape in shapes.items()
        }

        if mode == "full":
            leaves = {
                name: value.clone().requires_grad_()
                for name, value in generated.items()
            }
            loss, metadata = combined(
                leaves, current_teacher=teacher, stage=False)
        else:
            indices = torch.tensor([0, 2])
            leaves = {
                name: value.index_select(0, indices).clone().requires_grad_()
                for name, value in generated.items()
            }
            loss, metadata = combined(
                gradient_students=leaves,
                student_reference=generated,
                current_teacher=teacher,
                gradient_student_indices=indices,
                stage=False,
            )

        expected_terms = []
        for name, branch in combined.branches.items():
            if mode == "full":
                raw, _ = branch(
                    leaves[name], current_teacher=teacher[name], stage=False)
            else:
                raw, _ = branch(
                    gradient_students=leaves[name],
                    student_reference=generated[name],
                    current_teacher=teacher[name],
                    gradient_student_indices=indices,
                    stage=False,
                )
            normalized, denominator = normalize_mmd_with_detached_scale(
                raw, combined.normalization_eps)
            expected_terms.append(combined.branch_weights[name] * normalized)
            assert_same(
                metadata["branches"][name]["raw_mmd"],
                raw.detach(),
                f"{mode} {name} raw MMD",
            )
            assert_same(
                metadata["branches"][name]["normalization_denominator"],
                denominator.detach(),
                f"{mode} {name} normalization denominator",
            )

        expected = torch.stack(expected_terms).sum() / 3.0
        assert_same(loss, expected, f"{mode} combined MMD")
        observed_gradients = torch.autograd.grad(
            loss, tuple(leaves.values()), retain_graph=True)
        expected_gradients = torch.autograd.grad(
            expected, tuple(leaves.values()))
        for name, observed, reference in zip(
                leaves, observed_gradients, expected_gradients):
            assert_same(observed, reference, f"{mode} {name} feature gradient")
            assert torch.isfinite(observed).all() and observed.norm() > 0
        print(
            f"PASS {args.release} {mode}: weighted three-encoder MMD, "
            "metadata and feature gradients")

    eps = 1.0e-4
    for value in (-2.0e-4, 0.0, 1.0e-8, 0.1, 7.0):
        observed_input = torch.tensor(
            value, dtype=torch.float64, requires_grad=True)
        expected_input = observed_input.detach().clone().requires_grad_(True)
        observed, observed_denominator = normalize_mmd_with_detached_scale(
            observed_input, eps)
        expected_denominator = expected_input.detach() + eps
        expected = expected_input / expected_denominator
        assert_same(observed.detach(), expected.detach(), "normalization value")
        assert_same(
            observed_denominator, expected_denominator,
            "normalization denominator")
        observed_gradient, = torch.autograd.grad(observed, observed_input)
        expected_gradient, = torch.autograd.grad(expected, expected_input)
        assert_same(
            observed_gradient, expected_gradient, "normalization gradient")
    print(f"PASS {args.release}: detached MMD normalization behavior")


if __name__ == "__main__":
    main()