import numpy as np
import pytest

from gri_utils import rotate
from gri_utils.matrices import decompose
from gri_utils.matrices.decompose import is_symmetric


def test_2d_straight():
    mat = ((4, 0), (0, 9))
    vals, uvecs = decompose(mat)
    assert len(vals) == 2
    assert vals[0] == pytest.approx(9)
    assert vals[1] == pytest.approx(4)
    assert uvecs[0].tolist() == pytest.approx([0, 1])
    assert uvecs[1].tolist() == pytest.approx([1, 0])


def test_2d_rotate():
    mat = ((4, 0), (0, 9))
    mat = rotate.rotate_cw(np.pi / 4, mat)
    vals, uvecs = decompose(mat)
    assert len(vals) == 2
    assert vals[0] == pytest.approx(9)
    assert vals[1] == pytest.approx(4)
    assert uvecs[0].tolist() == pytest.approx([2**0.5 / 2, 2**0.5 / 2])
    # flips to up
    assert uvecs[1].tolist() == pytest.approx([-(2**0.5) / 2, 2**0.5 / 2])


def test_2d_batched():
    """Test decompose with a batch of 2x2 matrices."""
    # Stack of two matrices: identity scaled by 4 and 9
    mat1 = np.array([[4, 0], [0, 9]])
    mat2 = np.array([[1, 0], [0, 16]])
    mats = np.stack([mat1, mat2])  # (2, 2, 2)

    vals, uvecs = decompose(mats)

    # Check shapes
    assert vals.shape == (2, 2)
    assert uvecs.shape == (2, 2, 2)

    # First matrix: eigenvalues 9, 4
    assert vals[0].tolist() == pytest.approx([9, 4])
    assert uvecs[0, 0].tolist() == pytest.approx([0, 1])
    assert uvecs[0, 1].tolist() == pytest.approx([1, 0])

    # Second matrix: eigenvalues 16, 1
    assert vals[1].tolist() == pytest.approx([16, 1])
    assert uvecs[1, 0].tolist() == pytest.approx([0, 1])
    assert uvecs[1, 1].tolist() == pytest.approx([1, 0])


def test_3d_straight():
    mat = ((4, 0, 0), (0, 9, 0), (0, 0, 1))
    vals, uvecs = decompose(mat)
    assert len(vals) == 3
    assert vals[0] == pytest.approx(9)
    assert vals[1] == pytest.approx(4)
    assert vals[2] == pytest.approx(1)
    assert uvecs[0].tolist() == pytest.approx([0, 1, 0])
    assert uvecs[1].tolist() == pytest.approx([1, 0, 0])
    assert uvecs[2].tolist() == pytest.approx([0, 0, 1])


def test_mult():
    mat = ((9, 0, 0), (0, 4, 0), (0, 0, 25))
    vals, uvecs = decompose(mat)
    assert len(vals) == 3
    assert vals[0] == pytest.approx(25)
    assert vals[1] == pytest.approx(9)
    assert vals[2] == pytest.approx(4)
    assert uvecs[0].tolist() == pytest.approx([0, 0, 1])
    assert uvecs[1].tolist() == pytest.approx([1, 0, 0])
    assert uvecs[2].tolist() == pytest.approx([0, 1, 0])
    vecs = uvecs * (vals**0.5)[:, None]
    assert vecs[0].tolist() == pytest.approx([0, 0, 5])
    assert vecs[1].tolist() == pytest.approx([3, 0, 0])
    assert vecs[2].tolist() == pytest.approx([0, 2, 0])
    vecs = (uvecs.T * vals**0.5).T
    assert vecs[0].tolist() == pytest.approx([0, 0, 5])
    assert vecs[1].tolist() == pytest.approx([3, 0, 0])
    assert vecs[2].tolist() == pytest.approx([0, 2, 0])


def test_3d_rotate():
    mat = ((4, 0, 0), (0, 9, 0), (0, 0, 1))
    mat = rotate.rotate_cw(np.pi / 4, mat)
    mat = rotate.rotate_cw(np.pi / 4, mat, 1)
    vals, uvecs = decompose(mat)
    assert len(vals) == 3
    assert vals[0] == pytest.approx(9)
    assert vals[1] == pytest.approx(4)
    assert vals[2] == pytest.approx(1)
    assert uvecs[0].tolist() == pytest.approx([0.5, 2**0.5 / 2, 0.5])
    assert uvecs[1].tolist() == pytest.approx([0.5, -(2**0.5) / 2, 0.5])
    assert uvecs[2].tolist() == pytest.approx([-(2**0.5) / 2, 0, 2**0.5 / 2])


def test_3d_batched():
    """Test decompose with a batch of 3x3 matrices."""
    mat1 = np.array([[4, 0, 0], [0, 9, 0], [0, 0, 1]])
    mat2 = np.array([[25, 0, 0], [0, 16, 0], [0, 0, 9]])
    mats = np.stack([mat1, mat2])  # (2, 3, 3)

    vals, uvecs = decompose(mats)

    # Check shapes
    assert vals.shape == (2, 3)
    assert uvecs.shape == (2, 3, 3)

    # First matrix: eigenvalues 9, 4, 1
    assert vals[0].tolist() == pytest.approx([9, 4, 1])

    # Second matrix: eigenvalues 25, 16, 9
    assert vals[1].tolist() == pytest.approx([25, 16, 9])


def test_namedtuple_access():
    """Test that Decomposition NamedTuple works correctly."""
    mat = np.array([[4, 0], [0, 9]])
    result = decompose(mat)

    # Named access
    assert result.vals[0] == pytest.approx(9)
    assert result.uvecs[0].tolist() == pytest.approx([0, 1])

    # Index access (backward compatible)
    assert result[0][0] == pytest.approx(9)
    assert result[1][0].tolist() == pytest.approx([0, 1])

    # Unpacking (backward compatible)
    vals, _uvecs = result
    assert vals[0] == pytest.approx(9)


def test_is_symmetric():
    assert is_symmetric(np.array(((1, 0), (0, 3))))
    assert is_symmetric(np.array(((1, 5), (5, 3))))
    assert not is_symmetric(np.array(((1, 4), (5, 3))))
    assert not is_symmetric(np.array(((1, 2, 3), (4, 5, 6))))


def test_is_symmetric_batched():
    """Test is_symmetric with batched matrices."""
    sym1 = np.array([[1, 2], [2, 3]])
    sym2 = np.array([[4, 5], [5, 6]])
    nonsym = np.array([[1, 2], [3, 4]])

    # All symmetric
    mats = np.stack([sym1, sym2])
    result = is_symmetric(mats)
    assert isinstance(result, np.ndarray)
    assert result.tolist() == [True, True]

    # Mixed
    mats = np.stack([sym1, nonsym, sym2])
    result = is_symmetric(mats)
    assert isinstance(result, np.ndarray)
    assert result.tolist() == [True, False, True]
