# Copyright (C) 2025-2026 Changkai Zhang.
#
# This file is part of Nicole library.
#
# Nicole is free software: you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published
# by the Free Software Foundation, either version 3 of the License,
# or (at your option) any later version.
#
# Nicole is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Nicole. If not, see <https://www.gnu.org/licenses/>.


"""Tests for identity and isometry tensor construction."""

import math
import torch
import pytest

from nicole import Direction, Tensor, Index, Sector
from nicole import identity, isometry, isometry_n, contract
from nicole import U1Group, Z2Group, SU2Group, ProductGroup
from ..utils import assert_charge_neutral


# Identity tensor tests

def test_identity_basic():
    """Test basic identity tensor construction."""
    group = U1Group()
    idx = Index(Direction.OUT, group, sectors=(Sector(0, 3), Sector(1, 2)))
    ident = identity(idx, itags=("p", "p_dual"))

    # Identity should have matching charges on both indices
    for (q_left, q_right), block in ident.data.items():
        assert q_left == q_right
        assert torch.allclose(block, torch.eye(block.shape[0], dtype=block.dtype))


def test_identity_default_itags():
    """Test identity tensor with default itags."""
    group = U1Group()
    idx = Index(Direction.OUT, group, sectors=(Sector(0, 2),))
    ident = identity(idx)
    
    assert ident.itags[0] == "_init_"
    assert ident.itags[1] == "_init_"


def test_identity_dtype():
    """Test identity tensor with different dtypes."""
    group = U1Group()
    idx = Index(Direction.OUT, group, sectors=(Sector(0, 2),))
    
    # float32
    ident_f32 = identity(idx, dtype=torch.float32)
    assert ident_f32.dtype == torch.float32
    
    # complex128
    ident_c128 = identity(idx, dtype=torch.complex128)
    assert ident_c128.dtype == torch.complex128


def test_identity_directions():
    """Test that identity has flipped directions."""
    group = U1Group()
    idx = Index(Direction.OUT, group, sectors=(Sector(0, 2),))
    ident = identity(idx)
    
    assert ident.indices[0].direction == Direction.OUT
    assert ident.indices[1].direction == Direction.IN


def test_identity_blocks_are_identity_matrices():
    """Test that all identity blocks are identity matrices."""
    group = U1Group()
    idx = Index(Direction.OUT, group, sectors=(Sector(0, 2), Sector(1, 3), Sector(-1, 4)))
    ident = identity(idx)
    
    for (q_left, q_right), block in ident.data.items():
        assert q_left == q_right
        expected = torch.eye(block.shape[0], dtype=ident.dtype)
        assert torch.allclose(block, expected)


def test_identity_charge_neutral():
    """Test that identity tensor is charge neutral."""
    group = U1Group()
    idx = Index(Direction.OUT, group, sectors=(Sector(0, 2), Sector(1, 3)))
    ident = identity(idx)
    
    assert_charge_neutral(ident)


def test_identity_with_z2():
    """Test identity tensor with Z2 symmetry."""
    group = Z2Group()
    idx = Index(Direction.OUT, group, sectors=(Sector(0, 2), Sector(1, 3)))
    ident = identity(idx)
    
    assert set(ident.data.keys()) == {(0, 0), (1, 1)}
    for key, block in ident.data.items():
        assert torch.allclose(block, torch.eye(block.shape[0], dtype=block.dtype))


def test_identity_contraction_preserves_tensor():
    """Test that contracting with identity preserves the tensor."""
    group = U1Group()
    idx = Index(Direction.OUT, group, sectors=(Sector(0, 2), Sector(1, 3)))
    
    # Create tensor with two indices that can be connected by identity
    T = Tensor.random([idx, idx.flip()], seed=1, itags=["a", "b"])
    # Identity connects two indices with opposite directions
    ident = identity(idx, itags=("b", "c"))
    
    # Contract T's second index (IN, tag="b") with ident's first index (OUT, tag="b")
    # They have opposite directions and matching tags
    result = contract(T, ident)
    
    # Result should have two indices left: T's first and ident's second
    assert len(result.indices) == 2
    # Check that it's charge neutral
    assert_charge_neutral(result)


# Identity tensor tests - SU(2)

def test_identity_su2_basic():
    """Test identity tensor construction for SU(2)."""
    group = SU2Group()
    idx = Index(Direction.OUT, group, sectors=(Sector(1, 2),))
    
    ident = identity(idx, itags=("p", "p_dual"))
    
    # Should have intw populated
    assert ident.intw is not None
    assert len(ident.intw) == 1
    
    # Block shape should be (dim, dim, 1)
    key = (1, 1)
    assert key in ident.data
    assert ident.data[key].shape == (2, 2, 1)
    
    # Reduced tensor should be identity matrix with trailing dimension
    expected = torch.eye(2, dtype=torch.float64).unsqueeze(-1)
    assert torch.allclose(ident.data[key], expected)
    
    # Bridge should exist
    assert key in ident.intw
    bridge = ident.intw[key]
    assert bridge.num_components == 1
    assert bridge.om_dimension == 1  # 2 edges: no OM (n-2=0)
    
    # Weights should be √(2j+1)
    # Charge q=1 means 2j=1, so irrep_dim = 2j+1 = 2
    expected_weight = math.sqrt(2)
    assert math.isclose(bridge.weights[0, 0].item(), expected_weight, rel_tol=1e-12)


