import torch
import pytest
from ttglow import TensorTrain
from ttglow import linalg


class TestQR:
    def test_qr_reconstruction(self):
        """Test that Q * R reconstructs the original tensor."""
        torch.manual_seed(42)
        tt = TensorTrain.random([4, 5, 6], ranks=[3, 4])

        # Perform QR decomposition
        qr_tt, R = linalg.qr(tt)

        # Contract R back into the last core
        # R has shape (r_new, r_final) where r_final should be 1
        # Last core has shape (r_left, n, r_new)
        last_core = qr_tt.cores[-1]
        r_left, n, r_new = last_core.shape

        # Reshape last core to matrix and multiply by R
        # last_core_mat: (r_left * n, r_new), R: (r_new, r_final)
        last_core_mat = last_core.reshape(r_left * n, r_new)
        reconstructed_mat = last_core_mat @ R
        qr_tt.cores[-1] = reconstructed_mat.reshape(r_left, n, -1)
        qr_tt.ranks[-1] = reconstructed_mat.shape[1]

        # Compare full tensors
        T_original = tt.to_tensor()
        T_reconstructed = qr_tt.to_tensor()

        assert torch.allclose(T_original, T_reconstructed, atol=1e-10)

    def test_qr_left_orthogonality(self):
        """Test that each core (except last) is left-orthogonal."""
        torch.manual_seed(42)
        tt = TensorTrain.random([4, 5, 6], ranks=[3, 4])

        qr_tt, R = linalg.qr(tt)

        # Check each core except the last
        for k in range(qr_tt.d - 1):
            core = qr_tt.cores[k]  # (r_left, n, r_right)
            r_left, n, r_right = core.shape

            # Reshape to matrix: (r_left * n, r_right)
            core_mat = core.reshape(r_left * n, r_right)

            # Check Q^T @ Q = I
            gram = core_mat.T @ core_mat
            identity = torch.eye(r_right, dtype=core.dtype, device=core.device)

            assert torch.allclose(gram, identity, atol=1e-10), \
                f"Core {k} is not left-orthogonal"

    def test_qr_preserves_tensor(self):
        """Test that the full tensor is preserved after QR and reconstruction."""
        torch.manual_seed(123)
        dims = [3, 4, 5, 6]
        ranks = [2, 3, 2]
        tt = TensorTrain.random(dims, ranks)

        T_original = tt.to_tensor()

        # Perform QR
        qr_tt, R = linalg.qr(tt)

        # Reconstruct by contracting R back
        last_core = qr_tt.cores[-1]
        r_left, n, r_new = last_core.shape
        last_core_mat = last_core.reshape(r_left * n, r_new)
        reconstructed_mat = last_core_mat @ R
        qr_tt.cores[-1] = reconstructed_mat.reshape(r_left, n, -1)
        qr_tt.ranks[-1] = reconstructed_mat.shape[1]

        T_reconstructed = qr_tt.to_tensor()

        assert torch.allclose(T_original, T_reconstructed, atol=1e-10)

    def test_qr_rank1(self):
        """Test QR on rank-1 tensor train."""
        tt = TensorTrain.ones([3, 4, 5], value=2.0)

        qr_tt, R = linalg.qr(tt)

        # Check orthogonality
        for k in range(qr_tt.d - 1):
            core = qr_tt.cores[k]
            r_left, n, r_right = core.shape
            core_mat = core.reshape(r_left * n, r_right)
            gram = core_mat.T @ core_mat
            identity = torch.eye(r_right, dtype=core.dtype, device=core.device)
            assert torch.allclose(gram, identity, atol=1e-10)

    def test_qr_output_shapes(self):
        """Test that QR output has correct shapes."""
        torch.manual_seed(42)
        tt = TensorTrain.random([4, 5, 6], ranks=[3, 4])

        qr_tt, R = linalg.qr(tt)

        # Check that qr_tt is still a valid tensor train
        assert qr_tt.d == tt.d
        assert qr_tt.dims == tt.dims

        # R should be a matrix
        assert R.dim() == 2

        # First rank should always be 1
        assert qr_tt.ranks[0] == 1

    def test_qr_different_ranks(self):
        """Test QR with various rank configurations."""
        torch.manual_seed(42)

        test_cases = [
            ([5, 5, 5], [2, 2]),
            ([10, 10], [5]),
            ([3, 4, 5, 6], [2, 3, 2]),
        ]

        for dims, ranks in test_cases:
            tt = TensorTrain.random(dims, ranks)
            qr_tt, R = linalg.qr(tt)

            # Check orthogonality of all cores except last
            for k in range(qr_tt.d - 1):
                core = qr_tt.cores[k]
                r_left, n, r_right = core.shape
                core_mat = core.reshape(r_left * n, r_right)
                gram = core_mat.T @ core_mat
                identity = torch.eye(r_right, dtype=core.dtype, device=core.device)
                assert torch.allclose(gram, identity, atol=1e-10), \
                    f"Failed for dims={dims}, ranks={ranks}, core={k}"


