"""Tests for tensor comparison utility."""

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

from mlx_music.debug.compare import TensorComparison, compare_tensors


class TestCompareTensors:
    def test_identical_tensors(self):
        a = mx.ones((2, 3))
        b = mx.ones((2, 3))
        result = compare_tensors(a, b, name="test")
        assert result.cosine_sim > 0.999
        assert result.max_abs_diff < 1e-6
        assert result.passed

    def test_different_tensors(self):
        a = mx.ones((2, 3))
        b = mx.zeros((2, 3))
        result = compare_tensors(a, b, name="test")
        assert result.cosine_sim < 0.01
        assert not result.passed

    def test_close_tensors(self):
        a = mx.ones((2, 3))
        b = a + mx.random.normal((2, 3)) * 1e-4
        mx.eval(b)
        result = compare_tensors(a, b, name="test", atol=1e-3)
        assert result.passed

    def test_shape_mismatch(self):
        a = mx.ones((2, 3))
        b = mx.ones((3, 2))
        result = compare_tensors(a, b, name="test")
        assert not result.passed
        assert "shape" in result.note.lower()
