"""
Exact SO(3) isotypic projectors on (R^3)^{\otimes k} via the Casimir
operator J^2, instead of Monte-Carlo character averaging.

The spin-1 (vector) generators in the real Cartesian basis are
  (J_a)_{bc} = -i * epsilon_{abc}       (standard so(3) generators)
Built exactly with sympy, then verified to satisfy [J_a,J_b] = i eps_abc J_c
and J^2 = J_x^2+J_y^2+J_z^2 = 2*I_3 (i.e. j=1, j(j+1)=2) -- symbolically exact.

For k copies, total J_a = sum_{l=1}^k  I x ... x J_a^{(l)} x ... x I,
J^2_total is Hermitian on (C^3)^{\otimes k}; its eigenspaces are EXACTLY
the isotypic components (eigenvalue j(j+1)).  No integration needed.
"""
import numpy as np
import sympy as sp

i = sp.I
eps = lambda a,b,c: sp.LeviCivita(a,b,c)

def J_component(a):
    # a in {0,1,2} = x,y,z ; (J_a)_{bc} = -i * eps(a,b,c)
    M = sp.zeros(3,3)
    for b in range(3):
        for c in range(3):
            M[b,c] = -i*eps(a,b,c)
    return M

Jx, Jy, Jz = J_component(0), J_component(1), J_component(2)

# --- symbolic sanity checks ---
comm = Jx*Jy - Jy*Jx
print("[Jx,Jy] - i*Jz  == 0 ?", sp.simplify(comm - i*Jz) == sp.zeros(3,3))
J2_single = sp.simplify(Jx*Jx + Jy*Jy + Jz*Jz)
print("J^2 (single spin-1 leg), should be 2*I_3:")
sp.pprint(J2_single)

# convert to numpy (complex) for fast Kronecker-sum construction at larger k
Jx_np = np.array(Jx.tolist(), dtype=complex)
Jy_np = np.array(Jy.tolist(), dtype=complex)
Jz_np = np.array(Jz.tolist(), dtype=complex)

def total_J2(k):
    dim = 3**k
    Jtot = {a: np.zeros((dim,dim), dtype=complex) for a in range(3)}
    comps = [Jx_np, Jy_np, Jz_np]
    for leg in range(k):
        for a in range(3):
            mats = [np.eye(3, dtype=complex)]*k
            mats[leg] = comps[a]
            M = mats[0]
            for m in mats[1:]:
                M = np.kron(M, m)
            Jtot[a] += M
    return Jtot[0]@Jtot[0] + Jtot[1]@Jtot[1] + Jtot[2]@Jtot[2]

def exact_projectors(k, jmax):
    J2 = total_J2(k)
    assert np.abs(J2 - J2.conj().T).max() < 1e-10, "J^2 not Hermitian!"
    evals, evecs = np.linalg.eigh(J2)
    Ps = {}
    for j in range(jmax+1):
        target = j*(j+1)
        mask = np.abs(evals - target) < 1e-6
        if not np.any(mask):
            Ps[j] = np.zeros((3**k,3**k))
            continue
        V = evecs[:, mask]
        P = (V @ V.conj().T).real
        Ps[j] = P
    # sanity: eigenvalues actually cluster near integers j(j+1)
    return Ps, evals

print("\nBuilding exact projectors for k=2,3,4 via Casimir diagonalization...")
Ps2, ev2 = exact_projectors(2, 2)
Ps3, ev3 = exact_projectors(3, 3)
Ps4, ev4 = exact_projectors(4, 4)

for k,Ps,jmax in [(2,Ps2,2),(3,Ps3,3),(4,Ps4,4)]:
    print(f"\nk={k}:")
    for j in range(jmax+1):
        tr = np.trace(Ps[j]).real
        print(f"  j={j}: trace(P_j) = {tr:.10f}  (expect (2j+1)*m_j)")

np.savez("projectors_exact.npz",
         P2_0=Ps2[0],P2_1=Ps2[1],P2_2=Ps2[2],
         P3_0=Ps3[0],P3_1=Ps3[1],P3_2=Ps3[2],P3_3=Ps3[3],
         P4_0=Ps4[0],P4_1=Ps4[1],P4_2=Ps4[2],P4_3=Ps4[3],P4_4=Ps4[4])
print("\nsaved projectors_exact.npz")