# Copyright (C) 2025-2026 Changkai Zhang.
#
# This file is part of Nicole (TN) library.
#
# Nicole (TN) 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 (TN) 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 (TN). If not, see <https://www.gnu.org/licenses/>.


"""Integration tests for end-to-end workflows."""

import math
import torch

from nicole import Direction, Tensor, contract, decomp, identity, isometry, U1Group, Z2Group, permute, conj
from nicole import Index, Sector
from .utils import assert_charge_neutral, assert_blocks_equal


# Construction → Contract → SVD → Reconstruct workflows

def test_workflow_construct_contract_svd_reconstruct():
    """Test complete workflow: construct → contract → svd → reconstruct."""
    group = U1Group()
    
    # Construct tensors
    idx_a = Index(Direction.OUT, group, sectors=(Sector(0, 3), Sector(1, 2), Sector(-1, 2)))
    idx_b = Index(Direction.IN, group, sectors=(Sector(0, 3), Sector(1, 2), Sector(-1, 2)))
    idx_c = Index(Direction.OUT, group, sectors=(Sector(0, 3), Sector(1, 2), Sector(-1, 2)))
    
    A = Tensor.random([idx_a, idx_b], seed=1, itags=["a", "b"])
    B = Tensor.random([idx_b.dual(), idx_c], seed=2, itags=["b", "c"])
    
    # Contract
    C = contract(A, B)
    assert_charge_neutral(C)
    
    # SVD using decomp
    U, S, Vh = decomp(C, axes=0, mode="SVD")
    assert_charge_neutral(U)
    assert_charge_neutral(S)
    assert_charge_neutral(Vh)
    
    # Reconstruct
    S_Vh = contract(S, Vh, axes=(1, 0))
    reconstructed = contract(U, S_Vh, axes=(1, 0))
    
    # Verify reconstruction
    diff_norm = (C - reconstructed).norm()
    rel_error = diff_norm / C.norm()
    assert rel_error < 1e-12


def test_workflow_mps_like_contraction():
    """Test MPS-like tensor network contraction."""
    group = U1Group()
    
    # Create MPS-like tensors
    phys = Index(Direction.OUT, group, sectors=(Sector(0, 2), Sector(1, 1)))
    bond1 = Index(Direction.OUT, group, sectors=(Sector(0, 3), Sector(1, 2)))
    bond2 = Index(Direction.OUT, group, sectors=(Sector(0, 3), Sector(1, 2)))
    
    # Three-site MPS
    M1 = Tensor.random([phys, bond1], seed=1, itags=["p1", "b1"])
    M2 = Tensor.random([bond1.dual(), phys, bond2], seed=2, itags=["b1", "p2", "b2"])
    M3 = Tensor.random([bond2.dual(), phys], seed=3, itags=["b2", "p3"])
    
    # Contract all
    M12 = contract(M1, M2)
    M123 = contract(M12, M3)
    
    # Should have only physical indices left
    assert len(M123.indices) == 3
    assert all("p" in tag for tag in M123.itags)
    assert_charge_neutral(M123)


def test_workflow_identity_insertion():
    """Test inserting and removing identities."""
    group = U1Group()
    idx = Index(Direction.OUT, group, sectors=(Sector(0, 2), Sector(1, 3)))
    
    # Create tensor
    T = Tensor.random([idx, idx.dual()], seed=1, itags=["a", "b"])
    original_norm = T.norm()
    
    # Insert identity
    ident = identity(idx, itags=("b", "c"))
    T_with_id = contract(T, ident)
    
    # Norm should be preserved
    assert math.isclose(T_with_id.norm(), original_norm)


def test_workflow_fusion_contraction():
    """Test fusion using isometry and contraction."""
    group = U1Group()
    
    # Create indices
    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)))
    
    # Create isometry - by default fused has opposite direction (IN)
    iso = isometry(idx_a, idx_b, itags=("a", "b", "ab"))
    
    # Check structure of isometry
    assert len(iso.indices) == 3
    assert_charge_neutral(iso)
    
    # Verify isometry was created successfully
    assert iso.norm() > 0


def test_workflow_svd_truncation_and_contraction():
    """Test SVD with truncation followed by contraction."""
    group = U1Group()
    
    # Create large tensor
    idx1 = Index(Direction.OUT, group, sectors=(Sector(0, 10),))
    idx2 = Index(Direction.IN, group, sectors=(Sector(0, 10),))
    
    T = Tensor.random([idx1, idx2], seed=1, itags=["a", "b"])
    
    # SVD using decomp (UR mode for efficiency)
    U, R = decomp(T, axes=0, mode="UR")
    
    # Reconstruct
    reconstructed = contract(U, R, axes=(1, 0))
    
    # Should approximately recover original
    rel_error = (T - reconstructed).norm() / T.norm()
    assert rel_error < 1e-12


