import torch
import pytest
from ttglow import TensorTrain, dot, hadamard, scale, add


class TestTensorTrain:
    def test_creation(self):
        dims = [4, 5, 6]
        ranks = [2, 3]
        tt = TensorTrain(dims, ranks)

        assert tt.d == 3
        assert tt.dims == [4, 5, 6]
        assert tt.ranks == [1, 2, 3, 1]
        assert len(tt) == 3

    def test_core_shapes(self):
        dims = [4, 5, 6]
        ranks = [2, 3]
        tt = TensorTrain(dims, ranks)

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

    def test_random(self):
        dims = [4, 5, 6]
        ranks = [2, 3]
        tt = TensorTrain.random(dims, ranks)

        # Check cores are non-zero (random)
        for k in range(tt.d):
            assert tt[k].abs().sum() > 0

    def test_default_ranks(self):
        dims = [4, 5, 6]
        tt = TensorTrain(dims)  # No ranks specified

        assert tt.ranks == [1, 1, 1, 1]
        assert tt[0].shape == (1, 4, 1)
        assert tt[1].shape == (1, 5, 1)
        assert tt[2].shape == (1, 6, 1)

    def test_ones(self):
        tt = TensorTrain.ones([3, 4, 5])
        T = tt.to_tensor()

        assert T.shape == (3, 4, 5)
        assert torch.allclose(T, torch.ones(3, 4, 5, dtype=torch.float64))

    def test_ones_with_value(self):
        tt = TensorTrain.ones([3, 4, 5], value=3.14)
        T = tt.to_tensor()

        expected = 3.14 * torch.ones(3, 4, 5, dtype=torch.float64)
        assert torch.allclose(T, expected)

    def test_ones_ranks(self):
        tt = TensorTrain.ones([3, 4, 5])
        assert tt.ranks == [1, 1, 1, 1]  # All rank-1

    def test_zeros(self):
        tt = TensorTrain.zeros([3, 4, 5])
        full = tt.to_tensor()
        assert torch.allclose(full, torch.zeros(3, 4, 5, dtype=torch.float64))

    def test_zeros_with_ranks(self):
        tt = TensorTrain.zeros([3, 4, 5], ranks=[2, 3])
        assert tt.ranks == [1, 2, 3, 1]
        full = tt.to_tensor()
        assert torch.allclose(full, torch.zeros(3, 4, 5, dtype=torch.float64))

    def test_from_dense_exact(self):
        """Test from_dense produces exact representation."""
        torch.manual_seed(42)
        dense = torch.randn(4, 5, 6, dtype=torch.float64)
        tt = TensorTrain.from_dense(dense)

        reconstructed = tt.to_tensor()
        assert torch.allclose(dense, reconstructed, atol=1e-10)

    def test_from_dense_1d(self):
        """Test from_dense with 1D tensor."""
        dense = torch.randn(10, dtype=torch.float64)
        tt = TensorTrain.from_dense(dense)

        assert tt.d == 1
        assert tt.shape == (10,)
        assert torch.allclose(dense, tt.to_tensor(), atol=1e-10)

    def test_from_dense_2d(self):
        """Test from_dense with 2D tensor."""
        torch.manual_seed(42)
        dense = torch.randn(4, 5, dtype=torch.float64)
        tt = TensorTrain.from_dense(dense)

        assert tt.d == 2
        assert tt.shape == (4, 5)
        assert torch.allclose(dense, tt.to_tensor(), atol=1e-10)

    def test_from_dense_4d(self):
        """Test from_dense with 4D tensor."""
        torch.manual_seed(42)
        dense = torch.randn(3, 4, 5, 6, dtype=torch.float64)
        tt = TensorTrain.from_dense(dense)

        assert tt.d == 4
        assert tt.shape == (3, 4, 5, 6)
        assert torch.allclose(dense, tt.to_tensor(), atol=1e-10)

    def test_from_dense_max_rank(self):
        """Test from_dense with max_rank truncation."""
        torch.manual_seed(42)
        dense = torch.randn(8, 8, 8, dtype=torch.float64)
        tt = TensorTrain.from_dense(dense, max_rank=3)

        # Ranks should be at most 3
        for r in tt.ranks[1:-1]:
            assert r <= 3

        # Should still be a reasonable approximation
        reconstructed = tt.to_tensor()
        rel_error = (dense - reconstructed).norm() / dense.norm()
        assert rel_error < 0.9  # Rough approximation (heavily truncated)

    def test_from_dense_rel_tol(self):
        """Test from_dense with relative tolerance."""
        torch.manual_seed(42)
        # Create low-rank tensor
        a = torch.randn(4, 2, dtype=torch.float64)
        b = torch.randn(2, 5, dtype=torch.float64)
        c = torch.randn(2, 6, dtype=torch.float64)
        dense = torch.einsum('ir,rj,rk->ijk', a, b, c)

        # With tight tolerance, should recover exactly
        tt = TensorTrain.from_dense(dense, rel_tol=1e-10)
        assert torch.allclose(dense, tt.to_tensor(), atol=1e-8)

        # Ranks should be small (original was rank 2)
        assert max(tt.ranks[1:-1]) <= 4

    def test_from_dense_abs_tol(self):
        """Test from_dense with absolute tolerance."""
        torch.manual_seed(42)
        dense = torch.randn(4, 5, 6, dtype=torch.float64)
        tt = TensorTrain.from_dense(dense, abs_tol=0.5)

        # Some truncation should occur
        reconstructed = tt.to_tensor()
        # Should still preserve large-scale structure
        assert reconstructed.shape == dense.shape

    def test_from_dense_preserves_dtype(self):
        """Test from_dense preserves dtype."""
        dense32 = torch.randn(3, 4, 5, dtype=torch.float32)
        tt32 = TensorTrain.from_dense(dense32)
        assert tt32.dtype == torch.float32

        dense64 = torch.randn(3, 4, 5, dtype=torch.float64)
        tt64 = TensorTrain.from_dense(dense64)
        assert tt64.dtype == torch.float64

    def test_from_dense_scalar_raises(self):
        """Test from_dense raises for scalar input."""
        scalar = torch.tensor(5.0)
        with pytest.raises(ValueError):
            TensorTrain.from_dense(scalar)

    def test_setitem(self):
        tt = TensorTrain([4, 5], [2])
        new_core = torch.ones(1, 4, 2)
        tt[0] = new_core
        assert torch.allclose(tt[0], new_core)

    def test_setitem_wrong_shape(self):
        tt = TensorTrain([4, 5], [2])
        wrong_core = torch.ones(1, 4, 3)  # Wrong right rank
        with pytest.raises(ValueError):
            tt[0] = wrong_core

    def test_shape(self):
        tt = TensorTrain([4, 5, 6], [2, 3])
        assert tt.shape == (4, 5, 6)

    def test_ndim(self):
        tt = TensorTrain([4, 5, 6], [2, 3])
        assert tt.ndim == 3

    def test_numel(self):
        tt = TensorTrain([4, 5, 6], [2, 3])
        assert tt.numel == 4 * 5 * 6

    def test_numel_rank1(self):
        tt = TensorTrain([3, 4, 5])
        assert tt.numel == 3 * 4 * 5

    def test_clone(self):
        tt = TensorTrain.random([4, 5, 6], [2, 3])
        tt_clone = tt.clone()

        # Check it's a true copy
        for k in range(tt.d):
            assert torch.allclose(tt[k], tt_clone[k])

        # Modify original, clone should be unchanged
        tt[0] = torch.zeros_like(tt[0])
        assert not torch.allclose(tt[0], tt_clone[0])

    def test_detach(self):
        tt = TensorTrain.random([4, 5, 6], [2, 3])
        # Enable gradients on original
        for core in tt.cores:
            core.requires_grad_(True)

        tt_detached = tt.detach()

        # Check values are preserved
        for k in range(tt.d):
            assert torch.allclose(tt[k], tt_detached[k])

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

    def test_detach_preserves_shape(self):
        tt = TensorTrain.random([3, 4, 5], [2, 3])
        tt_detached = tt.detach()

        assert tt_detached.shape == tt.shape
        assert tt_detached.ranks == tt.ranks
        assert tt_detached.dtype == tt.dtype

    def test_to_dtype(self):
        tt = TensorTrain.random([4, 5, 6], [2, 3], dtype=torch.float32)
        tt_f64 = tt.to(dtype=torch.float64)

        assert tt.dtype == torch.float32
        assert tt_f64.dtype == torch.float64
        for core in tt_f64.cores:
            assert core.dtype == torch.float64

    def test_to_device_cpu(self):
        tt = TensorTrain.random([4, 5, 6], [2, 3], device="cpu")
        tt_cpu = tt.to(device="cpu")

        assert tt_cpu.device == "cpu"
        for core in tt_cpu.cores:
            assert core.device.type == "cpu"

    def test_tt_shape(self):
        tt = TensorTrain([4, 5, 6], [2, 3])
        shapes = tt.tt_shape()

        assert shapes == [(1, 4, 2), (2, 5, 3), (3, 6, 1)]

    def test_tt_ranks(self):
        tt = TensorTrain([4, 5, 6], [2, 3])
        ranks = tt.tt_ranks()

        assert ranks == (1, 2, 3, 1)

    def test_ncores(self):
        tt = TensorTrain([4, 5, 6], [2, 3])
        assert tt.ncores() == 3

    def test_tt_cores(self):
        tt = TensorTrain([4, 5, 6], [2, 3])
        cores = tt.tt_cores()

        assert len(cores) == 3
        assert cores[0].shape == (1, 4, 2)
        assert cores[1].shape == (2, 5, 3)
        assert cores[2].shape == (3, 6, 1)

    def test_storage_numel(self):
        tt = TensorTrain([4, 5, 6], [2, 3])
        # (1*4*2) + (2*5*3) + (3*6*1) = 8 + 30 + 18 = 56
        assert tt.storage_numel() == 56

    def test_compression_ratio(self):
        tt = TensorTrain([4, 5, 6], [2, 3])
        # logical: 4*5*6 = 120
        # storage: (1*4*2) + (2*5*3) + (3*6*1) = 56
        expected_ratio = 120.0 / 56.0
        assert abs(tt.compression_ratio() - expected_ratio) < 1e-10

    def test_to_dense(self):
        tt = TensorTrain.random([4, 5, 6], [2, 3])
        T1 = tt.to_dense()
        T2 = tt.to_tensor()

        assert torch.allclose(T1, T2)
        assert T1.shape == (4, 5, 6)

    def test_total_params(self):
        tt = TensorTrain([4, 5, 6], [2, 3])
        # (1*4*2) + (2*5*3) + (3*6*1) = 8 + 30 + 18 = 56
        assert tt.total_params() == 56

    def test_copy(self):
        tt = TensorTrain.random([4, 5, 6], [2, 3])
        tt_copy = tt.copy()

        # Check it's a true copy
        for k in range(tt.d):
            assert torch.allclose(tt[k], tt_copy[k])

        # Modify original, copy should be unchanged
        tt[0] = torch.zeros_like(tt[0])
        assert not torch.allclose(tt[0], tt_copy[0])

    def test_to_tensor_shape(self):
        tt = TensorTrain.random([4, 5, 6], [2, 3])
        T = tt.to_tensor()
        assert T.shape == (4, 5, 6)

    def test_to_tensor_rank1(self):
        # Rank-1 TT: full tensor is outer product of vectors
        tt = TensorTrain([3, 4, 5])  # Default rank-1
        tt[0] = torch.arange(3).reshape(1, 3, 1).double()
        tt[1] = torch.arange(4).reshape(1, 4, 1).double()
        tt[2] = torch.arange(5).reshape(1, 5, 1).double()

        T = tt.to_tensor()
        # T[i,j,k] = i * j * k
        assert T[2, 3, 4] == 2 * 3 * 4

    def test_to_tensor_norm(self):
        # Verify norm via full tensor matches dot product
        torch.manual_seed(123)
        tt = TensorTrain.random([3, 4, 5], [2, 3])

        T = tt.to_tensor()
        norm_full = torch.sqrt((T * T).sum()).item()
        norm_dot = dot(tt, tt)[-1].item() ** 0.5

        assert abs(norm_full - norm_dot) < 1e-12


