"""Tests for models/acestep/convert.py — DiT and VAE weight conversion."""

import numpy as np
import pytest

from mlx_music.models.acestep.convert import (
    convert_dit_weights,
    convert_vae_weights,
    fuse_weight_norm,
    is_dit_conv1d_weight,
    is_dit_conv_transpose_weight,
    remap_dit_key,
    reshape_snake_param,
    should_skip_dit_key,
)


# ---------------------------------------------------------------------------
# DiT key skipping
# ---------------------------------------------------------------------------


class TestShouldSkipDitKey:
    """Test skip logic for Lambda wrappers, rotary embeddings, tokenizer, etc."""

    def test_skip_lambda_proj_in_0(self):
        assert should_skip_dit_key("decoder.proj_in.0.some_param")

    def test_skip_lambda_proj_in_2(self):
        assert should_skip_dit_key("decoder.proj_in.2.some_param")

    def test_skip_lambda_proj_out_0(self):
        assert should_skip_dit_key("decoder.proj_out.0.some_param")

    def test_skip_lambda_proj_out_2(self):
        assert should_skip_dit_key("decoder.proj_out.2.some_param")

    def test_skip_decoder_rotary_emb(self):
        assert should_skip_dit_key("decoder.rotary_emb.inv_freq")

    def test_skip_lyric_encoder_rotary_emb(self):
        assert should_skip_dit_key("encoder.lyric_encoder.rotary_emb.inv_freq")

    def test_skip_timbre_encoder_rotary_emb(self):
        assert should_skip_dit_key("encoder.timbre_encoder.rotary_emb.inv_freq")

    def test_skip_tokenizer(self):
        assert should_skip_dit_key("tokenizer.audio_acoustic_proj.weight")
        assert should_skip_dit_key(
            "tokenizer.attention_pooler.layers.0.self_attn.q_proj.weight"
        )

    def test_skip_detokenizer(self):
        assert should_skip_dit_key("detokenizer.embed_tokens.weight")
        assert should_skip_dit_key("detokenizer.layers.0.self_attn.q_proj.weight")

    def test_keep_proj_in_conv(self):
        """The Conv1d at index 1 should NOT be skipped."""
        assert not should_skip_dit_key("decoder.proj_in.1.weight")
        assert not should_skip_dit_key("decoder.proj_in.1.bias")

    def test_keep_proj_out_conv(self):
        assert not should_skip_dit_key("decoder.proj_out.1.weight")
        assert not should_skip_dit_key("decoder.proj_out.1.bias")

    def test_keep_attention_weights(self):
        assert not should_skip_dit_key("decoder.layers.0.self_attn.q_proj.weight")
        assert not should_skip_dit_key("decoder.layers.23.cross_attn.o_proj.weight")

    def test_keep_time_embed(self):
        assert not should_skip_dit_key("decoder.time_embed.linear_1.weight")
        assert not should_skip_dit_key("decoder.time_embed_r.time_proj.weight")

    def test_keep_encoder_weights(self):
        assert not should_skip_dit_key("encoder.text_projector.weight")
        assert not should_skip_dit_key(
            "encoder.lyric_encoder.layers.0.self_attn.q_proj.weight"
        )
        assert not should_skip_dit_key("encoder.lyric_encoder.embed_tokens.weight")
        assert not should_skip_dit_key("encoder.timbre_encoder.special_token")

    def test_keep_null_condition_emb(self):
        assert not should_skip_dit_key("null_condition_emb")


# ---------------------------------------------------------------------------
# DiT key remapping
# ---------------------------------------------------------------------------


