"""
01_werner_qubit_symbolic.py

Exact symbolic (sympy) check: the two-qubit Werner state

    rho(p) = p |Psi-><Psi-| + (1-p) I/4,   |Psi-> = (|01>-|10>)/sqrt(2)

is invariant under U (x) U for every U in SU(2). Since the adjoint
representation of SU(2) on the traceless qubit Bloch space R^3 is
irreducible (single isotype), the correlation matrix is forced to be
proportional to the identity. We check this exactly and compare the
resulting nuclear-norm threshold to the exact PPT/separability threshold.

Result: BOTH give exactly p = 1/3 -- the order-1 correlation-matrix
criterion is exactly tight here (a low-dimensional special case, since
PPT=separable for 2x2 systems by the Horodecki theorem).

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

X = Matrix([[0, 1], [1, 0]])
Y = Matrix([[0, -I], [I, 0]])
Z = Matrix([[1, 0], [0, -1]])
I2 = eye(2)


def kron(A, B):
    mA, nA = A.shape
    mB, nB = B.shape
    out = zeros(mA * mB, nA * nB)
    for i in range(mA):
        for j in range(nA):
            out[i * mB:(i + 1) * mB, j * nB:(j + 1) * nB] = A[i, j] * B
    return out


def op_A(P):
    return kron(P, I2)


def op_B(P):
    return kron(I2, P)


p = symbols('p', real=True)

psi = zeros(4, 1)
psi[1, 0] = 1 / sqrt(2)
psi[2, 0] = -1 / sqrt(2)
rho_singlet = simplify(psi * psi.H)

rho_p = simplify(p * rho_singlet + (1 - p) * eye(4) / 4)
print("rho(p) =")
sp.pprint(rho_p)


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())))


plist = [X, Y, Z]
T = Matrix(3, 3, lambda i, j: entry(rho_p, [op_A(plist[i]), op_B(plist[j])]))
print("\nCorrelation matrix T(p) =")
sp.pprint(T)

G = simplify(T.T * T)
eigs = G.eigenvals()
singular_values = []
for ev, mult in eigs.items():
    singular_values += [simplify(sqrt(ev))] * mult
nuclear_norm = simplify(sum(singular_values))
print("\nNuclear norm ||M_A(rho(p))||_* =", nuclear_norm)
print("Shadow-map threshold (||.||_* = 1):", sp.solve(sp.Eq(nuclear_norm, 1), p))


def partial_transpose_B(M):
    Mpt = zeros(4, 4)
    for a in range(2):
        for b in range(2):
            for c in range(2):
                for dd in range(2):
                    i, j = a * 2 + b, c * 2 + dd
                    i2, j2 = a * 2 + dd, c * 2 + b
                    Mpt[i2, j2] = M[i, j]
    return Mpt


rho_pt = partial_transpose_B(rho_p)
print("\nEigenvalues of the partial transpose rho(p)^{T_B}:")
for e_ in rho_pt.eigenvals().keys():
    print("  ", simplify(e_), " = 0 at p =", sp.solve(sp.Eq(e_, 0), p))