"""
Tests for matmul (TTMatrix @ TTMatrix) operations.
"""
import pytest
import torch
from ttglow import TTMatrix
from ttglow import linalg


class TestMatmulFull:
    """Tests for matmul of two TTMatrices."""

    def test_matmul_identity_left(self):
        """I @ H = H."""
        torch.manual_seed(42)
        H = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])
        I = TTMatrix.identity([3, 4])

        result = linalg.matmul(I, H)

        T_H = H.to_tensor()
        T_result = result.to_tensor()

        assert T_result.shape == T_H.shape
        assert torch.allclose(T_result, T_H, atol=1e-10)

    def test_matmul_identity_right(self):
        """H @ I = H."""
        torch.manual_seed(42)
        H = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])
        I = TTMatrix.identity([3, 4])

        result = linalg.matmul(H, I)

        T_H = H.to_tensor()
        T_result = result.to_tensor()

        assert T_result.shape == T_H.shape
        assert torch.allclose(T_result, T_H, atol=1e-10)

    def test_matmul_against_dense(self):
        """matmul should match dense matrix multiplication."""
        torch.manual_seed(42)
        H1 = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])
        H2 = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])

        result = linalg.matmul(H1, H2)

        T_H1 = H1.to_tensor()
        T_H2 = H2.to_tensor()
        T_expected = T_H1 @ T_H2
        T_result = result.to_tensor()

        assert T_result.shape == T_expected.shape
        assert torch.allclose(T_result, T_expected, atol=1e-10)

    def test_matmul_rectangular(self):
        """matmul with rectangular operators."""
        torch.manual_seed(42)
        # H1: (3, 4) -> (2, 5)  i.e. maps 12-dim to 10-dim
        # H2: (5, 6) -> (3, 4)  i.e. maps 30-dim to 12-dim
        # H1 @ H2: (5, 6) -> (2, 5)  i.e. maps 30-dim to 10-dim
        H1 = TTMatrix.random([(2, 3), (5, 4)], ranks=[2])
        H2 = TTMatrix.random([(3, 5), (4, 6)], ranks=[2])

        result = linalg.matmul(H1, H2)

        T_H1 = H1.to_tensor()
        T_H2 = H2.to_tensor()
        T_expected = T_H1 @ T_H2
        T_result = result.to_tensor()

        assert result.row_dims == [2, 5]
        assert result.col_dims == [5, 6]
        assert T_result.shape == (10, 30)
        assert torch.allclose(T_result, T_expected, atol=1e-10)

    def test_matmul_three_sites(self):
        """matmul with 3-site operators."""
        torch.manual_seed(42)
        H1 = TTMatrix.random([(2, 2), (3, 3), (4, 4)], ranks=[2, 2])
        H2 = TTMatrix.random([(2, 2), (3, 3), (4, 4)], ranks=[3, 3])

        result = linalg.matmul(H1, H2)

        T_H1 = H1.to_tensor()
        T_H2 = H2.to_tensor()
        T_expected = T_H1 @ T_H2
        T_result = result.to_tensor()

        assert T_result.shape == T_expected.shape
        assert torch.allclose(T_result, T_expected, atol=1e-10)

    def test_matmul_returns_ttmatrix(self):
        """matmul should return a TTMatrix."""
        H1 = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])
        H2 = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])

        result = linalg.matmul(H1, H2)

        assert isinstance(result, TTMatrix)


class TestMatmulRanks:
    """Tests for rank behavior in matmul."""

    def test_matmul_rank_product(self):
        """Result ranks are products of input ranks."""
        H1 = TTMatrix.random([(3, 3), (4, 4), (5, 5)], ranks=[2, 3])
        H2 = TTMatrix.random([(3, 3), (4, 4), (5, 5)], ranks=[4, 5])

        result = linalg.matmul(H1, H2)

        # Internal ranks should be products: 2*4=8, 3*5=15
        assert result.ranks == [1, 8, 15, 1]

    def test_matmul_identity_preserves_rank(self):
        """Multiplying by identity increases rank minimally."""
        H = TTMatrix.random([(3, 3), (4, 4)], ranks=[5])
        I = TTMatrix.identity([3, 4])  # rank-1

        result = linalg.matmul(H, I)

        # I has rank 1, so result rank = H.rank * 1 = H.rank
        assert result.ranks == [1, 5, 1]


class TestMatmulErrors:
    """Tests for error handling in matmul."""

    def test_matmul_mismatched_dims_raises(self):
        """matmul with incompatible dimensions should raise."""
        H1 = TTMatrix.random([(3, 4), (5, 6)], ranks=[2])  # cols: 4, 6
        H2 = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])  # rows: 3, 4

        with pytest.raises(ValueError):
            linalg.matmul(H1, H2)

    def test_matmul_different_num_sites_raises(self):
        """matmul with different number of sites should raise."""
        H1 = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])
        H2 = TTMatrix.random([(3, 3), (4, 4), (5, 5)], ranks=[2, 2])

        with pytest.raises(ValueError):
            linalg.matmul(H1, H2)


class TestMatmulChain:
    """Tests for chaining matmul operations."""

    def test_matmul_associative(self):
        """(H1 @ H2) @ H3 = H1 @ (H2 @ H3)."""
        torch.manual_seed(42)
        H1 = TTMatrix.random([(2, 2), (3, 3)], ranks=[2])
        H2 = TTMatrix.random([(2, 2), (3, 3)], ranks=[2])
        H3 = TTMatrix.random([(2, 2), (3, 3)], ranks=[2])

        left = linalg.matmul(linalg.matmul(H1, H2), H3)
        right = linalg.matmul(H1, linalg.matmul(H2, H3))

        T_left = left.to_tensor()
        T_right = right.to_tensor()

        assert torch.allclose(T_left, T_right, atol=1e-10)

    def test_matmul_apply_consistency(self):
        """(H1 @ H2) @ psi = H1 @ (H2 @ psi)."""
        torch.manual_seed(42)
        from ttglow import TensorTrain

        H1 = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])
        H2 = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])
        psi = TensorTrain.random([3, 4], ranks=[2])

        # (H1 @ H2) @ psi
        H12 = linalg.matmul(H1, H2)
        result1 = linalg.apply(H12, psi)

        # H1 @ (H2 @ psi)
        H2_psi = linalg.apply(H2, psi)
        result2 = linalg.apply(H1, H2_psi)

        T1 = result1.to_tensor()
        T2 = result2.to_tensor()

        assert torch.allclose(T1, T2, atol=1e-10)