class TestRemapDitKey:
    """Test Sequential index removal for proj_in/proj_out."""

    def test_proj_in_weight(self):
        assert remap_dit_key("decoder.proj_in.1.weight") == "decoder.proj_in.weight"

    def test_proj_in_bias(self):
        assert remap_dit_key("decoder.proj_in.1.bias") == "decoder.proj_in.bias"

    def test_proj_out_weight(self):
        assert remap_dit_key("decoder.proj_out.1.weight") == "decoder.proj_out.weight"

    def test_proj_out_bias(self):
        assert remap_dit_key("decoder.proj_out.1.bias") == "decoder.proj_out.bias"

    def test_passthrough_attention(self):
        key = "decoder.layers.5.self_attn.q_proj.weight"
        assert remap_dit_key(key) == key

    def test_passthrough_mlp(self):
        key = "decoder.layers.0.mlp.gate_proj.weight"
        assert remap_dit_key(key) == key

    def test_passthrough_scale_shift(self):
        key = "decoder.layers.0.scale_shift_table"
        assert remap_dit_key(key) == key

    def test_passthrough_time_embed(self):
        key = "decoder.time_embed.linear_1.weight"
        assert remap_dit_key(key) == key

    def test_passthrough_encoder(self):
        key = "encoder.lyric_encoder.layers.0.self_attn.q_proj.weight"
        assert remap_dit_key(key) == key

    def test_passthrough_null_condition(self):
        key = "null_condition_emb"
        assert remap_dit_key(key) == key


# ---------------------------------------------------------------------------
# DiT Conv1d detection
# ---------------------------------------------------------------------------


class TestIsDitConvWeight:
    def test_proj_in_is_conv1d(self):
        assert is_dit_conv1d_weight("decoder.proj_in.1.weight")

    def test_proj_out_is_conv_transpose(self):
        assert is_dit_conv_transpose_weight("decoder.proj_out.1.weight")
        assert not is_dit_conv1d_weight("decoder.proj_out.1.weight")

    def test_proj_in_bias_not_conv(self):
        """Bias is 1D, should not be flagged as conv weight."""
        assert not is_dit_conv1d_weight("decoder.proj_in.1.bias")
        assert not is_dit_conv_transpose_weight("decoder.proj_in.1.bias")

    def test_linear_weight_not_conv(self):
        assert not is_dit_conv1d_weight("decoder.layers.0.self_attn.q_proj.weight")

    def test_time_embed_not_conv(self):
        assert not is_dit_conv1d_weight("decoder.time_embed.linear_1.weight")


# ---------------------------------------------------------------------------
# DiT full conversion
# ---------------------------------------------------------------------------