class TestDot:
    def test_dot_shapes(self):
        tt1 = TensorTrain.random([4, 5, 6], [2, 3])
        tt2 = TensorTrain.random([4, 5, 6], [2, 3])

        matrices = dot(tt1, tt2)

        assert len(matrices) == 4  # d+1 matrices
        assert matrices[0].shape == (1, 1)  # L_0
        assert matrices[1].shape == (2, 2)  # L_1
        assert matrices[2].shape == (3, 3)  # L_2
        assert matrices[3].shape == (1, 1)  # L_3 (inner product)

    def test_dot_self_positive(self):
        tt = TensorTrain.random([4, 5, 6], [2, 3])
        matrices = dot(tt, tt)

        # <tt, tt> should be positive
        inner_product = matrices[-1].item()
        assert inner_product > 0

    def test_dot_against_full_tensors(self):
        torch.manual_seed(42)
        dims = [3, 4, 5]
        ranks = [2, 3]
        tt1 = TensorTrain.random(dims, ranks)
        tt2 = TensorTrain.random(dims, ranks)

        # Compute via TT dot product
        inner_tt = dot(tt1, tt2)[-1].item()

        # Compute via full tensors
        T1 = tt1.to_tensor()
        T2 = tt2.to_tensor()
        inner_full = (T1 * T2).sum().item()

        assert abs(inner_tt - inner_full) < 1e-10

    def test_dot_different_ranks(self):
        tt1 = TensorTrain.random([4, 5, 6], [2, 3])
        tt2 = TensorTrain.random([4, 5, 6], [3, 4])

        matrices = dot(tt1, tt2)

        # Intermediate shapes are (r1_k, r2_k)
        assert matrices[0].shape == (1, 1)
        assert matrices[1].shape == (2, 3)
        assert matrices[2].shape == (3, 4)
        assert matrices[3].shape == (1, 1)

    def test_dot_mismatched_dims(self):
        tt1 = TensorTrain.random([4, 5, 6], [2, 3])
        tt2 = TensorTrain.random([4, 5, 7], [2, 3])  # Different last dim

        with pytest.raises(ValueError):
            dot(tt1, tt2)

    def test_dot_mismatched_d(self):
        tt1 = TensorTrain.random([4, 5, 6], [2, 3])
        tt2 = TensorTrain.random([4, 5], [2])  # Different d

        with pytest.raises(ValueError):
            dot(tt1, tt2)


