import unittest

import mlx.core as mx


class TestVoxCPMIntegration(unittest.TestCase):
    def test_sanitize_rope_generation(self):
        """Test strict loading logic for RoPE parameters."""
        from mlx_audio.tts.models.voxcpm import Model, ModelArgs
        from mlx_audio.tts.models.voxcpm.config import (
            AudioVAEConfig,
            DiTConfig,
            EncoderConfig,
            LMConfig,
        )

        args = ModelArgs(
            lm_config=LMConfig(
                num_hidden_layers=1,
                hidden_size=64,
                num_attention_heads=4,
                num_key_value_heads=2,
                intermediate_size=128,
            ),
            encoder_config=EncoderConfig(num_layers=1, hidden_dim=64),
            dit_config=DiTConfig(num_layers=1, hidden_dim=64),
            audio_vae_config=AudioVAEConfig(encoder_rates=[2], decoder_rates=[2]),
        )
        model = Model(args)

        # Test sanitizing an empty weight dict.
        # It should AUTO-POPULATE the RoPE parameters.
        weights = {}
        new_weights = model.sanitize(weights)

        # Check for RoPE keys
        expected_keys = [
            "base_lm.rope.inv_freq",
            "base_lm.rope.long_factor",
            "base_lm.rope.short_factor",
            "residual_lm.rope.inv_freq",
            "feat_encoder.encoder.rope.inv_freq",
        ]

        mx_array_type = type(mx.array(0))
        for k in expected_keys:
            self.assertIn(k, new_weights, f"Key {k} was not auto-generated by sanitize")
            self.assertIsInstance(new_weights[k], mx_array_type)

    def test_sanitize_audio_vae_prefix_strip(self):
        pass


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