"""Tests for kronadd function."""

import torch
import pytest
from ttglow.ttmatrix import TTMatrix, kronadd, kron, add


class TestKronadd:
    """Tests for kronadd operation."""

    def test_overlapping_sites(self):
        """Test kronadd with overlapping sites: ([0,1], [1,2])."""
        torch.manual_seed(42)

        # A acts on sites 0, 1 with dims [(2,2), (3,3)]
        A = TTMatrix.random([(2, 2), (3, 3)], ranks=[2])
        # B acts on sites 1, 2 with dims [(3,3), (4,4)]
        B = TTMatrix.random([(3, 3), (4, 4)], ranks=[3])

        # kronadd: result acts on sites 0, 1, 2
        result = kronadd(A, B, sites=([0, 1], [1, 2]))

        # Verify dimensions
        assert result.d == 3
        assert result.dims == [(2, 2), (3, 3), (4, 4)]

        # Verify by computing full tensors
        # result = (A ⊗ I_4) + (I_2 ⊗ B)
        A_dense = A.to_tensor()  # (6, 6)
        B_dense = B.to_tensor()  # (12, 12)

        I2 = torch.eye(2, dtype=A.dtype)
        I4 = torch.eye(4, dtype=A.dtype)

        # A ⊗ I_4: (2*3, 2*3) ⊗ (4, 4) = (24, 24)
        A_kron_I = torch.kron(A_dense, I4)
        # I_2 ⊗ B: (2, 2) ⊗ (3*4, 3*4) = (24, 24)
        I_kron_B = torch.kron(I2, B_dense)

        expected = A_kron_I + I_kron_B
        actual = result.to_tensor()

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

    def test_adjacent_sites_no_overlap(self):
        """Test kronadd with adjacent but non-overlapping sites: ([0,1], [2,3])."""
        torch.manual_seed(42)

        # A acts on sites 0, 1
        A = TTMatrix.random([(2, 2), (3, 3)], ranks=[2])
        # B acts on sites 2, 3
        B = TTMatrix.random([(4, 4), (5, 5)], ranks=[3])

        # kronadd: result acts on sites 0, 1, 2, 3
        result = kronadd(A, B, sites=([0, 1], [2, 3]))

        # Verify dimensions
        assert result.d == 4
        assert result.dims == [(2, 2), (3, 3), (4, 4), (5, 5)]

        # Verify by computing full tensors
        A_dense = A.to_tensor()  # (6, 6)
        B_dense = B.to_tensor()  # (20, 20)

        I_A = torch.eye(20, dtype=A.dtype)  # Identity for B's space
        I_B = torch.eye(6, dtype=A.dtype)   # Identity for A's space

        # A ⊗ I: (6, 6) ⊗ (20, 20) = (120, 120)
        A_kron_I = torch.kron(A_dense, I_A)
        # I ⊗ B: (6, 6) ⊗ (20, 20) = (120, 120)
        I_kron_B = torch.kron(I_B, B_dense)

        expected = A_kron_I + I_kron_B
        actual = result.to_tensor()

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

    def test_single_site_operators(self):
        """Test kronadd with single-site operators."""
        torch.manual_seed(42)

        # Single-site operators
        A = TTMatrix.random([(3, 3)], ranks=[])
        B = TTMatrix.random([(4, 4)], ranks=[])

        # Adjacent single sites
        result = kronadd(A, B, sites=([0], [1]))

        assert result.d == 2
        assert result.dims == [(3, 3), (4, 4)]

        # Verify
        A_dense = A.to_tensor()
        B_dense = B.to_tensor()

        expected = torch.kron(A_dense, torch.eye(4)) + torch.kron(torch.eye(3), B_dense)
        actual = result.to_tensor()

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

    def test_b_before_a(self):
        """Test kronadd with B's sites before A's sites."""
        torch.manual_seed(42)

        A = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])
        B = TTMatrix.random([(2, 2), (3, 3)], ranks=[3])

        # B on [0,1], A on [1,2]
        result = kronadd(A, B, sites=([1, 2], [0, 1]))

        assert result.d == 3
        assert result.dims == [(2, 2), (3, 3), (4, 4)]

        # Verify
        A_dense = A.to_tensor()  # (12, 12)
        B_dense = B.to_tensor()  # (6, 6)

        I2 = torch.eye(2, dtype=A.dtype)
        I4 = torch.eye(4, dtype=A.dtype)

        # I_2 ⊗ A
        I_kron_A = torch.kron(I2, A_dense)
        # B ⊗ I_4
        B_kron_I = torch.kron(B_dense, I4)

        expected = I_kron_A + B_kron_I
        actual = result.to_tensor()

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

    def test_full_overlap(self):
        """Test kronadd with fully overlapping sites (same as regular add)."""
        torch.manual_seed(42)

        A = TTMatrix.random([(2, 2), (3, 3)], ranks=[2])
        B = TTMatrix.random([(2, 2), (3, 3)], ranks=[3])

        # Same sites
        result = kronadd(A, B, sites=([0, 1], [0, 1]))

        # Should be equivalent to regular add
        expected = add(A, B)

        assert result.d == expected.d
        assert result.dims == expected.dims
        assert torch.allclose(result.to_tensor(), expected.to_tensor(), atol=1e-10)

    def test_error_gap_in_sites(self):
        """Test that kronadd raises error for gap between sites."""
        A = TTMatrix.random([(2, 2), (3, 3)], ranks=[2])
        B = TTMatrix.random([(4, 4), (5, 5)], ranks=[3])

        # Gap between [0,1] and [3,4]
        with pytest.raises(ValueError, match="contiguous union"):
            kronadd(A, B, sites=([0, 1], [3, 4]))

    def test_error_noncontiguous_sites(self):
        """Test that kronadd raises error for non-contiguous sites."""
        A = TTMatrix.random([(2, 2), (3, 3)], ranks=[2])
        B = TTMatrix.random([(4, 4), (5, 5)], ranks=[3])

        with pytest.raises(ValueError, match="must be contiguous"):
            kronadd(A, B, sites=([0, 2], [1, 3]))

    def test_error_sites_length_mismatch(self):
        """Test that kronadd raises error when sites length doesn't match operator."""
        A = TTMatrix.random([(2, 2), (3, 3)], ranks=[2])
        B = TTMatrix.random([(4, 4)], ranks=[])

        with pytest.raises(ValueError, match="must match"):
            kronadd(A, B, sites=([0, 1], [2, 3]))  # B has 1 site but 2 given

    def test_error_dimension_mismatch_at_overlap(self):
        """Test that kronadd raises error for dimension mismatch at overlap."""
        A = TTMatrix.random([(2, 2), (3, 3)], ranks=[2])
        B = TTMatrix.random([(4, 4), (5, 5)], ranks=[3])  # B[0] = (4,4) != A[1] = (3,3)

        with pytest.raises(ValueError, match="Dimension mismatch"):
            kronadd(A, B, sites=([0, 1], [1, 2]))

    def test_ranks_structure(self):
        """Test that result ranks have expected structure."""
        torch.manual_seed(42)

        A = TTMatrix.random([(2, 2), (3, 3)], ranks=[5])  # A has internal rank 5
        B = TTMatrix.random([(3, 3), (4, 4)], ranks=[7])  # B has internal rank 7

        result = kronadd(A, B, sites=([0, 1], [1, 2]))

        # Result ranks:
        # Bond 0-1: A active (rank 5), B identity (rank 1) -> 5 + 1 = 6
        # Bond 1-2: A terminated (rank 1), B active (rank 7) -> 1 + 7 = 8
        assert result.ranks == [1, 6, 8, 1]

    def test_chained_kronadd(self):
        """Test chaining multiple kronadd operations."""
        torch.manual_seed(42)

        # Build H = H_01 + H_12 + H_23
        H01 = TTMatrix.random([(2, 2), (2, 2)], ranks=[2])
        H12 = TTMatrix.random([(2, 2), (2, 2)], ranks=[2])
        H23 = TTMatrix.random([(2, 2), (2, 2)], ranks=[2])

        # Chain: first add H01 + H12, then add H23
        H_012 = kronadd(H01, H12, sites=([0, 1], [1, 2]))
        H_0123 = kronadd(H_012, H23, sites=([0, 1, 2], [2, 3]))

        assert H_0123.d == 4
        assert H_0123.dims == [(2, 2), (2, 2), (2, 2), (2, 2)]

        # Verify against explicit construction
        I2 = torch.eye(2)
        I4 = torch.eye(4)
        I8 = torch.eye(8)

        H01_dense = H01.to_tensor()
        H12_dense = H12.to_tensor()
        H23_dense = H23.to_tensor()

        # H_01 ⊗ I_4
        term1 = torch.kron(H01_dense, I4)
        # I_2 ⊗ H_12 ⊗ I_2
        term2 = torch.kron(torch.kron(I2, H12_dense), I2)
        # I_4 ⊗ H_23
        term3 = torch.kron(I4, H23_dense)

        expected = term1 + term2 + term3
        actual = H_0123.to_tensor()

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