class TestHadamard:
    def test_hadamard_against_full_tensors(self):
        torch.manual_seed(42)
        dims = [3, 4, 5]
        tt1 = TensorTrain.random(dims, ranks=[2, 3])
        tt2 = TensorTrain.random(dims, ranks=[3, 2])

        # Compute Hadamard product in TT format
        tt_result = hadamard(tt1, tt2)

        # Convert to full tensors and compare
        T1 = tt1.to_tensor()
        T2 = tt2.to_tensor()
        T_expected = T1 * T2

        T_result = tt_result.to_tensor()

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

    def test_hadamard_ranks(self):
        # Hadamard product multiplies ranks
        tt1 = TensorTrain.random([3, 4, 5], ranks=[2, 3])
        tt2 = TensorTrain.random([3, 4, 5], ranks=[4, 2])

        tt_result = hadamard(tt1, tt2)

        # Result ranks should be r1_k * r2_k
        assert tt_result.ranks == [1, 2 * 4, 3 * 2, 1]

    def test_hadamard_mismatched_dims(self):
        tt1 = TensorTrain.random([3, 4, 5], ranks=[2, 3])
        tt2 = TensorTrain.random([3, 4, 6], ranks=[2, 3])  # Different last dim

        with pytest.raises(ValueError):
            hadamard(tt1, tt2)