class TestConvertDitWeights:
    """Integration test for the full DiT conversion pipeline."""

    def _make_dit_weights(self):
        """Create a minimal mock DiT weight dictionary."""
        return {
            # Conv1d in Sequential (should be remapped + transposed)
            "decoder.proj_in.1.weight": np.random.randn(2048, 192, 2).astype(
                np.float32
            ),
            "decoder.proj_in.1.bias": np.random.randn(2048).astype(np.float32),
            # ConvTranspose1d in Sequential (should be remapped + transposed)
            # PyTorch ConvTranspose1d stores as [in_channels, out_channels, kernel_size]
            "decoder.proj_out.1.weight": np.random.randn(2048, 64, 2).astype(
                np.float32
            ),
            "decoder.proj_out.1.bias": np.random.randn(64).astype(np.float32),
            # Lambda wrappers (should be skipped)
            # (Lambda modules have no parameters, but just in case)
            # Rotary embedding (should be skipped)
            "decoder.rotary_emb.inv_freq": np.random.randn(64).astype(np.float32),
            "encoder.lyric_encoder.rotary_emb.inv_freq": np.random.randn(64).astype(
                np.float32
            ),
            # Tokenizer (should be skipped)
            "tokenizer.audio_acoustic_proj.weight": np.random.randn(2048, 64).astype(
                np.float32
            ),
            # Attention weights (should pass through unchanged)
            "decoder.layers.0.self_attn.q_proj.weight": np.random.randn(
                2048, 2048
            ).astype(np.float32),
            "decoder.layers.0.self_attn.q_norm.weight": np.random.randn(128).astype(
                np.float32
            ),
            # MLP weights (should pass through unchanged)
            "decoder.layers.0.mlp.gate_proj.weight": np.random.randn(6144, 2048).astype(
                np.float32
            ),
            # Time embed (should pass through unchanged)
            "decoder.time_embed.linear_1.weight": np.random.randn(2048, 256).astype(
                np.float32
            ),
            "decoder.time_embed.linear_1.bias": np.random.randn(2048).astype(
                np.float32
            ),
            # Scale-shift table (should pass through unchanged)
            "decoder.layers.0.scale_shift_table": np.random.randn(1, 6, 2048).astype(
                np.float32
            ),
            # Encoder weights
            "encoder.text_projector.weight": np.random.randn(2048, 1024).astype(
                np.float32
            ),
            "encoder.lyric_encoder.embed_tokens.weight": np.random.randn(
                2048, 1024
            ).astype(np.float32),
            # Null condition
            "null_condition_emb": np.random.randn(1, 1, 2048).astype(np.float32),
        }

    def test_skips_lambda_and_rotary_keys(self):
        weights = self._make_dit_weights()
        result = convert_dit_weights(weights)
        # Should NOT contain rotary or tokenizer keys
        for key in result:
            assert "rotary_emb" not in key
            assert "tokenizer" not in key
            assert "detokenizer" not in key

    def test_remaps_proj_in_key(self):
        weights = self._make_dit_weights()
        result = convert_dit_weights(weights)
        assert "decoder.proj_in.weight" in result
        assert "decoder.proj_in.bias" in result
        assert "decoder.proj_in.1.weight" not in result

    def test_remaps_proj_out_key(self):
        weights = self._make_dit_weights()
        result = convert_dit_weights(weights)
        assert "decoder.proj_out.weight" in result
        assert "decoder.proj_out.bias" in result
        assert "decoder.proj_out.1.weight" not in result

    def test_transposes_proj_in_weight(self):
        weights = self._make_dit_weights()
        result = convert_dit_weights(weights)
        # Original: [2048, 192, 2] -> transposed: [2048, 2, 192]
        w = result["decoder.proj_in.weight"]
        assert w.shape == (2048, 2, 192), f"Expected (2048, 2, 192), got {w.shape}"

    def test_transposes_proj_out_weight(self):
        weights = self._make_dit_weights()
        result = convert_dit_weights(weights)
        # ConvTranspose1d: PyTorch [I=2048, O=64, K=2] -> MLX [O=64, K=2, I=2048]
        w = result["decoder.proj_out.weight"]
        assert w.shape == (64, 2, 2048), f"Expected (64, 2, 2048), got {w.shape}"

    def test_does_not_transpose_bias(self):
        weights = self._make_dit_weights()
        result = convert_dit_weights(weights)
        # Bias should stay 1D
        assert result["decoder.proj_in.bias"].shape == (2048,)

    def test_passes_through_attention_weights(self):
        weights = self._make_dit_weights()
        result = convert_dit_weights(weights)
        key = "decoder.layers.0.self_attn.q_proj.weight"
        assert key in result
        assert result[key].shape == (2048, 2048)

    def test_passes_through_scale_shift_table(self):
        weights = self._make_dit_weights()
        result = convert_dit_weights(weights)
        key = "decoder.layers.0.scale_shift_table"
        assert key in result
        assert result[key].shape == (1, 6, 2048)

    def test_passes_through_null_condition_emb(self):
        weights = self._make_dit_weights()
        result = convert_dit_weights(weights)
        assert "null_condition_emb" in result
        assert result["null_condition_emb"].shape == (1, 1, 2048)


# ---------------------------------------------------------------------------
# VAE weight_norm fusion
# ---------------------------------------------------------------------------


