"""Tests for Wan2.2 I2V-14B support."""

import mlx.core as mx
import numpy as np
import pytest

from wan_test_helpers import _make_tiny_config


def _make_tiny_i2v_config():
    """Create a tiny I2V-14B config for testing."""
    config = _make_tiny_config()
    config.model_type = "i2v"
    config.in_dim = 9  # 4 noise + 4 image + 1 mask (scaled down from 16+16+4=36)
    config.out_dim = 4
    config.vae_z_dim = 4
    config.vae_stride = (4, 8, 8)
    config.dual_model = True
    config.boundary = 0.900
    config.sample_shift = 5.0
    config.sample_guide_scale = (3.5, 3.5)
    config.teacache_coefficients = None
    return config


class TestI2VConfig:
    """Test I2V-14B config preset."""

    def test_wan22_i2v_14b_preset(self):
        from mlx_video.models.wan.config import WanModelConfig

        config = WanModelConfig.wan22_i2v_14b()
        assert config.model_type == "i2v"
        assert config.in_dim == 36
        assert config.out_dim == 16
        assert config.dim == 5120
        assert config.num_layers == 40
        assert config.dual_model is True
        assert config.boundary == 0.900
        assert config.sample_shift == 5.0
        assert config.sample_guide_scale == (3.5, 3.5)
        assert config.vae_stride == (4, 8, 8)
        assert config.vae_z_dim == 16
        assert config.teacache_coefficients is None

    def test_i2v_vs_t2v_differences(self):
        from mlx_video.models.wan.config import WanModelConfig

        i2v = WanModelConfig.wan22_i2v_14b()
        t2v = WanModelConfig.wan22_t2v_14b()

        assert i2v.model_type == "i2v"
        assert t2v.model_type == "t2v"
        assert i2v.in_dim == 36 and t2v.in_dim == 16
        assert i2v.boundary == 0.900 and t2v.boundary == 0.875
        assert i2v.sample_shift == 5.0 and t2v.sample_shift == 12.0

    def test_i2v_serialization_roundtrip(self):
        from mlx_video.models.wan.config import WanModelConfig

        config = WanModelConfig.wan22_i2v_14b()
        d = config.to_dict()
        restored = WanModelConfig.from_dict(d)
        assert restored.model_type == "i2v"
        assert restored.in_dim == 36
        assert restored.boundary == 0.900


