import unittest
from unittest.mock import patch

import torch
from torch import nn

from h3_blackwell_runtime.lora import DynamicLoraMixin
from h3_blackwell_runtime.sampler import sample_video_res_multistep, sample_video_turbo, turbo_sigmas


class _Linear(DynamicLoraMixin, nn.Module):
    def __init__(self, weight):
        super().__init__()
        self.register_buffer("weight", weight)
        self.in_features = weight.shape[1]
        self.out_features = weight.shape[0]
        self._init_dynamic_lora()

    def forward(self, x):
        return self._apply_lora(x, torch.nn.functional.linear(x, self.weight))


class TurboLoraContracts(unittest.TestCase):
    def test_dynamic_lora_matches_unmerged_formula(self):
        base_weight = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
        down = torch.tensor([[2.0, -1.0]])
        up = torch.tensor([[3.0], [-2.0]])
        x = torch.tensor([[0.5, -1.0]])
        linear = _Linear(base_weight)
        linear.add_lora("turbo", down, up, alpha=0.5)
        linear.set_lora("turbo", strength=0.75)
        expected = torch.nn.functional.linear(x, base_weight)
        expected += 0.75 * 0.5 * torch.nn.functional.linear(torch.nn.functional.linear(x, down), up)
        torch.testing.assert_close(linear(x), expected)

    def test_disabled_lora_returns_exact_base_result(self):
        weight = torch.randn(3, 2)
        x = torch.randn(4, 2)
        linear = _Linear(weight)
        linear.add_lora("turbo", torch.randn(1, 2), torch.randn(3, 1), alpha=1.0)
        expected = torch.nn.functional.linear(x, weight)
        self.assertTrue(torch.equal(linear(x), expected))
        linear.set_lora("turbo", strength=0.0)
        self.assertTrue(torch.equal(linear(x), expected))

    def test_unknown_lora_is_rejected(self):
        linear = _Linear(torch.randn(3, 2))
        with self.assertRaisesRegex(ValueError, "not attached"):
            linear.set_lora("missing")

    def test_four_step_shift_six_schedule(self):
        actual = turbo_sigmas(4, 6.0, device="cpu")
        expected = torch.tensor([1.0, 0.9473684211, 0.8571428571, 0.6666666667, 0.0])
        torch.testing.assert_close(actual, expected)

    def test_eight_step_shift_twelve_schedule_has_nine_grid_points(self):
        actual = turbo_sigmas(8, 12.0, device="cpu")
        self.assertEqual(actual.shape, (9,))
        self.assertEqual(actual[0].item(), 1.0)
        self.assertEqual(actual[-1].item(), 0.0)
        self.assertTrue(bool(torch.all(actual[:-1] > actual[1:])))

    def test_turbo_sampler_uses_positive_dataward_updates_and_independent_clocks(self):
        video = torch.zeros(1, 1, 1, 1, 1)
        audio = torch.zeros(1, 1, 1, 1)
        calls = []

        def packer(text, current_video, current_audio, sigma, model_timesteps, **kwargs):
            calls.append((current_video.clone(), current_audio.clone(), model_timesteps.clone()))
            return (None, None, None, None, None, None)

        def model(*args):
            return torch.ones(1), torch.ones(1)

        with (
            patch("h3_blackwell_runtime.sampler.unpatchify_video", return_value=torch.full_like(video, 2.0)),
            patch("h3_blackwell_runtime.sampler._unpack_audio", return_value=torch.full_like(audio, 3.0)),
        ):
            result_video, result_audio = sample_video_turbo(
                model, packer, torch.empty(0), video, audio,
                steps=2, video_shift=6.0, audio_shift=3.0, return_audio=True,
            )

        torch.testing.assert_close(result_video, torch.full_like(video, 2.0))
        torch.testing.assert_close(result_audio, torch.full_like(audio, 3.0))
        torch.testing.assert_close(calls[1][0], torch.full_like(video, 2.0 / 7.0))
        torch.testing.assert_close(calls[1][1], torch.full_like(audio, 0.75))
        torch.testing.assert_close(calls[1][2], torch.tensor([1.0 / 7.0, 0.25]))

    def test_base_sampler_records_opt_in_audio_step_trace(self):
        video = torch.zeros(1, 1, 1, 1, 1)
        audio = torch.zeros(1, 32, 2, 1)
        trace = []

        def packer(*args, **kwargs):
            return (None, None, None, None, None, None)

        def model(*args):
            return torch.ones(1), torch.ones(1)

        with (
            patch("h3_blackwell_runtime.sampler.beta_sigmas", return_value=torch.tensor([1.0, 0.0])),
            patch("h3_blackwell_runtime.sampler.unpatchify_video", return_value=torch.zeros_like(video)),
            patch("h3_blackwell_runtime.sampler._unpack_audio", return_value=torch.ones_like(audio)),
        ):
            _video, final_audio = sample_video_res_multistep(
                model, packer, torch.empty(0), video, audio,
                steps=1, return_audio=True, audio_step_trace=trace,
            )

        self.assertEqual(len(trace), 1)
        self.assertEqual(trace[0]["step"], 1)
        torch.testing.assert_close(trace[0]["audio_before"], torch.zeros_like(audio))
        torch.testing.assert_close(trace[0]["audio_denoised"], torch.full_like(audio, 4.0))
        torch.testing.assert_close(trace[0]["audio_after"], torch.full_like(audio, 4.0))
        torch.testing.assert_close(final_audio, torch.ones_like(audio))


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