def test_identity_su2_multiple_sectors():
    """Test identity tensor with multiple SU(2) sectors."""
    group = SU2Group()
    idx = Index(Direction.OUT, group, sectors=(Sector(1, 2), Sector(2, 3), Sector(3, 4)))
    
    ident = identity(idx, itags=("a", "b"))
    
    # Should have 3 diagonal blocks
    assert len(ident.data) == 3
    assert len(ident.intw) == 3
    
    # Check each block
    for sector in idx.sectors:
        q = sector.charge
        dim = sector.dim
        key = (q, q)
        
        # Block shape
        assert ident.data[key].shape == (dim, dim, 1)
        
        # Reduced tensor is identity
        expected = torch.eye(dim, dtype=ident.dtype).unsqueeze(-1)
        assert torch.allclose(ident.data[key], expected)
        
        # Weights normalization: √(2j+1) where q = 2j
        bridge = ident.intw[key]
        irrep_dim = q + 1  # 2j + 1
        expected_weight = math.sqrt(irrep_dim)
        assert math.isclose(bridge.weights[0, 0].item(), expected_weight, rel_tol=1e-12)


def test_identity_su2_direction_in():
    """Test identity tensor with Direction.IN."""
    group = SU2Group()
    idx = Index(Direction.IN, group, sectors=(Sector(1, 2),))
    
    ident = identity(idx)
    
    # First index should be IN, second should be OUT (flipped)
    assert ident.indices[0].direction == Direction.IN
    assert ident.indices[1].direction == Direction.OUT
    
    # Bridge should have correct directions
    key = (1, 1)
    bridge = ident.intw[key]
    assert bridge.num_external == 2


def test_identity_su2_dtype_complex():
    """Test identity tensor with complex dtype for SU(2)."""
    group = SU2Group()
    idx = Index(Direction.OUT, group, sectors=(Sector(1, 2),))
    
    ident = identity(idx, dtype=torch.complex128)
    
    assert ident.dtype == torch.complex128
    key = (1, 1)
    assert ident.data[key].dtype == torch.complex128
    
    # Weights should also be complex
    assert ident.intw[key].weights.dtype == torch.complex128


def test_identity_su2_charge_neutral():
    """Test that SU(2) identity tensor is charge neutral."""
    group = SU2Group()
    idx = Index(Direction.OUT, group, sectors=(Sector(1, 2), Sector(2, 3)))
    
    ident = identity(idx)
    
    # Identity should be charge neutral
    assert_charge_neutral(ident)


def test_identity_su2_norm_single_sector():
    """Test norm of SU(2) identity with single sector."""
    group = SU2Group()
    idx = Index(Direction.OUT, group, sectors=(Sector(1, 2),))
    
    ident = identity(idx)
    
    # norm² = sector.dim * irrep_dim = 2 * 2 = 4
    # norm = 2
    assert math.isclose(ident.norm(), 2.0, rel_tol=1e-12)


def test_identity_su2_norm_multiple_sectors():
    """Test norm of SU(2) identity with multiple sectors."""
    group = SU2Group()
    idx = Index(Direction.OUT, group, sectors=(
        Sector(1, 2),   # spin-1/2, dim=2, irrep_dim=2
        Sector(2, 3),   # spin-1, dim=3, irrep_dim=3
        Sector(3, 4),   # spin-3/2, dim=4, irrep_dim=4
    ))
    
    ident = identity(idx)
    
    # norm² = Σ (sector.dim * irrep_dim) = 2*2 + 3*3 + 4*4 = 4 + 9 + 16 = 29
    expected_norm = math.sqrt(29)
    assert math.isclose(ident.norm(), expected_norm, rel_tol=1e-12)


def test_identity_su2_norm_spin_0():
    """Test norm of SU(2) identity for spin-0."""
    group = SU2Group()
    idx = Index(Direction.OUT, group, sectors=(Sector(0, 1),))
    
    ident = identity(idx)
    
    # norm² = sector.dim * irrep_dim = 1 * 1 = 1
    # norm = 1
    assert math.isclose(ident.norm(), 1.0, rel_tol=1e-12)


def test_identity_product_group_su2_norm():
    """Test norm of identity with ProductGroup containing SU(2)."""
    group = ProductGroup([U1Group(), SU2Group()])
    idx = Index(Direction.OUT, group, sectors=(
        Sector((0, 1), 2),   # U1=0, spin-1/2: dim=2, irrep_dim=1*2=2
        Sector((1, 2), 3),   # U1=1, spin-1: dim=3, irrep_dim=1*3=3
    ))
    
    ident = identity(idx)
    
    # norm² = 2*2 + 3*3 = 4 + 9 = 13
    expected_norm = math.sqrt(13)
    assert math.isclose(ident.norm(), expected_norm, rel_tol=1e-12)


def test_identity_product_group_with_su2():
    """Test identity with ProductGroup containing SU(2)."""
    group = ProductGroup([U1Group(), SU2Group()])
    idx = Index(Direction.OUT, group, sectors=(
        Sector((0, 1), 2),
        Sector((1, 2), 3),
    ))
    
    ident = identity(idx)
    
    # Should have intw (non-Abelian due to SU(2))
    assert ident.intw is not None
    assert len(ident.data) == 2
    
    # Check blocks
    for sector in idx.sectors:
        q = sector.charge
        dim = sector.dim
        key = (q, q)
        
        assert ident.data[key].shape == (dim, dim, 1)
        assert key in ident.intw
        
        # Weights normalization: √(irrep_dim)
        irrep_dim = group.irrep_dim(q)
        expected_weight = math.sqrt(irrep_dim)
        assert math.isclose(ident.intw[key].weights[0, 0].item(), expected_weight, rel_tol=1e-12)


# Isometry tensor tests

