"""Tests for TensorTrain reshape operations."""

import torch
import pytest
from ttglow import TensorTrain


class TestMergeModes:
    """Tests for merge_modes operation."""

    def test_merge_two_modes(self):
        """Test merging two adjacent modes."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 3, 4], ranks=[5, 6])
        tt2 = tt.merge_modes(0, 2)

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

        # Values should match
        expected = tt.to_tensor().reshape(6, 4)
        actual = tt2.to_tensor()
        assert torch.allclose(expected, actual)

    def test_merge_all_modes(self):
        """Test merging all modes into one."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 3, 4], ranks=[5, 6])
        tt2 = tt.merge_modes(0, 3)

        assert tt2.shape == (24,)
        assert tt2.d == 1

        expected = tt.to_tensor().reshape(24)
        actual = tt2.to_tensor()
        assert torch.allclose(expected, actual)

    def test_merge_last_two_modes(self):
        """Test merging last two modes."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 3, 4], ranks=[5, 6])
        tt2 = tt.merge_modes(1, 3)

        assert tt2.shape == (2, 12)
        assert tt2.d == 2

        expected = tt.to_tensor().reshape(2, 12)
        actual = tt2.to_tensor()
        assert torch.allclose(expected, actual)

    def test_merge_single_mode_returns_clone(self):
        """Test that merging a single mode returns a clone."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 3, 4], ranks=[5, 6])
        tt2 = tt.merge_modes(1, 2)

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

    def test_merge_invalid_range(self):
        """Test that invalid ranges raise errors."""
        tt = TensorTrain.random([2, 3, 4], ranks=[5, 6])

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

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

        with pytest.raises(ValueError):
            tt.merge_modes(2, 1)


class TestSplitMode:
    """Tests for split_mode operation."""

    def test_split_into_two(self):
        """Test splitting one mode into two."""
        torch.manual_seed(42)
        tt = TensorTrain.random([6, 4], ranks=[5])
        tt2 = tt.split_mode(0, [2, 3])

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

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

    def test_split_into_three(self):
        """Test splitting one mode into three."""
        torch.manual_seed(42)
        tt = TensorTrain.random([24], ranks=[])
        tt2 = tt.split_mode(0, [2, 3, 4])

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

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

    def test_split_middle_mode(self):
        """Test splitting a middle mode."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 12, 5], ranks=[3, 4])
        tt2 = tt.split_mode(1, [3, 4])

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

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

    def test_split_single_factor_returns_clone(self):
        """Test that splitting with single factor returns clone."""
        torch.manual_seed(42)
        tt = TensorTrain.random([6, 4], ranks=[5])
        tt2 = tt.split_mode(0, [6])

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

    def test_split_invalid_shape(self):
        """Test that invalid shape raise errors."""
        tt = TensorTrain.random([6, 4], ranks=[5])

        with pytest.raises(ValueError, match="doesn't match"):
            tt.split_mode(0, [2, 2])  # 2*2=4 != 6

    def test_split_invalid_mode(self):
        """Test that invalid mode raises error."""
        tt = TensorTrain.random([6, 4], ranks=[5])

        with pytest.raises(ValueError):
            tt.split_mode(-1, [2, 3])

        with pytest.raises(ValueError):
            tt.split_mode(5, [2, 3])


class TestReshape:
    """Tests for reshape operation."""

    def test_reshape_merge(self):
        """Test reshape that requires merging."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 3, 4], ranks=[5, 6])
        tt2 = tt.reshape([6, 4])

        assert tt2.shape == (6, 4)
        expected = tt.to_tensor().reshape(6, 4)
        assert torch.allclose(expected, tt2.to_tensor())

    def test_reshape_split(self):
        """Test reshape that requires splitting."""
        torch.manual_seed(42)
        tt = TensorTrain.random([6, 4], ranks=[5])
        tt2 = tt.reshape([2, 3, 4])

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

    def test_reshape_to_single_mode(self):
        """Test reshape to single mode."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 3, 4], ranks=[5, 6])
        tt2 = tt.reshape([24])

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

    def test_reshape_identity(self):
        """Test reshape to same shape."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 3, 4], ranks=[5, 6])
        tt2 = tt.reshape([2, 3, 4])

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

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

        with pytest.raises(ValueError, match="Cannot reshape"):
            tt.reshape([10, 3])  # 30 != 24

    def test_reshape_requires_permutation(self):
        """Test that reshape requiring permutation raises error."""
        tt = TensorTrain.random([2, 3, 4], ranks=[5, 6])

        with pytest.raises(ValueError, match="no valid"):
            tt.reshape([4, 6])  # would need permutation

        with pytest.raises(ValueError, match="no valid"):
            tt.reshape([3, 8])  # would need permutation


