"""cluster.py -- the full 15x15 bigraduated cluster shadow map M_S for a 2-qubit source
cluster S (within 4 qubits), and helpers for building the Pauli tensor of a general
(possibly mixed) density matrix."""
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 = [I2, X, Y, Z]


def cluster_map(C, s0, s1):
    """C: (4,4,4,4) Pauli correlation tensor. S = {s0,s1} (source cluster), complement =
    the other two qubits. Returns the normalized 15x15 matrix M_S(rho), normalization
    1/sqrt((d_S-1)(d_Sc-1)) = 1/sqrt(3*3) = 1/3 for two-qubit clusters in 4 qubits."""
    others = [k for k in range(4) if k not in (s0, s1)]
    c0, c1 = others
    rows = [(i, j) for i in range(4) for j in range(4) if (i, j) != (0, 0)]
    cols = [(i, j) for i in range(4) for j in range(4) if (i, j) != (0, 0)]
    M = np.zeros((15, 15))
    for ri, (ia, ib) in enumerate(rows):
        for ci, (ic, idd) in enumerate(cols):
            idx = [0, 0, 0, 0]
            idx[s0] = ia
            idx[s1] = ib
            idx[c0] = ic
            idx[c1] = idd
            M[ri, ci] = C[tuple(idx)]
    return M / 3.0


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


def full_tensor_mixed(rho):
    """Pauli-tensor extraction for a general (possibly mixed) 16x16 density matrix rho,
    via direct trace. Slower than core2.full_tensor (which needs a pure-state vector)
    but works for explicit mixtures (e.g. the Smolin state)."""
    def kron4(a, b, c, d):
        return np.kron(np.kron(a, b), np.kron(c, d))
    C = np.zeros((4, 4, 4, 4))
    for i0 in range(4):
        for i1 in range(4):
            for i2 in range(4):
                for i3 in range(4):
                    op = kron4(_PAULI[i0], _PAULI[i1], _PAULI[i2], _PAULI[i3])
                    C[i0, i1, i2, i3] = np.real(np.trace(rho @ op))
    return C