def test_isometry_basic():
    """Test basic isometry tensor construction."""
    group = U1Group()
    idx_a = Index(Direction.OUT, group, sectors=(Sector(0, 2), Sector(1, 1)))
    idx_b = Index(Direction.OUT, group, sectors=(Sector(0, 1), Sector(-1, 2)))
    
    fused_tensor = isometry(idx_a, idx_b, itags=("a", "b", "ab"))
    
    assert len(fused_tensor.indices) == 3
    assert fused_tensor.itags == ("a", "b", "ab")


def test_isometry_fused_charges():
    """Test that isometry has correct fused charges."""
    group = U1Group()
    idx_a = Index(Direction.OUT, group, sectors=(Sector(0, 2), Sector(1, 1)))
    idx_b = Index(Direction.OUT, group, sectors=(Sector(0, 1), Sector(-1, 2)))
    
    fused_tensor = isometry(idx_a, idx_b, itags=("a", "b", "ab"))

    # Check fused index metadata
    fused_index = fused_tensor.indices[2]
    fused_charges = {sector.charge for sector in fused_index.sectors}
    expected_charges = {idx_a.group.fuse_unique(sa.charge, sb.charge)
        for sa in idx_a.sectors for sb in idx_b.sectors}
    assert fused_charges == expected_charges


def test_isometry_structure():
    """Test that isometry has correct selection structure."""
    group = U1Group()
    idx_a = Index(Direction.OUT, group, sectors=(Sector(0, 2), Sector(1, 1)))
    idx_b = Index(Direction.OUT, group, sectors=(Sector(0, 1), Sector(-1, 2)))
    
    fused_tensor = isometry(idx_a, idx_b, itags=("a", "b", "ab"))

    # Check selection structure: each product multiplet maps to at most one fused multiplet
    for key, block in fused_tensor.data.items():
        reshaped = block.reshape(block.shape[0] * block.shape[1], block.shape[2])
        mask = torch.abs(reshaped) > 1e-12
        # Each product basis row contributes to at most one fused column
        assert torch.all(mask.sum(axis=1) <= 1)
        # Each fused column receives either zero or one contributions (no superposition for U(1))
        assert torch.all((mask.sum(axis=0) == 0) | (mask.sum(axis=0) == 1))
        # Values should be exactly 0 or 1
        assert torch.all((torch.abs(reshaped) < 1e-12) | (torch.abs(reshaped - 1.0) < 1e-12))


def test_isometry_default_itags():
    """Test isometry with default itags."""
    group = U1Group()
    idx_a = Index(Direction.OUT, group, sectors=(Sector(0, 2),))
    idx_b = Index(Direction.OUT, group, sectors=(Sector(0, 1),))
    
    fused_tensor = isometry(idx_a, idx_b)
    
    assert fused_tensor.itags == ("_init_", "_init_", "_init_")


def test_isometry_dtype():
    """Test isometry with different dtypes."""
    group = U1Group()
    idx_a = Index(Direction.OUT, group, sectors=(Sector(0, 2),))
    idx_b = Index(Direction.OUT, group, sectors=(Sector(0, 1),))
    
    # float32
    iso_f32 = isometry(idx_a, idx_b, dtype=torch.float32)
    assert iso_f32.dtype == torch.float32
    
    # complex128
    iso_c128 = isometry(idx_a, idx_b, dtype=torch.complex128)
    assert iso_c128.dtype == torch.complex128


def test_isometry_fused_direction():
    """Test isometry with custom fused direction."""
    group = U1Group()
    idx_a = Index(Direction.OUT, group, sectors=(Sector(0, 2),))
    idx_b = Index(Direction.OUT, group, sectors=(Sector(0, 1),))
    
    # Default should be dual of first
    iso_default = isometry(idx_a, idx_b)
    assert iso_default.indices[2].direction == Direction.IN
    
    # Custom direction
    iso_custom = isometry(idx_a, idx_b, fused_direction=Direction.OUT)
    assert iso_custom.indices[2].direction == Direction.OUT


def test_isometry_charge_neutral():
    """Test that isometry tensor is charge neutral."""
    group = U1Group()
    idx_a = Index(Direction.OUT, group, sectors=(Sector(0, 2), Sector(1, 1)))
    idx_b = Index(Direction.OUT, group, sectors=(Sector(0, 1), Sector(-1, 2)))
    
    fused_tensor = isometry(idx_a, idx_b)
    
    assert_charge_neutral(fused_tensor)


def test_isometry_dimensions():
    """Test that isometry has correct dimensions."""
    group = U1Group()
    idx_a = Index(Direction.OUT, group, sectors=(Sector(0, 2), Sector(1, 3)))
    idx_b = Index(Direction.OUT, group, sectors=(Sector(0, 4), Sector(-1, 5)))
    
    fused_tensor = isometry(idx_a, idx_b)
    
    # Fused dimension should be sum of products
    fused_index = fused_tensor.indices[2]
    # Total fused dim = 2*4 + 2*5 + 3*4 + 3*5 = 8 + 10 + 12 + 15 = 45
    assert fused_index.dim == 2*4 + 2*5 + 3*4 + 3*5


def test_isometry_different_groups_raises():
    """Test that isometry raises error for different groups."""
    u1 = U1Group()
    z2 = Z2Group()
    idx_u1 = Index(Direction.OUT, u1, sectors=(Sector(0, 2),))
    idx_z2 = Index(Direction.OUT, z2, sectors=(Sector(0, 2),))
    
    with pytest.raises(ValueError, match="same symmetry group"):
        isometry(idx_u1, idx_z2)