class TestModelYParameter:
    """Test y parameter channel concatenation in WanModel."""

    def test_forward_without_y(self):
        """Standard T2V forward pass (no y) still works."""
        from mlx_video.models.wan.model import WanModel

        config = _make_tiny_config()
        model = WanModel(config)

        C, F, H, W = config.in_dim, 1, 4, 4
        pt, ph, pw = config.patch_size
        seq_len = (F // pt) * (H // ph) * (W // pw)

        x_list = [mx.random.normal((C, F, H, W))]
        t = mx.array([500.0])
        context = [mx.random.normal((6, config.text_dim))]

        out = model(x_list, t, context, seq_len)
        mx.eval(out[0])
        assert out[0].shape == (C, F, H, W)

    def test_forward_with_y(self):
        """I2V forward pass with y channel concatenation."""
        from mlx_video.models.wan.model import WanModel

        config = _make_tiny_i2v_config()
        model = WanModel(config)

        C_noise = 4  # noise channels
        C_y = 5  # mask (1) + image (4)
        F, H, W = 1, 4, 4
        pt, ph, pw = config.patch_size
        seq_len = (F // pt) * (H // ph) * (W // pw)

        x_list = [mx.random.normal((C_noise, F, H, W))]
        y_list = [mx.random.normal((C_y, F, H, W))]
        t = mx.array([500.0])
        context = [mx.random.normal((6, config.text_dim))]

        out = model(x_list, t, context, seq_len, y=y_list)
        mx.eval(out[0])
        # Output should match noise channels (out_dim), not concatenated in_dim
        assert out[0].shape == (config.out_dim, F, H, W)

    def test_y_none_is_noop(self):
        """Passing y=None should be identical to not passing y."""
        from mlx_video.models.wan.model import WanModel

        config = _make_tiny_config()
        model = WanModel(config)

        C, F, H, W = config.in_dim, 1, 4, 4
        pt, ph, pw = config.patch_size
        seq_len = (F // pt) * (H // ph) * (W // pw)

        mx.random.seed(42)
        x = mx.random.normal((C, F, H, W))
        t = mx.array([500.0])
        ctx = [mx.random.normal((6, config.text_dim))]

        out1 = model([x], t, ctx, seq_len)[0]
        out2 = model([x], t, ctx, seq_len, y=None)[0]
        mx.eval(out1, out2)
        assert mx.allclose(out1, out2, atol=1e-5).item()

    def test_batched_cfg_with_y(self):
        """Batched CFG (B=2) with y should work."""
        from mlx_video.models.wan.model import WanModel

        config = _make_tiny_i2v_config()
        model = WanModel(config)

        C_noise, C_y = 4, 5
        F, H, W = 1, 4, 4
        pt, ph, pw = config.patch_size
        seq_len = (F // pt) * (H // ph) * (W // pw)

        latents = mx.random.normal((C_noise, F, H, W))
        y = mx.random.normal((C_y, F, H, W))
        t = mx.array([500.0, 500.0])
        ctx = [mx.random.normal((6, config.text_dim)), mx.random.normal((6, config.text_dim))]

        out = model([latents, latents], t, ctx, seq_len, y=[y, y])
        mx.eval(out[0], out[1])
        assert len(out) == 2
        assert out[0].shape == (config.out_dim, F, H, W)
        assert out[1].shape == (config.out_dim, F, H, W)


class TestVAEEncoder:
    """Test Wan2.1 VAE encoder."""

    def test_encoder3d_instantiation(self):
        from mlx_video.models.wan.vae import Encoder3d

        enc = Encoder3d(dim=32, z_dim=8)  # z_dim=8 (will output 8ch, but WanVAE wraps with z*2)
        assert enc.conv1 is not None
        assert len(enc.downsamples) > 0
        assert len(enc.middle) == 3

    def test_encoder3d_output_shape(self):
        """Encoder should downsample spatially by 8x and temporally by 4x."""
        from mlx_video.models.wan.vae import Encoder3d

        enc = Encoder3d(dim=32, z_dim=8)
        # Random input: [B=1, 3, T=5, H=32, W=32]
        x = mx.random.normal((1, 3, 5, 32, 32))
        out = enc(x)
        mx.eval(out)
        # With default dim_mult=[1,2,4,4] and temporal_downsample=[True,True,False]:
        # Spatial: 32 -> 16 -> 8 -> 4 (3 spatial downsamples)
        # Temporal: 5 -> 3 -> 2 (2 temporal downsamples: downsample3d stride 2)
        assert out.shape[0] == 1
        assert out.shape[1] == 8  # z_dim
        assert out.shape[3] == 32 // 8  # spatial /8
        assert out.shape[4] == 32 // 8

    def test_wan_vae_encode(self):
        """WanVAE with encoder=True should produce normalized latents."""
        from mlx_video.models.wan.vae import WanVAE

        vae = WanVAE(z_dim=16, encoder=True)
        # Input: [B=1, 3, T=5, H=32, W=32]
        x = mx.random.normal((1, 3, 5, 32, 32))
        z = vae.encode(x)
        mx.eval(z)
        assert z.shape[0] == 1
        assert z.shape[1] == 16  # z_dim

    def test_wan_vae_encoder_flag(self):
        """WanVAE without encoder flag should not have encoder attribute."""
        from mlx_video.models.wan.vae import WanVAE

        vae_no_enc = WanVAE(z_dim=4, encoder=False)
        assert not hasattr(vae_no_enc, 'encoder')

        vae_enc = WanVAE(z_dim=4, encoder=True)
        assert hasattr(vae_enc, 'encoder')


class TestResampleDownsample:
    """Test downsample modes in Resample."""

    def test_downsample2d(self):
        from mlx_video.models.wan.vae import Resample

        r = Resample(dim=16, mode="downsample2d")
        x = mx.random.normal((1, 16, 2, 8, 8))
        out = r(x)
        mx.eval(out)
        # Spatial /2, temporal unchanged, channels same
        assert out.shape == (1, 16, 2, 4, 4)

    def test_downsample3d(self):
        from mlx_video.models.wan.vae import Resample

        r = Resample(dim=16, mode="downsample3d")
        x = mx.random.normal((1, 16, 4, 8, 8))
        out = r(x)
        mx.eval(out)
        # Spatial /2, temporal /2, channels same
        assert out.shape == (1, 16, 2, 4, 4)

    def test_upsample2d_still_works(self):
        from mlx_video.models.wan.vae import Resample

        r = Resample(dim=16, mode="upsample2d")
        x = mx.random.normal((1, 16, 2, 4, 4))
        out = r(x)
        mx.eval(out)
        assert out.shape == (1, 8, 2, 8, 8)

    def test_upsample3d_still_works(self):
        from mlx_video.models.wan.vae import Resample

        r = Resample(dim=16, mode="upsample3d")
        x = mx.random.normal((1, 16, 2, 4, 4))
        out = r(x)
        mx.eval(out)
        assert out.shape == (1, 8, 4, 8, 8)


class TestI2VMaskConstruction:
    """Test mask construction for I2V-14B."""

    def test_mask_shape(self):
        """I2V-14B mask should have 4 channels with correct temporal structure."""
        num_frames = 81
        h_latent, w_latent = 10, 18  # example latent dims
        t_latent = (num_frames - 1) // 4 + 1  # = 21

        # Build mask following reference logic
        msk = mx.ones((1, num_frames, h_latent, w_latent))
        msk = mx.concatenate([msk[:, :1], mx.zeros((1, num_frames - 1, h_latent, w_latent))], axis=1)
        msk = mx.concatenate([mx.repeat(msk[:, :1], 4, axis=1), msk[:, 1:]], axis=1)
        msk = msk.reshape(1, msk.shape[1] // 4, 4, h_latent, w_latent)
        msk = msk.transpose(0, 2, 1, 3, 4)[0]  # [4, T_lat, H_lat, W_lat]

        assert msk.shape == (4, t_latent, h_latent, w_latent)

    def test_mask_values(self):
        """First temporal position should be 1, rest 0."""
        num_frames = 9
        h_latent, w_latent = 4, 4
        t_latent = (num_frames - 1) // 4 + 1  # = 3

        msk = mx.ones((1, num_frames, h_latent, w_latent))
        msk = mx.concatenate([msk[:, :1], mx.zeros((1, num_frames - 1, h_latent, w_latent))], axis=1)
        msk = mx.concatenate([mx.repeat(msk[:, :1], 4, axis=1), msk[:, 1:]], axis=1)
        msk = msk.reshape(1, msk.shape[1] // 4, 4, h_latent, w_latent)
        msk = msk.transpose(0, 2, 1, 3, 4)[0]

        mx.eval(msk)
        # First temporal position: all 4 channels should be 1
        assert mx.all(msk[:, 0] == 1.0).item()
        # Rest: all should be 0
        assert mx.all(msk[:, 1:] == 0.0).item()

    def test_y_tensor_shape(self):
        """y = concat([mask_4ch, encoded_video_16ch]) should be 20 channels."""
        mask = mx.zeros((4, 5, 10, 18))
        encoded = mx.zeros((16, 5, 10, 18))
        y = mx.concatenate([mask, encoded], axis=0)
        assert y.shape == (20, 5, 10, 18)


# ---------------------------------------------------------------------------
# Integration: I2V end-to-end pipeline
# ---------------------------------------------------------------------------


class TestI2VEndToEndPipeline:
    """Full I2V pipeline: image → preprocess → VAE encode → y tensor → denoise → VAE decode."""

    def test_full_i2v_pipeline(self):
        """End-to-end I2V: synthetic image → VAE encode → build y → denoise → VAE decode."""
        from mlx_video.models.wan.model import WanModel
        from mlx_video.models.wan.scheduler import FlowMatchEulerScheduler
        from mlx_video.models.wan.vae import WanVAE

        mx.random.seed(0)

        # --- Tiny I2V model config (z_dim=16 to match VAE normalization stats) ---
        config = _make_tiny_i2v_config()
        config.vae_z_dim = 16
        config.out_dim = 16  # must match VAE z_dim for decode
        config.in_dim = 16 + 4 + 16  # noise(out_dim=16) + mask(4) + image(z_dim=16) = 36
        model = WanModel(config)

        # --- Tiny VAE (with encoder) ---
        vae = WanVAE(z_dim=config.vae_z_dim, encoder=True)

        # --- Synthetic image: [B=1, 3, T=1, H=32, W=32] in [-1, 1] ---
        height, width = 32, 32
        num_frames = 5  # small temporal extent
        img = mx.random.uniform(-1, 1, (1, 3, 1, height, width))

        # Build video: first frame = image, rest = zeros -> [1, 3, F, H, W]
        video = mx.concatenate([
            img,
            mx.zeros((1, 3, num_frames - 1, height, width)),
        ], axis=2)

        # --- VAE encode ---
        z_video = vae.encode(video)  # [1, z_dim, T_lat, H_lat, W_lat]
        mx.eval(z_video)
        assert z_video.ndim == 5
        assert z_video.shape[1] == config.vae_z_dim

        z_video = z_video[0]  # [z_dim, T_lat, H_lat, W_lat]
        t_latent = z_video.shape[1]
        h_latent = z_video.shape[2]
        w_latent = z_video.shape[3]

        # --- Build I2V mask (4 channels) ---
        msk = mx.ones((1, num_frames, h_latent, w_latent))
        msk = mx.concatenate([msk[:, :1], mx.zeros((1, num_frames - 1, h_latent, w_latent))], axis=1)
        msk = mx.concatenate([mx.repeat(msk[:, :1], 4, axis=1), msk[:, 1:]], axis=1)
        msk = msk.reshape(1, msk.shape[1] // 4, 4, h_latent, w_latent)
        msk = msk.transpose(0, 2, 1, 3, 4)[0]  # [4, T_lat, H_lat, W_lat]

        # --- Build y tensor: [mask(4ch) + encoded(z_dim ch)] ---
        y_i2v = mx.concatenate([msk, z_video], axis=0)
        mx.eval(y_i2v)
        assert y_i2v.shape[0] == 4 + config.vae_z_dim

        # --- Denoising loop (2 steps) ---
        C_noise = config.out_dim  # noise channels
        pt, ph, pw = config.patch_size
        seq_len = (t_latent // pt) * (h_latent // ph) * (w_latent // pw)

        sched = FlowMatchEulerScheduler()
        num_steps = 2
        sched.set_timesteps(num_steps, shift=config.sample_shift)

        latents = mx.random.normal((C_noise, t_latent, h_latent, w_latent))
        context = mx.random.normal((4, config.text_dim))

        for i in range(num_steps):
            t_val = sched.timesteps[i].item()
            pred = model(
                [latents],
                mx.array([t_val]),
                [context],
                seq_len,
                y=[y_i2v],
            )[0]
            latents = sched.step(pred[None], t_val, latents[None]).squeeze(0)
            mx.eval(latents)

        assert latents.shape == (C_noise, t_latent, h_latent, w_latent)
        assert not mx.any(mx.isnan(latents)).item(), "NaN in denoised latents"
        assert not mx.any(mx.isinf(latents)).item(), "Inf in denoised latents"

        # --- VAE decode ---
        decoded = vae.decode(latents[None])  # [1, 3, T_out, H_out, W_out]
        mx.eval(decoded)
        assert decoded.ndim == 5
        assert decoded.shape[0] == 1
        assert decoded.shape[1] == 3  # RGB output
        assert not mx.any(mx.isnan(decoded)).item(), "NaN in decoded video"
        assert not mx.any(mx.isinf(decoded)).item(), "Inf in decoded video"
        # VAE decode clips to [-1, 1]
        assert float(decoded.max()) <= 1.0
        assert float(decoded.min()) >= -1.0


class TestDualModelSwitching:
    """Test dual-model selection logic: high_noise vs low_noise based on boundary."""

    def test_model_selection_by_timestep(self):
        """Verify high_noise model used for timesteps >= boundary, low_noise otherwise."""
        from mlx_video.models.wan.model import WanModel
        from mlx_video.models.wan.scheduler import FlowMatchEulerScheduler

        mx.random.seed(1)
        config = _make_tiny_i2v_config()
        assert config.dual_model is True

        high_noise_model = WanModel(config)
        low_noise_model = WanModel(config)

        boundary = config.boundary * config.num_train_timesteps  # 0.9 * 1000 = 900

        C_noise = config.out_dim  # 4
        C_y = config.in_dim - config.out_dim  # 9 - 4 = 5
        F, H, W = 1, 4, 4
        pt, ph, pw = config.patch_size
        seq_len = (F // pt) * (H // ph) * (W // pw)

        sched = FlowMatchEulerScheduler()
        num_steps = 5
        sched.set_timesteps(num_steps, shift=config.sample_shift)

        guide_scale = config.sample_guide_scale  # (3.5, 3.5)
        assert isinstance(guide_scale, tuple) and len(guide_scale) == 2

        latents = mx.random.normal((C_noise, F, H, W))
        y_i2v = mx.random.normal((C_y, F, H, W))
        context = mx.random.normal((4, config.text_dim))

        high_used_steps = []
        low_used_steps = []

        timestep_list = sched.timesteps.tolist()
        for i in range(num_steps):
            timestep_val = timestep_list[i]

            if timestep_val >= boundary:
                model = high_noise_model
                gs = guide_scale[1]
                high_used_steps.append(i)
            else:
                model = low_noise_model
                gs = guide_scale[0]
                low_used_steps.append(i)

            # CFG pass: cond + uncond
            preds = model(
                [latents, latents],
                mx.array([timestep_val, timestep_val]),
                [context, context],
                seq_len,
                y=[y_i2v, y_i2v],
            )
            noise_pred_cond, noise_pred_uncond = preds[0], preds[1]
            noise_pred = noise_pred_uncond + gs * (noise_pred_cond - noise_pred_uncond)

            latents = sched.step(noise_pred[None], timestep_val, latents[None]).squeeze(0)
            mx.eval(latents)

        # With shift=5.0, early timesteps should be high (>=900), later ones low
        assert len(high_used_steps) > 0, "High-noise model was never selected"
        assert len(low_used_steps) > 0, "Low-noise model was never selected"
        # High-noise steps should come before low-noise steps (timesteps decrease)
        if high_used_steps and low_used_steps:
            assert max(high_used_steps) < min(low_used_steps) or \
                   min(high_used_steps) < max(low_used_steps), \
                   "Model switching should happen during the loop"

        assert latents.shape == (C_noise, F, H, W)
        assert not mx.any(mx.isnan(latents)).item()

    def test_guide_scale_tuple_applied_per_model(self):
        """Verify (low_gs, high_gs) tuple applies different scales per model."""
        from mlx_video.models.wan.model import WanModel
        from mlx_video.models.wan.scheduler import FlowMatchEulerScheduler

        mx.random.seed(2)
        config = _make_tiny_i2v_config()
        config.sample_guide_scale = (2.0, 5.0)  # distinct values

        model = WanModel(config)
        boundary = config.boundary * config.num_train_timesteps

        C_noise = config.out_dim
        F, H, W = 1, 4, 4
        pt, ph, pw = config.patch_size
        seq_len = (F // pt) * (H // ph) * (W // pw)

        sched = FlowMatchEulerScheduler()
        sched.set_timesteps(5, shift=config.sample_shift)

        latents = mx.random.normal((C_noise, F, H, W))
        context = mx.random.normal((4, config.text_dim))
        guide_scale = config.sample_guide_scale
        C_y = config.in_dim - config.out_dim  # y channels
        y_i2v = mx.random.normal((C_y, F, H, W))

        # Track which guide scale was used at each step
        gs_per_step = []

        timestep_list = sched.timesteps.tolist()
        for i in range(5):
            timestep_val = timestep_list[i]

            if timestep_val >= boundary:
                gs = guide_scale[1]  # high_gs = 5.0
            else:
                gs = guide_scale[0]  # low_gs = 2.0
            gs_per_step.append(gs)

            pred = model(
                [latents, latents],
                mx.array([timestep_val, timestep_val]),
                [context, context],
                seq_len,
                y=[y_i2v, y_i2v],
            )
            noise_pred = pred[1] + gs * (pred[0] - pred[1])
            latents = sched.step(noise_pred[None], timestep_val, latents[None]).squeeze(0)
            mx.eval(latents)

        # Verify both guide scales were used
        assert 5.0 in gs_per_step, "High guide scale (5.0) was never used"
        assert 2.0 in gs_per_step, "Low guide scale (2.0) was never used"
        # High gs should appear first (high timesteps come first)
        first_high = gs_per_step.index(5.0)
        last_low = len(gs_per_step) - 1 - gs_per_step[::-1].index(2.0)
        assert first_high < last_low, "High gs steps should precede low gs steps"

    def test_single_model_fallback_with_tuple_guide_scale(self):
        """When dual_model=False, guide_scale tuple should use first element."""
        from mlx_video.models.wan.model import WanModel
        from mlx_video.models.wan.scheduler import FlowMatchEulerScheduler

        mx.random.seed(3)
        config = _make_tiny_config()
        config.dual_model = False
        config.sample_guide_scale = (3.0, 5.0)

        model = WanModel(config)
        guide_scale = config.sample_guide_scale

        C, F, H, W = config.in_dim, 1, 4, 4
        pt, ph, pw = config.patch_size
        seq_len = (F // pt) * (H // ph) * (W // pw)

        sched = FlowMatchEulerScheduler()
        sched.set_timesteps(3, shift=3.0)

        latents = mx.random.normal((C, F, H, W))
        context = mx.random.normal((4, config.text_dim))

        # Mimic generate_wan.py single-model logic:
        # gs = guide_scale if isinstance(guide_scale, (int, float)) else guide_scale[0]
        gs = guide_scale if isinstance(guide_scale, (int, float)) else guide_scale[0]
        assert gs == 3.0, "Single model should use first element of guide_scale tuple"

        for i in range(3):
            t_val = sched.timesteps[i].item()
            pred = model(
                [latents, latents],
                mx.array([t_val, t_val]),
                [context, context],
                seq_len,
            )
            noise_pred = pred[1] + gs * (pred[0] - pred[1])
            latents = sched.step(noise_pred[None], t_val, latents[None]).squeeze(0)
            mx.eval(latents)

        assert latents.shape == (C, F, H, W)
        assert not mx.any(mx.isnan(latents)).item()