class TestQRStep:
    def test_qr_step_orthogonality(self):
        """Test that qr_step produces orthogonal Q."""
        torch.manual_seed(42)
        core = torch.randn(3, 5, 4, dtype=torch.float64)

        Q_core, R = linalg.qr_step(core)

        # Check shapes
        assert Q_core.shape[0] == 3  # r_left preserved
        assert Q_core.shape[1] == 5  # n preserved
        assert Q_core.shape[2] <= 4  # r_right can reduce

        # Check orthogonality: Q^T @ Q = I
        r_left, n, r_right = Q_core.shape
        Q_mat = Q_core.reshape(r_left * n, r_right)
        gram = Q_mat.T @ Q_mat
        identity = torch.eye(r_right, dtype=Q_core.dtype)

        assert torch.allclose(gram, identity, atol=1e-10)

    def test_qr_step_reconstruction(self):
        """Test that Q @ R reconstructs the original core matrix."""
        torch.manual_seed(42)
        core = torch.randn(3, 5, 4, dtype=torch.float64)

        Q_core, R = linalg.qr_step(core)

        # Reconstruct
        r_left, n, r_right = core.shape
        core_mat = core.reshape(r_left * n, r_right)

        Q_mat = Q_core.reshape(r_left * n, -1)
        reconstructed = Q_mat @ R

        assert torch.allclose(core_mat, reconstructed, atol=1e-10)