def test_isometry_z2():
    """Test isometry with Z2 symmetry."""
    group = Z2Group()
    idx_a = Index(Direction.OUT, group, sectors=(Sector(0, 2), Sector(1, 1)))
    idx_b = Index(Direction.OUT, group, sectors=(Sector(0, 1), Sector(1, 2)))
    
    fused_tensor = isometry(idx_a, idx_b)
    
    # Z2: 0+0=0, 0+1=1, 1+0=1, 1+1=0
    fused_index = fused_tensor.indices[2]
    fused_charges = {sector.charge for sector in fused_index.sectors}
    assert fused_charges == {0, 1}


def test_isometry_positive_norm():
    """Test that isometry is successfully constructed with non-zero norm."""
    group = U1Group()
    idx_a = Index(Direction.OUT, group, sectors=(Sector(0, 2), Sector(1, 1)))
    idx_b = Index(Direction.OUT, group, sectors=(Sector(0, 1), Sector(-1, 1)))

    iso = isometry(idx_a, idx_b, itags=("a", "b", "ab"))

    assert iso.norm() > 0


def test_isometry_orthonormality():
    """Test that isometry columns are orthonormal."""
    group = U1Group()
    idx_a = Index(Direction.OUT, group, sectors=(Sector(0, 2),))
    idx_b = Index(Direction.OUT, group, sectors=(Sector(0, 1),))
    
    iso = isometry(idx_a, idx_b)
    
    # For each block, columns should be orthonormal
    for key, block in iso.data.items():
        reshaped = block.reshape(block.shape[0] * block.shape[1], block.shape[2])
        # Compute Gram matrix
        gram = reshaped.T @ reshaped
        # Should be identity (columns are orthonormal)
        assert torch.allclose(gram, torch.eye(gram.shape[0], dtype=gram.dtype), atol=1e-12)


def test_identity_and_isometry_consistent():
    """Test that identity is consistent with isometry for single index."""
    group = U1Group()
    idx = Index(Direction.OUT, group, sectors=(Sector(0, 2), Sector(1, 3)))
    
    # Identity should behave like isometry with one index
    ident = identity(idx, itags=("a", "b"))
    
    # Check structure
    assert len(ident.indices) == 2
    assert ident.indices[0] == idx
    assert ident.indices[1] == idx.flip()


# Isometry tensor tests - SU(2)

def test_isometry_su2_basic():
    """Test basic SU(2) isometry construction."""
    group = SU2Group()
    idx1 = Index(Direction.IN, group, sectors=(Sector(1, 2),))  # spin-1/2
    idx2 = Index(Direction.IN, group, sectors=(Sector(1, 2),))  # spin-1/2
    
    iso = isometry(idx1, idx2, itags=("a", "b", "fused"))
    
    # Should have 3 indices
    assert len(iso.indices) == 3
    assert iso.itags == ("a", "b", "fused")
    
    # Should have intw populated
    assert iso.intw is not None
    
    # Fusion of 1/2 ⊗ 1/2 → 0, 2 (spin-0 and spin-1)
    # Should have 2 blocks: (1,1,0) and (1,1,2)
    assert (1, 1, 0) in iso.data
    assert (1, 1, 2) in iso.data
    
    # Check block shapes: (dim1, dim2, dim_fused, 1)
    # The fused dimension is from combine_indices, which accumulates dims
    # Just check trailing dimension is 1
    assert iso.data[(1, 1, 0)].shape[-1] == 1
    assert iso.data[(1, 1, 2)].shape[-1] == 1
    assert iso.data[(1, 1, 0)].shape[:2] == (2, 2)
    assert iso.data[(1, 1, 2)].shape[:2] == (2, 2)
    
    # Check bridges exist
    assert (1, 1, 0) in iso.intw
    assert (1, 1, 2) in iso.intw
    
    # Weights should be √(irrep_dim(qf))
    # For qf=0: irrep_dim=1, weight=1
    assert math.isclose(iso.intw[(1, 1, 0)].weights[0, 0].item(), 1.0, rel_tol=1e-12)
    # For qf=2: irrep_dim=3, weight=√3
    assert math.isclose(iso.intw[(1, 1, 2)].weights[0, 0].item(), math.sqrt(3), rel_tol=1e-12)


def test_isometry_su2_multiple_sectors():
    """Test SU(2) isometry with multiple sectors and fusion channels."""
    group = SU2Group()
    idx1 = Index(Direction.OUT, group, sectors=(Sector(1, 2), Sector(2, 3)))
    idx2 = Index(Direction.OUT, group, sectors=(Sector(1, 2),))
    
    iso = isometry(idx1, idx2)
    
    # Should have multiple blocks from different fusion channels
    # 1⊗1 → 0,2
    # 2⊗1 → 1,3
    assert len(iso.data) >= 4
    assert iso.intw is not None
    
    # All blocks should have trailing dimension of 1
    for key, block in iso.data.items():
        assert block.shape[-1] == 1
        assert key in iso.intw


def test_isometry_su2_charge_neutral():
    """Test that SU(2) isometry is charge neutral."""
    group = SU2Group()
    idx1 = Index(Direction.IN, group, sectors=(Sector(1, 2),))
    idx2 = Index(Direction.IN, group, sectors=(Sector(1, 2),))
    
    iso = isometry(idx1, idx2)
    
    assert_charge_neutral(iso)


def test_isometry_su2_dtype_complex():
    """Test SU(2) isometry with complex dtype."""
    group = SU2Group()
    idx1 = Index(Direction.OUT, group, sectors=(Sector(1, 2),))
    idx2 = Index(Direction.OUT, group, sectors=(Sector(1, 2),))
    
    iso = isometry(idx1, idx2, dtype=torch.complex128)
    
    assert iso.dtype == torch.complex128
    for key, block in iso.data.items():
        assert block.dtype == torch.complex128
        assert iso.intw[key].weights.dtype == torch.complex128


