"""core2.py -- fast, vectorized computation of the 4-qubit Pauli/Bloch correlation
tensor and the single-party combined shadow map M_a, plus the source-aggregated
functional Phi_sym.

C[i0,i1,i2,i3] = tr(rho * sigma_i0 x sigma_i1 x sigma_i2 x sigma_i3), i_k in {0,1,2,3}
(0 = identity, 1,2,3 = X,Y,Z), for a PURE state psi (rho = |psi><psi|).
"""
import numpy as np

I2 = np.eye(2, dtype=complex)
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)
PAULI = np.stack([I2, X, Y, Z])  # shape (4,2,2), index (pauli_label, row, col)


def full_tensor(psi):
    """psi: length-16 complex state vector, qubit order (A,B,C,D).
    Returns the (4,4,4,4) real Pauli correlation tensor."""
    t = psi.reshape(2, 2, 2, 2)
    tc = np.conj(t)
    C = np.einsum('abcd,iae,jbf,kcg,ldh,efgh->ijkl',
                   tc, PAULI, PAULI, PAULI, PAULI, t, optimize=True)
    return C.real


# all (j0,j1,j2) index combinations excluding the all-identity (0,0,0) target sector
_COLS = np.array([(j0, j1, j2) for j0 in range(4) for j1 in range(4) for j2 in range(4)
                   if (j0, j1, j2) != (0, 0, 0)])


def shadow_map_singleparty(C, a):
    """Normalized combined single-party shadow map M_a(rho): 3 x 63 real matrix,
    normalization 1/sqrt(2^(n-1)-1) = 1/sqrt(7) for n=4 qubits."""
    others = [k for k in range(4) if k != a]
    rows = np.array([1, 2, 3])
    n_cols = len(_COLS)
    Ridx = np.repeat(rows, n_cols)
    J0 = np.tile(_COLS[:, 0], 3)
    J1 = np.tile(_COLS[:, 1], 3)
    J2 = np.tile(_COLS[:, 2], 3)
    idx_full = [None] * 4
    idx_full[a] = Ridx
    idx_full[others[0]] = J0
    idx_full[others[1]] = J1
    idx_full[others[2]] = J2
    vals = C[idx_full[0], idx_full[1], idx_full[2], idx_full[3]]
    M = vals.reshape(3, n_cols)
    return M / np.sqrt(7.0)


def nuclear_norm(M):
    return np.linalg.svd(M, compute_uv=False).sum()


def phi_sym(psi):
    """Phi_sym(rho) = (1/4) sum_a ||M_a(rho)||_*, for a pure 4-qubit state psi.
    Returns (mean_value, [values per party A,B,C,D])."""
    C = full_tensor(psi)
    vals = [nuclear_norm(shadow_map_singleparty(C, a)) for a in range(4)]
    return float(np.mean(vals)), vals