# Multi-operation workflows

def test_workflow_permute_contract_permute():
    """Test permutation before and after contraction."""
    group = U1Group()
    
    idx_a = Index(Direction.OUT, group, sectors=(Sector(0, 2), Sector(1, 2)))
    idx_b = Index(Direction.OUT, group, sectors=(Sector(0, 2), Sector(-1, 2)))
    idx_c = Index(Direction.IN, group, sectors=(Sector(0, 2), Sector(1, 2)))
    
    A = Tensor.random([idx_a, idx_b], seed=1, itags=["a", "b"])
    B = Tensor.random([idx_b.dual(), idx_c], seed=2, itags=["b", "c"])
    
    # Permute A
    A_perm = permute(A, [1, 0])
    
    # Contract (should still work)
    C = contract(A_perm, B)
    
    # Permute result
    C_perm = permute(C, [1, 0])
    
    assert_charge_neutral(C_perm)


def test_workflow_arithmetic_operations_chain():
    """Test chain of arithmetic operations."""
    group = U1Group()
    idx = Index(Direction.OUT, group, sectors=(Sector(0, 2), Sector(1, 3)))
    
    A = Tensor.random([idx, idx.flip()], seed=1, itags=["a", "b"])
    B = Tensor.random([idx, idx.flip()], seed=2, itags=["a", "b"])
    C = Tensor.random([idx, idx.flip()], seed=3, itags=["a", "b"])
    
    # Complex expression: 2*A + 3*B - 0.5*C
    result = 2 * A + 3 * B - 0.5 * C
    
    # Verify
    for key in result.data:
        expected = 2 * A.data[key] + 3 * B.data[key] - 0.5 * C.data[key]
        assert torch.allclose(result.data[key], expected)


def test_workflow_conjugation_and_contraction():
    """Test conjugation combined with contraction."""
    group = U1Group()
    
    idx_a = Index(Direction.OUT, group, sectors=(Sector(0, 2), Sector(1, 2)))
    idx_b = Index(Direction.IN, group, sectors=(Sector(0, 2), Sector(1, 2)))
    
    T = Tensor.random([idx_a, idx_b], dtype=torch.complex128, seed=1, itags=["a", "b"])
    
    # Conjugate flips directions: OUT->IN, IN->OUT
    T_conj = conj(T)
    
    # T_conj now has (IN, OUT) directions
    # Create another tensor with matching structure for contraction
    T2 = Tensor.random([idx_a.flip(), idx_b.flip()], dtype=torch.complex128, seed=2, itags=["a", "b"])
    # T2 has (IN, OUT) directions, matching T_conj
    
    # For contraction to work, we need opposite directions and matching tags
    # T_conj: (IN, OUT) with tags ["a", "b"]
    # T2: (IN, OUT) with tags ["a", "b"]
    # These have same directions, so won't contract automatically
    
    # Instead, let's just verify conjugation works
    assert T_conj.dtype.is_complex
    assert T_conj.indices[0].direction == idx_a.direction.reverse()
    assert T_conj.indices[1].direction == idx_b.direction.reverse()


def test_workflow_retag_and_contract():
    """Test retagging before contraction."""
    group = U1Group()
    
    idx = Index(Direction.OUT, group, sectors=(Sector(0, 2),))
    
    A = Tensor.random([idx, idx.dual()], seed=1, itags=["old1", "old2"])
    B = Tensor.random([idx.dual(), idx], seed=2, itags=["x", "y"])
    
    # Retag A to match B
    A.retag(["x", "z"])
    
    # Now contract should work
    result = contract(A, B)
    
    assert_charge_neutral(result)


# Complex workflows

def test_workflow_tensor_network_contraction_order():
    """Test that contraction order doesn't affect result."""
    group = U1Group()
    
    idx1 = Index(Direction.OUT, group, sectors=(Sector(0, 2),))
    idx2 = Index(Direction.OUT, group, sectors=(Sector(0, 2),))
    idx3 = Index(Direction.OUT, group, sectors=(Sector(0, 2),))
    
    A = Tensor.random([idx1, idx2.dual()], seed=1, itags=["a", "b"])
    B = Tensor.random([idx2, idx3.dual()], seed=2, itags=["b", "c"])
    C = Tensor.random([idx3, idx1.dual()], seed=3, itags=["c", "a"])
    
    # Contract in different orders - result will be scalar
    result1 = contract(contract(A, B), C)
    result2 = contract(A, contract(B, C))
    
    # Should give same result (up to numerical precision)
    assert math.isclose(result1.norm(), result2.norm())
    assert result1.is_scalar()
    assert result2.is_scalar()