def test_isometry_su2_fused_direction():
    """Test SU(2) isometry with specified fused direction."""
    group = SU2Group()
    idx1 = Index(Direction.IN, group, sectors=(Sector(1, 2),))
    idx2 = Index(Direction.IN, group, sectors=(Sector(1, 2),))
    
    # Specify fused direction as OUT
    iso = isometry(idx1, idx2, fused_direction=Direction.OUT)
    
    # Fused index should have OUT direction
    assert iso.indices[2].direction == Direction.OUT


def test_isometry_su2_norm():
    """Test norm of SU(2) isometry."""
    group = SU2Group()
    idx1 = Index(Direction.IN, group, sectors=(Sector(1, 2),))
    idx2 = Index(Direction.IN, group, sectors=(Sector(1, 2),))
    
    iso = isometry(idx1, idx2)
    
    # Calculate expected norm: norm² = Σ_blocks (sector_dims × irrep_dim(qf))
    # For each block (qa, qb, qf), contribution is:
    # dim_a × dim_b × dim_fused × irrep_dim(qf)
    # But with identity-like structure, effectively: contribution per block element
    
    # Let me compute manually for fusion 1⊗1→0,2:
    # Block (1,1,0): shape (2,2,1,1), weight=√1=1, reduced is identity-like
    # Block (1,1,2): shape (2,2,3,1), weight=√3, reduced is identity-like
    
    # Actually, with identity-like reduced tensor:
    # norm² = Σ (number of 1's in reduced) × weight²
    # For each block: da*db ones, each contributes weight²
    # (1,1,0): 2*2=4 ones, weight²=1, contribution = 4
    # (1,1,2): 2*2=4 ones, weight²=3, contribution = 12
    # Total norm² = 4 + 12 = 16, norm = 4
    expected_norm = 4.0
    assert math.isclose(iso.norm(), expected_norm, rel_tol=1e-12)


def test_isometry_product_group_with_su2():
    """Test isometry with ProductGroup containing SU(2)."""
    group = ProductGroup([U1Group(), SU2Group()])
    idx1 = Index(Direction.OUT, group, sectors=(Sector((0, 1), 2),))
    idx2 = Index(Direction.OUT, group, sectors=(Sector((0, 1), 2),))
    
    iso = isometry(idx1, idx2)
    
    # Should have intw
    assert iso.intw is not None
    
    # U1 part: 0⊗0→0, SU(2) part: 1⊗1→0,2
    # Should have blocks for (0,1)⊗(0,1)→(0,0) and (0,1)⊗(0,1)→(0,2)
    assert ((0, 1), (0, 1), (0, 0)) in iso.data
    assert ((0, 1), (0, 1), (0, 2)) in iso.data


def test_isometry_su2_multiple_pairs_same_fused():
    """Test SU(2) isometry where multiple sector pairs fuse to same charge.
    
    This tests the offset mechanism: 1/2⊗1/2→0 and 1⊗1→0 both fuse to spin-0.
    The fused dimension should accumulate contributions from both pairs.
    """
    group = SU2Group()
    # Create indices with spin-1/2 and spin-1 sectors
    idx1 = Index(Direction.IN, group, sectors=(Sector(1, 2), Sector(2, 3)))
    idx2 = Index(Direction.IN, group, sectors=(Sector(1, 2), Sector(2, 3)))
    
    iso = isometry(idx1, idx2)
    
    # Check that blocks exist for both pairs fusing to spin-0
    # 1⊗1→0 (spin-1/2 ⊗ spin-1/2 → spin-0)
    assert (1, 1, 0) in iso.data
    assert (1, 1, 0) in iso.intw
    
    # 2⊗2→0 (spin-1 ⊗ spin-1 → spin-0)
    assert (2, 2, 0) in iso.data
    assert (2, 2, 0) in iso.intw
    
    # Verify weights are √(irrep_dim(0)) = √1 = 1 for both
    assert math.isclose(iso.intw[(1, 1, 0)].weights[0, 0].item(), 1.0, rel_tol=1e-12)
    assert math.isclose(iso.intw[(2, 2, 0)].weights[0, 0].item(), 1.0, rel_tol=1e-12)
    
    # Check fused dimension includes contributions from both pairs
    # The fused index for charge 0 should have accumulated dimension from:
    # - (1,1,0): contributes 2*2 = 4
    # - (2,2,0): contributes 3*3 = 9
    # Total dimension for fused charge 0 should be at least 4+9 = 13
    fused_idx = iso.indices[2]
    for sector in fused_idx.sectors:
        if sector.charge == 0:
            assert sector.dim >= 13, f"Expected fused dim >= 13 for charge 0, got {sector.dim}"
    
    # Verify block shapes reflect the accumulated dimension
    # Both blocks should map to the same fused dimension
    fused_dim_0 = iso.data[(1, 1, 0)].shape[2]
    assert iso.data[(2, 2, 0)].shape[2] == fused_dim_0


def test_isometry_su2_norm_multiple_fused_charges():
    """Test norm calculation with multiple pairs fusing to various charges."""
    group = SU2Group()
    idx1 = Index(Direction.OUT, group, sectors=(Sector(1, 2), Sector(2, 3)))
    idx2 = Index(Direction.OUT, group, sectors=(Sector(1, 2),))
    
    iso = isometry(idx1, idx2)
    
    # Calculate expected norm manually
    # norm² = Σ_blocks (number of 1's in reduced tensor) × weight²
    # For identity-like reduced tensor: da*db ones per block
    
    expected_norm_sq = 0.0
    for key, block in iso.data.items():
        qa, qb, qf = key
        da, db = block.shape[0], block.shape[1]
        num_ones = da * db
        weight = iso.intw[key].weights[0, 0].item()
        expected_norm_sq += num_ones * (weight ** 2)
    
    expected_norm = math.sqrt(expected_norm_sq)
    actual_norm = iso.norm()
    
    assert math.isclose(actual_norm, expected_norm, rel_tol=1e-10)


