"""Tests for TensorTrain swap_adjacent and transpose operations."""

import torch
import pytest
from ttglow import TensorTrain


class TestSwapAdjacent:
    """Tests for swap_adjacent operation."""

    def test_swap_first_two_modes(self):
        """Test swapping modes 0 and 1."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 3, 4], ranks=[5, 6])
        tt2 = tt.swap_adjacent(0)

        assert tt2.shape == (3, 2, 4)
        assert tt2.d == 3

        # Values should match permuted tensor
        expected = tt.to_tensor().permute(1, 0, 2)
        actual = tt2.to_tensor()
        assert torch.allclose(expected, actual, atol=1e-10)

    def test_swap_last_two_modes(self):
        """Test swapping modes d-2 and d-1."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 3, 4], ranks=[5, 6])
        tt2 = tt.swap_adjacent(1)

        assert tt2.shape == (2, 4, 3)
        assert tt2.d == 3

        expected = tt.to_tensor().permute(0, 2, 1)
        actual = tt2.to_tensor()
        assert torch.allclose(expected, actual, atol=1e-10)

    def test_swap_middle_modes(self):
        """Test swapping middle modes in a 4-mode TT."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 3, 4, 5], ranks=[3, 4, 5])
        tt2 = tt.swap_adjacent(1)

        assert tt2.shape == (2, 4, 3, 5)

        expected = tt.to_tensor().permute(0, 2, 1, 3)
        actual = tt2.to_tensor()
        assert torch.allclose(expected, actual, atol=1e-10)

    def test_swap_with_max_rank(self):
        """Test that max_rank truncation works."""
        torch.manual_seed(42)
        tt = TensorTrain.random([4, 5, 6], ranks=[10, 10])
        tt2 = tt.swap_adjacent(0, max_rank=3)

        assert tt2.shape == (5, 4, 6)
        assert tt2.ranks[1] <= 3

    def test_swap_with_rel_tol(self):
        """Test that rel_tol truncation works."""
        torch.manual_seed(42)
        tt = TensorTrain.random([4, 5, 6], ranks=[10, 10])
        tt2 = tt.swap_adjacent(0, rel_tol=0.1)

        assert tt2.shape == (5, 4, 6)
        # Rank should be reduced (exact amount depends on singular values)

    def test_swap_invalid_index_negative(self):
        """Test that negative index raises error."""
        tt = TensorTrain.random([2, 3, 4], ranks=[5, 6])

        with pytest.raises(ValueError):
            tt.swap_adjacent(-1)

    def test_swap_invalid_index_too_large(self):
        """Test that index >= d-1 raises error."""
        tt = TensorTrain.random([2, 3, 4], ranks=[5, 6])

        with pytest.raises(ValueError):
            tt.swap_adjacent(2)  # d=3, so max valid index is 1

    def test_swap_preserves_dtype(self):
        """Test that swap preserves dtype."""
        tt = TensorTrain.random([2, 3, 4], ranks=[5, 6], dtype=torch.float32)
        tt2 = tt.swap_adjacent(0)

        assert tt2.dtype == torch.float32

    def test_double_swap_identity(self):
        """Test that swapping twice returns to original."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 3, 4], ranks=[5, 6])
        tt2 = tt.swap_adjacent(0).swap_adjacent(0)

        # Shape should be back to original
        assert tt2.shape == tt.shape

        # Values should match (up to numerical precision)
        assert torch.allclose(tt.to_tensor(), tt2.to_tensor(), atol=1e-10)


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

    def test_transpose_adjacent(self):
        """Test transpose of adjacent modes (single swap)."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 3, 4], ranks=[5, 6])
        tt2 = tt.transpose(0, 1)

        assert tt2.shape == (3, 2, 4)

        expected = tt.to_tensor().permute(1, 0, 2)
        actual = tt2.to_tensor()
        assert torch.allclose(expected, actual, atol=1e-10)

    def test_transpose_non_adjacent(self):
        """Test transpose of non-adjacent modes."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 3, 4, 5], ranks=[3, 4, 5])
        tt2 = tt.transpose(0, 2)

        assert tt2.shape == (4, 3, 2, 5)

        expected = tt.to_tensor().permute(2, 1, 0, 3)
        actual = tt2.to_tensor()
        assert torch.allclose(expected, actual, atol=1e-10)

    def test_transpose_first_and_last(self):
        """Test transpose of first and last modes."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 3, 4, 5], ranks=[3, 4, 5])
        tt2 = tt.transpose(0, 3)

        assert tt2.shape == (5, 3, 4, 2)

        expected = tt.to_tensor().permute(3, 1, 2, 0)
        actual = tt2.to_tensor()
        assert torch.allclose(expected, actual, atol=1e-10)

    def test_transpose_reversed_args(self):
        """Test that transpose(i, j) == transpose(j, i)."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 3, 4, 5], ranks=[3, 4, 5])

        tt2 = tt.transpose(1, 3)
        tt3 = tt.transpose(3, 1)

        assert tt2.shape == tt3.shape
        assert torch.allclose(tt2.to_tensor(), tt3.to_tensor(), atol=1e-10)

    def test_transpose_same_dim(self):
        """Test that transpose(i, i) returns clone."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 3, 4], ranks=[5, 6])
        tt2 = tt.transpose(1, 1)

        assert tt2.shape == tt.shape
        assert torch.allclose(tt.to_tensor(), tt2.to_tensor())

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

        with pytest.raises(ValueError):
            tt.transpose(-1, 1)

        with pytest.raises(ValueError):
            tt.transpose(3, 1)

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

        with pytest.raises(ValueError):
            tt.transpose(0, -1)

        with pytest.raises(ValueError):
            tt.transpose(0, 5)

    def test_transpose_5d(self):
        """Test transpose on 5-dimensional TT."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 3, 4, 5, 6], ranks=[2, 3, 4, 5])
        tt2 = tt.transpose(1, 4)

        assert tt2.shape == (2, 6, 4, 5, 3)

        expected = tt.to_tensor().permute(0, 4, 2, 3, 1)
        actual = tt2.to_tensor()
        assert torch.allclose(expected, actual, atol=1e-10)


class TestRoundTrip:
    """Tests for round-trip transpose operations."""

    def test_transpose_twice_identity(self):
        """Test that transposing twice returns to original."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 3, 4, 5], ranks=[3, 4, 5])
        tt2 = tt.transpose(0, 2).transpose(0, 2)

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

    def test_three_transposes(self):
        """Test composition of three transposes."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 3, 4, 5], ranks=[3, 4, 5])

        # Shape (2, 3, 4, 5)
        # transpose(0,1) -> (3, 2, 4, 5)
        # transpose(2,3) -> (3, 2, 5, 4)
        # transpose(0,2) -> (5, 2, 3, 4)
        tt2 = tt.transpose(0, 1).transpose(2, 3).transpose(0, 2)

        assert tt2.shape == (5, 2, 3, 4)

        expected = tt.to_tensor().permute(3, 0, 1, 2)
        actual = tt2.to_tensor()
        assert torch.allclose(expected, actual, atol=1e-10)