def test_workflow_build_mpo_and_apply():
    """Test building an MPO-like operator and applying it."""
    group = U1Group()
    
    # Create a simple MPO-MPS contraction workflow
    # MPS state has physical index as OUT direction
    phys_mps = Index(Direction.OUT, group, sectors=(Sector(0, 2), Sector(1, 1)))
    # MPO has physical_in as IN to match with MPS OUT direction
    phys_mpo = Index(Direction.IN, group, sectors=(Sector(0, 2), Sector(1, 1)))
    bond = Index(Direction.OUT, group, sectors=(Sector(0, 3),))
    
    # MPO: operator - phys_out OUT, phys_in IN, bond OUT
    W = Tensor.random([phys_mpo.flip(), phys_mpo, bond], seed=1, itags=["p_out", "p", "b"])
    
    # MPS: state vector - phys OUT, bond OUT
    psi = Tensor.random([phys_mps, bond], seed=2, itags=["p", "b"])
    
    # Apply W to psi (contract on physical index "p")
    # W has "p" with IN direction, psi has "p" with OUT direction
    result = contract(W, psi)
    
    # Result should have physical_out and bond indices (2 bonds actually)
    assert len(result.indices) == 3
    assert_charge_neutral(result)
    assert result.norm() > 0


def test_workflow_z2_tensors():
    """Test workflow with Z2 symmetry."""
    group = Z2Group()
    
    idx1 = Index(Direction.OUT, group, sectors=(Sector(0, 2), Sector(1, 1)))
    idx2 = Index(Direction.IN, group, sectors=(Sector(0, 2), Sector(1, 1)))
    
    A = Tensor.random([idx1, idx2], seed=1, itags=["a", "b"])
    B = Tensor.random([idx2.dual(), idx1], seed=2, itags=["b", "c"])
    
    # Contract
    C = contract(A, B)
    
    # SVD using decomp (UR mode for efficiency)
    U, R = decomp(C, axes=0, mode="UR")
    
    # Reconstruct
    reconstructed = contract(U, R, axes=(1, 0))
    
    # Verify
    rel_error = (C - reconstructed).norm() / C.norm()
    assert rel_error < 1e-12


def test_workflow_norm_preservation():
    """Test that norm is preserved through various operations."""
    group = U1Group()
    
    idx = Index(Direction.OUT, group, sectors=(Sector(0, 2), Sector(1, 3)))
    
    T = Tensor.random([idx, idx.dual()], seed=1, itags=["a", "b"])
    original_norm = T.norm()
    
    # Permute
    T_perm = permute(T, [1, 0])
    assert math.isclose(T_perm.norm(), original_norm)
    
    # Transpose back
    T_back = permute(T_perm, [1, 0])
    assert math.isclose(T_back.norm(), original_norm)
    
    # Conjugate (real tensor, norm should be same)
    T_conj = conj(T)
    assert math.isclose(T_conj.norm(), original_norm)


def test_workflow_large_tensor_operations():
    """Test operations on larger tensors."""
    group = U1Group()
    
    # Create larger indices
    charges = [(i, 2) for i in range(5)]
    idx1 = Index(Direction.OUT, group, sectors=tuple(Sector(c, d) for c, d in charges))
    idx2 = Index(Direction.IN, group, sectors=tuple(Sector(c, d) for c, d in charges))
    idx3 = Index(Direction.OUT, group, sectors=tuple(Sector(c, d) for c, d in charges))
    
    # Create tensors
    A = Tensor.random([idx1, idx2], seed=1, itags=["a", "b"])
    B = Tensor.random([idx2.dual(), idx3], seed=2, itags=["b", "c"])
    
    # Contract
    C = contract(A, B)
    
    # Should complete successfully
    assert C.norm() > 0
    assert_charge_neutral(C)


def test_workflow_mixed_dtypes():
    """Test workflow with dtype promotion."""
    group = U1Group()
    idx = Index(Direction.OUT, group, sectors=(Sector(0, 2),))
    
    # float32 tensor
    A = Tensor.random([idx, idx.flip()], dtype=torch.float32, seed=1, itags=["a", "b"])
    
    # float64 tensor
    B = Tensor.random([idx, idx.flip()], dtype=torch.float64, seed=2, itags=["a", "b"])
    
    # Add (should promote to float64)
    C = A + B
    assert C.dtype == torch.float64
    
    # Multiply by complex (should promote to complex)
    D = C * (1 + 1j)
    assert D.dtype.is_complex

