"""
Tests for apply (TTMatrix @ TensorTrain) operations.
"""
import pytest
import torch
from ttglow import TTMatrix, TensorTrain
from ttglow import linalg


class TestApplyFull:
    """Tests for apply when TTMatrix acts on all dimensions."""

    def test_apply_identity(self):
        """Apply identity operator should return same tensor."""
        torch.manual_seed(42)
        psi = TensorTrain.random([3, 4, 5], ranks=[2, 2])
        I = TTMatrix.identity([3, 4, 5])

        result = linalg.apply(I, psi)

        T_psi = psi.to_tensor()
        T_result = result.to_tensor()
        assert torch.allclose(T_result, T_psi, atol=1e-10)

    def test_apply_full_against_dense(self):
        """Apply should match dense matrix-vector multiplication."""
        torch.manual_seed(42)
        psi = TensorTrain.random([3, 4], ranks=[2])
        H = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])

        result = linalg.apply(H, psi)

        # Dense computation
        T_psi = psi.to_tensor().flatten()
        T_H = H.to_tensor()
        T_expected = (T_H @ T_psi).reshape(3, 4)
        T_result = result.to_tensor()

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

    def test_apply_rectangular(self):
        """Apply rectangular operator (changes dimensions)."""
        torch.manual_seed(42)
        # Operator maps (3, 4) -> (2, 5)
        psi = TensorTrain.random([3, 4], ranks=[2])
        H = TTMatrix.random([(2, 3), (5, 4)], ranks=[2])

        result = linalg.apply(H, psi)

        # Dense computation
        T_psi = psi.to_tensor().flatten()
        T_H = H.to_tensor()
        T_expected = (T_H @ T_psi).reshape(2, 5)
        T_result = result.to_tensor()

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

    def test_apply_returns_tensortrain(self):
        """Apply should return a TensorTrain."""
        psi = TensorTrain.random([3, 4], ranks=[2])
        H = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])

        result = linalg.apply(H, psi)

        assert isinstance(result, TensorTrain)


class TestApplyPartial:
    """Tests for apply with dims parameter (partial application)."""

    def test_apply_partial_first_dim(self):
        """Apply single-site operator to first dimension."""
        torch.manual_seed(42)
        psi = TensorTrain.random([3, 4, 5], ranks=[2, 2])
        H = TTMatrix.random([(2, 3)], ranks=[])  # Maps 3 -> 2

        result = linalg.apply(H, psi, dims=[0])

        # Dense computation
        T_psi = psi.to_tensor()  # (3, 4, 5)
        T_H = H.to_tensor()      # (2, 3)
        # Apply H to first index: result[i,j,k] = sum_a H[i,a] * psi[a,j,k]
        T_expected = torch.einsum('ia,ajk->ijk', T_H, T_psi)
        T_result = result.to_tensor()

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

    def test_apply_partial_middle_dim(self):
        """Apply single-site operator to middle dimension."""
        torch.manual_seed(42)
        psi = TensorTrain.random([3, 4, 5], ranks=[2, 2])
        H = TTMatrix.random([(6, 4)], ranks=[])  # Maps 4 -> 6

        result = linalg.apply(H, psi, dims=[1])

        # Dense computation
        T_psi = psi.to_tensor()  # (3, 4, 5)
        T_H = H.to_tensor()      # (6, 4)
        T_expected = torch.einsum('ib,ajk->aik', T_H, T_psi.permute(1, 0, 2)).permute(1, 0, 2)
        # Simpler: T_expected[a,i,k] = sum_j H[i,j] * psi[a,j,k]
        T_expected = torch.einsum('ij,ajk->aik', T_H, T_psi)
        T_result = result.to_tensor()

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

    def test_apply_partial_last_dim(self):
        """Apply single-site operator to last dimension."""
        torch.manual_seed(42)
        psi = TensorTrain.random([3, 4, 5], ranks=[2, 2])
        H = TTMatrix.random([(7, 5)], ranks=[])  # Maps 5 -> 7

        result = linalg.apply(H, psi, dims=[2])

        # Dense computation
        T_psi = psi.to_tensor()  # (3, 4, 5)
        T_H = H.to_tensor()      # (7, 5)
        # result[i,j,k] = sum_l H[k,l] * psi[i,j,l]
        T_expected = torch.einsum('kl,ijl->ijk', T_H, T_psi)
        T_result = result.to_tensor()

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

    def test_apply_partial_two_sites(self):
        """Apply two-site operator to subset of dimensions."""
        torch.manual_seed(42)
        psi = TensorTrain.random([3, 4, 5, 6], ranks=[2, 2, 2])
        H = TTMatrix.random([(4, 4), (5, 5)], ranks=[2])  # Square 2-site op

        result = linalg.apply(H, psi, dims=[1, 2])

        # Dense computation
        T_psi = psi.to_tensor()  # (3, 4, 5, 6)
        T_H = H.to_tensor()      # (20, 20)

        # Reshape psi: (3, 4*5, 6) then apply H to middle
        T_psi_mid = T_psi.reshape(3, 20, 6)
        T_result_mid = torch.einsum('ij,ajb->aib', T_H, T_psi_mid)
        T_expected = T_result_mid.reshape(3, 4, 5, 6)

        T_result = result.to_tensor()

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

    def test_apply_partial_identity_single_site(self):
        """Apply identity to single site should not change tensor."""
        torch.manual_seed(42)
        psi = TensorTrain.random([3, 4, 5], ranks=[2, 2])
        I = TTMatrix.identity([4])  # Identity on dim 4

        result = linalg.apply(I, psi, dims=[1])

        T_psi = psi.to_tensor()
        T_result = result.to_tensor()

        assert T_result.shape == (3, 4, 5)
        assert torch.allclose(T_result, T_psi, atol=1e-10)

    @pytest.mark.skip(reason="Non-contiguous dims not yet fully supported")
    def test_apply_partial_non_contiguous(self):
        """Apply two-site operator to non-contiguous dimensions."""
        torch.manual_seed(42)
        psi = TensorTrain.random([3, 4, 5, 6], ranks=[2, 2, 2])
        H = TTMatrix.random([(3, 3), (5, 5)], ranks=[2])  # Acts on dims 0, 2

        result = linalg.apply(H, psi, dims=[0, 2])

        # Dense computation - more complex due to non-contiguous
        T_psi = psi.to_tensor()  # (3, 4, 5, 6)
        T_H = H.to_tensor()      # (15, 15)

        # Permute to bring dims 0,2 together: (3, 5, 4, 6)
        T_psi_perm = T_psi.permute(0, 2, 1, 3)
        # Reshape: (15, 4, 6)
        T_psi_flat = T_psi_perm.reshape(15, 4, 6)
        # Apply H: (15, 15) @ (15, 4, 6) -> (15, 4, 6)
        T_result_flat = torch.einsum('ij,jkl->ikl', T_H, T_psi_flat)
        # Reshape back: (3, 5, 4, 6)
        T_result_perm = T_result_flat.reshape(3, 5, 4, 6)
        # Permute back: (3, 4, 5, 6)
        T_expected = T_result_perm.permute(0, 2, 1, 3)

        T_result = result.to_tensor()

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