class TestScale:
    def test_scale_against_full_tensor(self):
        torch.manual_seed(42)
        tt = TensorTrain.random([3, 4, 5], ranks=[2, 3])
        alpha = 2.5

        tt_scaled = scale(tt, alpha)

        T = tt.to_tensor()
        T_scaled = tt_scaled.to_tensor()

        assert torch.allclose(T_scaled, alpha * T, atol=1e-12)

    def test_scale_negative(self):
        torch.manual_seed(42)
        tt = TensorTrain.random([3, 4, 5], ranks=[2, 3])
        alpha = -3.0

        tt_scaled = scale(tt, alpha)

        T = tt.to_tensor()
        T_scaled = tt_scaled.to_tensor()

        assert torch.allclose(T_scaled, alpha * T, atol=1e-12)

    def test_scale_zero(self):
        tt = TensorTrain.random([3, 4, 5], ranks=[2, 3])

        tt_scaled = scale(tt, 0.0)
        T_scaled = tt_scaled.to_tensor()

        assert torch.allclose(T_scaled, torch.zeros_like(T_scaled), atol=1e-12)

    def test_scale_preserves_ranks(self):
        tt = TensorTrain.random([3, 4, 5], ranks=[2, 3])

        tt_scaled = scale(tt, 5.0)

        assert tt_scaled.ranks == tt.ranks

    def test_scale_does_not_modify_original(self):
        torch.manual_seed(42)
        tt = TensorTrain.random([3, 4, 5], ranks=[2, 3])
        original_core0 = tt[0].clone()

        scale(tt, 10.0)

        assert torch.allclose(tt[0], original_core0)