class TestNorm:
    def test_norm_against_full_tensor(self):
        """Test that norm matches the full tensor norm."""
        torch.manual_seed(42)
        tt = TensorTrain.random([4, 5, 6], ranks=[3, 4])

        # Compute norm via TT
        norm_tt = linalg.norm(tt)

        # Compute norm via full tensor
        T = tt.to_tensor()
        norm_full = torch.linalg.norm(T).item()

        assert abs(norm_tt - norm_full) < 1e-10

    def test_norm_squared_against_full_tensor(self):
        """Test that norm_squared matches the full tensor norm squared."""
        torch.manual_seed(42)
        tt = TensorTrain.random([4, 5, 6], ranks=[3, 4])

        # Compute norm squared via TT
        norm_sq_tt = linalg.norm_squared(tt)

        # Compute norm squared via full tensor
        T = tt.to_tensor()
        norm_sq_full = (T * T).sum().item()

        assert abs(norm_sq_tt - norm_sq_full) < 1e-10

    def test_norm_relationship(self):
        """Test that norm = sqrt(norm_squared)."""
        torch.manual_seed(42)
        tt = TensorTrain.random([3, 4, 5], ranks=[2, 3])

        norm_val = linalg.norm(tt)
        norm_sq_val = linalg.norm_squared(tt)

        assert abs(norm_val ** 2 - norm_sq_val) < 1e-10

    def test_norm_positive(self):
        """Test that norm is always positive for non-zero tensors."""
        torch.manual_seed(42)
        tt = TensorTrain.random([4, 5, 6], ranks=[3, 4])

        norm_val = linalg.norm(tt)

        assert norm_val > 0

    def test_norm_zero_tensor(self):
        """Test that norm is zero for zero tensor."""
        tt = TensorTrain([3, 4, 5], ranks=[2, 3])
        # Cores are initialized to zeros

        norm_val = linalg.norm(tt)

        assert abs(norm_val) < 1e-10

    def test_norm_ones(self):
        """Test norm of ones tensor."""
        dims = [3, 4, 5]
        tt = TensorTrain.ones(dims, value=1.0)

        norm_val = linalg.norm(tt)

        # Total number of elements
        n_elements = 3 * 4 * 5
        expected_norm = (n_elements ** 0.5)

        assert abs(norm_val - expected_norm) < 1e-10

    def test_norm_scaling(self):
        """Test that norm(alpha * tt) = |alpha| * norm(tt)."""
        torch.manual_seed(42)
        tt = TensorTrain.random([3, 4, 5], ranks=[2, 3])
        alpha = 2.5

        from ttglow import scale
        tt_scaled = scale(tt, alpha)

        norm_original = linalg.norm(tt)
        norm_scaled = linalg.norm(tt_scaled)

        assert abs(norm_scaled - abs(alpha) * norm_original) < 1e-10

    def test_norm_different_sizes(self):
        """Test norm on various tensor sizes."""
        torch.manual_seed(42)

        test_cases = [
            ([5, 5, 5], [2, 2]),
            ([10, 10], [5]),
            ([3, 4, 5, 6], [2, 3, 2]),
        ]

        for dims, ranks in test_cases:
            tt = TensorTrain.random(dims, ranks)
            T = tt.to_tensor()

            norm_tt = linalg.norm(tt)
            norm_full = torch.linalg.norm(T).item()

            assert abs(norm_tt - norm_full) < 1e-10, \
                f"Failed for dims={dims}, ranks={ranks}"


class TestQRDirectional:
    def test_qr_right_to_left_orthogonality(self):
        """Test that right-to-left QR produces right-orthogonal cores."""
        torch.manual_seed(42)
        tt = TensorTrain.random([4, 5, 6], ranks=[3, 4])

        qr_tt, L = linalg.qr(tt, direction='right-to-left')

        # Check each core except the first (leftmost)
        for k in range(1, qr_tt.d):
            core = qr_tt.cores[k]  # (r_left, n, r_right)
            r_left, n, r_right = core.shape

            # Reshape to matrix: (r_left, n * r_right)
            core_mat = core.reshape(r_left, n * r_right)

            # Check Q @ Q^T = I (right-orthogonal)
            gram = core_mat @ core_mat.T
            identity = torch.eye(r_left, dtype=core.dtype, device=core.device)

            assert torch.allclose(gram, identity, atol=1e-10), \
                f"Core {k} is not right-orthogonal"

    def test_qr_right_to_left_reconstruction(self):
        """Test that right-to-left QR can reconstruct original tensor."""
        torch.manual_seed(42)
        tt = TensorTrain.random([4, 5, 6], ranks=[3, 4])

        T_original = tt.to_tensor()

        qr_tt, L = linalg.qr(tt, direction='right-to-left')

        # Contract L back into the first core
        first_core = qr_tt.cores[0]
        r_left, n, r_right = first_core.shape
        first_core_mat = first_core.reshape(r_left, n * r_right)
        reconstructed_mat = L @ first_core_mat
        qr_tt.cores[0] = reconstructed_mat.reshape(-1, n, r_right)
        qr_tt.ranks[0] = reconstructed_mat.shape[0]

        T_reconstructed = qr_tt.to_tensor()

        assert torch.allclose(T_original, T_reconstructed, atol=1e-10)

    def test_qr_both_directions_preserve_tensor(self):
        """Test that both directions preserve tensor values."""
        torch.manual_seed(123)
        tt = TensorTrain.random([3, 4, 5], ranks=[2, 3])

        T_original = tt.to_tensor()

        # Test left-to-right
        qr_lr, R = linalg.qr(tt, direction='left-to-right')
        last_core = qr_lr.cores[-1]
        r_left, n, r_right = last_core.shape
        last_core_mat = last_core.reshape(r_left * n, r_right)
        reconstructed = (last_core_mat @ R).reshape(r_left, n, -1)
        qr_lr.cores[-1] = reconstructed
        qr_lr.ranks[-1] = reconstructed.shape[2]
        T_lr = qr_lr.to_tensor()

        assert torch.allclose(T_original, T_lr, atol=1e-10)

        # Test right-to-left
        qr_rl, L = linalg.qr(tt, direction='right-to-left')
        first_core = qr_rl.cores[0]
        r_left, n, r_right = first_core.shape
        first_core_mat = first_core.reshape(r_left, n * r_right)
        reconstructed = (L @ first_core_mat).reshape(-1, n, r_right)
        qr_rl.cores[0] = reconstructed
        qr_rl.ranks[0] = reconstructed.shape[0]
        T_rl = qr_rl.to_tensor()

        assert torch.allclose(T_original, T_rl, atol=1e-10)

    def test_qr_invalid_direction(self):
        """Test that invalid direction raises error."""
        tt = TensorTrain.random([3, 4, 5], ranks=[2, 3])

        with pytest.raises(ValueError, match="Invalid direction"):
            linalg.qr(tt, direction='invalid')