class TestFuseWeightNorm:
    """Test weight_norm fusion: W = g * (v / ||v||)."""

    def test_basic_fusion_identity(self):
        """When v is already unit-norm per channel, g directly scales it."""
        # v has unit norm per output channel
        v = np.array([[[1.0, 0.0, 0.0]]], dtype=np.float32)  # [1, 1, 3], norm=1
        g = np.array([[[2.0]]], dtype=np.float32)  # [1, 1, 1]
        result = fuse_weight_norm(g, v)
        expected = np.array([[[2.0, 0.0, 0.0]]], dtype=np.float32)
        np.testing.assert_allclose(result, expected, atol=1e-6)

    def test_fusion_normalizes_direction(self):
        """The direction vector should be normalized before scaling."""
        # v = [3, 4] has norm 5
        v = np.array([[[3.0, 4.0]]], dtype=np.float32)  # [1, 1, 2]
        g = np.array([[[5.0]]], dtype=np.float32)  # [1, 1, 1]
        result = fuse_weight_norm(g, v)
        # Expected: 5 * [3/5, 4/5] = [3, 4]
        expected = np.array([[[3.0, 4.0]]], dtype=np.float32)
        np.testing.assert_allclose(result, expected, atol=1e-6)

    def test_fusion_multiple_output_channels(self):
        """Test with multiple output channels (axis 0)."""
        # 2 output channels, 3 input, kernel 1
        v = np.array(
            [
                [[3.0, 4.0, 0.0]],  # norm = 5
                [[0.0, 0.0, 1.0]],  # norm = 1
            ],
            dtype=np.float32,
        )  # [2, 1, 3]
        g = np.array([[[10.0]], [[2.0]]], dtype=np.float32)  # [2, 1, 1]
        result = fuse_weight_norm(g, v)
        # Channel 0: 10 * [3/5, 4/5, 0] = [6, 8, 0]
        # Channel 1: 2 * [0, 0, 1] = [0, 0, 2]
        expected = np.array(
            [
                [[6.0, 8.0, 0.0]],
                [[0.0, 0.0, 2.0]],
            ],
            dtype=np.float32,
        )
        np.testing.assert_allclose(result, expected, atol=1e-6)

    def test_fusion_1d_g(self):
        """Test with weight_g as 1D vector [O] (common in saved state dicts)."""
        v = np.random.randn(4, 3, 5).astype(np.float32)
        g = np.random.randn(4).astype(np.float32)  # [O]
        result = fuse_weight_norm(g, v)
        assert result.shape == v.shape

        # Verify per-channel: result[i] = g[i] * v[i] / ||v[i]||
        for i in range(4):
            v_norm = np.sqrt(np.sum(v[i] ** 2))
            expected_i = g[i] * v[i] / v_norm
            np.testing.assert_allclose(result[i], expected_i, atol=1e-5)

    def test_fusion_preserves_shape(self):
        """Output should have same shape as weight_v."""
        v = np.random.randn(16, 8, 7).astype(np.float32)
        g = np.random.randn(16, 1, 1).astype(np.float32)
        result = fuse_weight_norm(g, v)
        assert result.shape == (16, 8, 7)


# ---------------------------------------------------------------------------
# Snake1d parameter reshaping
# ---------------------------------------------------------------------------


class TestReshapeSnakeParam:
    def test_squeeze_3d_to_1d(self):
        """PyTorch [1, C, 1] -> MLX [C]."""
        value = np.random.randn(1, 64, 1).astype(np.float32)
        result = reshape_snake_param(value)
        assert result.shape == (64,)
        np.testing.assert_allclose(result, value.squeeze())

    def test_already_1d(self):
        """If already 1D, squeeze is a no-op."""
        value = np.random.randn(64).astype(np.float32)
        result = reshape_snake_param(value)
        assert result.shape == (64,)

    def test_large_channel_dim(self):
        value = np.random.randn(1, 2048, 1).astype(np.float32)
        result = reshape_snake_param(value)
        assert result.shape == (2048,)


