"""
common.py

Shared numpy utilities for the Tiles-state / DPS-hierarchy scripts (03-08).
Convention: qutrits (d=3), Gell-Mann generators scaled so that
tr(sigma_i sigma_j) = d * delta_ij = 3 * delta_ij, matching the paper's own
convention (see the Tiles benchmark in symmetric_shadow_maps_formal.tex).
"""
import numpy as np

d = 3

# --- Gell-Mann matrices, paper convention ---
_lam = [None] * 8
_lam[0] = np.array([[0, 1, 0], [1, 0, 0], [0, 0, 0]], dtype=complex)
_lam[1] = np.array([[0, -1j, 0], [1j, 0, 0], [0, 0, 0]], dtype=complex)
_lam[2] = np.array([[1, 0, 0], [0, -1, 0], [0, 0, 0]], dtype=complex)
_lam[3] = np.array([[0, 0, 1], [0, 0, 0], [1, 0, 0]], dtype=complex)
_lam[4] = np.array([[0, 0, -1j], [0, 0, 0], [1j, 0, 0]], dtype=complex)
_lam[5] = np.array([[0, 0, 0], [0, 0, 1], [0, 1, 0]], dtype=complex)
_lam[6] = np.array([[0, 0, 0], [0, 0, -1j], [0, 1j, 0]], dtype=complex)
_lam[7] = (1 / np.sqrt(3)) * np.array([[1, 0, 0], [0, 1, 0], [0, 0, -2]], dtype=complex)
GELLMANN = [np.sqrt(3 / 2) * L for L in _lam]

I3 = np.eye(3, dtype=complex)
I9 = np.eye(9, dtype=complex)


def opA(P):
    return np.kron(P, I3)


def opB(P):
    return np.kron(I3, P)


def e(i):
    v = np.zeros(3)
    v[i] = 1
    return v


# --- The Tiles UPB bound-entangled state (Bennett, DiVincenzo, Mor, Shor,
#     Smolin, Terhal 1999), as used in the paper's own qutrit benchmark ---
def _build_tiles():
    sqrt2, sqrt3 = np.sqrt(2), np.sqrt(3)
    upb = [
        np.kron(e(0), (e(0) - e(1)) / sqrt2),
        np.kron(e(2), (e(1) - e(2)) / sqrt2),
        np.kron((e(0) - e(1)) / sqrt2, e(2)),
        np.kron((e(1) - e(2)) / sqrt2, e(0)),
        np.kron((e(0) + e(1) + e(2)) / sqrt3, (e(0) + e(1) + e(2)) / sqrt3),
    ]
    P_UPB = sum(np.outer(v, v) for v in upb)
    return ((np.eye(9) - P_UPB) / 4).astype(complex)


RHO_TILES = _build_tiles()


def noisy_tiles(p):
    """rho(p) = p * rho_Tiles + (1-p) * I/9"""
    return p * RHO_TILES + (1 - p) * I9 / 9


# --- Qutrit Werner state (antisymmetric-subspace family), Werner 1989 ---
def _swap_matrix(dim=3):
    V = np.zeros((dim * dim, dim * dim))
    for a in range(dim):
        for b in range(dim):
            V[b * dim + a, a * dim + b] = 1
    return V


SWAP_3 = _swap_matrix(3)
P_ANTI = (np.eye(9) - SWAP_3) / 2
DIM_ANTI = np.trace(P_ANTI).real  # = 3


def werner_qutrit(p):
    """rho(p) = p * P_anti/3 + (1-p) * I/9. Known exact separability
    threshold: p = 1/(d+1) = 1/4 (Werner 1989)."""
    return p * P_ANTI / DIM_ANTI + (1 - p) * I9 / 9


# --- Correlation matrix / shadow-map criterion (paper Section "tensor
#     viewpoint" / Tiles benchmark) ---
def correlation_matrix(rho):
    T = np.zeros((8, 8))
    for i in range(8):
        for j in range(8):
            T[i, j] = np.trace(rho @ opA(GELLMANN[i]) @ opB(GELLMANN[j])).real
    return T


def shadow_map_nuclear_norm(rho):
    """||M_A(rho)||_*, normalization sqrt((d_A-1)(d_B-1)) = 2 for qutrits.
    Separable states satisfy this <= 1 (Theorem "cut-bound" in the note)."""
    T = correlation_matrix(rho)
    return np.linalg.svd(T / 2.0, compute_uv=False).sum()


# --- Plain PPT (Peres-Horodecki) check ---
def plain_ppt_min_eig(rho, dim=3):
    rho_pt = np.zeros((dim * dim, dim * dim), dtype=complex)
    for a in range(dim):
        for b in range(dim):
            for ap in range(dim):
                for bp in range(dim):
                    i, j = a * dim + b, ap * dim + bp
                    i2, j2 = a * dim + bp, ap * dim + b
                    rho_pt[i2, j2] = rho[i, j]
    return np.linalg.eigvalsh(rho_pt).min()


def plain_ppt_feasible(rho, dim=3, tol=1e-9):
    return plain_ppt_min_eig(rho, dim) >= -tol


# --- numpy partial traces, used only by the operator-Sinkhorn filter ---
def partial_trace_B_np(X, dim=3):
    T = X.reshape(dim, dim, dim, dim)
    return np.einsum('ikjk->ij', T)


def partial_trace_A_np(X, dim=3):
    T = X.reshape(dim, dim, dim, dim)
    return np.einsum('kikj->ij', T)


def _inv_sqrt_psd(M, eps=1e-12):
    w, v = np.linalg.eigh(M)
    w = np.clip(w, eps, None)
    return (v * (w ** -0.5)) @ v.conj().T


def operator_sinkhorn(rho, dim=3, max_iter=3000, tol=1e-11, verbose=False):
    """Local-filtering (SLOCC) normal-form algorithm: alternately rescale
    each side by (reduced state)^{-1/2} until both marginals are maximally
    mixed. Standard algorithm (Verstraete-Dehaene-DeMoor 2001/2003); the
    resulting fixed point is the Leinaas-Myrheim-Ovrum (2006) normal form."""
    X = rho.copy() / np.trace(rho).real
    devA = devB = None
    for it in range(max_iter):
        rhoA = partial_trace_B_np(X, dim)
        rhoA /= np.trace(rhoA).real
        devA = np.linalg.norm(rhoA - np.eye(dim) / dim)
        FA = np.kron(_inv_sqrt_psd(rhoA), np.eye(dim))
        X = FA @ X @ FA.conj().T
        X /= np.trace(X).real

        rhoB = partial_trace_A_np(X, dim)
        rhoB /= np.trace(rhoB).real
        devB = np.linalg.norm(rhoB - np.eye(dim) / dim)
        FB = np.kron(np.eye(dim), _inv_sqrt_psd(rhoB))
        X = FB @ X @ FB.conj().T
        X /= np.trace(X).real

        if devA < tol and devB < tol:
            if verbose:
                print(f"  Sinkhorn converged after {it + 1} iterations")
            break
    else:
        if verbose:
            print(f"  Sinkhorn did NOT fully converge in {max_iter} iters "
                  f"(devA={devA:.2e}, devB={devB:.2e})")
    return X