def test_isometry_su2_three_sectors_comprehensive():
    """Comprehensive test with 3 sectors creating many fusion channels."""
    group = SU2Group()
    # Spin-1/2, spin-1, and spin-3/2
    idx1 = Index(Direction.OUT, group, sectors=(
        Sector(1, 2),   # spin-1/2
        Sector(2, 3),   # spin-1
        Sector(3, 4),   # spin-3/2
    ))
    idx2 = Index(Direction.OUT, group, sectors=(Sector(1, 2),))  # spin-1/2
    
    iso = isometry(idx1, idx2)
    
    # Verify charge conservation
    assert_charge_neutral(iso)
    
    # Check that multiple blocks exist
    # 1⊗1 → 0,2
    # 2⊗1 → 1,3
    # 3⊗1 → 2,4
    assert len(iso.data) >= 6
    assert len(iso.intw) == len(iso.data)
    
    # Verify all weights follow √(irrep_dim(qf)) pattern
    for key, bridge in iso.intw.items():
        qa, qb, qf = key
        expected_weight = math.sqrt(group.irrep_dim(qf))
        assert math.isclose(bridge.weights[0, 0].item(), expected_weight, rel_tol=1e-12)
    
    # Verify norm is computed correctly
    expected_norm_sq = 0.0
    for key, block in iso.data.items():
        qa, qb, qf = key
        # Get sector dims from original indices
        da = next(s.dim for s in idx1.sectors if s.charge == qa)
        db = next(s.dim for s in idx2.sectors if s.charge == qb)
        num_ones = da * db
        weight_sq = iso.intw[key].weights[0, 0].item() ** 2
        expected_norm_sq += num_ones * weight_sq
    
    expected_norm = math.sqrt(expected_norm_sq)
    assert math.isclose(iso.norm(), expected_norm, rel_tol=1e-10)


# Isometry_n tensor tests

def test_isometry_n_basic():
    """Test basic isometry_n construction with 3 indices."""
    group = U1Group()
    idx1 = Index(Direction.OUT, group, sectors=(Sector(0, 2), Sector(1, 1)))
    idx2 = Index(Direction.OUT, group, sectors=(Sector(0, 1), Sector(-1, 2)))
    idx3 = Index(Direction.OUT, group, sectors=(Sector(0, 3), Sector(1, 1)))
    
    iso = isometry_n([idx1, idx2, idx3], itags=("a", "b", "c", "fused"))
    
    # Should have 4 indices (3 unfused + 1 fused)
    assert len(iso.indices) == 4
    assert iso.itags == ("a", "b", "c", "fused")


def test_isometry_n_minimum_indices():
    """Test isometry_n with minimum 2 indices."""
    group = U1Group()
    idx1 = Index(Direction.OUT, group, sectors=(Sector(0, 2),))
    idx2 = Index(Direction.OUT, group, sectors=(Sector(0, 1),))
    
    iso = isometry_n([idx1, idx2])
    
    # Should have 3 indices (2 unfused + 1 fused)
    assert len(iso.indices) == 3
    # All default tags
    assert all(tag == "_init_" for tag in iso.itags)


def test_isometry_n_too_few_indices_raises():
    """Test that isometry_n raises error with fewer than 2 indices."""
    group = U1Group()
    idx = Index(Direction.OUT, group, sectors=(Sector(0, 2),))
    
    with pytest.raises(ValueError, match="at least 2 indices"):
        isometry_n([idx])
    
    with pytest.raises(ValueError, match="at least 2 indices"):
        isometry_n([])


def test_isometry_n_direction_opposite():
    """Test that isometry_n creates indices with opposite directions."""
    group = U1Group()
    idx1 = Index(Direction.OUT, group, sectors=(Sector(0, 2),))
    idx2 = Index(Direction.IN, group, sectors=(Sector(0, 1),))
    idx3 = Index(Direction.OUT, group, sectors=(Sector(0, 3),))
    
    iso = isometry_n([idx1, idx2, idx3])
    
    # First 3 indices should have opposite directions to inputs
    assert iso.indices[0].direction == Direction.IN  # Opposite of idx1 (OUT)
    assert iso.indices[1].direction == Direction.OUT  # Opposite of idx2 (IN)
    assert iso.indices[2].direction == Direction.IN  # Opposite of idx3 (OUT)
    # Last index should be the fused index with default direction OUT
    assert iso.indices[3].direction == Direction.OUT


def test_isometry_n_fused_direction_parameter():
    """Test that fused direction parameter works correctly."""
    group = U1Group()
    idx1 = Index(Direction.OUT, group, sectors=(Sector(0, 2),))
    idx2 = Index(Direction.OUT, group, sectors=(Sector(0, 1),))
    idx3 = Index(Direction.OUT, group, sectors=(Sector(0, 3),))
    
    # Test with Direction.OUT (default)
    iso_out = isometry_n([idx1, idx2, idx3], direction=Direction.OUT)
    assert iso_out.indices[3].direction == Direction.OUT
    
    # Test with Direction.IN
    iso_in = isometry_n([idx1, idx2, idx3], direction=Direction.IN)
    assert iso_in.indices[3].direction == Direction.IN