class TestSVDDirectional:
    def test_svd_right_to_left_orthogonality(self):
        """Test that right-to-left SVD produces right-orthogonal cores."""
        torch.manual_seed(42)
        tt = TensorTrain.random([4, 5, 6], ranks=[3, 4])

        svd_tt = linalg.svd(tt, direction='right-to-left', max_rank=10)

        # Check each core except the first is right-orthogonal
        for k in range(1, svd_tt.d):
            core = svd_tt.cores[k]
            r_left, n, r_right = core.shape
            core_mat = core.reshape(r_left, n * r_right)
            gram = core_mat @ core_mat.T
            identity = torch.eye(r_left, dtype=core.dtype)

            assert torch.allclose(gram, identity, atol=1e-10), \
                f"Core {k} is not right-orthogonal"

    def test_svd_both_directions_preserve_tensor(self):
        """Test that SVD in both directions preserves tensor."""
        torch.manual_seed(42)
        tt = TensorTrain.random([3, 4, 5], ranks=[2, 3])

        T_original = tt.to_tensor()

        # Test left-to-right
        svd_lr = linalg.svd(tt, direction='left-to-right')
        T_lr = svd_lr.to_tensor()
        assert torch.allclose(T_original, T_lr, atol=1e-10)

        # Test right-to-left
        svd_rl = linalg.svd(tt, direction='right-to-left')
        T_rl = svd_rl.to_tensor()
        assert torch.allclose(T_original, T_rl, atol=1e-10)

    def test_svd_direction_with_rank_truncation(self):
        """Test SVD with rank truncation in both directions."""
        torch.manual_seed(42)
        tt = TensorTrain.random([5, 5, 5], ranks=[10, 10])

        # Truncate with both directions
        svd_lr = linalg.svd(tt, max_rank=3, direction='left-to-right')
        svd_rl = linalg.svd(tt, max_rank=3, direction='right-to-left')

        # Check ranks are reduced
        assert all(r <= 3 for r in svd_lr.ranks)
        assert all(r <= 3 for r in svd_rl.ranks)

        # Both should approximate original (may differ due to truncation)
        T_original = tt.to_tensor()
        T_lr = svd_lr.to_tensor()
        T_rl = svd_rl.to_tensor()

        # Check they're reasonable approximations
        # (Can't check exact norm without from_tensor, but can check tensor norms)
        error_lr = torch.linalg.norm(T_original - T_lr).item()
        error_rl = torch.linalg.norm(T_original - T_rl).item()

        # Errors should be bounded (truncation introduces some error)
        # With aggressive truncation (rank 10 -> 3), errors can be significant
        assert error_lr < 100.0, f"L-R error too large: {error_lr}"
        assert error_rl < 100.0, f"R-L error too large: {error_rl}"