class TestFlatten:
    """Tests for flatten operation."""

    def test_flatten_all(self):
        """Test flattening all dimensions."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 3, 4], ranks=[5, 6])
        tt2 = tt.flatten()

        assert tt2.shape == (24,)
        assert tt2.d == 1

        expected = tt.to_tensor().flatten()
        actual = tt2.to_tensor()
        assert torch.allclose(expected, actual)

    def test_flatten_first_two(self):
        """Test flattening first two dimensions."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 3, 4, 5], ranks=[2, 3, 4])
        tt2 = tt.flatten(0, 1)

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

        expected = tt.to_tensor().flatten(0, 1)
        actual = tt2.to_tensor()
        assert torch.allclose(expected, actual)

    def test_flatten_last_two(self):
        """Test flattening last two dimensions."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 3, 4, 5], ranks=[2, 3, 4])
        tt2 = tt.flatten(2, 3)

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

        expected = tt.to_tensor().flatten(2, 3)
        actual = tt2.to_tensor()
        assert torch.allclose(expected, actual)

    def test_flatten_middle(self):
        """Test flattening middle dimensions."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 3, 4, 5], ranks=[2, 3, 4])
        tt2 = tt.flatten(1, 2)

        assert tt2.shape == (2, 12, 5)

        expected = tt.to_tensor().flatten(1, 2)
        actual = tt2.to_tensor()
        assert torch.allclose(expected, actual)

    def test_flatten_negative_index(self):
        """Test flatten with negative indexing."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 3, 4, 5], ranks=[2, 3, 4])

        # flatten(1, -1) should flatten from dim 1 to last dim
        tt2 = tt.flatten(1, -1)

        assert tt2.shape == (2, 60)

        expected = tt.to_tensor().flatten(1, -1)
        actual = tt2.to_tensor()
        assert torch.allclose(expected, actual)

    def test_flatten_single_dim(self):
        """Test flatten of single dimension (no-op)."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 3, 4], ranks=[5, 6])
        tt2 = tt.flatten(1, 1)

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

    def test_flatten_invalid_dims(self):
        """Test that invalid dimensions raise errors."""
        tt = TensorTrain.random([2, 3, 4], ranks=[5, 6])

        with pytest.raises(ValueError):
            tt.flatten(5, 6)

        with pytest.raises(ValueError):
            tt.flatten(2, 1)  # start > end

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


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

    def test_merge_then_split(self):
        """Test merge followed by split recovers original."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 3, 4], ranks=[5, 6])
        tt2 = tt.merge_modes(0, 2).split_mode(0, [2, 3])

        # Note: ranks may differ but values should match
        assert tt2.shape == (2, 3, 4)
        assert torch.allclose(tt.to_tensor(), tt2.to_tensor(), atol=1e-10)

    def test_split_then_merge(self):
        """Test split followed by merge recovers original."""
        torch.manual_seed(42)
        tt = TensorTrain.random([6, 4], ranks=[5])
        tt2 = tt.split_mode(0, [2, 3]).merge_modes(0, 2)

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