# ---------------------------------------------------------------------------
# VAE full conversion
# ---------------------------------------------------------------------------


class TestConvertVaeWeights:
    """Integration test for the full VAE conversion pipeline."""

    def _make_vae_weights(self):
        """Create a minimal mock VAE weight dictionary.

        Mimics the diffusers AutoencoderOobleck state_dict structure:
        all Conv1d layers have weight_g + weight_v pairs (weight_norm).
        """
        return {
            # Encoder conv1 (weight_norm)
            "encoder.conv1.weight_g": np.random.randn(128, 1, 1).astype(np.float32),
            "encoder.conv1.weight_v": np.random.randn(128, 2, 7).astype(np.float32),
            "encoder.conv1.bias": np.random.randn(128).astype(np.float32),
            # Encoder block residual conv (weight_norm)
            "encoder.block.0.res_unit1.conv1.weight_g": np.random.randn(
                128, 1, 1
            ).astype(np.float32),
            "encoder.block.0.res_unit1.conv1.weight_v": np.random.randn(
                128, 128, 7
            ).astype(np.float32),
            "encoder.block.0.res_unit1.conv1.bias": np.random.randn(128).astype(
                np.float32
            ),
            # Decoder conv_t1 (ConvTranspose1d, weight_norm)
            # PyTorch ConvTranspose1d stores as [in_channels, out_channels, kernel_size]
            "decoder.block.0.conv_t1.weight_g": np.random.randn(256, 1, 1).astype(
                np.float32
            ),
            "decoder.block.0.conv_t1.weight_v": np.random.randn(256, 128, 20).astype(
                np.float32
            ),
            "decoder.block.0.conv_t1.bias": np.random.randn(128).astype(np.float32),
            # Decoder output conv (weight_norm, no bias)
            "decoder.conv2.weight_g": np.random.randn(2, 1, 1).astype(np.float32),
            "decoder.conv2.weight_v": np.random.randn(2, 128, 7).astype(np.float32),
            # Snake1d alpha/beta (PyTorch [1, C, 1] format)
            "encoder.block.0.res_unit1.snake1.alpha": np.random.randn(1, 128, 1).astype(
                np.float32
            ),
            "encoder.block.0.res_unit1.snake1.beta": np.random.randn(1, 128, 1).astype(
                np.float32
            ),
            "decoder.snake1.alpha": np.random.randn(1, 128, 1).astype(np.float32),
            "decoder.snake1.beta": np.random.randn(1, 128, 1).astype(np.float32),
        }

    def test_fuses_weight_norm_to_single_weight(self):
        """weight_g + weight_v should become a single .weight key."""
        weights = self._make_vae_weights()
        result = convert_vae_weights(weights)
        # Should have .weight but NOT .weight_g or .weight_v
        assert "encoder.conv1.weight" in result
        assert "encoder.conv1.weight_g" not in result
        assert "encoder.conv1.weight_v" not in result

    def test_fused_weight_is_transposed(self):
        """Fused weight should be transposed [O, I, K] -> [O, K, I]."""
        weights = self._make_vae_weights()
        result = convert_vae_weights(weights)
        # Original encoder.conv1: [128, 2, 7] -> transposed: [128, 7, 2]
        w = result["encoder.conv1.weight"]
        assert w.shape == (128, 7, 2), f"Expected (128, 7, 2), got {w.shape}"

    def test_residual_conv_transposed(self):
        weights = self._make_vae_weights()
        result = convert_vae_weights(weights)
        # res_unit1.conv1: [128, 128, 7] -> [128, 7, 128]
        w = result["encoder.block.0.res_unit1.conv1.weight"]
        assert w.shape == (128, 7, 128)

    def test_conv_transpose_1d_transposed(self):
        weights = self._make_vae_weights()
        result = convert_vae_weights(weights)
        # conv_t1: PyTorch [I=256, O=128, K=20] -> MLX [O=128, K=20, I=256]
        w = result["decoder.block.0.conv_t1.weight"]
        assert w.shape == (128, 20, 256), f"Expected (128, 20, 256), got {w.shape}"

    def test_bias_preserved(self):
        weights = self._make_vae_weights()
        result = convert_vae_weights(weights)
        assert "encoder.conv1.bias" in result
        assert result["encoder.conv1.bias"].shape == (128,)

    def test_snake_alpha_reshaped(self):
        """Snake1d alpha should be [1, C, 1] -> [C]."""
        weights = self._make_vae_weights()
        result = convert_vae_weights(weights)
        key = "encoder.block.0.res_unit1.snake1.alpha"
        assert key in result
        assert result[key].shape == (128,)

    def test_snake_beta_reshaped(self):
        weights = self._make_vae_weights()
        result = convert_vae_weights(weights)
        key = "encoder.block.0.res_unit1.snake1.beta"
        assert key in result
        assert result[key].shape == (128,)

    def test_no_output_conv_bias_when_none(self):
        """decoder.conv2 has no bias — should not appear in output."""
        weights = self._make_vae_weights()
        result = convert_vae_weights(weights)
        assert "decoder.conv2.bias" not in result

    def test_fusion_math_correctness(self):
        """Verify the fused weight value matches manual computation."""
        # Use simple values for hand-verification
        weights = {
            "test.weight_g": np.array([[[5.0]]], dtype=np.float32),  # [1, 1, 1]
            "test.weight_v": np.array([[[3.0, 4.0]]], dtype=np.float32),  # [1, 1, 2]
        }
        result = convert_vae_weights(weights)
        # v_norm = sqrt(3^2 + 4^2) = 5
        # fused = 5 * [3/5, 4/5] = [3, 4]
        # then transposed: [1, 1, 2] -> [1, 2, 1]
        w = np.array(result["test.weight"])
        expected = np.array(
            [[[3.0], [4.0]]], dtype=np.float32
        )  # [1, 2, 1] after transpose
        np.testing.assert_allclose(w, expected, atol=1e-5)

    def test_mismatched_weight_norm_raises(self):
        """If weight_g exists without weight_v (or vice versa), should raise."""
        weights = {
            "broken.weight_g": np.array([[[1.0]]], dtype=np.float32),
            # No matching weight_v
        }
        with pytest.raises(ValueError, match="weight_v"):
            convert_vae_weights(weights)