class TestSum:
    def test_sum_all_against_full_tensor(self):
        """Test that sum over all dimensions matches full tensor sum."""
        torch.manual_seed(42)
        tt = TensorTrain.random([4, 5, 6], ranks=[2, 2])

        # Convert to dense tensor
        T = tt.to_tensor()

        # Sum via full tensor
        sum_full = T.sum().item()

        # Sum via TT
        sum_tt = linalg.sum(tt)

        assert abs(sum_tt - sum_full) < 1e-10

    def test_sum_single_dim_against_full_tensor(self):
        """Test that sum over a single dimension matches full tensor."""
        torch.manual_seed(42)
        tt = TensorTrain.random([4, 5, 6], ranks=[2, 2])

        T = tt.to_tensor()

        # Test sum over each dimension
        for dim in range(3):
            # Sum via full tensor
            T_summed = T.sum(dim=dim)

            # Sum via TT
            tt_summed = linalg.sum(tt, dim=dim)
            T_tt_summed = tt_summed.to_tensor()

            assert torch.allclose(T_tt_summed, T_summed, atol=1e-10), \
                f"Sum over dim={dim} failed"

    def test_sum_preserves_other_dims(self):
        """Test that summing over one dim preserves shape of others."""
        torch.manual_seed(42)
        tt = TensorTrain.random([4, 5, 6], ranks=[2, 2])

        # Sum over middle dimension
        tt_summed = linalg.sum(tt, dim=1)

        # Should have shape (4, 6)
        assert tt_summed.shape == (4, 6)

    def test_sum_first_dim(self):
        """Test sum over first dimension."""
        torch.manual_seed(42)
        tt = TensorTrain.random([4, 5, 6], ranks=[2, 2])

        T = tt.to_tensor()
        T_summed = T.sum(dim=0)

        tt_summed = linalg.sum(tt, dim=0)
        T_tt_summed = tt_summed.to_tensor()

        assert T_tt_summed.shape == (5, 6)
        assert torch.allclose(T_tt_summed, T_summed, atol=1e-10)

    def test_sum_last_dim(self):
        """Test sum over last dimension."""
        torch.manual_seed(42)
        tt = TensorTrain.random([4, 5, 6], ranks=[2, 2])

        T = tt.to_tensor()
        T_summed = T.sum(dim=2)

        tt_summed = linalg.sum(tt, dim=2)
        T_tt_summed = tt_summed.to_tensor()

        assert T_tt_summed.shape == (4, 5)
        assert torch.allclose(T_tt_summed, T_summed, atol=1e-10)

    def test_sum_ones_tensor(self):
        """Test sum on a tensor of ones."""
        dims = [3, 4, 5]
        tt = TensorTrain.ones(dims, value=1.0)

        # Sum all: should equal numel
        sum_all = linalg.sum(tt)
        assert abs(sum_all - 3 * 4 * 5) < 1e-10

        # Sum over dim 0: should give tensor of 3's with shape (4, 5)
        tt_sum0 = linalg.sum(tt, dim=0)
        assert tt_sum0.shape == (4, 5)
        T_sum0 = tt_sum0.to_tensor()
        assert torch.allclose(T_sum0, 3.0 * torch.ones(4, 5, dtype=torch.float64), atol=1e-10)


