#!/usr/bin/env python3
"""CPU regression tests for the paper mean-only color-match float-domain shift."""

from __future__ import annotations

import sys
import unittest
from pathlib import Path

import torch

DELIVERY_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(DELIVERY_ROOT))

from benchmark.verify_paper_reproduction import (
    MEAN_ONLY_PROTOCOL as VERIFIER_MEAN_ONLY_PROTOCOL,
)
from scripts.infer import (
    COLOR_MATCH_FORMULA,
    MEAN_ONLY_PROTOCOL,
    MODEL_CHOICES,
    match_color_mean,
)


class MatchColorMeanTests(unittest.TestCase):
    def test_shifts_each_channel_by_its_own_mean_difference(self) -> None:
        sr = torch.tensor(
            [[[[0.25, 0.5], [0.75, 1.0]], [[0.125, 0.25], [0.5, 0.875]], [[0.0, 0.5], [0.25, 0.5]]]],
            dtype=torch.float32,
        )
        lr = torch.tensor(
            [[[[0.75, 0.75], [0.75, 0.75]], [[0.25, 0.25], [0.25, 0.25]], [[0.5, 0.5], [0.5, 0.5]]]],
            dtype=torch.float32,
        )

        shifted = match_color_mean(sr, lr)

        expected = torch.tensor(
            [[[[0.375, 0.625], [0.875, 1.0]], [[0.0, 0.0625], [0.3125, 0.6875]], [[0.1875, 0.6875], [0.4375, 0.6875]]]],
            dtype=torch.float32,
        )
        torch.testing.assert_close(shifted, expected, rtol=0.0, atol=0.0)

    def test_zero_shift_leaves_values_untouched(self) -> None:
        sr = torch.tensor([[[[0.25, 0.5], [0.5, 0.75]]]], dtype=torch.float32)
        lr = torch.full((1, 1, 2, 2), 0.5, dtype=torch.float32)

        torch.testing.assert_close(match_color_mean(sr, lr), sr, rtol=0.0, atol=0.0)

    def test_unclipped_shift_moves_every_channel_mean_to_the_lr_mean(self) -> None:
        sr = torch.arange(12, dtype=torch.float32).reshape(1, 1, 3, 4) / 16.0
        lr = torch.full((1, 1, 3, 4), 0.5, dtype=torch.float32)

        shifted = match_color_mean(sr, lr)

        torch.testing.assert_close(shifted.mean(dim=(-2, -1)), lr.mean(dim=(-2, -1)), rtol=0.0, atol=0.0)
        difference = shifted - sr
        torch.testing.assert_close(
            difference, difference.flatten()[0].expand_as(difference), rtol=0.0, atol=0.0
        )

    def test_computes_in_float32_regardless_of_input_dtype(self) -> None:
        sr = torch.tensor([[[[0.25, 0.5], [0.5, 0.75]]]], dtype=torch.bfloat16)
        lr = torch.tensor([[[[0.5, 0.5], [0.5, 0.5]]]], dtype=torch.bfloat16)

        self.assertEqual(match_color_mean(sr, lr).dtype, torch.float32)

    def test_does_not_mutate_its_inputs(self) -> None:
        sr = torch.tensor([[[[0.25, 0.75], [0.25, 0.75]]]], dtype=torch.float32)
        lr = torch.tensor([[[[0.0, 0.0], [1.0, 1.0]]]], dtype=torch.float32)
        sr_copy, lr_copy = sr.clone(), lr.clone()

        match_color_mean(sr, lr)

        torch.testing.assert_close(sr, sr_copy, rtol=0.0, atol=0.0)
        torch.testing.assert_close(lr, lr_copy, rtol=0.0, atol=0.0)


class ProtocolConstantTests(unittest.TestCase):
    def test_inference_protocol_constant_matches_the_verifier(self) -> None:
        self.assertEqual(MEAN_ONLY_PROTOCOL, VERIFIER_MEAN_ONLY_PROTOCOL)

    def test_documented_formula_is_the_per_channel_mean_shift(self) -> None:
        self.assertEqual(COLOR_MATCH_FORMULA, "clamp(sr - mean(sr,H,W) + mean(lr,H,W), 0, 1)")

    def test_documented_model_choices(self) -> None:
        self.assertEqual(MODEL_CHOICES, ("raesr", "raesr-d"))


class FrontendWeightsTests(unittest.TestCase):
    def test_released_frontend_weights_match_the_projector_contract(self) -> None:
        from raesr.runtime.frontends.raesr_d import PatchTokenProjector

        weights = DELIVERY_ROOT / "weights" / "raesr_d_frontend.pt"
        self.assertTrue(weights.is_file(), f"missing frontend weights: {weights}")
        state_dict = torch.load(weights, map_location="cpu", weights_only=True)
        projector = PatchTokenProjector()
        reference = projector.state_dict()
        self.assertEqual(sorted(state_dict), sorted(reference))
        for key, tensor in reference.items():
            self.assertEqual(tuple(state_dict[key].shape), tuple(tensor.shape), key)
        projector.load_state_dict(state_dict, strict=True)

    def test_frontend_state_dict_contains_no_python_objects(self) -> None:
        weights = DELIVERY_ROOT / "weights" / "raesr_d_frontend.pt"
        self.assertTrue(weights.is_file(), f"missing frontend weights: {weights}")
        state_dict = torch.load(weights, map_location="cpu", weights_only=True)
        self.assertTrue(all(torch.is_tensor(value) for value in state_dict.values()))


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