"""
SDP-based sharpening of the biseparable threshold for min(||M_AB||_*, ||M_AC||_*, ||M_AD||_*)
on 4 qubits, via the PPT-mixture relaxation (Jungnitsch-Moroder-Guehne 2011 style SDP).

WHY NOT A ONE-SHOT SDP
-----------------------
||M_S(rho)||_* is CONVEX in rho. Maximizing a convex function over a convex set is itself
a non-convex problem -- no SDP solver can do this directly.

THE FIX -- alternating (Frank-Wolfe) scheme using the dual (support-function) form of the
nuclear norm:
    ||X||_* = max_{||O||_op <= 1} tr(O^T X)
For FIXED O_AB, O_AC, O_AD (each with operator norm <= 1), min_S tr(O_S^T M_S(rho)) is a
min of THREE LINEAR functions of rho, hence CONCAVE, hence
    max_{rho in PPT-mixtures} min_S tr(O_S^T M_S(rho))
IS a legitimate concave-maximization problem -> a genuine SDP.

Loop:
  1) fix O's -> solve the SDP -> get rho*
  2) at rho*, compute the TRUE nuclear norms and their exact dual witnesses
     O_S = U_S V_S^T (from the SVD of M_S(rho*)) -> update O's
  3) repeat

This is a heuristic (finds a local stationary point of a genuinely non-convex problem),
but it searches the FULL convex PPT-mixture body (a strict superset of the biseparable
states), not just a hand-picked family of pure-state mixtures -- a much stronger stress
test of the conjectured biseparable supremum (~3.0) than black-box optimization over a
parametrized ansatz.

Requires: pip install cvxpy numpy scipy
Every piece of linear algebra here (Pauli-tensor extraction, partial-transpose
permutation, the fast coefficient-matrix construction) was verified against a slow
reference implementation in pure numpy before being translated to cvxpy; only the
cvxpy Problem-building/solving itself is unverified in the sandbox this was written in
(no cvxpy, no network there).
"""
import numpy as np
import cvxpy as cp

# ----------------------------------------------------------------------
# 1. Pauli tensor operators, index k = i0*64 + i1*16 + i2*4 + i3
# ----------------------------------------------------------------------
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 kron4(a, b, c, d):
    return np.kron(np.kron(a, b), np.kron(c, d))


PAULI_OPS = np.zeros((256, 16, 16), dtype=complex)
for i0 in range(4):
    for i1 in range(4):
        for i2 in range(4):
            for i3 in range(4):
                k = i0 * 64 + i1 * 16 + i2 * 4 + i3
                PAULI_OPS[k] = kron4(PAULI[i0], PAULI[i1], PAULI[i2], PAULI[i3])

CLUSTERS = [('AB', 0, 1), ('AC', 0, 2), ('AD', 0, 3)]


def build_coeff_matrix(s0, s1):
    """(225,256) complex matrix Coeff s.t., for rho flattened row-major (vec[a*16+b]=rho[a,b]),
    (Coeff @ vec).reshape(15,15).real / 3.0 == the normalized bigraduated cluster map M_S,
    S = {s0,s1}, complement the other two qubits. Verified against a slow trace-based
    reference (max abs diff ~2e-17)."""
    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)]
    Coeff = np.zeros((225, 256), dtype=complex)
    entry = 0
    for (ia, ib) in rows:
        for (ic, idd) in cols:
            idx = [0, 0, 0, 0]
            idx[s0] = ia
            idx[s1] = ib
            idx[c0] = ic
            idx[c1] = idd
            k = idx[0] * 64 + idx[1] * 16 + idx[2] * 4 + idx[3]
            # trace(rho @ P_k) = sum_{a,b} rho[a,b] P_k[b,a]; row-major vec_rho[a*16+b]=rho[a,b]
            Coeff[entry, :] = PAULI_OPS[k].T.flatten()
            entry += 1
    return Coeff


COEFF = {name: build_coeff_matrix(s0, s1) for name, s0, s1 in CLUSTERS}


# ----------------------------------------------------------------------
# 2. Partial transpose as an explicit (256,256) permutation matrix
# ----------------------------------------------------------------------
def partial_transpose_perm(T, n=4):
    perm = np.zeros(4 ** n, dtype=int)
    for row in range(2 ** n):
        rbits = [(row >> (n - 1 - i)) & 1 for i in range(n)]
        for col in range(2 ** n):
            cbits = [(col >> (n - 1 - i)) & 1 for i in range(n)]
            new_row = [cbits[i] if i in T else rbits[i] for i in range(n)]
            new_col = [rbits[i] if i in T else cbits[i] for i in range(n)]
            new_row_idx = sum(b << (n - 1 - i) for i, b in enumerate(new_row))
            new_col_idx = sum(b << (n - 1 - i) for i, b in enumerate(new_col))
            perm[new_row_idx * (2 ** n) + new_col_idx] = row * (2 ** n) + col
    return perm


def perm_matrix(T, n=4):
    perm = partial_transpose_perm(T, n)
    P = np.zeros((4 ** n, 4 ** n))
    for new_idx, old_idx in enumerate(perm):
        P[new_idx, old_idx] = 1.0
    return P