class TestContract:
    """Tests for partial contraction of two tensor trains."""

    def test_contract_single_dim_middle(self):
        """Contract over a single middle dimension."""
        torch.manual_seed(42)
        tt1 = TensorTrain.random([4, 5, 6], ranks=[2, 2])
        tt2 = TensorTrain.random([4, 5, 6], ranks=[2, 2])

        # Contract over dimension 1 (size 5)
        tt_result = linalg.contract(tt1, tt2, dims=[1])

        # Compare with full tensors
        T1 = tt1.to_tensor()
        T2 = tt2.to_tensor()
        # einsum: contract j index, keep i and k
        T_expected = torch.einsum('ijk,ijk->ik', T1, T2)
        T_result = tt_result.to_tensor()

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

    def test_contract_single_dim_first(self):
        """Contract over the first dimension."""
        torch.manual_seed(42)
        tt1 = TensorTrain.random([4, 5, 6], ranks=[2, 2])
        tt2 = TensorTrain.random([4, 5, 6], ranks=[2, 2])

        # Contract over dimension 0
        tt_result = linalg.contract(tt1, tt2, dims=[0])

        T1 = tt1.to_tensor()
        T2 = tt2.to_tensor()
        T_expected = torch.einsum('ijk,ijk->jk', T1, T2)
        T_result = tt_result.to_tensor()

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

    def test_contract_single_dim_last(self):
        """Contract over the last dimension."""
        torch.manual_seed(42)
        tt1 = TensorTrain.random([4, 5, 6], ranks=[2, 2])
        tt2 = TensorTrain.random([4, 5, 6], ranks=[2, 2])

        # Contract over dimension 2
        tt_result = linalg.contract(tt1, tt2, dims=[2])

        T1 = tt1.to_tensor()
        T2 = tt2.to_tensor()
        T_expected = torch.einsum('ijk,ijk->ij', T1, T2)
        T_result = tt_result.to_tensor()

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

    def test_contract_multiple_dims(self):
        """Contract over multiple dimensions."""
        torch.manual_seed(42)
        tt1 = TensorTrain.random([4, 5, 6, 7], ranks=[2, 2, 2])
        tt2 = TensorTrain.random([4, 5, 6, 7], ranks=[2, 2, 2])

        # Contract over dimensions 1 and 3
        tt_result = linalg.contract(tt1, tt2, dims=[1, 3])

        T1 = tt1.to_tensor()
        T2 = tt2.to_tensor()
        # Contract j and l, keep i and k
        T_expected = torch.einsum('ijkl,ijkl->ik', T1, T2)
        T_result = tt_result.to_tensor()

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

    def test_contract_contiguous_dims(self):
        """Contract over contiguous dimensions (common for density matrices)."""
        torch.manual_seed(42)
        tt1 = TensorTrain.random([3, 4, 5, 6], ranks=[2, 2, 2])
        tt2 = TensorTrain.random([3, 4, 5, 6], ranks=[2, 2, 2])

        # Contract over dimensions 2 and 3 (trace out "right half")
        tt_result = linalg.contract(tt1, tt2, dims=[2, 3])

        T1 = tt1.to_tensor()
        T2 = tt2.to_tensor()
        T_expected = torch.einsum('ijkl,ijkl->ij', T1, T2)
        T_result = tt_result.to_tensor()

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

    def test_contract_all_dims_equals_dot(self):
        """Contracting all dims should equal dot product."""
        torch.manual_seed(42)
        tt1 = TensorTrain.random([4, 5, 6], ranks=[2, 2])
        tt2 = TensorTrain.random([4, 5, 6], ranks=[2, 2])

        # Contract all dimensions
        result = linalg.contract(tt1, tt2, dims=[0, 1, 2])

        # Compare with dot
        from ttglow import dot
        dot_result = dot(tt1, tt2)[-1].item()

        assert abs(result - dot_result) < 1e-10

    def test_contract_no_dims_equals_hadamard(self):
        """Contracting no dims should equal hadamard product."""
        torch.manual_seed(42)
        tt1 = TensorTrain.random([4, 5, 6], ranks=[2, 2])
        tt2 = TensorTrain.random([4, 5, 6], ranks=[2, 2])

        # Contract no dimensions
        tt_result = linalg.contract(tt1, tt2, dims=[])

        # Compare with hadamard
        from ttglow import hadamard
        tt_hadamard = hadamard(tt1, tt2)

        T_result = tt_result.to_tensor()
        T_hadamard = tt_hadamard.to_tensor()

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

    def test_contract_different_ranks(self):
        """Contract tensors with different TT ranks."""
        torch.manual_seed(42)
        tt1 = TensorTrain.random([4, 5, 6], ranks=[2, 3])
        tt2 = TensorTrain.random([4, 5, 6], ranks=[3, 2])

        tt_result = linalg.contract(tt1, tt2, dims=[1])

        T1 = tt1.to_tensor()
        T2 = tt2.to_tensor()
        T_expected = torch.einsum('ijk,ijk->ik', T1, T2)
        T_result = tt_result.to_tensor()

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

    def test_contract_self_for_norm(self):
        """Contract tensor with itself over all dims gives norm squared."""
        torch.manual_seed(42)
        tt = TensorTrain.random([4, 5, 6], ranks=[2, 2])

        result = linalg.contract(tt, tt, dims=[0, 1, 2])

        # Compare with norm_squared
        norm_sq = linalg.norm_squared(tt)

        assert abs(result - norm_sq) < 1e-10

    def test_contract_mismatched_dims_raises(self):
        """Contracting tensors with mismatched dimensions should raise."""
        tt1 = TensorTrain.random([4, 5, 6], ranks=[2, 2])
        tt2 = TensorTrain.random([4, 5, 7], ranks=[2, 2])

        with pytest.raises(ValueError):
            linalg.contract(tt1, tt2, dims=[1])

    def test_contract_invalid_dim_raises(self):
        """Contracting over invalid dimension should raise."""
        tt1 = TensorTrain.random([4, 5, 6], ranks=[2, 2])
        tt2 = TensorTrain.random([4, 5, 6], ranks=[2, 2])

        with pytest.raises(ValueError):
            linalg.contract(tt1, tt2, dims=[3])  # dim 3 doesn't exist


