import numpy as np

np.set_printoptions(precision=5, suppress=True)

# ---------- Gell-Mann generators for d=3, normalized so Tr(sigma_i sigma_j) = 3*delta_ij ----------
i_ = 1j
lam = [None]*9
lam[1] = np.array([[0,1,0],[1,0,0],[0,0,0]], dtype=complex)
lam[2] = np.array([[0,-i_,0],[i_,0,0],[0,0,0]], dtype=complex)
lam[3] = np.array([[1,0,0],[0,-1,0],[0,0,0]], dtype=complex)
lam[4] = np.array([[0,0,1],[0,0,0],[1,0,0]], dtype=complex)
lam[5] = np.array([[0,0,-i_],[0,0,0],[i_,0,0]], dtype=complex)
lam[6] = np.array([[0,0,0],[0,0,1],[0,1,0]], dtype=complex)
lam[7] = np.array([[0,0,0],[0,0,-i_],[0,i_,0]], dtype=complex)
lam[8] = (1/np.sqrt(3))*np.array([[1,0,0],[0,1,0],[0,0,-2]], dtype=complex)

# check standard normalization Tr(lam_a lam_b) = 2 delta_ab
for a in range(1,9):
    for b in range(1,9):
        val = np.trace(lam[a]@lam[b])
        if a==b and not np.isclose(val,2):
            print("WARN std norm", a,b,val)
        if a!=b and not np.isclose(val,0):
            print("WARN std orth", a,b,val)

sigma = [None] + [np.sqrt(3/2)*lam[k] for k in range(1,9)]  # d=3 -> Tr(sigma_i sigma_j)=3 delta_ij

# sanity check
for a in range(1,9):
    for b in range(1,9):
        val = np.trace(sigma[a]@sigma[b]).real
        expected = 3.0 if a==b else 0.0
        assert abs(val-expected) < 1e-9, (a,b,val)
print("Generator normalization OK: Tr(sigma_i sigma_j) = 3 delta_ij")

# ---------- Tiles UPB (Bennett, DiVincenzo, Mor, Shor, Smolin, Terhal 1999) ----------
e0 = np.array([1,0,0], dtype=complex)
e1 = np.array([0,1,0], dtype=complex)
e2 = np.array([0,0,1], dtype=complex)

def nrm(v):
    return v/np.linalg.norm(v)

psi = []
psi.append(np.kron(e0, nrm(e0-e1)))
psi.append(np.kron(e2, nrm(e1-e2)))
psi.append(np.kron(nrm(e0-e1), e2))
psi.append(np.kron(nrm(e1-e2), e0))
psi.append(np.kron(nrm(e0+e1+e2), nrm(e0+e1+e2)))

# check orthonormality
G = np.array([[np.vdot(p,q) for q in psi] for p in psi])
print("\nGram matrix of the 5 UPB vectors (should be identity):")
print(np.round(G,6))

P = sum(np.outer(p, p.conj()) for p in psi)
I9 = np.eye(9, dtype=complex)
rho = (I9 - P)/4.0

print("\nTr(rho) =", np.trace(rho).real, " (should be 1)")
print("rho is Hermitian:", np.allclose(rho, rho.conj().T))
eigvals_rho = np.linalg.eigvalsh(rho)
print("eigenvalues of rho (should be >=0, rank 4 nonzero):", np.round(eigvals_rho,5))

# ---------- PPT check ----------
def partial_transpose_B(rho, dA=3, dB=3):
    r = rho.reshape(dA,dB,dA,dB)
    rpt = r.transpose(0,3,2,1)
    return rpt.reshape(dA*dB, dA*dB)

rho_pt = partial_transpose_B(rho)
eig_pt = np.linalg.eigvalsh(rho_pt)
print("\nEigenvalues of partial transpose (PPT check):")
print(np.round(eig_pt,6))
print("min eigenvalue of PT:", eig_pt.min(), " -> PPT" if eig_pt.min() > -1e-9 else " -> NPT (entangled via ordinary PPT already)")

# ---------- correlation tensor / shadow map ----------
T = np.zeros((8,8))
for a in range(1,9):
    for b in range(1,9):
        op = np.kron(sigma[a], sigma[b])
        T[a-1,b-1] = np.trace(rho @ op).real

s = np.linalg.svd(T, compute_uv=False)
nuclear_T = s.sum()
dA, dB = 3, 3
norm_const = np.sqrt((dA-1)*(dB-1))
M_norm = nuclear_T / norm_const

print("\nSingular values of correlation tensor T:", np.round(s,5))
print("Nuclear norm ||T||_* =", nuclear_T)
print("Normalization constant sqrt((dA-1)(dB-1)) =", norm_const)
print("Shadow-map value ||M_A(rho)||_* =", M_norm, " (separable bound: <= 1)")

# ---------- CCNR / realignment criterion for comparison ----------
def realign(rho, dA=3, dB=3):
    r = rho.reshape(dA,dB,dA,dB)
    # standard realignment: R_{(i mu),(j nu)} = rho_{ij,mu nu}
    R = r.transpose(0,2,1,3).reshape(dA*dA, dB*dB)
    return R

R = realign(rho)
s_R = np.linalg.svd(R, compute_uv=False)
ccnr = s_R.sum()
print("\nCCNR (realignment) trace norm:", ccnr, " (separable bound: <= 1)")