"""Tests for core/convert_utils.py"""

import numpy as np
import pytest

from mlx_music.core.convert_utils import (
    is_conv1d_weight,
    transpose_conv1d_weight,
)


class TestTransposeConv1dWeight:
    def test_basic_transpose(self):
        """Conv1d: [O, I, K] -> [O, K, I]"""
        w = np.random.randn(4, 3, 5).astype(np.float32)
        result = transpose_conv1d_weight(w)
        assert result.shape == (4, 5, 3)
        # Verify values: result[o, k, i] == w[o, i, k]
        assert np.allclose(result[0, 0, :], w[0, :, 0])
        assert np.allclose(result[2, 3, :], w[2, :, 3])

    def test_preserves_out_channels(self):
        """Out channels (axis 0) should be unchanged."""
        w = np.random.randn(16, 8, 3).astype(np.float32)
        result = transpose_conv1d_weight(w)
        assert result.shape[0] == 16

    def test_rejects_non_3d(self):
        """Should raise on 2D (Linear) or 1D (bias) tensors."""
        with pytest.raises(ValueError):
            transpose_conv1d_weight(np.zeros((4, 3)))
        with pytest.raises(ValueError):
            transpose_conv1d_weight(np.zeros((4,)))

    def test_roundtrip(self):
        """Transposing twice should NOT return to original (it's axes 1,2 swap)."""
        w = np.random.randn(4, 3, 5).astype(np.float32)
        t1 = transpose_conv1d_weight(w)
        t2 = transpose_conv1d_weight(t1)
        # Double swap of axes 1,2 returns to original
        assert np.allclose(t2, w)


class TestIsConv1dWeight:
    def test_3d_weight(self):
        assert is_conv1d_weight("decoder.proj_in.1.weight", (2048, 192, 2))

    def test_1d_bias_rejected(self):
        assert not is_conv1d_weight("decoder.proj_in.1.bias", (2048,))

    def test_2d_linear_rejected(self):
        assert not is_conv1d_weight("layers.0.self_attn.q_proj.weight", (2048, 2048))

    def test_vae_conv_detected(self):
        assert is_conv1d_weight("encoder.conv1.weight", (128, 2, 7))
        assert is_conv1d_weight("decoder.block.0.res_unit1.conv1.weight", (128, 128, 7))