class TestTensordot:
    """Tests for tensordot (generalized contraction keeping dims from both tensors)."""

    def test_tensordot_single_dim(self):
        """Contract last dim of tt1 with first dim of tt2."""
        torch.manual_seed(42)
        # tt1: shape (4, 5), tt2: shape (5, 6)
        # Contract tt1's dim 1 with tt2's dim 0 (both size 5)
        # Result shape: (4, 6)
        tt1 = TensorTrain.random([4, 5], ranks=[2])
        tt2 = TensorTrain.random([5, 6], ranks=[2])

        tt_result = linalg.tensordot(tt1, tt2, dims=1)

        T1 = tt1.to_tensor()
        T2 = tt2.to_tensor()
        T_expected = torch.tensordot(T1, T2, dims=1)
        T_result = tt_result.to_tensor()

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

    def test_tensordot_multiple_dims(self):
        """Contract last 2 dims of tt1 with first 2 dims of tt2."""
        torch.manual_seed(42)
        # tt1: shape (3, 4, 5), tt2: shape (4, 5, 6)
        # Contract tt1's dims [1,2] with tt2's dims [0,1]
        # Result shape: (3, 6)
        tt1 = TensorTrain.random([3, 4, 5], ranks=[2, 2])
        tt2 = TensorTrain.random([4, 5, 6], ranks=[2, 2])

        tt_result = linalg.tensordot(tt1, tt2, dims=2)

        T1 = tt1.to_tensor()
        T2 = tt2.to_tensor()
        T_expected = torch.tensordot(T1, T2, dims=2)
        T_result = tt_result.to_tensor()

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

    def test_tensordot_density_matrix(self):
        """Use tensordot to compute reduced density matrix."""
        torch.manual_seed(42)
        # State psi with shape (n_system, n_environment)
        # Reduced density: rho[i, j] = sum_k psi[i, k] * psi[j, k]
        psi = TensorTrain.random([4, 5], ranks=[3])

        # Contract psi's last dim with psi's last dim
        # tensordot(psi, psi, dims=([1], [1])) should give shape (4, 4)
        rho = linalg.tensordot(psi, psi, dims=([1], [1]))

        T_psi = psi.to_tensor()
        T_rho_expected = torch.einsum('ik,jk->ij', T_psi, T_psi)
        T_rho = rho.to_tensor()

        assert T_rho.shape == (4, 4)
        assert torch.allclose(T_rho, T_rho_expected, atol=1e-10)

    def test_tensordot_density_matrix_multisite(self):
        """Reduced density matrix tracing out multiple sites."""
        torch.manual_seed(42)
        # psi with shape (n0, n1, n2, n3) - trace out sites 2,3
        # rho[i,j,l,k] = sum_{m,n} psi[i,j,m,n] * psi[k,l,m,n]
        # Note: Due to TT structure constraints, tt2's kept dims appear in
        # reversed order. This is mathematically equivalent to:
        # rho[i,j,k,l] = sum_{m,n} psi[i,j,m,n] * psi[l,k,m,n]
        psi = TensorTrain.random([3, 4, 5, 6], ranks=[2, 2, 2])

        # Contract last 2 dims of psi with last 2 dims of psi
        rho = linalg.tensordot(psi, psi, dims=([2, 3], [2, 3]))

        T_psi = psi.to_tensor()
        # Expected: rho[i,j,l,k] due to reversed tt2 dims in TT format
        T_rho_expected = torch.einsum('ijkl,mnkl->ijnm', T_psi, T_psi)
        T_rho = rho.to_tensor()

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

    def test_tensordot_outer_product(self):
        """dims=0 should give outer product."""
        torch.manual_seed(42)
        tt1 = TensorTrain.random([3, 4], ranks=[2])
        tt2 = TensorTrain.random([5, 6], ranks=[2])

        tt_result = linalg.tensordot(tt1, tt2, dims=0)

        T1 = tt1.to_tensor()
        T2 = tt2.to_tensor()
        T_expected = torch.tensordot(T1, T2, dims=0)
        T_result = tt_result.to_tensor()

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

    def test_tensordot_full_contraction(self):
        """Contracting all dims should give scalar (like dot)."""
        torch.manual_seed(42)
        tt1 = TensorTrain.random([4, 5, 6], ranks=[2, 2])
        tt2 = TensorTrain.random([4, 5, 6], ranks=[2, 2])

        result = linalg.tensordot(tt1, tt2, dims=3)

        # Compare with dot
        from ttglow import dot
        dot_result = dot(tt1, tt2)[-1].item()

        assert abs(result - dot_result) < 1e-10

    def test_tensordot_different_ranks(self):
        """Tensordot with different TT ranks."""
        torch.manual_seed(42)
        tt1 = TensorTrain.random([3, 4, 5], ranks=[2, 4])
        tt2 = TensorTrain.random([5, 6], ranks=[3])

        tt_result = linalg.tensordot(tt1, tt2, dims=1)

        T1 = tt1.to_tensor()
        T2 = tt2.to_tensor()
        T_expected = torch.tensordot(T1, T2, dims=1)
        T_result = tt_result.to_tensor()

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

    def test_tensordot_explicit_dims(self):
        """Test with explicit dimension specification."""
        torch.manual_seed(42)
        tt1 = TensorTrain.random([3, 4, 5], ranks=[2, 2])
        tt2 = TensorTrain.random([5, 6, 7], ranks=[2, 2])

        # Contract tt1's dim 2 with tt2's dim 0
        tt_result = linalg.tensordot(tt1, tt2, dims=([2], [0]))

        T1 = tt1.to_tensor()
        T2 = tt2.to_tensor()
        T_expected = torch.tensordot(T1, T2, dims=([2], [0]))
        T_result = tt_result.to_tensor()

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

    def test_tensordot_mismatched_dims_raises(self):
        """Mismatched contraction dimensions should raise."""
        tt1 = TensorTrain.random([3, 4], ranks=[2])
        tt2 = TensorTrain.random([5, 6], ranks=[2])

        with pytest.raises(ValueError):
            linalg.tensordot(tt1, tt2, dims=1)  # tt1's dim 1 (4) != tt2's dim 0 (5)
