"""Tests for TensorTrain cat and stack operations."""

import torch
import pytest
from ttglow import TensorTrain
from ttglow.tensortrain import cat, stack


class TestCat:
    """Tests for cat operation."""

    def test_cat_along_first_dim(self):
        """Test concatenation along first dimension."""
        torch.manual_seed(42)
        tt1 = TensorTrain.random([2, 3, 4], ranks=[5, 6])
        tt2 = TensorTrain.random([5, 3, 4], ranks=[3, 4])

        result = cat([tt1, tt2], dim=0)

        assert result.shape == (7, 3, 4)
        assert result.d == 3

        # Verify against dense
        expected = torch.cat([tt1.to_tensor(), tt2.to_tensor()], dim=0)
        actual = result.to_tensor()
        assert torch.allclose(expected, actual, atol=1e-10)

    def test_cat_along_middle_dim(self):
        """Test concatenation along middle dimension."""
        torch.manual_seed(42)
        tt1 = TensorTrain.random([2, 3, 4], ranks=[5, 6])
        tt2 = TensorTrain.random([2, 5, 4], ranks=[3, 4])

        result = cat([tt1, tt2], dim=1)

        assert result.shape == (2, 8, 4)

        expected = torch.cat([tt1.to_tensor(), tt2.to_tensor()], dim=1)
        actual = result.to_tensor()
        assert torch.allclose(expected, actual, atol=1e-10)

    def test_cat_along_last_dim(self):
        """Test concatenation along last dimension."""
        torch.manual_seed(42)
        tt1 = TensorTrain.random([2, 3, 4], ranks=[5, 6])
        tt2 = TensorTrain.random([2, 3, 7], ranks=[3, 4])

        result = cat([tt1, tt2], dim=2)

        assert result.shape == (2, 3, 11)

        expected = torch.cat([tt1.to_tensor(), tt2.to_tensor()], dim=2)
        actual = result.to_tensor()
        assert torch.allclose(expected, actual, atol=1e-10)

    def test_cat_negative_dim(self):
        """Test concatenation with negative dimension."""
        torch.manual_seed(42)
        tt1 = TensorTrain.random([2, 3, 4], ranks=[5, 6])
        tt2 = TensorTrain.random([2, 3, 5], ranks=[3, 4])

        result = cat([tt1, tt2], dim=-1)

        assert result.shape == (2, 3, 9)

        expected = torch.cat([tt1.to_tensor(), tt2.to_tensor()], dim=-1)
        actual = result.to_tensor()
        assert torch.allclose(expected, actual, atol=1e-10)

    def test_cat_multiple_tensors(self):
        """Test concatenation of more than two tensors."""
        torch.manual_seed(42)
        tt1 = TensorTrain.random([2, 3, 4], ranks=[2, 3])
        tt2 = TensorTrain.random([3, 3, 4], ranks=[2, 3])
        tt3 = TensorTrain.random([4, 3, 4], ranks=[2, 3])

        result = cat([tt1, tt2, tt3], dim=0)

        assert result.shape == (9, 3, 4)

        expected = torch.cat([tt1.to_tensor(), tt2.to_tensor(), tt3.to_tensor()], dim=0)
        actual = result.to_tensor()
        assert torch.allclose(expected, actual, atol=1e-10)

    def test_cat_single_tensor(self):
        """Test cat with single tensor returns clone."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 3, 4], ranks=[5, 6])

        result = cat([tt], dim=0)

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

    def test_cat_empty_raises(self):
        """Test that empty list raises error."""
        with pytest.raises(ValueError):
            cat([], dim=0)

    def test_cat_dimension_mismatch_raises(self):
        """Test that mismatched dimensions raise error."""
        tt1 = TensorTrain.random([2, 3, 4], ranks=[5, 6])
        tt2 = TensorTrain.random([2, 5, 4], ranks=[3, 4])

        with pytest.raises(ValueError, match="Dimension mismatch"):
            cat([tt1, tt2], dim=0)  # dim 1 doesn't match

    def test_cat_2d_tensors(self):
        """Test cat on 2D tensor trains."""
        torch.manual_seed(42)
        tt1 = TensorTrain.random([3, 4], ranks=[5])
        tt2 = TensorTrain.random([5, 4], ranks=[3])

        result = cat([tt1, tt2], dim=0)

        assert result.shape == (8, 4)

        expected = torch.cat([tt1.to_tensor(), tt2.to_tensor()], dim=0)
        actual = result.to_tensor()
        assert torch.allclose(expected, actual, atol=1e-10)

    def test_cat_4d_tensors(self):
        """Test cat on 4D tensor trains."""
        torch.manual_seed(42)
        tt1 = TensorTrain.random([2, 3, 4, 5], ranks=[2, 3, 4])
        tt2 = TensorTrain.random([2, 3, 6, 5], ranks=[2, 3, 4])

        result = cat([tt1, tt2], dim=2)

        assert result.shape == (2, 3, 10, 5)

        expected = torch.cat([tt1.to_tensor(), tt2.to_tensor()], dim=2)
        actual = result.to_tensor()
        assert torch.allclose(expected, actual, atol=1e-10)


class TestStack:
    """Tests for stack operation."""

    def test_stack_along_first_dim(self):
        """Test stacking along first dimension."""
        torch.manual_seed(42)
        tt1 = TensorTrain.random([2, 3, 4], ranks=[5, 6])
        tt2 = TensorTrain.random([2, 3, 4], ranks=[3, 4])

        result = stack([tt1, tt2], dim=0)

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

        expected = torch.stack([tt1.to_tensor(), tt2.to_tensor()], dim=0)
        actual = result.to_tensor()
        assert torch.allclose(expected, actual, atol=1e-10)

    def test_stack_along_middle_dim(self):
        """Test stacking along middle dimension."""
        torch.manual_seed(42)
        tt1 = TensorTrain.random([2, 3, 4], ranks=[5, 6])
        tt2 = TensorTrain.random([2, 3, 4], ranks=[3, 4])

        result = stack([tt1, tt2], dim=1)

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

        expected = torch.stack([tt1.to_tensor(), tt2.to_tensor()], dim=1)
        actual = result.to_tensor()
        assert torch.allclose(expected, actual, atol=1e-10)

    def test_stack_along_last_dim(self):
        """Test stacking along last dimension."""
        torch.manual_seed(42)
        tt1 = TensorTrain.random([2, 3, 4], ranks=[5, 6])
        tt2 = TensorTrain.random([2, 3, 4], ranks=[3, 4])

        result = stack([tt1, tt2], dim=3)

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

        expected = torch.stack([tt1.to_tensor(), tt2.to_tensor()], dim=3)
        actual = result.to_tensor()
        assert torch.allclose(expected, actual, atol=1e-10)

    def test_stack_negative_dim(self):
        """Test stacking with negative dimension."""
        torch.manual_seed(42)
        tt1 = TensorTrain.random([2, 3, 4], ranks=[5, 6])
        tt2 = TensorTrain.random([2, 3, 4], ranks=[3, 4])

        result = stack([tt1, tt2], dim=-1)

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

        expected = torch.stack([tt1.to_tensor(), tt2.to_tensor()], dim=-1)
        actual = result.to_tensor()
        assert torch.allclose(expected, actual, atol=1e-10)

    def test_stack_multiple_tensors(self):
        """Test stacking more than two tensors."""
        torch.manual_seed(42)
        tt1 = TensorTrain.random([2, 3, 4], ranks=[2, 3])
        tt2 = TensorTrain.random([2, 3, 4], ranks=[2, 3])
        tt3 = TensorTrain.random([2, 3, 4], ranks=[2, 3])

        result = stack([tt1, tt2, tt3], dim=0)

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

        expected = torch.stack([tt1.to_tensor(), tt2.to_tensor(), tt3.to_tensor()], dim=0)
        actual = result.to_tensor()
        assert torch.allclose(expected, actual, atol=1e-10)

    def test_stack_single_tensor(self):
        """Test stack with single tensor."""
        torch.manual_seed(42)
        tt = TensorTrain.random([2, 3, 4], ranks=[5, 6])

        result = stack([tt], dim=0)

        assert result.shape == (1, 2, 3, 4)

        expected = torch.stack([tt.to_tensor()], dim=0)
        actual = result.to_tensor()
        assert torch.allclose(expected, actual, atol=1e-10)

    def test_stack_empty_raises(self):
        """Test that empty list raises error."""
        with pytest.raises(ValueError):
            stack([], dim=0)

    def test_stack_shape_mismatch_raises(self):
        """Test that mismatched shapes raise error."""
        tt1 = TensorTrain.random([2, 3, 4], ranks=[5, 6])
        tt2 = TensorTrain.random([2, 5, 4], ranks=[3, 4])

        with pytest.raises(ValueError, match="same shape"):
            stack([tt1, tt2], dim=0)

    def test_stack_2d_tensors(self):
        """Test stack on 2D tensor trains."""
        torch.manual_seed(42)
        tt1 = TensorTrain.random([3, 4], ranks=[5])
        tt2 = TensorTrain.random([3, 4], ranks=[3])

        result = stack([tt1, tt2], dim=0)

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

        expected = torch.stack([tt1.to_tensor(), tt2.to_tensor()], dim=0)
        actual = result.to_tensor()
        assert torch.allclose(expected, actual, atol=1e-10)
