"""
General-r check of Proposition coherence-templates: for r G-fixed states
phi_1,...,phi_r and ANY density matrix rho = sum_{a,b} c_{ab} |phi_a><phi_b|
(c a Hermitian PSD matrix, not necessarily rank-1/pure), the full
correlation tensor is T(rho) = sum_{a,b} c_{ab} T_{ab}, with T_{ab} fixed
(independent of c), REGARDLESS of which cut is subsequently taken.

Real-tensor count: T_{ba} = conj(T_{ab}) (since sigma is Hermitian), so the
independent REAL data is {T_{aa}}_{a=1}^r (each already real) together with
{Re(T_{ab}), Im(T_{ab})}_{a<b} -- total r + 2*binom(r,2) = r^2 real tensors,
matching the real dimension of the space of r x r Hermitian matrices.
(Corrects an earlier mis-stated count of r(r+1)/2 in the TODO comment.)

Tested here for r=3, using three different perfect matchings of six qubits
into singlets as the three G-fixed basis states, and a genuinely MIXED
(not pure/rank-1) random density matrix c -- a strictly more general test
than the r=2 pure-superposition case checked earlier.
"""
import numpy as np

def s(a,b):
    if (a,b)==(0,1): return 1/np.sqrt(2)
    if (a,b)==(1,0): return -1/np.sqrt(2)
    return 0.0
def build_pairing(pairs):
    psi = np.zeros((2,)*6, dtype=complex)
    for idx in np.ndindex(2,2,2,2,2,2):
        val=1.0
        for (p_,q_) in pairs:
            val *= s(idx[p_], idx[q_])
            if val==0: break
        psi[idx]=val
    return psi.reshape(64)

# three different perfect matchings of {A,B,C,D,E,F} = {0,1,2,3,4,5}
phi1 = build_pairing([(0,3),(1,4),(2,5)])   # (A,D)(B,E)(C,F)
phi2 = build_pairing([(0,4),(1,5),(2,3)])   # (A,E)(B,F)(C,D)
phi3 = build_pairing([(0,5),(1,3),(2,4)])   # (A,F)(B,D)(C,E)
Phi = np.stack([phi1,phi2,phi3], axis=1)    # 64x3

G = Phi.conj().T @ Phi   # Gram matrix
print("Gram matrix (should be Hermitian, diag=1, off-diag |.|<1):")
print(np.round(G,4))
print("condition number:", np.linalg.cond(G))

# --- Pauli machinery ---
X=np.array([[0,1],[1,0]],dtype=complex); Y=np.array([[0,-1j],[1j,0]],dtype=complex); Z=np.array([[1,0],[0,-1]],dtype=complex)
paulis=[X,Y,Z]
def apply_leg(psi,axis,P):
    p2=np.moveaxis(psi,axis,0); out=np.tensordot(P,p2,axes=([1],[0])); return np.moveaxis(out,0,axis)
def corr_tensor(bra,ket):
    bra6, ket6 = bra.reshape((2,)*6), ket.reshape((2,)*6)
    c=np.zeros((3,3,3,3,3,3),dtype=complex)
    for iA in range(3):
     for iB in range(3):
      for iC in range(3):
       for iD in range(3):
        for iE in range(3):
         for iF in range(3):
            k=ket6
            for ax,ii in zip(range(6),(iA,iB,iC,iD,iE,iF)):
                k=apply_leg(k,ax,paulis[ii])
            c[iA,iB,iC,iD,iE,iF]=np.vdot(bra6,k)
    return c

print("\ncomputing all T_ab (a,b=1,2,3), 9 tensors total...")
phis = [phi1,phi2,phi3]
T = {}
for a in range(3):
    for b in range(3):
        T[(a,b)] = corr_tensor(phis[a],phis[b])
        print(f"  T[{a+1},{b+1}] done, max imag part={np.abs(T[(a,b)].imag).max():.2e}" if a==b else
              f"  T[{a+1},{b+1}] done")

# check T_ba = conj(T_ab)
for a in range(3):
    for b in range(3):
        err = np.abs(T[(a,b)] - np.conj(T[(b,a)])).max()
        assert err < 1e-10, (a,b,err)
print("T_ba = conj(T_ab) verified for all pairs.")

# --- random genuinely MIXED c (PSD, not rank 1) ---
rng = np.random.default_rng(42)
W = rng.normal(size=(3,3)) + 1j*rng.normal(size=(3,3))
c_raw = W @ W.conj().T     # Hermitian PSD, generically full rank
print("\nrandom c_raw (Hermitian PSD, rank =", np.linalg.matrix_rank(c_raw), "):")
print(np.round(c_raw,3))

rho_raw = Phi @ c_raw @ Phi.conj().T     # 64x64
tr = np.trace(rho_raw).real
rho = rho_raw/tr
c = c_raw/tr
print(f"\ntrace(rho_raw)={tr:.6f}; normalized rho has trace {np.trace(rho).real:.10f}")
evals_rho = np.linalg.eigvalsh(rho)
print("eigenvalues of rho (should be >=0, sum=1):", np.round(evals_rho[np.abs(evals_rho)>1e-9],6))

# --- brute-force TRUE correlation tensor of rho ---
def corr_tensor_rho(rho):
    rho6 = rho.reshape((2,)*12)  # not directly useful; do it via trace instead
    c=np.zeros((3,3,3,3,3,3),dtype=complex)
    for iA in range(3):
     for iB in range(3):
      for iC in range(3):
       for iD in range(3):
        for iE in range(3):
         for iF in range(3):
            O = paulis[iA]
            for ii in (iB,iC,iD,iE,iF):
                O = np.kron(O, paulis[ii])
            c[iA,iB,iC,iD,iE,iF] = np.trace(rho @ O)
    return c
T_true = corr_tensor_rho(rho).real

# --- predicted via T(rho) = sum_ab c_ab T_ab ---
T_pred = np.zeros((3,3,3,3,3,3), dtype=complex)
for a in range(3):
    for b in range(3):
        T_pred += c[a,b] * T[(a,b)]
T_pred = T_pred.real

err = np.abs(T_true - T_pred).max()
print(f"\nmax|T_true - T_pred| (full 6-index tensor, r=3, genuinely mixed rho): {err:.2e}")

# --- verify at BOTH cuts via simple reshape, no new contraction ---
M1_true, M1_pred = T_true.reshape(27,27), T_pred.reshape(27,27)
M2_true, M2_pred = T_true.reshape(9,81), T_pred.reshape(9,81)
print(f"cut ABC|DEF: max matrix error = {np.abs(M1_true-M1_pred).max():.2e}, "
      f"||M||_* true={np.linalg.svd(M1_true,compute_uv=False).sum():.4f} "
      f"pred={np.linalg.svd(M1_pred,compute_uv=False).sum():.4f}")
print(f"cut AB|CDEF: max matrix error = {np.abs(M2_true-M2_pred).max():.2e}, "
      f"||M||_* true={np.linalg.svd(M2_true,compute_uv=False).sum():.4f} "
      f"pred={np.linalg.svd(M2_pred,compute_uv=False).sum():.4f}")