class TestAdd:
    def test_add_against_full_tensors(self):
        torch.manual_seed(42)
        dims = [3, 4, 5]
        tt1 = TensorTrain.random(dims, ranks=[2, 3])
        tt2 = TensorTrain.random(dims, ranks=[3, 2])

        tt_sum = add(tt1, tt2)

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

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

    def test_add_ranks(self):
        # Addition sums the ranks (except boundaries)
        tt1 = TensorTrain.random([3, 4, 5], ranks=[2, 3])
        tt2 = TensorTrain.random([3, 4, 5], ranks=[4, 2])

        tt_sum = add(tt1, tt2)

        # Result ranks should be r1_k + r2_k
        assert tt_sum.ranks == [1, 2 + 4, 3 + 2, 1]

    def test_add_with_scalar(self):
        # TT + alpha using ones
        torch.manual_seed(42)
        tt = TensorTrain.random([3, 4, 5], ranks=[2, 3])
        alpha = 2.5

        tt_alpha = TensorTrain.ones(tt.dims, value=alpha)
        tt_sum = add(tt, tt_alpha)

        T = tt.to_tensor()
        T_sum = tt_sum.to_tensor()

        assert torch.allclose(T_sum, T + alpha, atol=1e-12)

    def test_add_commutative(self):
        torch.manual_seed(42)
        tt1 = TensorTrain.random([3, 4, 5], ranks=[2, 3])
        tt2 = TensorTrain.random([3, 4, 5], ranks=[3, 2])

        sum1 = add(tt1, tt2).to_tensor()
        sum2 = add(tt2, tt1).to_tensor()

        assert torch.allclose(sum1, sum2, atol=1e-12)

    def test_add_mismatched_dims(self):
        tt1 = TensorTrain.random([3, 4, 5], ranks=[2, 3])
        tt2 = TensorTrain.random([3, 4, 6], ranks=[2, 3])

        with pytest.raises(ValueError):
            add(tt1, tt2)

    def test_add_does_not_modify_original(self):
        torch.manual_seed(42)
        tt1 = TensorTrain.random([3, 4, 5], ranks=[2, 3])
        tt2 = TensorTrain.random([3, 4, 5], ranks=[3, 2])
        original_core0 = tt1[0].clone()

        add(tt1, tt2)

        assert torch.allclose(tt1[0], original_core0)


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

    def test_add_operator(self):
        """tt1 + tt2 uses add."""
        torch.manual_seed(42)
        tt1 = TensorTrain.random([3, 4, 5], ranks=[2, 3])
        tt2 = TensorTrain.random([3, 4, 5], ranks=[3, 2])

        result = tt1 + tt2

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

        assert torch.allclose(T_result, T1 + T2, atol=1e-12)

    def test_sub_operator(self):
        """tt1 - tt2 uses add and scale."""
        torch.manual_seed(42)
        tt1 = TensorTrain.random([3, 4, 5], ranks=[2, 3])
        tt2 = TensorTrain.random([3, 4, 5], ranks=[3, 2])

        result = tt1 - tt2

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

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

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

        result = tt * 2.5

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

        assert torch.allclose(T_result, 2.5 * T, atol=1e-12)

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

        result = 3.0 * tt

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

        assert torch.allclose(T_result, 3.0 * T, atol=1e-12)

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

        result = -tt

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

        assert torch.allclose(T_result, -T, atol=1e-12)

    def test_chained_operations(self):
        """Chained operator expressions work correctly."""
        torch.manual_seed(42)
        tt1 = TensorTrain.random([3, 4, 5], ranks=[2, 2])
        tt2 = TensorTrain.random([3, 4, 5], ranks=[2, 2])
        tt3 = TensorTrain.random([3, 4, 5], ranks=[2, 2])

        # (tt1 + tt2) - 2 * tt3
        result = (tt1 + tt2) - 2.0 * tt3

        T1 = tt1.to_tensor()
        T2 = tt2.to_tensor()
        T3 = tt3.to_tensor()
        T_expected = (T1 + T2) - 2.0 * T3
        T_result = result.to_tensor()

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


