#!/usr/bin/env python3
"""
halo_bound.py -- numerical check of the qubit "quantum halo" bound (Chapter 12).

CLAIM.  If a two-qubit state rho_AB is m-extendible on B (B's partner is shared symmetrically with m
recipients), its negativity is at most
        N_max(m) = 1 / (m + 2 sqrt(m) - 1)            (= 1/(2m) when rho_A = I/2).
CHECK.  Maximise the negativity of rho_AB1 over (1+m)-qubit states sigma_{A B1..Bm}, symmetrised over
permutations of B1..Bm so every rho_ABk is the same marginal.  For two qubits the partial transpose has at
most one negative eigenvalue, so N = max(0, -lambda_min(rho^Gamma)); the smooth objective lambda_min is
minimised with L-BFGS-B from many random starts.  This is a lower bound on the true maximum: it confirms
that the claimed value is reached, and no start exceeded it.  It is not a proof.
Also checks the maximally mixed case by restricting to sigma with rho_A = I/2 (penalty method).
numpy + scipy only.  No files are written.
"""
import os
for _v in ("OMP_NUM_THREADS", "OPENBLAS_NUM_THREADS", "MKL_NUM_THREADS", "VECLIB_MAXIMUM_THREADS"):
    os.environ.setdefault(_v, "1")
import itertools
import warnings
import numpy as np
from scipy.optimize import minimize

warnings.filterwarnings("ignore", message=r".*encountered in matmul")   # spurious numpy 2.0 + Accelerate warnings


def perm_op(n, perm):
    """Permutation matrix sending qubit k to position perm[k] (qubit 0 = most significant)."""
    d = 2 ** n
    P = np.zeros((d, d))
    for b in range(d):
        bits = [(b >> (n - 1 - k)) & 1 for k in range(n)]
        nb = [0] * n
        for k in range(n):
            nb[perm[k]] = bits[k]
        P[int(''.join(map(str, nb)), 2), b] = 1
    return P


def model(m):
    n, d = m + 1, 2 ** (m + 1)
    perms = [perm_op(n, [0] + list(p)) for p in itertools.permutations(range(1, n))]

    def sigma(x):
        G = (x[:d * d] + 1j * x[d * d:]).reshape(d, d)
        s = G @ G.conj().T
        s /= np.trace(s).real
        return sum(P @ s @ P.T for P in perms) / len(perms)

    def rho_ab(s):
        return np.einsum('ajbj->ab', s.reshape(4, 2 ** (m - 1), 4, 2 ** (m - 1)))

    return d, sigma, rho_ab


def lam_min_pt(r):
    pt = r.reshape(2, 2, 2, 2).transpose(0, 3, 2, 1).reshape(4, 4)
    return float(np.linalg.eigvalsh((pt + pt.conj().T) / 2)[0])


def rho_a(r):
    return np.einsum('ajbj->ab', r.reshape(2, 2, 2, 2))


def best_negativity(m, starts, mixed_a=False, seed=11):
    d, sigma, rho_ab = model(m)
    rng = np.random.default_rng(seed)
    best, best_ra = 0.0, None
    for _ in range(starts):
        def f(x):
            r = rho_ab(sigma(x))
            v = lam_min_pt(r)
            if mixed_a:                                     # push rho_A to I/2
                v += 50.0 * float(np.abs(rho_a(r) - np.eye(2) / 2).sum()) ** 2
            return v
        res = minimize(f, rng.normal(size=2 * d * d), method='L-BFGS-B', options=dict(maxiter=4000, maxfun=400000))
        r = rho_ab(sigma(res.x))
        if mixed_a and np.abs(rho_a(r) - np.eye(2) / 2).max() > 1e-3:
            continue
        neg = max(0.0, -lam_min_pt(r))
        if neg > best:
            best, best_ra = neg, np.linalg.eigvalsh(rho_a(r))
    return best, best_ra


def main():
    print("halo_bound.py -- max negativity of m-extendible two-qubit states (numerical lower bound on the maximum)")
    print(f"{'m':>3} {'numerical max':>14} {'claimed 1/(m+2sqrt(m)-1)':>25} {'rel. diff':>10} {'rho_A at optimum':>18} | {'rho_A=I/2: numerical':>21} {'claimed 1/(2m)':>15}")
    ok = True
    for m, starts in ((1, 16), (2, 16), (3, 8)):
        n, ra = best_negativity(m, starts)
        nm, _ = best_negativity(m, starts, mixed_a=True, seed=5)
        claim, claim_mixed = 1 / (m + 2 * np.sqrt(m) - 1), 1 / (2 * m)
        rel = abs(n - claim) / claim
        ok &= rel < 1e-4 and n <= claim * (1 + 1e-6) and abs(nm - claim_mixed) / claim_mixed < 1e-3
        print(f"{m:>3} {n:>14.6f} {claim:>25.6f} {rel:>10.1e} {np.array2string(np.round(ra, 4)):>18} | {nm:>21.6f} {claim_mixed:>15.6f}")
    print("claimed values reached, and no random start exceeded them" if ok else "CHECK FAILED")


if __name__ == "__main__":
    main()