class TestApplyDimsFormats:
    """Tests for different dims parameter formats."""

    def test_apply_dims_none_full(self):
        """dims=None with matching sizes applies to all."""
        torch.manual_seed(42)
        psi = TensorTrain.random([3, 4], ranks=[2])
        H = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])

        result1 = linalg.apply(H, psi, dims=None)
        result2 = linalg.apply(H, psi, dims=[0, 1])

        assert torch.allclose(result1.to_tensor(), result2.to_tensor(), atol=1e-10)

    def test_apply_dims_none_partial(self):
        """dims=None with fewer matrix dims applies to first dims."""
        torch.manual_seed(42)
        psi = TensorTrain.random([3, 4, 5], ranks=[2, 2])
        H = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])  # Only 2 sites

        result1 = linalg.apply(H, psi, dims=None)
        result2 = linalg.apply(H, psi, dims=[0, 1])

        assert torch.allclose(result1.to_tensor(), result2.to_tensor(), atol=1e-10)

    @pytest.mark.skip(reason="Explicit dim pairing not yet fully supported")
    def test_apply_dims_tuple(self):
        """dims as tuple of two lists for explicit pairing."""
        torch.manual_seed(42)
        psi = TensorTrain.random([3, 4, 5], ranks=[2, 2])
        # H has dims in different order than target
        H = TTMatrix.random([(5, 5), (3, 3)], ranks=[2])

        # H's dim 0 (size 5) maps to psi's dim 2
        # H's dim 1 (size 3) maps to psi's dim 0
        result = linalg.apply(H, psi, dims=([0, 1], [2, 0]))

        # Dense: apply H (acting on dims 2,0 of psi in that order)
        T_psi = psi.to_tensor()  # (3, 4, 5)
        T_H = H.to_tensor()      # (15, 15)

        # Permute psi to (5, 3, 4) to bring target dims first
        T_psi_perm = T_psi.permute(2, 0, 1)
        T_psi_flat = T_psi_perm.reshape(15, 4)
        T_result_flat = T_H @ T_psi_flat
        T_result_perm = T_result_flat.reshape(5, 3, 4)
        T_expected = T_result_perm.permute(1, 2, 0)

        T_result = result.to_tensor()

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


class TestApplyErrors:
    """Tests for error handling in apply."""

    def test_apply_mismatched_dims_raises(self):
        """Applying to wrong-sized dimension should raise error."""
        psi = TensorTrain.random([3, 4, 5], ranks=[2, 2])
        H = TTMatrix.random([(2, 2)], ranks=[])  # Expects dim 2, not 3

        with pytest.raises(ValueError):
            linalg.apply(H, psi, dims=[0])  # psi[0] has size 3, H expects 2

    def test_apply_invalid_dim_raises(self):
        """Applying to non-existent dimension should raise error."""
        psi = TensorTrain.random([3, 4, 5], ranks=[2, 2])
        H = TTMatrix.random([(3, 3)], ranks=[])

        with pytest.raises(ValueError):
            linalg.apply(H, psi, dims=[5])  # psi only has dims 0,1,2

    def test_apply_wrong_number_of_dims_raises(self):
        """Number of dims must match TTMatrix sites."""
        psi = TensorTrain.random([3, 4, 5], ranks=[2, 2])
        H = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])  # 2 sites

        with pytest.raises(ValueError):
            linalg.apply(H, psi, dims=[0])  # Only 1 dim, but H has 2 sites


class TestApplyRanks:
    """Tests for rank behavior in apply."""

    def test_apply_rank_product(self):
        """Result ranks are products of input ranks (worst case)."""
        psi = TensorTrain.random([3, 4, 5], ranks=[2, 3])
        H = TTMatrix.random([(3, 3), (4, 4), (5, 5)], ranks=[2, 2])

        result = linalg.apply(H, psi)

        # Result ranks should be at most product of corresponding ranks
        # But exact structure depends on implementation
        assert isinstance(result, TensorTrain)
        assert result.shape == (3, 4, 5)
