"""
Numerical check of the 'one invariant tensor, many cuts' claim
for the 6-qubit singlet-network example states.

Qubits ordered A,B,C,D,E,F -> tensor axes 0..5.
S1 = ABC | DEF   (3|3 cut)
S2 = AB  | CDEF  (2|4 cut)

We build the two 'basis' states
  |psi1> = singlets (A,D)(B,E)(C,F)
  |psi2> = singlets (A,E)(B,F)(C,D)
and the family
  |Xi(alpha)> = (cos(a) psi1 + sin(a) psi2) / norm

For rho_Xi = |Xi><Xi|, linearity in rho gives EXACTLY
  T(Xi) = [ cos^2(a) T1 + sin^2(a) T2 + cos(a)sin(a) C12 ] / N2
where T1 = <psi1|O|psi1>, T2 = <psi2|O|psi2>, C12 = <psi1|O|psi2> + <psi2|O|psi1>,
N2 = <Xi_raw|Xi_raw>.

Key point: T1, T2, C12 do NOT depend on alpha or on the cut.
Once computed ONCE (full 6-leg tensors), every cut's shadow-map block
for every alpha is obtained by (i) taking this fixed linear combination
of THREE fixed numbers times three fixed tensors, and (ii) a plain
numpy .reshape() -- no further contraction over the 64-dim Hilbert space.

This script verifies that against two independent brute-force
quantum simulations (cut S1 and cut S2, both done from scratch).
"""

import numpy as np

# ---------- Pauli matrices ----------
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]

# ---------- build the two basis states (6-qubit amplitude tensors) ----------
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

psi1 = build_pairing([(0,3),(1,4),(2,5)])   # (A,D)(B,E)(C,F)
psi2 = build_pairing([(0,4),(1,5),(2,3)])   # (A,E)(B,F)(C,D)
overlap = np.vdot(psi1, psi2)
print("overlap <psi1|psi2> =", overlap)

def apply_pauli_leg(psi, axis, P):
    psi2 = np.moveaxis(psi, axis, 0)
    out = np.tensordot(P, psi2, axes=([1],[0]))
    return np.moveaxis(out, 0, axis)

def corr_tensor(bra, ket):
    """<bra| sigma_i1 x ... x sigma_i6 |ket>, all six legs active (i in {0,1,2}=x,y,z)."""
    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):
            ket_ = ket
            for axis,ii in zip(range(6),(iA,iB,iC,iD,iE,iF)):
                ket_ = apply_pauli_leg(ket_, axis, paulis[ii])
            c[iA,iB,iC,iD,iE,iF] = np.vdot(bra, ket_)
    return c

print("computing T1 = <psi1|O|psi1> ...")
T1 = corr_tensor(psi1, psi1)
print("computing T2 = <psi2|O|psi2> ...")
T2 = corr_tensor(psi2, psi2)
print("computing cross term <psi1|O|psi2> ...")
X12 = corr_tensor(psi1, psi2)
X21 = corr_tensor(psi2, psi1)
C12 = X12 + X21

print("max imag part T1,T2,C12:",
      np.abs(T1.imag).max(), np.abs(T2.imag).max(), np.abs(C12.imag).max())
T1, T2, C12 = T1.real, T2.real, C12.real

np.save("T1.npy", T1); np.save("T2.npy", T2); np.save("C12.npy", C12)

# ---------- ground truth: brute-force Xi(alpha) for a couple of alphas, both cuts ----------
alpha_list = [np.pi/5, 0.9, -0.3]

def build_Xi(alpha):
    raw = np.cos(alpha)*psi1 + np.sin(alpha)*psi2
    n = np.linalg.norm(raw)
    return raw/n, n**2

def predict_from_basis(alpha):
    N2 = 1 + np.sin(2*alpha)*overlap.real
    return (np.cos(alpha)**2*T1 + np.sin(alpha)**2*T2 + np.cos(alpha)*np.sin(alpha)*C12) / N2

max_err_cut1 = 0.0
max_err_cut2 = 0.0
for a in alpha_list:
    Xi, N2_check = build_Xi(a)
    Tgt = corr_tensor(Xi, Xi).real          # brute-force ground truth, full simulation
    Tpred = predict_from_basis(a)           # from the 3 fixed tensors, no new simulation

    err = np.abs(Tgt - Tpred).max()
    print(f"alpha={a:+.4f}: max|T_bruteforce - T_predicted| = {err:.3e}   (N2 check: {N2_check:.6f})")

    # cut 1: ABC|DEF  (27x27)
    M1_true = Tgt.reshape(27,27)
    M1_pred = Tpred.reshape(27,27)
    e1 = np.abs(M1_true - M1_pred).max()
    max_err_cut1 = max(max_err_cut1, e1)

    # cut 2: AB|CDEF  (9x81)
    M2_true = Tgt.reshape(9,81)
    M2_pred = Tpred.reshape(9,81)
    e2 = np.abs(M2_true - M2_pred).max()
    max_err_cut2 = max(max_err_cut2, e2)

    print(f"   cut ABC|DEF : max matrix error = {e1:.3e},  ||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 = {e2:.3e},  ||M||_* true={np.linalg.svd(M2_true,compute_uv=False).sum():.4f} "
          f"pred={np.linalg.svd(M2_pred,compute_uv=False).sum():.4f}")

print()
print(f"WORST CASE over all tested alpha, both cuts: {max(max_err_cut1, max_err_cut2):.3e}")