def test_isometry_n_different_groups_raises():
    """Test that isometry_n raises error for different groups."""
    u1 = U1Group()
    z2 = Z2Group()
    idx_u1 = Index(Direction.OUT, u1, sectors=(Sector(0, 2),))
    idx_z2 = Index(Direction.OUT, z2, sectors=(Sector(0, 2),))
    
    with pytest.raises(ValueError, match="same symmetry group"):
        isometry_n([idx_u1, idx_z2])


def test_isometry_n_itags_validation():
    """Test that itags length validation works."""
    group = U1Group()
    idx1 = Index(Direction.OUT, group, sectors=(Sector(0, 2),))
    idx2 = Index(Direction.OUT, group, sectors=(Sector(0, 1),))
    idx3 = Index(Direction.OUT, group, sectors=(Sector(0, 3),))
    
    # Too few tags
    with pytest.raises(ValueError, match="itags must have length 4"):
        isometry_n([idx1, idx2, idx3], itags=("a", "b"))
    
    # Too many tags
    with pytest.raises(ValueError, match="itags must have length 4"):
        isometry_n([idx1, idx2, idx3], itags=("a", "b", "c", "d", "e"))
    
    # Correct number of tags should work
    iso = isometry_n([idx1, idx2, idx3], itags=("a", "b", "c", "fused"))
    assert iso.itags == ("a", "b", "c", "fused")


def test_isometry_n_charge_neutral():
    """Test that isometry_n tensor is charge neutral."""
    group = U1Group()
    idx1 = Index(Direction.OUT, group, sectors=(Sector(0, 2), Sector(1, 1)))
    idx2 = Index(Direction.OUT, group, sectors=(Sector(0, 1), Sector(-1, 2)))
    idx3 = Index(Direction.IN, group, sectors=(Sector(0, 3), Sector(1, 1)))
    
    iso = isometry_n([idx1, idx2, idx3])
    
    assert_charge_neutral(iso)


def test_isometry_n_dtype():
    """Test isometry_n with different dtypes."""
    group = U1Group()
    idx1 = Index(Direction.OUT, group, sectors=(Sector(0, 2),))
    idx2 = Index(Direction.OUT, group, sectors=(Sector(0, 1),))
    
    # float32
    iso_f32 = isometry_n([idx1, idx2], dtype=torch.float32)
    assert iso_f32.dtype == torch.float32
    
    # complex128
    iso_c128 = isometry_n([idx1, idx2], dtype=torch.complex128)
    assert iso_c128.dtype == torch.complex128


def test_isometry_n_z2_symmetry():
    """Test isometry_n with Z2 symmetry."""
    group = Z2Group()
    idx1 = Index(Direction.OUT, group, sectors=(Sector(0, 2), Sector(1, 1)))
    idx2 = Index(Direction.OUT, group, sectors=(Sector(0, 1), Sector(1, 2)))
    idx3 = Index(Direction.OUT, group, sectors=(Sector(0, 3), Sector(1, 1)))
    
    iso = isometry_n([idx1, idx2, idx3])
    
    # Check structure
    assert len(iso.indices) == 4
    assert_charge_neutral(iso)
    
    # Fused index should have both Z2 charges
    fused_index = iso.indices[3]
    fused_charges = {sector.charge for sector in fused_index.sectors}
    assert 0 in fused_charges or 1 in fused_charges  # Should have at least one charge


def test_isometry_n_product_group():
    """Test isometry_n with ProductGroup."""
    u1 = U1Group()
    z2 = Z2Group()
    group = ProductGroup([u1, z2])
    
    idx1 = Index(Direction.OUT, group, sectors=(Sector((0, 0), 2), Sector((1, 1), 1)))
    idx2 = Index(Direction.OUT, group, sectors=(Sector((0, 0), 1), Sector((-1, 1), 2)))
    
    iso = isometry_n([idx1, idx2])
    
    # Check structure
    assert len(iso.indices) == 3
    assert_charge_neutral(iso)


def test_isometry_n_contraction_with_tensor():
    """Test that isometry_n can contract with a tensor to fuse indices."""
    group = U1Group()
    idx1 = Index(Direction.OUT, group, sectors=(Sector(0, 2), Sector(1, 1)))
    idx2 = Index(Direction.OUT, group, sectors=(Sector(0, 1), Sector(-1, 1)))
    idx3 = Index(Direction.OUT, group, sectors=(Sector(0, 2), Sector(1, 1)))
    idx_extra = Index(Direction.OUT, group, sectors=(Sector(0, 2),))
    
    # Create a random tensor with 4 indices (extra one to avoid 1-index tensor after contraction)
    T = Tensor.random([idx1, idx2, idx3, idx_extra], seed=42, itags=["a", "b", "c", "extra"])
    
    # Create isometry to fuse first 3 indices
    iso = isometry_n([idx1, idx2, idx3], itags=["a", "b", "c", "fused"])
    
    # Contract - should automatically match on opposite directions and tags
    result = contract(T, iso)
    
    # Result should have the fused index and the extra index
    assert len(result.indices) == 2
    assert "fused" in result.itags
    assert "extra" in result.itags
    
    # Should be charge neutral
    assert_charge_neutral(result)
    
    # Check dimension makes sense
    # Fused dimension should be product of all input dimensions
    expected_total_dim = idx1.dim * idx2.dim * idx3.dim
    fused_idx = result.indices[0] if result.itags[0] == "fused" else result.indices[1]
    # Result fused index dimension should be <= expected (due to symmetry constraints)
    assert fused_idx.dim <= expected_total_dim


