import numpy as np
from itertools import permutations, product
from math import comb

def dicke3(k):
    psi = np.zeros(8, dtype=complex); n=0
    for bits in range(8):
        if bin(bits).count("1")==k: psi[bits]=1.0; n+=1
    return psi/np.sqrt(n)
D3 = {k: dicke3(k) for k in range(4)}
Q = np.zeros(64, dtype=complex)
for k in range(4):
    Q += ((-1)**k) * np.kron(D3[k], D3[3-k])
Q /= np.linalg.norm(Q)

# --- correlation tensor (Cartesian, all six legs, A,B,C,D,E,F) ---
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]
Qt = Q.reshape((2,)*6)
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):
    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=ket
            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(bra,k)
    return c
T = corr_tensor(Qt,Qt).real
M = T.reshape(27,27)   # cut ABC|DEF
print("nuclear norm of raw M (cut ABC|DEF):", np.linalg.svd(M,compute_uv=False).sum())

# --- exact Casimir (from before) for k=3 legs ---
def J_component(a):
    eps = np.zeros((3,3,3))
    eps[0,1,2]=eps[1,2,0]=eps[2,0,1]=1
    eps[0,2,1]=eps[2,1,0]=eps[1,0,2]=-1
    Mm = np.zeros((3,3), dtype=complex)
    for b in range(3):
        for c in range(3):
            Mm[b,c] = -1j*eps[a,b,c]
    return Mm
Jx,Jy,Jz = J_component(0),J_component(1),J_component(2)
def total_J2(k):
    dim=3**k
    comps=[Jx,Jy,Jz]
    Jtot=[np.zeros((dim,dim),dtype=complex) for _ in range(3)]
    for leg in range(k):
        for a in range(3):
            mats=[np.eye(3,dtype=complex)]*k
            mats[leg]=comps[a]
            Mm=mats[0]
            for mm in mats[1:]: Mm=np.kron(Mm,mm)
            Jtot[a]+=Mm
    return Jtot[0]@Jtot[0]+Jtot[1]@Jtot[1]+Jtot[2]@Jtot[2]
J2_3 = total_J2(3)

# --- multinomial "type" basis u_alpha for Sym^3(C^3): dimension binom(3+2,2)=10 ---
types = [(a,b,c) for a in range(4) for b in range(4) for c in range(4) if a+b+c==3]
print("\ntypes (a_x,a_y,a_z):", types, " count:", len(types))

def type_vector(alpha):
    ax,ay,az = alpha
    letters = ['x']*ax+['y']*ay+['z']*az   # length 3
    idxmap = {'x':0,'y':1,'z':2}
    seen = set()
    vec = np.zeros(27, dtype=complex)
    count = 0
    for perm in set(permutations(letters)):
        idx = tuple(idxmap[l] for l in perm)
        flat = idx[0]*9+idx[1]*3+idx[2]
        vec[flat] = 1.0
        count += 1
    vec /= np.linalg.norm(vec)
    return vec

U = np.zeros((27,10), dtype=complex)
for i,alpha in enumerate(types):
    U[:,i] = type_vector(alpha)

print("orthonormality check (U^T U should be I_10), max dev:", np.abs(U.conj().T@U - np.eye(10)).max())

# --- restrict Casimir to the 10-dim symmetric subspace ---
J2_sym = U.conj().T @ J2_3 @ U
evals_sym = np.linalg.eigvalsh(J2_sym)
print("\nEigenvalues of J^2 restricted to Sym^3(C^3) (10-dim):")
print(np.round(np.sort(evals_sym),6))
print("expected: j=1 (val=2, x3) and j=3 (val=12, x7) -- j=0,2 should be ABSENT")

# --- restrict shadow-map block M to the symmetric subspace on BOTH sides ---
M_sym = U.conj().T @ M @ U   # 10x10 (reduced, S_3-symmetric on both ABC and DEF)
print("\nnuclear norm of M restricted to Sym^3 x Sym^3 (10x10):", np.linalg.svd(M_sym,compute_uv=False).sum())

# isotypic projectors within the 10-dim space, from J2_sym eigenvectors
evals, evecs = np.linalg.eigh(J2_sym)
for jtarget, label in [(1,'j=1 (val=2)'), (3,'j=3 (val=12)')]:
    target = jtarget*(jtarget+1)
    mask = np.abs(evals-target)<1e-6
    print(f"{label}: multiplicity found = {mask.sum()} (expect {2*jtarget+1})")

for j0 in (0,2):
    target=j0*(j0+1)
    mask=np.abs(evals-target)<1e-6
    print(f"j={j0}: multiplicity found = {mask.sum()} (expect 0)")

# --- verify ALL signal lives in the symmetric x symmetric block ---
Proj_sym_27 = U @ U.conj().T   # 27x27 projector onto Sym^3 within full space
M_outside = M - Proj_sym_27 @ M @ Proj_sym_27
print("\nnuclear norm of M OUTSIDE the Sym^3 x Sym^3 block:",
      np.linalg.svd(M_outside, compute_uv=False).sum(), " (should be ~0)")

# --- extract the actual scalar A_1, A_3 values within the multiplicity-free channels ---
for jtarget in (1,3):
    target = jtarget*(jtarget+1)
    mask = np.abs(evals-target)<1e-6
    P = evecs[:,mask] @ evecs[:,mask].conj().T   # 10x10 projector
    block = P @ M_sym @ P
    nn = np.linalg.svd(block, compute_uv=False).sum()
    A_j = nn/(2*jtarget+1)
    print(f"j={jtarget}: ||A_j||_* (now a genuine SCALAR, multiplicity 1) = {A_j:.6f}")

print(f"\ncheck: 3*A_1 + 7*A_3 = {3*2.5+7*(15-3*2.5)/7 if False else ''}")