#!/usr/bin/env python3
"""
record_key.py -- records of the noise vs records of the pointer (Prop. 2 of the note).

Private qubit Q, maximally entangled with a reference R. A random Pauli frame X^a Z^b
(a, b uniform) is applied to Q. The environment writes the fine record Z = (x, a, b); the
public pointer bit x is independent of (a, b) and is omitted below (it only conditions).
A fragment F carries l of the frame bits: l = 0 (x only), 1 (a), 2 (a and b).

Checks:
  I(R : F | X)      = 0 for every l      (a one-time-pad key says nothing about the message)
  I(R : F | Q, X)   = l bits             (but it helps whoever holds the system Q)
  average fidelity of recovering the private state from (Q, F): exactly 1/2, 2/3, 1, computed from the
  entanglement fidelity F_e of the decoded channel, F_avg = (2 F_e + 1)/3.  Optimality: for l <= 1 the
  state of R Q given F is separable (PPT, two qubits), so no recovery reaches F_e > 1/2 (F_avg > 2/3);
  for l = 0 it is a product state, so every recovery gives exactly F_avg = 1/2.
No files are written. Runs in well under a second.
"""
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 numpy as np

I2 = np.eye(2); X = np.array([[0, 1], [1, 0]], complex); Z = np.diag([1, -1]).astype(complex)
Y = np.array([[0, -1j], [1j, 0]])
phi = np.array([1, 0, 0, 1], complex) / np.sqrt(2)          # |Phi+> on R (x) Q
PHI = np.outer(phi, phi.conj())


def S(rho):
    w = np.linalg.eigvalsh((rho + rho.conj().T) / 2)
    w = w[w > 1e-14]
    return float(-(w * np.log2(w)).sum())


def frame(a, b):
    return np.linalg.matrix_power(X, a) @ np.linalg.matrix_power(Z, b)


def joint_state(l):
    """rho on R (x) Q (x) F, with F a classical register holding the leaked frame bits."""
    dF = 2 ** l
    rho = np.zeros((4 * dF, 4 * dF), complex)
    for a, b in itertools.product((0, 1), (0, 1)):
        U = np.kron(I2, frame(a, b))
        rq = U @ PHI @ U.conj().T
        f = [a, b][:l]
        k = int(''.join(map(str, f)), 2) if l else 0
        e = np.zeros((dF, dF)); e[k, k] = 1
        rho += 0.25 * np.kron(rq, e)
    return rho, dF


def ptrace(rho, dims, keep):
    n = len(dims)
    r = rho.reshape(dims + dims)
    idx = list(range(n))
    for i in reversed([i for i in idx if i not in keep]):
        r = np.trace(r, axis1=i, axis2=i + r.ndim // 2)
    d = int(np.prod([dims[i] for i in keep]))
    return r.reshape(d, d)


def mutual_cond(rho, dims, A, B, C):
    """I(A : B | C) = S(AC) + S(BC) - S(ABC) - S(C)."""
    s = lambda keep: S(ptrace(rho, dims, sorted(keep))) if keep else 0.0
    return s(A + C) + s(B + C) - s(A + B + C) - s(C)


def decoded_fidelity(l):
    """Exact average fidelity of 'undo the known frame bits, output Q', via F_avg = (2 F_e + 1)/3,
    F_e = sum_k |Tr K_k / 2|^2 over the Kraus operators K_ab = (1/2) V_ab X^a Z^b of the decoded channel."""
    Fe = 0.0
    for a, b in itertools.product((0, 1), (0, 1)):
        ka = a if l >= 1 else 0; kb = b if l >= 2 else 0
        K = 0.5 * frame(ka, kb).conj().T @ frame(a, b)
        Fe += abs(np.trace(K) / 2) ** 2
    return (2 * Fe + 1) / 3, Fe


def conditional_ppt_min(l):
    """Smallest eigenvalue of the partial transpose (on R) of the R Q state given each value of F.
    >= 0 means separable (two qubits), which caps the singlet fraction, hence F_e, at 1/2."""
    worst = np.inf
    for f in itertools.product((0, 1), repeat=l):
        rho = np.zeros((4, 4), complex)
        for a, b in itertools.product((0, 1), (0, 1)):
            if tuple([a, b][:l]) != f:
                continue
            U = np.kron(I2, frame(a, b))
            rho += U @ PHI @ U.conj().T
        rho /= np.trace(rho)
        pt = rho.reshape(2, 2, 2, 2).transpose(2, 1, 0, 3).reshape(4, 4)
        worst = min(worst, float(np.linalg.eigvalsh((pt + pt.conj().T) / 2).min()))
    return worst


def main():
    print("records of the noise vs records of the pointer (Prop. 2)")
    print(f"{'leaked bits l':>14} {'I(R:F|X)':>10} {'I(R:F|Q,X)':>12} {'F_e':>7} {'avg fidelity':>13} {'expected':>9} {'min PT eig | F':>15}")
    exp_f = {0: 1 / 2, 1: 2 / 3, 2: 1.0}
    ok = True
    for l in (0, 1, 2):
        rho, dF = joint_state(l)
        dims = [2, 2, dF]
        i_rf = mutual_cond(rho, dims, [0], [2], []) if l else 0.0
        i_rf_q = mutual_cond(rho, dims, [0], [2], [1]) if l else 0.0
        f, fe = decoded_fidelity(l)
        ppt = conditional_ppt_min(l)
        print(f"{l:>14d} {i_rf:>10.6f} {i_rf_q:>12.6f} {fe:>7.4f} {f:>13.6f} {exp_f[l]:>9.6f} {ppt:>15.2e}")
        ok &= abs(i_rf) < 1e-9 and abs(i_rf_q - l) < 1e-9 and abs(f - exp_f[l]) < 1e-12
        if l <= 1:
            ok &= ppt > -1e-12            # separable given F: no recovery beats F_e = 1/2
    print("l <= 1: the R Q state given F is PPT, hence separable, so the fidelities above are optimal.")
    print("all checks pass" if ok else "CHECK FAILED")


if __name__ == "__main__":
    main()