import unittest

import torch

from diffusion_cli.text_encoder import TextConditioning
from diffusion_cli.zimage_model import (
    NextDiT,
    ZImageConfig,
    ZImageModel,
    detectZImageConfig,
    normalizeZImageStateDict,
)


class ZImageModelTest(unittest.TestCase):
    def testStateDictNormalizationStripsModelPrefix(self):
        state_dict = normalizeZImageStateDict(
            {
                "model.diffusion_model.x_embedder.weight": "a",
                "diffusion_model.final_layer.linear.weight": "b",
                "layers.0.ffn_norm1.weight": "c",
            }
        )

        self.assertEqual(state_dict["x_embedder.weight"], "a")
        self.assertEqual(state_dict["final_layer.linear.weight"], "b")
        self.assertEqual(state_dict["layers.0.ffn_norm1.weight"], "c")

    def testDetectConfigReadsZImageShapes(self):
        state_dict = {
            "model.diffusion_model.cap_embedder.1.weight": torch.empty(
                3840,
                2560,
            ),
            "model.diffusion_model.x_embedder.weight": torch.empty(3840, 64),
            "model.diffusion_model.cap_pad_token": torch.empty(1, 3840),
            "model.diffusion_model.layers.0.ffn_norm1.weight": torch.empty(
                3840,
            ),
            "model.diffusion_model.layers.1.ffn_norm1.weight": torch.empty(
                3840,
            ),
            "model.diffusion_model.context_refiner.0.ffn_norm1.weight":
                torch.empty(3840),
        }

        config = detectZImageConfig(state_dict)

        self.assertEqual(config.dim, 3840)
        self.assertEqual(config.cap_feat_dim, 2560)
        self.assertEqual(config.n_layers, 2)
        self.assertEqual(config.n_refiner_layers, 1)
        self.assertEqual(config.patch_size, 2)
        self.assertEqual(config.pad_tokens_multiple, 32)

    def testTinyForwardMatchesLatentShapeAndIsFinite(self):
        torch.manual_seed(1)
        config = ZImageConfig(
            dim=258,
            cap_feat_dim=8,
            n_layers=1,
            n_refiner_layers=1,
            n_heads=3,
            n_kv_heads=3,
            patch_size=2,
            in_channels=4,
            axes_dims=(30, 28, 28),
            pad_tokens_multiple=0,
        )
        model = NextDiT(config, dtype=torch.float32)
        latent = torch.randn(1, 4, 4, 4)
        context = torch.randn(1, 3, 8)
        timesteps = torch.tensor([0.5])

        output = model(latent, timesteps, context, 3)

        self.assertEqual(output.shape, latent.shape)
        self.assertTrue(torch.all(torch.isfinite(output)))

    def testWrapperValidatesFiniteOutput(self):
        class IdentityModel(torch.nn.Module):
            def forward(
                self,
                latent,
                timesteps,
                context,
                num_tokens,
                attention_mask,
            ):
                return latent + timesteps.reshape(-1, 1, 1, 1) * 0

        wrapper = ZImageModel(
            model_path=None,
            device=torch.device("cpu"),
            dtype=torch.float32,
            model=IdentityModel(),
        )
        conditioning = TextConditioning(
            hidden_states=torch.zeros(1, 2, 8),
            attention_mask=torch.ones(1, 2),
            token_count=2,
        )
        latent = torch.zeros(1, 4, 4, 4)

        output = wrapper.forward(latent, torch.tensor([0.5]), conditioning)

        self.assertEqual(output.shape, latent.shape)
        self.assertTrue(torch.all(torch.isfinite(output)))

    def testWrapperConvertsFlowOutputToDenoisedPrediction(self):
        class ConstantVelocityModel(torch.nn.Module):
            def forward(
                self,
                latent,
                timesteps,
                context,
                num_tokens,
                attention_mask,
            ):
                del timesteps, context, num_tokens, attention_mask
                return torch.ones_like(latent) * 2.0

        wrapper = ZImageModel(
            model_path=None,
            device=torch.device("cpu"),
            dtype=torch.float32,
            model=ConstantVelocityModel(),
        )
        conditioning = TextConditioning(
            hidden_states=torch.zeros(1, 2, 8),
            attention_mask=torch.ones(1, 2),
            token_count=2,
        )
        latent = torch.ones(1, 4, 2, 2) * 3.0

        output = wrapper.forward(latent, torch.tensor([0.25]), conditioning)

        self.assertTrue(torch.equal(output, torch.ones_like(latent) * 2.5))


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