"""
02_werner_qutrit_symbolic.py

Same check as 01_werner_qubit_symbolic.py, generalized to d=3 (qutrits),
using the sqrt(3/2)-scaled Gell-Mann convention fixed in the paper's own
Tiles example (tr(sigma_i sigma_j) = d delta_ij = 3 delta_ij).

State family:  rho(p) = p * P_anti/dim(P_anti) + (1-p) * I/9
(the natural qutrit "Werner state" built from the antisymmetric subspace
of C^3 x C^3, dimension 3), invariant under U(x)U for all U in U(3).

Result (the interesting part): the order-1 shadow-map/correlation-matrix
criterion gives p_c = 1/2, but the TRUE separability threshold (Werner
1989, p_sep = 1/(d+1)) is p_c = 1/4. Unlike the qubit case, the criterion
is here only a valid but NOT tight sufficient condition -- symmetry forces
"concentration" of the signal (single isotype => correlation matrix
proportional to identity) but not "sharpening" of the threshold itself.

Requires: sympy.  Runtime: under a minute.
"""
import sympy as sp
from sympy import sqrt, I, simplify, Matrix, eye, zeros, re, symbols, Rational

d = 3

lam = [None] * 8
lam[0] = Matrix([[0, 1, 0], [1, 0, 0], [0, 0, 0]])
lam[1] = Matrix([[0, -I, 0], [I, 0, 0], [0, 0, 0]])
lam[2] = Matrix([[1, 0, 0], [0, -1, 0], [0, 0, 0]])
lam[3] = Matrix([[0, 0, 1], [0, 0, 0], [1, 0, 0]])
lam[4] = Matrix([[0, 0, -I], [0, 0, 0], [I, 0, 0]])
lam[5] = Matrix([[0, 0, 0], [0, 0, 1], [0, 1, 0]])
lam[6] = Matrix([[0, 0, 0], [0, 0, -I], [0, I, 0]])
lam[7] = (1 / sqrt(3)) * Matrix([[1, 0, 0], [0, 1, 0], [0, 0, -2]])

c = sqrt(Rational(3, 2))
sigma = [simplify(c * L) for L in lam]
for i in range(8):
    for j in range(8):
        val = simplify((sigma[i] * sigma[j]).trace())
        assert val == (3 if i == j else 0), (i, j, val)

I3 = eye(3)


def op_A(P):
    return sp.Matrix(sp.kronecker_product(P, I3))


def op_B(P):
    return sp.Matrix(sp.kronecker_product(I3, P))


V = zeros(9, 9)
for a in range(3):
    for b in range(3):
        V[b * 3 + a, a * 3 + b] = 1

I9 = eye(9)
P_anti = simplify((I9 - V) / 2)
dim_anti = simplify(P_anti.trace())  # = 3

p = symbols('p', real=True)
rho_p = simplify(p * P_anti / dim_anti + (1 - p) * I9 / 9)


def entry(rho, ops):
    M = None
    for op in ops:
        M = op if M is None else M * op
    return simplify(re(simplify((rho * M).trace())))


T = Matrix(8, 8, lambda i, j: entry(rho_p, [op_A(sigma[i]), op_B(sigma[j])]))
print("Correlation matrix T(p) (should be proportional to I_8):")
sp.pprint(T)

norm_const = 1 / sqrt(4)  # (d_a-1)(d_bar_a-1) = 2*2 = 4
Mn = simplify(norm_const * T)
nuclear_norm = simplify(8 * sp.Abs(Mn[0, 0]))
print("\nShadow-map nuclear norm:", nuclear_norm)
print("Shadow-map threshold:", sp.solve(sp.Eq(nuclear_norm, 1), p))


def partial_transpose_B_d(M, dim):
    Mpt = zeros(dim * dim, dim * dim)
    for a in range(dim):
        for b in range(dim):
            for cc in range(dim):
                for dd in range(dim):
                    i, j = a * dim + b, cc * dim + dd
                    i2, j2 = a * dim + dd, cc * dim + b
                    Mpt[i2, j2] = M[i, j]
    return Mpt


rho_pt = partial_transpose_B_d(rho_p, 3)
print("\nPartial-transpose eigenvalues (PPT / true-separability threshold):")
for ev in rho_pt.eigenvals().keys():
    print("  ", simplify(ev), " = 0 at p =", sp.solve(sp.Eq(ev, 0), p))

print("\nExpected: shadow-map threshold p=1/2 (NOT tight);"
      " true threshold (Werner 1989, p_sep=1/(d+1)) p=1/4.")