class TestIndexing:
    """Tests for element access and slicing."""

    def test_element_access_basic(self):
        """tt[i, j, k] returns correct element as 0-d tensor."""
        torch.manual_seed(42)
        tt = TensorTrain.random([3, 4, 5], ranks=[2, 3])
        full = tt.to_tensor()

        # Check return type is tensor (for autograd compatibility)
        elem = tt[0, 0, 0]
        assert isinstance(elem, torch.Tensor)
        assert elem.ndim == 0  # 0-d tensor

        # Check .item() returns Python scalar (like torch tensors)
        scalar = tt[0, 0, 0].item()
        assert isinstance(scalar, float)
        assert scalar == full[0, 0, 0].item()

        # Check several elements match
        assert torch.allclose(tt[0, 0, 0], full[0, 0, 0])
        assert torch.allclose(tt[1, 2, 3], full[1, 2, 3])
        assert torch.allclose(tt[2, 3, 4], full[2, 3, 4])

    def test_element_access_all_elements(self):
        """All elements match full tensor."""
        torch.manual_seed(42)
        tt = TensorTrain.random([3, 4, 5], ranks=[2, 3])
        full = tt.to_tensor()

        for i in range(3):
            for j in range(4):
                for k in range(5):
                    assert torch.allclose(tt[i, j, k], full[i, j, k])

    def test_element_access_2d(self):
        """Element access works for 2D TT."""
        torch.manual_seed(42)
        tt = TensorTrain.random([4, 5], ranks=[3])
        full = tt.to_tensor()

        for i in range(4):
            for j in range(5):
                assert torch.allclose(tt[i, j], full[i, j])

    def test_core_access_preserved(self):
        """tt[k] still returns k-th core (backward compatibility)."""
        tt = TensorTrain.random([3, 4, 5], ranks=[2, 3])

        assert tt[0].shape == (1, 3, 2)
        assert tt[1].shape == (2, 4, 3)
        assert tt[2].shape == (3, 5, 1)

    def test_slice_keep_all(self):
        """tt[:, :, :] returns equivalent TT."""
        torch.manual_seed(42)
        tt = TensorTrain.random([3, 4, 5], ranks=[2, 3])
        full = tt.to_tensor()

        sliced = tt[:, :, :]
        sliced_full = sliced.to_tensor()

        assert sliced.shape == tt.shape
        assert torch.allclose(sliced_full, full, atol=1e-12)

    def test_slice_contract_middle(self):
        """tt[:, j, :] contracts middle dimension."""
        torch.manual_seed(42)
        tt = TensorTrain.random([3, 4, 5], ranks=[2, 3])
        full = tt.to_tensor()

        j = 2
        sliced = tt[:, j, :]

        assert sliced.shape == (3, 5)
        assert torch.allclose(sliced.to_tensor(), full[:, j, :], atol=1e-12)

    def test_slice_contract_first(self):
        """tt[i, :, :] contracts first dimension."""
        torch.manual_seed(42)
        tt = TensorTrain.random([3, 4, 5], ranks=[2, 3])
        full = tt.to_tensor()

        i = 1
        sliced = tt[i, :, :]

        assert sliced.shape == (4, 5)
        assert torch.allclose(sliced.to_tensor(), full[i, :, :], atol=1e-12)

    def test_slice_contract_last(self):
        """tt[:, :, k] contracts last dimension."""
        torch.manual_seed(42)
        tt = TensorTrain.random([3, 4, 5], ranks=[2, 3])
        full = tt.to_tensor()

        k = 3
        sliced = tt[:, :, k]

        assert sliced.shape == (3, 4)
        assert torch.allclose(sliced.to_tensor(), full[:, :, k], atol=1e-12)

    def test_slice_contract_two(self):
        """tt[i, :, k] contracts first and last dimensions."""
        torch.manual_seed(42)
        tt = TensorTrain.random([3, 4, 5], ranks=[2, 3])
        full = tt.to_tensor()

        i, k = 1, 3
        sliced = tt[i, :, k]

        assert sliced.shape == (4,)
        assert torch.allclose(sliced.to_tensor(), full[i, :, k], atol=1e-12)

    def test_slice_range(self):
        """tt[:, 1:3, :] slices a range."""
        torch.manual_seed(42)
        tt = TensorTrain.random([3, 4, 5], ranks=[2, 3])
        full = tt.to_tensor()

        sliced = tt[:, 1:3, :]

        assert sliced.shape == (3, 2, 5)
        assert torch.allclose(sliced.to_tensor(), full[:, 1:3, :], atol=1e-12)

    def test_slice_ellipsis(self):
        """tt[..., k] uses ellipsis."""
        torch.manual_seed(42)
        tt = TensorTrain.random([3, 4, 5], ranks=[2, 3])
        full = tt.to_tensor()

        k = 2
        sliced = tt[..., k]

        assert sliced.shape == (3, 4)
        assert torch.allclose(sliced.to_tensor(), full[..., k], atol=1e-12)

    def test_slice_4d(self):
        """Slicing works on 4D TT."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 3, 4, 5], ranks=[2, 2, 2])
        full = tt.to_tensor()

        sliced = tt[:, 1, :, 3]

        assert sliced.shape == (2, 4)
        assert torch.allclose(sliced.to_tensor(), full[:, 1, :, 3], atol=1e-12)

    def test_too_many_indices_raises(self):
        """Too many indices raises IndexError."""
        tt = TensorTrain.random([3, 4, 5], ranks=[2, 3])

        with pytest.raises(IndexError):
            _ = tt[0, 1, 2, 3]  # 4 indices for 3D TT

