import unittest

import torch

from h3_blackwell_runtime.latent_upscaler import H3LatentResizer3D, _checkpoint_config, upscale_h3_latent
from h3_blackwell_runtime.runtime import normalize_upscale


class LatentUpscalerContracts(unittest.TestCase):
    def test_request_upscale_normalization(self):
        self.assertIsNone(normalize_upscale(None))
        self.assertIsNone(normalize_upscale("none"))
        self.assertIsNone(normalize_upscale(1))
        self.assertEqual(normalize_upscale("2"), 2.0)
        with self.assertRaisesRegex(ValueError, "greater than 1.0"):
            normalize_upscale(0.5)
        with self.assertRaisesRegex(ValueError, "at most 4.0"):
            normalize_upscale(4.1)
        with self.assertRaisesRegex(ValueError, "must be a number"):
            normalize_upscale({"scale": 2})

    def test_3d_model_preserves_time_and_scales_space(self):
        model = H3LatentResizer3D(in_blocks=1, out_blocks=1, channels=32, dropout=0, temporal_every=0).eval()
        latent = torch.randn(1, 24, 3, 2, 4)
        output = upscale_h3_latent(model, latent, scale=2)
        self.assertEqual(output.shape, (1, 24, 3, 4, 8))

    def test_invalid_latent_shape_is_rejected(self):
        model = H3LatentResizer3D(in_blocks=1, out_blocks=1, channels=32, temporal_every=0).eval()
        with self.assertRaisesRegex(ValueError, r"\[B,24,T,H,W\]"):
            upscale_h3_latent(model, torch.randn(1, 16, 3, 2, 4))

    def test_scale_one_preserves_latent_exactly(self):
        model = H3LatentResizer3D(in_blocks=1, out_blocks=1, channels=32, temporal_every=0).eval()
        latent = torch.randn(1, 24, 3, 2, 4)
        self.assertTrue(torch.equal(upscale_h3_latent(model, latent, scale=1), latent))

    def test_checkpoint_architecture_detection_matches_module_layout(self):
        model = H3LatentResizer3D(in_blocks=3, out_blocks=2, channels=32, temporal_every=2, temporal_kernel=3)
        config = _checkpoint_config(model.state_dict())
        self.assertEqual(config, {
            "in_channels": 24,
            "in_blocks": 3,
            "out_blocks": 2,
            "channels": 32,
            "temporal_every": 2,
            "temporal_kernel": 3,
        })


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