"""
Tests for TTMatrix (TT-format linear operators).
"""
import pytest
import torch
from ttglow import TTMatrix, TensorTrain
from ttglow import linalg
from ttglow import ttmatrix


class TestTTMatrixCreation:
    """Tests for TTMatrix creation and basic properties."""

    def test_creation_basic(self):
        """Test basic TTMatrix creation."""
        # 3 sites, each with 4x4 local matrix
        dims = [(4, 4), (4, 4), (4, 4)]
        ranks = [2, 2]
        ttm = TTMatrix(dims, ranks)

        assert ttm.d == 3
        assert ttm.row_dims == [4, 4, 4]
        assert ttm.col_dims == [4, 4, 4]
        assert ttm.ranks == [1, 2, 2, 1]

    def test_creation_rectangular(self):
        """Test TTMatrix with rectangular local matrices."""
        # Maps (3, 4, 5) -> (2, 3, 4)
        dims = [(2, 3), (3, 4), (4, 5)]
        ranks = [2, 3]
        ttm = TTMatrix(dims, ranks)

        assert ttm.d == 3
        assert ttm.row_dims == [2, 3, 4]
        assert ttm.col_dims == [3, 4, 5]
        assert ttm.shape == (2 * 3 * 4, 3 * 4 * 5)

    def test_core_shapes(self):
        """Test that cores have correct 4D shapes."""
        dims = [(3, 4), (5, 6), (7, 8)]
        ranks = [2, 3]
        ttm = TTMatrix(dims, ranks)

        assert ttm[0].shape == (1, 3, 4, 2)
        assert ttm[1].shape == (2, 5, 6, 3)
        assert ttm[2].shape == (3, 7, 8, 1)

    def test_default_ranks(self):
        """Test that default ranks are all ones."""
        dims = [(4, 4), (4, 4), (4, 4)]
        ttm = TTMatrix(dims)

        assert ttm.ranks == [1, 1, 1, 1]

    def test_wrong_ranks_length_raises(self):
        """Test that wrong number of ranks raises error."""
        dims = [(4, 4), (4, 4), (4, 4)]
        with pytest.raises(ValueError):
            TTMatrix(dims, ranks=[2, 2, 2])  # Should be 2 ranks, not 3


class TestTTMatrixRandom:
    """Tests for TTMatrix.random factory method."""

    def test_random_creation(self):
        """Test random TTMatrix creation."""
        torch.manual_seed(42)
        dims = [(3, 4), (5, 6)]
        ranks = [2]
        ttm = TTMatrix.random(dims, ranks)

        assert ttm.d == 2
        assert ttm[0].shape == (1, 3, 4, 2)
        assert ttm[1].shape == (2, 5, 6, 1)
        # Check it's not all zeros
        assert ttm[0].abs().sum() > 0

    def test_random_dtype(self):
        """Test random TTMatrix with specific dtype."""
        ttm = TTMatrix.random([(3, 3), (4, 4)], [2], dtype=torch.float32)
        assert ttm.dtype == torch.float32
        assert ttm[0].dtype == torch.float32


class TestTTMatrixIdentity:
    """Tests for TTMatrix.identity factory method."""

    def test_identity_creation(self):
        """Test identity TTMatrix creation."""
        dims = [3, 4, 5]
        ttm = TTMatrix.identity(dims)

        assert ttm.d == 3
        assert ttm.row_dims == [3, 4, 5]
        assert ttm.col_dims == [3, 4, 5]
        # Identity should be rank-1
        assert ttm.ranks == [1, 1, 1, 1]

    def test_identity_is_identity(self):
        """Test that identity TTMatrix represents identity operator."""
        dims = [3, 4, 5]
        ttm = TTMatrix.identity(dims)

        # Convert to dense and check it's identity
        T = ttm.to_tensor()
        n_total = 3 * 4 * 5
        expected = torch.eye(n_total, dtype=ttm.dtype)

        assert T.shape == (n_total, n_total)
        assert torch.allclose(T, expected)

    def test_identity_different_dims(self):
        """Test identity with different dimensions per site."""
        dims = [2, 3, 4]
        ttm = TTMatrix.identity(dims)
        T = ttm.to_tensor()

        n_total = 2 * 3 * 4
        assert torch.allclose(T, torch.eye(n_total, dtype=ttm.dtype))