# The 7 bipartitions of {A,B,C,D}={0,1,2,3}; PT taken w.r.t. the listed (smaller) side.
# PPT is equivalent for either side of a bipartition, so this choice is arbitrary but fixed.
BIPARTITIONS = [
    ('AB|CD', {0, 1}), ('AC|BD', {0, 2}), ('AD|BC', {0, 3}),
    ('A|BCD', {0}), ('B|ACD', {1}), ('C|ABD', {2}), ('D|ABC', {3}),
]
PT_MATRIX = {name: perm_matrix(side) for name, side in BIPARTITIONS}


def cvxpy_flatten_rowmajor(rho_expr):
    """16x16 cvxpy expression -> length-256 cvxpy expression, row-major."""
    return cp.hstack([rho_expr[i, :] for i in range(16)])


def cvxpy_partial_transpose(rho_expr, name):
    vec = cvxpy_flatten_rowmajor(rho_expr)
    pt_vec = PT_MATRIX[name] @ vec
    return cp.reshape(pt_vec, (16, 16), order='C')  # MUST match row-major PT_MATRIX construction


def cvxpy_cluster_maps(rho_expr):
    vec = cvxpy_flatten_rowmajor(rho_expr)
    out = {}
    for name, s0, s1 in CLUSTERS:
        flat = COEFF[name] @ vec / 3.0
        out[name] = cp.real(cp.reshape(flat, (15, 15), order='C'))  # MUST match row-major COEFF construction
    return out


# ----------------------------------------------------------------------
# 3. PPT-mixture SDP + one alternating step
# ----------------------------------------------------------------------
def solve_fixed_witness_step(O, solver=cp.SCS, verbose=False):
    """O: dict name(in {'AB','AC','AD'}) -> 15x15 real array with operator norm <= 1.
    Returns (rho_value, sdp_optimal_t, dict of numeric M_S values)."""
    rho_gammas = {}
    constraints = []
    for name, side in BIPARTITIONS:
        r = cp.Variable((16, 16), hermitian=True)
        constraints.append(r >> 0)                              # PSD
        constraints.append(cvxpy_partial_transpose(r, name) >> 0)  # PPT across this cut
        rho_gammas[name] = r

    rho = sum(rho_gammas.values())
    constraints.append(cp.real(cp.trace(rho)) == 1)

    M = cvxpy_cluster_maps(rho)
    t = cp.Variable()
    for name, _, _ in CLUSTERS:
        constraints.append(t <= cp.sum(cp.multiply(O[name], M[name])))

    prob = cp.Problem(cp.Maximize(t), constraints)
    prob.solve(solver=solver, verbose=verbose)

    if rho.value is None:
        raise RuntimeError(f"SDP did not solve to a usable solution (status={prob.status}).")

    M_vals = {name: M[name].value for name, _, _ in CLUSTERS}
    return rho.value, prob.value, M_vals


def true_norms_and_witnesses(M_vals):
    """Exact nuclear norms of the numeric M_S matrices, plus their optimal dual witnesses
    O_S = U_S V_S^T (operator norm exactly 1, achieves tr(O_S^T M_S) = ||M_S||_*)."""
    norms, O_opt = {}, {}
    for name, M in M_vals.items():
        U, s, Vt = np.linalg.svd(M)
        norms[name] = s.sum()
        O_opt[name] = U @ Vt
    return norms, O_opt


# ----------------------------------------------------------------------
# 4. Alternating search with multiple random restarts
# ----------------------------------------------------------------------
def alternating_search(n_restarts=5, n_iters=15, seed0=0, verbose=True):
    best_min, best_rho = -np.inf, None
    for r in range(n_restarts):
        rng = np.random.default_rng(seed0 + r)
        O = {}
        for name, _, _ in CLUSTERS:
            A = rng.normal(size=(15, 15))
            O[name] = A / np.linalg.norm(A, ord=2)  # operator norm 1
        if verbose:
            print(f"--- restart {r} ---")
        for it in range(n_iters):
            rho_val, t_val, M_vals = solve_fixed_witness_step(O)
            norms, O = true_norms_and_witnesses(M_vals)
            cur_min = min(norms.values())
            if verbose:
                nice = {k: round(v, 4) for k, v in norms.items()}
                print(f"  iter {it:2d}: SDP t={t_val:.4f}  true norms={nice}  min={cur_min:.4f}")
            if cur_min > best_min:
                best_min, best_rho = cur_min, rho_val
        if verbose:
            print()
    return best_min, best_rho


if __name__ == "__main__":
    print("Proven upper bound (pure-state extreme points + convexity of the sum): 11/3 =", 11 / 3)
    print("Conjectured true biseparable / PPT-mixture supremum: ~3.0")
    print()
    best_min, best_rho = alternating_search(n_restarts=5, n_iters=15)
    print("=" * 60)
    print("Best min(||M_AB||_*, ||M_AC||_*, ||M_AD||_*) found over PPT-mixtures:", best_min)
    print("- if this stays at/near 3.0 across restarts -> strong evidence 3.0 is exact")
    print("- if it clearly exceeds 3.0 -> best_rho is a concrete witness state to inspect")