"""Tests for TensorTrain squeeze and unsqueeze operations."""

import torch
import pytest
from ttglow import TensorTrain


class TestUnsqueeze:
    """Tests for unsqueeze operation."""

    def test_unsqueeze_at_beginning(self):
        """Test unsqueeze at position 0."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 3, 4], ranks=[5, 6])
        tt2 = tt.unsqueeze(0)

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

        # Values should match with unsqueezed dense tensor
        expected = tt.to_tensor().unsqueeze(0)
        actual = tt2.to_tensor()
        assert torch.allclose(expected, actual, atol=1e-10)

    def test_unsqueeze_at_end(self):
        """Test unsqueeze at last position."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 3, 4], ranks=[5, 6])
        tt2 = tt.unsqueeze(3)

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

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

    def test_unsqueeze_in_middle(self):
        """Test unsqueeze in the middle."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 3, 4], ranks=[5, 6])
        tt2 = tt.unsqueeze(1)

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

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

    def test_unsqueeze_negative_index(self):
        """Test unsqueeze with negative index."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 3, 4], ranks=[5, 6])

        # -1 means after the last dim
        tt2 = tt.unsqueeze(-1)
        assert tt2.shape == (2, 3, 4, 1)

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

    def test_unsqueeze_multiple(self):
        """Test multiple unsqueeze operations."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 3], ranks=[5])
        tt2 = tt.unsqueeze(0).unsqueeze(2).unsqueeze(4)

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

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

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

        with pytest.raises(ValueError):
            tt.unsqueeze(5)  # d=3, valid range is 0-3

        with pytest.raises(ValueError):
            tt.unsqueeze(-5)

    def test_unsqueeze_preserves_dtype(self):
        """Test that unsqueeze preserves dtype."""
        tt = TensorTrain.random([2, 3, 4], ranks=[5, 6], dtype=torch.float32)
        tt2 = tt.unsqueeze(1)

        assert tt2.dtype == torch.float32


class TestSqueeze:
    """Tests for squeeze operation."""

    def test_squeeze_at_beginning(self):
        """Test squeeze at position 0."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 3, 4], ranks=[5, 6])
        tt_expanded = tt.unsqueeze(0)  # (1, 2, 3, 4)
        tt2 = tt_expanded.squeeze(0)

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

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

    def test_squeeze_at_end(self):
        """Test squeeze at last position."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 3, 4], ranks=[5, 6])
        tt_expanded = tt.unsqueeze(3)  # (2, 3, 4, 1)
        tt2 = tt_expanded.squeeze(3)

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

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

    def test_squeeze_in_middle(self):
        """Test squeeze in the middle."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 3, 4], ranks=[5, 6])
        tt_expanded = tt.unsqueeze(2)  # (2, 3, 1, 4)
        tt2 = tt_expanded.squeeze(2)

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

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

    def test_squeeze_negative_index(self):
        """Test squeeze with negative index."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 3, 4], ranks=[5, 6])
        tt_expanded = tt.unsqueeze(-1)  # (2, 3, 4, 1)
        tt2 = tt_expanded.squeeze(-1)

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

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

    def test_squeeze_all(self):
        """Test squeeze without dim argument removes all singletons."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 3], ranks=[5])
        tt_expanded = tt.unsqueeze(0).unsqueeze(2).unsqueeze(4)  # (1, 2, 1, 3, 1)
        tt2 = tt_expanded.squeeze()

        assert tt2.shape == (2, 3)

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

    def test_squeeze_non_singleton_raises(self):
        """Test that squeezing non-singleton dim raises error."""
        tt = TensorTrain.random([2, 3, 4], ranks=[5, 6])

        with pytest.raises(ValueError, match="size is"):
            tt.squeeze(0)  # dim 0 has size 2, not 1

    def test_squeeze_invalid_dim(self):
        """Test that invalid dim raises error."""
        tt = TensorTrain.random([1, 3, 1], ranks=[2, 3])

        with pytest.raises(ValueError):
            tt.squeeze(5)

        with pytest.raises(ValueError):
            tt.squeeze(-5)

    def test_squeeze_preserves_non_singletons(self):
        """Test squeeze() preserves non-singleton dimensions."""
        torch.manual_seed(42)
        tt = TensorTrain.random([1, 3, 1, 4, 1], ranks=[2, 3, 4, 5])
        tt2 = tt.squeeze()

        assert tt2.shape == (3, 4)

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


class TestRoundTrip:
    """Tests for round-trip squeeze/unsqueeze operations."""

    def test_unsqueeze_squeeze_identity(self):
        """Test that unsqueeze followed by squeeze is identity."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 3, 4], ranks=[5, 6])

        for dim in range(4):
            tt2 = tt.unsqueeze(dim).squeeze(dim)
            assert tt2.shape == tt.shape
            assert torch.allclose(tt.to_tensor(), tt2.to_tensor(), atol=1e-10)

    def test_squeeze_unsqueeze_identity(self):
        """Test that squeeze followed by unsqueeze at same position."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 1, 4], ranks=[5, 6])

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

    def test_multiple_unsqueeze_squeeze(self):
        """Test complex sequence of unsqueeze and squeeze."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 3, 4], ranks=[5, 6])

        # Add singletons at various positions
        tt2 = tt.unsqueeze(0).unsqueeze(2).unsqueeze(-1)
        assert tt2.shape == (1, 2, 1, 3, 4, 1)

        # Remove them all
        tt3 = tt2.squeeze()
        assert tt3.shape == (2, 3, 4)

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