class TestTTMatrixLocal:
    """Tests for TTMatrix.local factory method."""

    def test_local_creation(self):
        """Test basic local operator creation."""
        matrix = torch.randn(3, 4)
        ttm = TTMatrix.local(matrix)

        assert ttm.d == 1
        assert ttm.row_dims == [3]
        assert ttm.col_dims == [4]
        assert ttm.ranks == [1, 1]

    def test_local_preserves_matrix(self):
        """Local operator should preserve the input matrix."""
        matrix = torch.randn(3, 4)
        ttm = TTMatrix.local(matrix)

        T = ttm.to_tensor()
        assert torch.allclose(T, matrix)

    def test_local_square(self):
        """Local operator with square matrix."""
        matrix = torch.randn(5, 5)
        ttm = TTMatrix.local(matrix)

        assert ttm.shape == (5, 5)
        assert torch.allclose(ttm.to_tensor(), matrix)

    def test_local_preserves_dtype(self):
        """Local operator preserves input dtype."""
        matrix = torch.randn(3, 3, dtype=torch.float32)
        ttm = TTMatrix.local(matrix)

        assert ttm.dtype == torch.float32

    def test_local_complex(self):
        """Local operator with complex matrix."""
        matrix = torch.randn(3, 3, dtype=torch.complex128)
        ttm = TTMatrix.local(matrix)

        assert ttm.dtype == torch.complex128
        assert torch.allclose(ttm.to_tensor(), matrix)

    def test_local_with_on_syntax(self):
        """Local operator works with .on() syntax."""
        torch.manual_seed(42)
        psi = TensorTrain.random([3, 4, 5], ranks=[2, 2])

        # Create a local operator for the middle site (match dtype with psi)
        matrix = torch.randn(6, 4, dtype=torch.float64)  # Maps dim 4 -> 6
        op = TTMatrix.local(matrix)

        result = op.on(1) @ psi

        # Verify against dense computation
        T_psi = psi.to_tensor()
        T_expected = torch.einsum('ij,ajk->aik', matrix, 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_local_non_2d_raises(self):
        """Local with non-2D tensor should raise."""
        tensor_1d = torch.randn(5)
        tensor_3d = torch.randn(2, 3, 4)

        with pytest.raises(ValueError):
            TTMatrix.local(tensor_1d)

        with pytest.raises(ValueError):
            TTMatrix.local(tensor_3d)


class TestTTMatrixToTensor:
    """Tests for TTMatrix.to_tensor conversion."""

    def test_to_tensor_shape(self):
        """Test that to_tensor returns correct shape."""
        dims = [(3, 4), (5, 6)]
        ttm = TTMatrix.random(dims, [2])

        T = ttm.to_tensor()
        assert T.shape == (3 * 5, 4 * 6)

    def test_to_tensor_identity(self):
        """Test to_tensor for identity matrix."""
        ttm = TTMatrix.identity([2, 3])
        T = ttm.to_tensor()

        assert T.shape == (6, 6)
        assert torch.allclose(T, torch.eye(6, dtype=ttm.dtype))

    def test_to_dense_alias(self):
        """Test that to_dense is alias for to_tensor."""
        ttm = TTMatrix.random([(3, 3), (4, 4)], [2])
        assert torch.allclose(ttm.to_tensor(), ttm.to_dense())


class TestTTMatrixOperations:
    """Tests for basic TTMatrix operations."""

    def test_copy(self):
        """Test TTMatrix copy."""
        ttm = TTMatrix.random([(3, 3), (4, 4)], [2])
        ttm_copy = ttm.copy()

        # Modify original
        ttm[0] = torch.zeros_like(ttm[0])

        # Copy should be unchanged
        assert ttm_copy[0].abs().sum() > 0

    def test_clone(self):
        """Test TTMatrix clone (torch-like API)."""
        ttm = TTMatrix.random([(3, 3), (4, 4)], [2])
        ttm_clone = ttm.clone()

        assert torch.allclose(ttm.to_tensor(), ttm_clone.to_tensor())

    def test_to_dtype(self):
        """Test dtype conversion."""
        ttm = TTMatrix.random([(3, 3), (4, 4)], [2], dtype=torch.float64)
        ttm32 = ttm.to(dtype=torch.float32)

        assert ttm32.dtype == torch.float32
        assert ttm32[0].dtype == torch.float32

    def test_detach(self):
        """Test detach from computation graph."""
        ttm = TTMatrix.random([(3, 3), (4, 4)], [2])
        # Enable gradients on original
        for core in ttm.cores:
            core.requires_grad_(True)

        ttm_detached = ttm.detach()

        # Check values are preserved
        assert torch.allclose(ttm.to_tensor(), ttm_detached.to_tensor())

        # Check detached cores don't require gradients
        for core in ttm_detached.cores:
            assert not core.requires_grad

    def test_setitem(self):
        """Test setting cores."""
        ttm = TTMatrix([(3, 4), (5, 6)], [2])
        new_core = torch.randn(1, 3, 4, 2, dtype=ttm.dtype)
        ttm[0] = new_core

        assert torch.allclose(ttm[0], new_core)

    def test_setitem_wrong_shape_raises(self):
        """Test that setting wrong shape raises error."""
        ttm = TTMatrix([(3, 4), (5, 6)], [2])
        wrong_core = torch.randn(1, 3, 5, 2, dtype=ttm.dtype)  # Wrong col dim

        with pytest.raises(ValueError):
            ttm[0] = wrong_core


class TestTTMatrixProperties:
    """Tests for TTMatrix properties."""

    def test_tt_shape(self):
        """Test tt_shape returns list of core shapes."""
        dims = [(3, 4), (5, 6), (7, 8)]
        ranks = [2, 3]
        ttm = TTMatrix(dims, ranks)

        shapes = ttm.tt_shape()
        assert shapes == [(1, 3, 4, 2), (2, 5, 6, 3), (3, 7, 8, 1)]

    def test_tt_ranks(self):
        """Test tt_ranks returns full rank tuple."""
        dims = [(3, 3), (4, 4)]
        ranks = [5]
        ttm = TTMatrix(dims, ranks)

        assert ttm.tt_ranks() == (1, 5, 1)

    def test_total_params(self):
        """Test total parameter count."""
        dims = [(3, 4), (5, 6)]
        ranks = [2]
        ttm = TTMatrix(dims, ranks)

        # Core 0: 1 * 3 * 4 * 2 = 24
        # Core 1: 2 * 5 * 6 * 1 = 60
        assert ttm.total_params() == 24 + 60

    def test_storage_numel(self):
        """Test storage_numel (alias for total_params)."""
        ttm = TTMatrix.random([(3, 4), (5, 6)], [2])
        assert ttm.storage_numel() == ttm.total_params()

    def test_numel(self):
        """Test logical element count (full matrix size)."""
        dims = [(3, 4), (5, 6)]
        ttm = TTMatrix(dims, [2])

        # Full matrix would be (3*5) x (4*6) = 15 x 24 = 360
        assert ttm.numel == 15 * 24

    def test_len(self):
        """Test len returns number of cores."""
        ttm = TTMatrix([(3, 3), (4, 4), (5, 5)], [2, 2])
        assert len(ttm) == 3

    def test_ncores(self):
        """Test ncores method."""
        ttm = TTMatrix([(3, 3), (4, 4), (5, 5)], [2, 2])
        assert ttm.ncores() == 3


class TestTranspose:
    """Tests for TTMatrix transpose."""

    def test_transpose_dims(self):
        """Transpose swaps row and column dims."""
        ttm = TTMatrix.random([(3, 4), (5, 6)], [2])
        ttm_t = linalg.transpose(ttm)

        assert ttm_t.row_dims == [4, 6]
        assert ttm_t.col_dims == [3, 5]
        assert ttm_t.shape == (24, 15)

    def test_transpose_against_dense(self):
        """Transpose should match dense transpose."""
        torch.manual_seed(42)
        ttm = TTMatrix.random([(3, 4), (5, 6)], [2])

        T_ttm = ttm.to_tensor()
        T_transpose = linalg.transpose(ttm).to_tensor()

        assert torch.allclose(T_transpose, T_ttm.T, atol=1e-10)

    def test_transpose_square(self):
        """Transpose of square matrix."""
        torch.manual_seed(42)
        ttm = TTMatrix.random([(3, 3), (4, 4)], [2])

        T_ttm = ttm.to_tensor()
        T_transpose = linalg.transpose(ttm).to_tensor()

        assert T_transpose.shape == T_ttm.shape
        assert torch.allclose(T_transpose, T_ttm.T, atol=1e-10)

    def test_transpose_involutory(self):
        """Transpose of transpose equals original."""
        torch.manual_seed(42)
        ttm = TTMatrix.random([(3, 4), (5, 6)], [2])

        ttm_tt = linalg.transpose(linalg.transpose(ttm))

        T_orig = ttm.to_tensor()
        T_double = ttm_tt.to_tensor()

        assert torch.allclose(T_double, T_orig, atol=1e-10)

    def test_transpose_preserves_ranks(self):
        """Transpose preserves TT ranks."""
        ttm = TTMatrix.random([(3, 4), (5, 6), (7, 8)], [2, 3])
        ttm_t = linalg.transpose(ttm)

        assert ttm_t.ranks == ttm.ranks


class TestTrace:
    """Tests for TTMatrix trace."""

    def test_trace_against_dense(self):
        """Trace should match dense trace."""
        torch.manual_seed(42)
        ttm = TTMatrix.random([(3, 3), (4, 4)], [2])

        T_ttm = ttm.to_tensor()
        tr_dense = torch.trace(T_ttm)
        tr_tt = linalg.trace(ttm)

        assert abs(tr_tt - tr_dense.item()) < 1e-10

    def test_trace_identity(self):
        """Trace of identity is dimension."""
        dims = [3, 4, 5]
        I = TTMatrix.identity(dims)

        tr = linalg.trace(I)

        # Identity trace = product of dims = 3*4*5 = 60
        assert abs(tr - 60.0) < 1e-10

    def test_trace_three_sites(self):
        """Trace with 3-site matrix."""
        torch.manual_seed(42)
        ttm = TTMatrix.random([(2, 2), (3, 3), (4, 4)], [2, 2])

        T_ttm = ttm.to_tensor()
        tr_dense = torch.trace(T_ttm)
        tr_tt = linalg.trace(ttm)

        assert abs(tr_tt - tr_dense.item()) < 1e-10

    def test_trace_non_square_raises(self):
        """Trace of non-square matrix should raise."""
        ttm = TTMatrix.random([(3, 4), (5, 5)], [2])

        with pytest.raises(ValueError):
            linalg.trace(ttm)

    def test_trace_single_site(self):
        """Trace of single-site matrix."""
        torch.manual_seed(42)
        ttm = TTMatrix.random([(5, 5)], [])

        T_ttm = ttm.to_tensor()
        tr_dense = torch.trace(T_ttm)
        tr_tt = linalg.trace(ttm)

        assert abs(tr_tt - tr_dense.item()) < 1e-10


class TestToTensorTrain:
    """Tests for TTMatrix <-> TensorTrain conversion."""

    def test_to_tensor_train_dims(self):
        """to_tensor_train merges row and col dims."""
        ttm = TTMatrix.random([(3, 4), (5, 6)], [2])
        tt = ttm.to_tensor_train()

        assert tt.dims == [12, 30]
        assert tt.ranks == [1, 2, 1]

    def test_to_tensor_train_preserves_core_data(self):
        """to_tensor_train preserves the underlying core data."""
        torch.manual_seed(42)
        ttm = TTMatrix.random([(3, 4), (5, 6)], [2])

        tt = ttm.to_tensor_train()

        # Core data should be reshaped but identical
        for k in range(ttm.d):
            ttm_core = ttm.cores[k]  # (r_l, n_row, n_col, r_r)
            tt_core = tt.cores[k]    # (r_l, n_row*n_col, r_r)

            r_l, n_row, n_col, r_r = ttm_core.shape
            assert torch.allclose(tt_core, ttm_core.reshape(r_l, n_row * n_col, r_r), atol=1e-10)

    def test_to_tensor_train_three_sites(self):
        """to_tensor_train with 3-site matrix."""
        ttm = TTMatrix.random([(2, 3), (4, 5), (6, 7)], [2, 3])
        tt = ttm.to_tensor_train()

        assert tt.dims == [6, 20, 42]
        assert tt.ranks == [1, 2, 3, 1]

    def test_from_tensor_train_basic(self):
        """from_tensor_train splits dims correctly."""
        tt = TensorTrain.random([12, 30], [2])
        ttm = TTMatrix.from_tensor_train(tt, dims=[(3, 4), (5, 6)])

        assert ttm.row_dims == [3, 5]
        assert ttm.col_dims == [4, 6]
        assert ttm.ranks == [1, 2, 1]

    def test_from_tensor_train_preserves_core_data(self):
        """from_tensor_train preserves the underlying core data."""
        torch.manual_seed(42)
        tt = TensorTrain.random([12, 30], [2])
        dims = [(3, 4), (5, 6)]

        ttm = TTMatrix.from_tensor_train(tt, dims=dims)

        # Core data should be reshaped but identical
        for k in range(tt.d):
            tt_core = tt.cores[k]    # (r_l, n, r_r)
            ttm_core = ttm.cores[k]  # (r_l, n_row, n_col, r_r)

            r_l, n, r_r = tt_core.shape
            n_row, n_col = dims[k]
            assert torch.allclose(ttm_core, tt_core.reshape(r_l, n_row, n_col, r_r), atol=1e-10)

    def test_roundtrip_ttmatrix_to_tt_to_ttmatrix(self):
        """TTMatrix -> TensorTrain -> TTMatrix preserves data."""
        torch.manual_seed(42)
        original_dims = [(3, 4), (5, 6)]
        ttm1 = TTMatrix.random(original_dims, [2])

        tt = ttm1.to_tensor_train()
        ttm2 = TTMatrix.from_tensor_train(tt, dims=original_dims)

        T1 = ttm1.to_tensor()
        T2 = ttm2.to_tensor()

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

    def test_roundtrip_tt_to_ttmatrix_to_tt(self):
        """TensorTrain -> TTMatrix -> TensorTrain preserves data."""
        torch.manual_seed(42)
        tt1 = TensorTrain.random([12, 30], [2])

        ttm = TTMatrix.from_tensor_train(tt1, dims=[(3, 4), (5, 6)])
        tt2 = ttm.to_tensor_train()

        T1 = tt1.to_tensor()
        T2 = tt2.to_tensor()

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

    def test_from_tensor_train_wrong_num_dims_raises(self):
        """from_tensor_train with wrong number of dims raises."""
        tt = TensorTrain.random([12, 30], [2])

        with pytest.raises(ValueError):
            TTMatrix.from_tensor_train(tt, dims=[(3, 4)])  # Only 1 dim, need 2

    def test_from_tensor_train_wrong_product_raises(self):
        """from_tensor_train with wrong dim product raises."""
        tt = TensorTrain.random([12, 30], [2])

        with pytest.raises(ValueError):
            # 3*5=15 != 12
            TTMatrix.from_tensor_train(tt, dims=[(3, 5), (5, 6)])


class TestTTMatrixAdd:
    """Tests for TTMatrix addition."""

    def test_add_against_dense(self):
        """add should match dense addition."""
        torch.manual_seed(42)
        H1 = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])
        H2 = TTMatrix.random([(3, 3), (4, 4)], ranks=[3])

        H_sum = ttmatrix.add(H1, H2)

        T1 = H1.to_tensor()
        T2 = H2.to_tensor()
        T_sum = H_sum.to_tensor()

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

    def test_add_ranks(self):
        """add ranks are sums 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])

        H_sum = ttmatrix.add(H1, H2)

        # Internal ranks should be sums: 2+4=6, 3+5=8
        assert H_sum.ranks == [1, 6, 8, 1]

    def test_add_mismatched_dims_raises(self):
        """add with mismatched dims should raise."""
        H1 = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])
        H2 = TTMatrix.random([(3, 3), (5, 5)], ranks=[2])

        with pytest.raises(ValueError):
            ttmatrix.add(H1, H2)


class TestTTMatrixScale:
    """Tests for TTMatrix scaling."""

    def test_scale_against_dense(self):
        """scale should match dense scaling."""
        torch.manual_seed(42)
        H = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])

        H_scaled = ttmatrix.scale(H, 2.5)

        T = H.to_tensor()
        T_scaled = H_scaled.to_tensor()

        assert torch.allclose(T_scaled, 2.5 * T, atol=1e-10)

    def test_scale_negative(self):
        """scale with negative scalar."""
        torch.manual_seed(42)
        H = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])

        H_scaled = ttmatrix.scale(H, -1.0)

        T = H.to_tensor()
        T_scaled = H_scaled.to_tensor()

        assert torch.allclose(T_scaled, -T, atol=1e-10)

    def test_scale_preserves_ranks(self):
        """scale preserves TT ranks."""
        H = TTMatrix.random([(3, 3), (4, 4)], ranks=[5])
        H_scaled = ttmatrix.scale(H, 3.0)

        assert H_scaled.ranks == H.ranks


class TestTTMatrixNorm:
    """Tests for TTMatrix norm."""

    def test_norm_against_dense(self):
        """norm should match dense Frobenius norm."""
        torch.manual_seed(42)
        H = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])

        n = ttmatrix.norm(H)

        T = H.to_tensor()
        n_dense = torch.linalg.norm(T, 'fro')

        assert abs(n - n_dense.item()) < 1e-10

    def test_norm_squared_against_dense(self):
        """norm_squared should match dense squared norm."""
        torch.manual_seed(42)
        H = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])

        ns = ttmatrix.norm_squared(H)

        T = H.to_tensor()
        ns_dense = (T ** 2).sum()

        assert abs(ns - ns_dense.item()) < 1e-10

    def test_norm_positive(self):
        """norm should be positive for non-zero matrix."""
        H = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])
        assert ttmatrix.norm(H) > 0

    def test_norm_identity(self):
        """norm of identity is sqrt(dimension)."""
        dims = [3, 4]
        I = TTMatrix.identity(dims)

        n = ttmatrix.norm(I)

        # Identity has 1s on diagonal, dimension = 3*4 = 12
        assert abs(n - 12 ** 0.5) < 1e-10


class TestTTMatrixRound:
    """Tests for TTMatrix rounding (rank truncation)."""

    def test_round_reduces_rank(self):
        """round should reduce rank."""
        H1 = TTMatrix.random([(3, 3), (4, 4)], ranks=[5])
        H2 = TTMatrix.random([(3, 3), (4, 4)], ranks=[5])
        H_sum = ttmatrix.add(H1, H2)  # rank 10

        H_rounded = ttmatrix.round(H_sum, max_rank=3)

        assert H_rounded.ranks[1] <= 3

    def test_round_preserves_approximately(self):
        """round should approximately preserve the matrix."""
        torch.manual_seed(42)
        H1 = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])
        H2 = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])
        H_sum = ttmatrix.add(H1, H2)  # rank 4

        H_rounded = ttmatrix.round(H_sum, max_rank=4)  # no actual truncation

        T_sum = H_sum.to_tensor()
        T_rounded = H_rounded.to_tensor()

        assert torch.allclose(T_rounded, T_sum, atol=1e-10)

    def test_round_with_tolerance(self):
        """round with relative tolerance."""
        torch.manual_seed(42)
        H = TTMatrix.random([(3, 3), (4, 4), (5, 5)], ranks=[10, 10])

        H_rounded = ttmatrix.round(H, rel_tol=0.1)

        # Should have reduced rank
        assert H_rounded.ranks[1] <= 10
        assert H_rounded.ranks[2] <= 10


class TestAdjoint:
    """Tests for TTMatrix adjoint (conjugate transpose)."""

    def test_adjoint_dims(self):
        """adjoint swaps row and col dims."""
        ttm = TTMatrix.random([(3, 4), (5, 6)], [2])
        adj = ttmatrix.adjoint(ttm)

        assert adj.row_dims == [4, 6]
        assert adj.col_dims == [3, 5]
        assert adj.shape == (24, 15)

    def test_adjoint_against_dense_real(self):
        """adjoint of real matrix equals transpose."""
        torch.manual_seed(42)
        ttm = TTMatrix.random([(3, 4), (5, 6)], [2])

        adj = ttmatrix.adjoint(ttm)

        T = ttm.to_tensor()
        T_adj = adj.to_tensor()

        assert torch.allclose(T_adj, T.T, atol=1e-10)

    def test_adjoint_against_dense_complex(self):
        """adjoint of complex matrix is conjugate transpose."""
        torch.manual_seed(42)
        ttm = TTMatrix.random([(3, 4), (5, 6)], [2], dtype=torch.complex128)

        adj = ttmatrix.adjoint(ttm)

        T = ttm.to_tensor()
        T_adj = adj.to_tensor()

        # Conjugate transpose
        assert torch.allclose(T_adj, T.T.conj(), atol=1e-10)

    def test_adjoint_involutory_real(self):
        """adjoint(adjoint(H)) = H for real matrices."""
        torch.manual_seed(42)
        ttm = TTMatrix.random([(3, 3), (4, 4)], [2])

        adj_adj = ttmatrix.adjoint(ttmatrix.adjoint(ttm))

        T_orig = ttm.to_tensor()
        T_adj_adj = adj_adj.to_tensor()

        assert torch.allclose(T_adj_adj, T_orig, atol=1e-10)

    def test_adjoint_involutory_complex(self):
        """adjoint(adjoint(H)) = H for complex matrices."""
        torch.manual_seed(42)
        ttm = TTMatrix.random([(3, 3), (4, 4)], [2], dtype=torch.complex128)

        adj_adj = ttmatrix.adjoint(ttmatrix.adjoint(ttm))

        T_orig = ttm.to_tensor()
        T_adj_adj = adj_adj.to_tensor()

        assert torch.allclose(T_adj_adj, T_orig, atol=1e-10)

    def test_adjoint_preserves_ranks(self):
        """adjoint preserves TT ranks."""
        ttm = TTMatrix.random([(3, 4), (5, 6)], [7])
        adj = ttmatrix.adjoint(ttm)

        assert adj.ranks == ttm.ranks


class TestOperatorOverloading:
    """Tests for TTMatrix operator overloading."""

    def test_matmul_operator_ttmatrix(self):
        """H1 @ H2 uses matmul."""
        torch.manual_seed(42)
        H1 = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])
        H2 = TTMatrix.random([(3, 3), (4, 4)], ranks=[3])

        result = H1 @ H2

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

        assert torch.allclose(T_result, T_H1 @ T_H2, atol=1e-10)

    def test_matmul_operator_tensortrain(self):
        """H @ psi uses apply."""
        torch.manual_seed(42)
        H = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])
        psi = TensorTrain.random([3, 4], ranks=[2])

        result = H @ psi

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

        assert isinstance(result, TensorTrain)
        assert torch.allclose(T_result, T_expected, atol=1e-10)

    def test_add_operator(self):
        """H1 + H2 uses add."""
        torch.manual_seed(42)
        H1 = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])
        H2 = TTMatrix.random([(3, 3), (4, 4)], ranks=[3])

        result = H1 + H2

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

        assert torch.allclose(T_result, T_H1 + T_H2, atol=1e-10)

    def test_sub_operator(self):
        """H1 - H2 uses add and scale."""
        torch.manual_seed(42)
        H1 = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])
        H2 = TTMatrix.random([(3, 3), (4, 4)], ranks=[3])

        result = H1 - H2

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

        assert torch.allclose(T_result, T_H1 - T_H2, atol=1e-10)

    def test_mul_operator(self):
        """H * scalar uses scale."""
        torch.manual_seed(42)
        H = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])

        result = H * 2.5

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

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

    def test_rmul_operator(self):
        """scalar * H uses scale."""
        torch.manual_seed(42)
        H = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])

        result = 3.0 * H

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

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

    def test_neg_operator(self):
        """-H uses scale with -1."""
        torch.manual_seed(42)
        H = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])

        result = -H

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

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

    def test_chained_operations(self):
        """Chained operator expressions work correctly."""
        torch.manual_seed(42)
        H1 = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])
        H2 = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])
        H3 = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])

        # (H1 + H2) @ H3 - 2 * H1
        result = (H1 + H2) @ H3 - 2.0 * H1

        T1 = H1.to_tensor()
        T2 = H2.to_tensor()
        T3 = H3.to_tensor()
        T_expected = (T1 + T2) @ T3 - 2.0 * T1
        T_result = result.to_tensor()

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


class TestLocatedOp:
    """Tests for H.on(sites) @ psi syntax."""

    def test_on_single_site_first(self):
        """H.on(0) @ psi applies to first site."""
        torch.manual_seed(42)
        psi = TensorTrain.random([3, 4, 5], ranks=[2, 2])
        H = TTMatrix.random([(2, 3)], ranks=[])  # Maps 3 -> 2

        result = H.on(0) @ psi

        # Dense computation
        T_psi = psi.to_tensor()  # (3, 4, 5)
        T_H = H.to_tensor()      # (2, 3)
        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_on_single_site_middle(self):
        """H.on(1) @ psi applies to middle site."""
        torch.manual_seed(42)
        psi = TensorTrain.random([3, 4, 5], ranks=[2, 2])
        H = TTMatrix.random([(6, 4)], ranks=[])  # Maps 4 -> 6

        result = H.on(1) @ psi

        # Dense computation
        T_psi = psi.to_tensor()  # (3, 4, 5)
        T_H = H.to_tensor()      # (6, 4)
        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_on_single_site_last(self):
        """H.on(2) @ psi applies to last site."""
        torch.manual_seed(42)
        psi = TensorTrain.random([3, 4, 5], ranks=[2, 2])
        H = TTMatrix.random([(7, 5)], ranks=[])  # Maps 5 -> 7

        result = H.on(2) @ psi

        # Dense computation
        T_psi = psi.to_tensor()  # (3, 4, 5)
        T_H = H.to_tensor()      # (7, 5)
        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_on_two_sites(self):
        """H.on(1, 2) @ psi applies to two contiguous sites."""
        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 = H.on(1, 2) @ psi

        # 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_on_equivalent_to_apply(self):
        """H.on(sites) @ psi gives same result as apply(H, psi, dims=sites)."""
        torch.manual_seed(42)
        psi = TensorTrain.random([3, 4, 5, 6], ranks=[2, 2, 2])
        H = TTMatrix.random([(4, 4), (5, 5)], ranks=[2])

        result_on = H.on(1, 2) @ psi
        result_apply = linalg.apply(H, psi, dims=[1, 2])

        assert torch.allclose(result_on.to_tensor(), result_apply.to_tensor(), atol=1e-10)

    def test_on_wrong_num_sites_raises(self):
        """H.on(wrong_count) raises ValueError."""
        H = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])  # 2-site operator

        with pytest.raises(ValueError):
            H.on(1)  # Only 1 site, but H has 2 sites

    def test_on_full_apply(self):
        """H.on(0, 1, 2) @ psi is equivalent to H @ psi."""
        torch.manual_seed(42)
        psi = TensorTrain.random([3, 4, 5], ranks=[2, 2])
        H = TTMatrix.random([(3, 3), (4, 4), (5, 5)], ranks=[2, 2])

        result_on = H.on(0, 1, 2) @ psi
        result_direct = H @ psi

        assert torch.allclose(result_on.to_tensor(), result_direct.to_tensor(), atol=1e-10)

    def test_on_repr(self):
        """LocatedOp has a useful repr."""
        from ttglow.ttmatrix import LocatedOp

        H = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])
        loc_op = H.on(1, 2)

        assert isinstance(loc_op, LocatedOp)
        assert "LocatedOp" in repr(loc_op)
        assert "(1, 2)" in repr(loc_op)


class TestKron:
    """Tests for Kronecker product of TTMatrices."""

    def test_kron_dims(self):
        """kron concatenates dimensions."""
        H1 = TTMatrix.random([(3, 4), (5, 6)], ranks=[2])
        H2 = TTMatrix.random([(7, 8)], ranks=[])

        result = ttmatrix.kron(H1, H2)

        assert result.d == 3
        assert result.dims == [(3, 4), (5, 6), (7, 8)]
        assert result.row_dims == [3, 5, 7]
        assert result.col_dims == [4, 6, 8]

    def test_kron_ranks(self):
        """kron has junction rank 1."""
        H1 = TTMatrix.random([(3, 3), (4, 4)], ranks=[5])
        H2 = TTMatrix.random([(5, 5), (6, 6)], ranks=[7])

        result = ttmatrix.kron(H1, H2)

        # ranks = [1, 5, 1, 7, 1]
        assert result.ranks == [1, 5, 1, 7, 1]

    def test_kron_against_dense(self):
        """kron matches torch.kron for dense matrices."""
        torch.manual_seed(42)
        H1 = TTMatrix.random([(3, 4), (5, 6)], ranks=[2])
        H2 = TTMatrix.random([(7, 8), (9, 10)], ranks=[3])

        result = ttmatrix.kron(H1, H2)

        T1 = H1.to_tensor()
        T2 = H2.to_tensor()
        T_expected = torch.kron(T1, T2)
        T_result = result.to_tensor()

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

    def test_kron_single_site_operators(self):
        """kron of single-site operators."""
        torch.manual_seed(42)
        A = TTMatrix.local(torch.randn(3, 3, dtype=torch.float64))
        B = TTMatrix.local(torch.randn(4, 4, dtype=torch.float64))

        result = ttmatrix.kron(A, B)

        assert result.d == 2
        assert result.dims == [(3, 3), (4, 4)]
        assert result.ranks == [1, 1, 1]

        T_A = A.to_tensor()
        T_B = B.to_tensor()
        T_result = result.to_tensor()

        assert torch.allclose(T_result, torch.kron(T_A, T_B), atol=1e-10)

    def test_kron_with_identity(self):
        """H ⊗ I extends operator to more sites."""
        torch.manual_seed(42)
        H = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])
        I = TTMatrix.identity([5])

        HI = ttmatrix.kron(H, I)
        IH = ttmatrix.kron(I, H)

        assert HI.d == 3
        assert HI.dims == [(3, 3), (4, 4), (5, 5)]

        assert IH.d == 3
        assert IH.dims == [(5, 5), (3, 3), (4, 4)]

        # Verify against dense
        T_H = H.to_tensor()
        T_I = torch.eye(5, dtype=H.dtype)

        assert torch.allclose(HI.to_tensor(), torch.kron(T_H, T_I), atol=1e-10)
        assert torch.allclose(IH.to_tensor(), torch.kron(T_I, T_H), atol=1e-10)

    def test_kron_apply_to_subset(self):
        """kron result can be applied with .on() syntax."""
        torch.manual_seed(42)
        psi = TensorTrain.random([2, 3, 4, 5, 6], ranks=[2, 2, 2, 2])

        A = TTMatrix.local(torch.randn(3, 3, dtype=torch.float64))
        B = TTMatrix.local(torch.randn(4, 4, dtype=torch.float64))
        AB = ttmatrix.kron(A, B)

        # Apply A⊗B to sites 1 and 2
        result = AB.on(1, 2) @ psi

        # Verify: should be same as applying A to site 1 and B to site 2
        T_psi = psi.to_tensor()
        T_A = A.to_tensor()
        T_B = B.to_tensor()
        # Apply A to dim 1, B to dim 2
        T_expected = torch.einsum('ij,ajklm->aiklm', T_A, T_psi)
        T_expected = torch.einsum('jk,aiklm->aijlm', T_B, T_expected)
        T_result = result.to_tensor()

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

    def test_kron_chain(self):
        """Multiple kron operations."""
        torch.manual_seed(42)
        A = TTMatrix.local(torch.randn(2, 2, dtype=torch.float64))
        B = TTMatrix.local(torch.randn(3, 3, dtype=torch.float64))
        C = TTMatrix.local(torch.randn(4, 4, dtype=torch.float64))

        ABC = ttmatrix.kron(ttmatrix.kron(A, B), C)

        assert ABC.d == 3
        assert ABC.dims == [(2, 2), (3, 3), (4, 4)]

        T_expected = torch.kron(torch.kron(A.to_tensor(), B.to_tensor()), C.to_tensor())
        assert torch.allclose(ABC.to_tensor(), T_expected, atol=1e-10)

    def test_kron_variadic(self):
        """kron accepts multiple arguments."""
        torch.manual_seed(42)
        A = TTMatrix.local(torch.randn(2, 2, dtype=torch.float64))
        B = TTMatrix.local(torch.randn(3, 3, dtype=torch.float64))
        C = TTMatrix.local(torch.randn(4, 4, dtype=torch.float64))
        D = TTMatrix.local(torch.randn(5, 5, dtype=torch.float64))

        ABCD = ttmatrix.kron(A, B, C, D)

        assert ABCD.d == 4
        assert ABCD.dims == [(2, 2), (3, 3), (4, 4), (5, 5)]

        T_expected = torch.kron(torch.kron(torch.kron(
            A.to_tensor(), B.to_tensor()), C.to_tensor()), D.to_tensor())
        assert torch.allclose(ABCD.to_tensor(), T_expected, atol=1e-10)

    def test_kron_operator(self):
        """& operator for Kronecker product."""
        torch.manual_seed(42)
        A = TTMatrix.local(torch.randn(2, 2, dtype=torch.float64))
        B = TTMatrix.local(torch.randn(3, 3, dtype=torch.float64))

        result = A & B

        T_expected = torch.kron(A.to_tensor(), B.to_tensor())
        assert torch.allclose(result.to_tensor(), T_expected, atol=1e-10)

    def test_kron_operator_chain(self):
        """Chained & operator: A & B & C & D."""
        torch.manual_seed(42)
        A = TTMatrix.local(torch.randn(2, 2, dtype=torch.float64))
        B = TTMatrix.local(torch.randn(3, 3, dtype=torch.float64))
        C = TTMatrix.local(torch.randn(4, 4, dtype=torch.float64))
        D = TTMatrix.local(torch.randn(5, 5, dtype=torch.float64))

        result = A & B & C & D

        assert result.d == 4
        assert result.dims == [(2, 2), (3, 3), (4, 4), (5, 5)]

        T_expected = torch.kron(torch.kron(torch.kron(
            A.to_tensor(), B.to_tensor()), C.to_tensor()), D.to_tensor())
        assert torch.allclose(result.to_tensor(), T_expected, atol=1e-10)

    def test_kron_operator_with_on(self):
        """(X & Z & X & Z).on(0, 1, 2, 3) @ psi."""
        torch.manual_seed(42)
        psi = TensorTrain.random([2, 2, 2, 2], ranks=[2, 2, 2])

        X = TTMatrix.local(torch.tensor([[0., 1.], [1., 0.]], dtype=torch.float64))
        Z = TTMatrix.local(torch.tensor([[1., 0.], [0., -1.]], dtype=torch.float64))

        op = X & Z & X & Z
        result = op.on(0, 1, 2, 3) @ psi

        # Verify against dense
        T_psi = psi.to_tensor()
        T_op = op.to_tensor()
        T_expected = (T_op @ T_psi.flatten()).reshape(2, 2, 2, 2)
        T_result = result.to_tensor()

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

    def test_kron_less_than_two_raises(self):
        """kron with less than 2 args raises."""
        A = TTMatrix.local(torch.randn(2, 2))

        with pytest.raises(ValueError):
            ttmatrix.kron(A)


class TestTrunc:
    """Tests for Trunc truncation operator."""

    def test_trunc_tensortrain(self):
        """Trunc @ TensorTrain truncates ranks."""
        from ttglow.ttmatrix import Trunc

        torch.manual_seed(42)
        # Create a high-rank TT by adding two TTs
        tt1 = TensorTrain.random([3, 4, 5], ranks=[5, 5])
        tt2 = TensorTrain.random([3, 4, 5], ranks=[5, 5])
        from ttglow.tensortrain import add
        tt_sum = add(tt1, tt2)  # ranks become [10, 10]

        result = Trunc(3) @ tt_sum

        assert result.ranks[1] <= 3
        assert result.ranks[2] <= 3

    def test_trunc_ttmatrix(self):
        """Trunc @ TTMatrix truncates ranks."""
        from ttglow.ttmatrix import Trunc

        torch.manual_seed(42)
        H1 = TTMatrix.random([(3, 3), (4, 4)], ranks=[5])
        H2 = TTMatrix.random([(3, 3), (4, 4)], ranks=[5])
        H_sum = ttmatrix.add(H1, H2)  # rank becomes 10

        result = Trunc(3) @ H_sum

        assert result.ranks[1] <= 3

    def test_trunc_in_circuit(self):
        """Trunc works in circuit-style expressions."""
        from ttglow.ttmatrix import Trunc

        torch.manual_seed(42)
        psi = TensorTrain.random([2, 2, 2, 2], ranks=[2, 2, 2])

        X = TTMatrix.local(torch.tensor([[0., 1.], [1., 0.]], dtype=torch.float64))
        Z = TTMatrix.local(torch.tensor([[1., 0.], [0., -1.]], dtype=torch.float64))

        # Circuit with truncation
        result = Trunc(5) @ X.on(0) @ Trunc(5) @ Z.on(1) @ psi

        assert isinstance(result, TensorTrain)
        assert result.shape == (2, 2, 2, 2)
        # Ranks should be bounded
        for r in result.ranks[1:-1]:
            assert r <= 5

    def test_trunc_chained_operators(self):
        """Full circuit: Trunc @ U @ Trunc @ U @ psi."""
        from ttglow.ttmatrix import Trunc

        torch.manual_seed(42)
        psi = TensorTrain.random([2, 2, 2, 2], ranks=[2, 2, 2])

        # Some 2-site operator
        U = TTMatrix.random([(2, 2), (2, 2)], ranks=[3])

        # Apply U twice with truncation between
        result = Trunc(4) @ U.on(2, 3) @ Trunc(4) @ U.on(0, 1) @ psi

        assert isinstance(result, TensorTrain)
        assert result.shape == (2, 2, 2, 2)

    def test_trunc_with_tolerance(self):
        """Trunc with relative tolerance."""
        from ttglow.ttmatrix import Trunc

        torch.manual_seed(42)
        tt = TensorTrain.random([3, 4, 5], ranks=[10, 10])

        result = Trunc(rel_tol=0.1) @ tt

        # Should have reduced rank
        assert isinstance(result, TensorTrain)

    def test_trunc_repr(self):
        """Trunc has useful repr."""
        from ttglow.ttmatrix import Trunc

        t1 = Trunc(10)
        assert "max_rank=10" in repr(t1)

        t2 = Trunc(rel_tol=1e-5)
        assert "rel_tol" in repr(t2)

        t3 = Trunc(max_rank=5, abs_tol=1e-10)
        assert "max_rank=5" in repr(t3)
        assert "abs_tol" in repr(t3)

    def test_trunc_preserves_approximately(self):
        """Truncation should approximately preserve the tensor."""
        from ttglow.ttmatrix import Trunc

        torch.manual_seed(42)
        tt = TensorTrain.random([3, 4, 5], ranks=[3, 3])

        # Truncate with high max_rank (should preserve exactly)
        result = Trunc(10) @ tt

        T_orig = tt.to_tensor()
        T_result = result.to_tensor()

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