# ---------------------------------------------------------------------------
# Integration: key name correctness against MLX model tree
# ---------------------------------------------------------------------------


class TestKeyNameAlignment:
    """Verify that converted keys match expected MLX model parameter paths.

    These keys are derived from the MLX model classes in:
      - dit.py (AceStepConditionGenerationModel, AceStepDiTModel, etc.)
      - vae.py (AutoencoderOobleck, OobleckEncoder, OobleckDecoder, etc.)
    """

    # DiT keys that should be present after conversion
    EXPECTED_DIT_KEYS = [
        # Patch embed/unembed (after remapping from Sequential .1.)
        "decoder.proj_in.weight",
        "decoder.proj_in.bias",
        "decoder.proj_out.weight",
        "decoder.proj_out.bias",
        # DiT layer attention
        "decoder.layers.0.self_attn.q_proj.weight",
        "decoder.layers.0.self_attn.k_proj.weight",
        "decoder.layers.0.self_attn.v_proj.weight",
        "decoder.layers.0.self_attn.o_proj.weight",
        "decoder.layers.0.self_attn.q_norm.weight",
        "decoder.layers.0.self_attn.k_norm.weight",
        # DiT layer cross-attention
        "decoder.layers.0.cross_attn.q_proj.weight",
        "decoder.layers.0.cross_attn.k_proj.weight",
        "decoder.layers.0.cross_attn.v_proj.weight",
        "decoder.layers.0.cross_attn.o_proj.weight",
        # DiT layer norms
        "decoder.layers.0.self_attn_norm.weight",
        "decoder.layers.0.cross_attn_norm.weight",
        "decoder.layers.0.mlp_norm.weight",
        # DiT layer MLP
        "decoder.layers.0.mlp.gate_proj.weight",
        "decoder.layers.0.mlp.up_proj.weight",
        "decoder.layers.0.mlp.down_proj.weight",
        # DiT layer AdaLN
        "decoder.layers.0.scale_shift_table",
        # Timestep embeddings
        "decoder.time_embed.linear_1.weight",
        "decoder.time_embed.linear_1.bias",
        "decoder.time_embed.linear_2.weight",
        "decoder.time_embed.linear_2.bias",
        "decoder.time_embed.time_proj.weight",
        "decoder.time_embed.time_proj.bias",
        # Output AdaLN
        "decoder.norm_out.weight",
        "decoder.scale_shift_table",
        # Condition encoder
        "encoder.text_projector.weight",
        "encoder.lyric_encoder.embed_tokens.weight",
        "encoder.lyric_encoder.embed_tokens.bias",
        "encoder.lyric_encoder.norm.weight",
        "encoder.lyric_encoder.layers.0.self_attn.q_proj.weight",
        "encoder.lyric_encoder.layers.0.input_layernorm.weight",
        "encoder.lyric_encoder.layers.0.post_attention_layernorm.weight",
        "encoder.lyric_encoder.layers.0.mlp.gate_proj.weight",
        # Timbre encoder
        "encoder.timbre_encoder.embed_tokens.weight",
        "encoder.timbre_encoder.embed_tokens.bias",
        "encoder.timbre_encoder.norm.weight",
        "encoder.timbre_encoder.special_token",
        "encoder.timbre_encoder.layers.0.self_attn.q_proj.weight",
        # Null condition
        "null_condition_emb",
    ]

    # VAE keys that should be present after conversion
    EXPECTED_VAE_KEYS = [
        # Encoder
        "encoder.conv1.weight",
        "encoder.conv1.bias",
        "encoder.conv2.weight",
        "encoder.conv2.bias",
        "encoder.snake1.alpha",
        "encoder.snake1.beta",
        "encoder.block.0.res_unit1.conv1.weight",
        "encoder.block.0.res_unit1.conv1.bias",
        "encoder.block.0.res_unit1.snake1.alpha",
        "encoder.block.0.res_unit1.snake1.beta",
        "encoder.block.0.conv1.weight",
        "encoder.block.0.conv1.bias",
        # Decoder
        "decoder.conv1.weight",
        "decoder.conv1.bias",
        "decoder.conv2.weight",
        "decoder.snake1.alpha",
        "decoder.snake1.beta",
        "decoder.block.0.conv_t1.weight",
        "decoder.block.0.conv_t1.bias",
        "decoder.block.0.res_unit1.conv1.weight",
    ]

    def test_dit_key_names_documented(self):
        """Sanity check that expected DiT key names are valid."""
        # This test documents the expected key structure rather than testing code.
        # It will fail if someone accidentally changes the key list.
        assert len(self.EXPECTED_DIT_KEYS) > 30, "Expected many DiT keys"
        for key in self.EXPECTED_DIT_KEYS:
            # Basic sanity: no double dots, no leading/trailing dots
            assert ".." not in key, f"Invalid key: {key}"
            assert not key.startswith("."), f"Invalid key: {key}"
            assert not key.endswith("."), f"Invalid key: {key}"

    def test_vae_key_names_documented(self):
        """Sanity check that expected VAE key names are valid."""
        assert len(self.EXPECTED_VAE_KEYS) > 10, "Expected many VAE keys"
        for key in self.EXPECTED_VAE_KEYS:
            assert ".." not in key
            assert not key.startswith(".")