def test_isometry_n_dimension_sorting():
    """Test that isometry_n fuses indices in order of increasing dimension."""
    group = U1Group()
    # Create indices with different dimensions: 6, 2, 4
    idx1 = Index(Direction.OUT, group, sectors=(Sector(0, 6),))  # dim=6
    idx2 = Index(Direction.OUT, group, sectors=(Sector(0, 2),))  # dim=2
    idx3 = Index(Direction.OUT, group, sectors=(Sector(0, 4),))  # dim=4
    
    iso = isometry_n([idx1, idx2, idx3])
    
    # The isometry should work correctly regardless of input order
    # We verify by checking it's charge neutral and has correct structure
    assert len(iso.indices) == 4
    assert_charge_neutral(iso)
    
    # Check that the first 3 indices match (in original order, with opposite directions)
    assert iso.indices[0].direction == Direction.IN  # Opposite of idx1
    assert iso.indices[1].direction == Direction.IN  # Opposite of idx2
    assert iso.indices[2].direction == Direction.IN  # Opposite of idx3


def test_isometry_n_4th_order():
    """Test isometry_n with 4 indices."""
    group = U1Group()
    idx1 = Index(Direction.OUT, group, sectors=(Sector(0, 2),))
    idx2 = Index(Direction.OUT, group, sectors=(Sector(0, 1),))
    idx3 = Index(Direction.IN, group, sectors=(Sector(0, 3),))
    idx4 = Index(Direction.OUT, group, sectors=(Sector(0, 2),))
    
    iso = isometry_n([idx1, idx2, idx3, idx4], itags=["a", "b", "c", "d", "fused"])
    
    # Should have 5 indices (4 unfused + 1 fused)
    assert len(iso.indices) == 5
    assert iso.itags == ("a", "b", "c", "d", "fused")
    
    # Check directions are opposite
    assert iso.indices[0].direction == Direction.IN
    assert iso.indices[1].direction == Direction.IN
    assert iso.indices[2].direction == Direction.OUT
    assert iso.indices[3].direction == Direction.IN
    assert iso.indices[4].direction == Direction.OUT  # fused


def test_isometry_n_orthonormality():
    """Test that isometry_n has orthonormal columns when viewed as a matrix."""
    group = U1Group()
    idx1 = Index(Direction.OUT, group, sectors=(Sector(0, 2),))
    idx2 = Index(Direction.OUT, group, sectors=(Sector(0, 1),))
    
    iso = isometry_n([idx1, idx2])
    
    # For each block, when reshaped to matrix form, columns should be orthonormal
    for key, block in iso.data.items():
        # Reshape to (product of first n dims, last dim)
        matrix_form = block.reshape(-1, block.shape[-1])
        
        # Compute Gram matrix (should be identity)
        gram = matrix_form.T @ matrix_form
        
        # Should be identity matrix (columns are orthonormal)
        assert torch.allclose(gram, torch.eye(gram.shape[0], dtype=gram.dtype), atol=1e-10)


def test_isometry_n_5th_order():
    """Test isometry_n with 5 indices to verify scalability."""
    group = U1Group()
    indices = [
        Index(Direction.OUT, group, sectors=(Sector(0, 2),)),
        Index(Direction.IN, group, sectors=(Sector(0, 1),)),
        Index(Direction.OUT, group, sectors=(Sector(0, 3),)),
        Index(Direction.OUT, group, sectors=(Sector(0, 2),)),
        Index(Direction.IN, group, sectors=(Sector(0, 1),)),
    ]
    
    iso = isometry_n(indices)
    
    # Should have 6 indices (5 unfused + 1 fused)
    assert len(iso.indices) == 6
    
    # Check directions are opposite to inputs
    for i in range(5):
        assert iso.indices[i].direction == indices[i].direction.reverse()
    
    # Fused index should have default OUT direction
    assert iso.indices[5].direction == Direction.OUT
    
    # Should be charge neutral
    assert_charge_neutral(iso)


def test_isometry_n_mixed_sectors():
    """Test isometry_n with indices having multiple sectors."""
    group = U1Group()
    idx1 = Index(Direction.OUT, group, sectors=(Sector(0, 2), Sector(1, 3), Sector(-1, 1)))
    idx2 = Index(Direction.OUT, group, sectors=(Sector(0, 1), Sector(1, 2), Sector(-1, 1)))
    idx3 = Index(Direction.IN, group, sectors=(Sector(0, 3), Sector(1, 1)))
    
    iso = isometry_n([idx1, idx2, idx3])
    
    # Should have 4 indices
    assert len(iso.indices) == 4
    
    # Should be charge neutral
    assert_charge_neutral(iso)
    
    # Fused index should have some sectors
    fused_index = iso.indices[3]
    assert len(fused_index.sectors) > 0


def test_isometry_n_consistency_with_isometry():
    """Test that isometry_n with 2 indices is consistent with isometry."""
    group = U1Group()
    idx1 = Index(Direction.OUT, group, sectors=(Sector(0, 2), Sector(1, 1)))
    idx2 = Index(Direction.OUT, group, sectors=(Sector(0, 1), Sector(-1, 2)))
    
    # Create with isometry_n
    iso_n = isometry_n([idx1, idx2], itags=("a", "b", "fused"))
    
    # Create with regular isometry (with flipped indices for opposite directions)
    iso = isometry(idx1.flip(), idx2.flip(), itags=("a", "b", "fused"))
    
    # They should have same structure
    assert len(iso_n.indices) == len(iso.indices)
    assert iso_n.itags == iso.itags
    
    # Directions should match
    for i in range(len(iso_n.indices)):
        assert iso_n.indices[i].direction == iso.indices[i].direction
    
    # Both should be charge neutral
    assert_charge_neutral(iso_n)
    assert_charge_neutral(iso)
