#!/usr/bin/env python3
"""
N6 v3 — THE INVOICE, EXECUTED IN THE TOY UNIVERSE
===================================================
Monogamy Gravity — the BMV experiment run exactly, inside the lattice.

CLAIM UNDER TEST (Folio I's invoice / Postulate 4):
    Vacuum-mediated entanglement between two bodies is never free: the
    field's which-path record exacts an irreducible dephasing fee, and
    the fee's structure is set by the vacuum's soft (massless) sector.

SETUP (exact, no approximations):
    Two "bodies" = branch qubits sigma_a, sigma_b = +/-1, each branch
    sourcing the chain linearly: H_int = eps*(sigma_a q_xa + sigma_b q_xb).
    Each of the 4 branches displaces every chain mode coherently:
        alpha_k(t) = c_k(t) j_k,   |c_k|^2 = (1-cos w_k t)/w_k^3
    Branch phases: theta_sigma(t) = sum_k (j_k^2/2w_k^3)(w_k t - sin w_k t)
    -> the cross term gives the NEWTONIAN ENTANGLING PHASE, secular rate
       2 eps^2 [K^-1]_ab  (the lattice Yukawa/Coulomb potential), while
    branch-distinguishability of the field gives the FEE:
        R_{ss'}(t) = exp(-1/2 sum_k |c_k|^2 (Dj_k)^2)
    The 4x4 two-qubit density matrix is exact:
        rho_{s,s'} = (1/4) R_{ss'} e^{i(theta_s - theta_s')}
    Measured: E_N(a:b) (log-negativity), V_a (qubit-a visibility),
    fee D = -ln V_a; fee-per-nat F = D/E_N.

GATE (written before the run):
    If the massless-sector arm transacts entanglement with NO
    accompanying fee (V_a = 1 while E_N grows), the fee claim of
    Postulate 4 is falsified: gravity-like mediation would be free.

ARMS:
    gapped m = 0.1 (no soft sector) vs near-critical m = 1e-3 (soft
    sector present); separations d; box sizes N (does the soft fee know
    the size of the room?).

Plain numpy; sums over exact lattice modes.
"""

import numpy as np
import json
import warnings
warnings.filterwarnings("ignore", category=RuntimeWarning)


def modes(N, m):
    K = np.zeros((N, N))
    for i in range(N):
        K[i, i] = m * m + 2.0
        K[i, (i + 1) % N] = -1.0
        K[(i + 1) % N, i] = -1.0
    w2, U = np.linalg.eigh(K)
    return np.sqrt(np.clip(w2, 1e-30, None)), U


def two_qubit_rho(t, w, U, xa, xb, eps):
    """Exact 4x4 density matrix of the branch qubits at time t."""
    Ua, Ub = U[xa, :], U[xb, :]
    c2 = (1 - np.cos(w * t)) / w ** 3                  # |c_k(t)|^2
    ph = (w * t - np.sin(w * t)) / (2 * w ** 3)        # phase kernel
    sigmas = [(+1, +1), (+1, -1), (-1, +1), (-1, -1)]
    j = [eps * (sa * Ua + sb * Ub) for sa, sb in sigmas]
    theta = [float(np.sum(ph * jk ** 2)) for jk in j]
    rho = np.zeros((4, 4), complex)
    for p in range(4):
        for q in range(4):
            dj = j[p] - j[q]
            R = np.exp(-0.5 * float(np.sum(c2 * dj ** 2)))
            rho[p, q] = 0.25 * R * np.exp(1j * (theta[p] - theta[q]))
    return rho


def logneg(rho):
    """Log-negativity across the a|b cut (basis order ++, +-, -+, --)."""
    rt = rho.reshape(2, 2, 2, 2).transpose(0, 3, 2, 1).reshape(4, 4)  # T_b
    ev = np.linalg.eigvalsh((rt + rt.conj().T) / 2)
    return float(np.log(np.sum(np.abs(ev))))


def visibility(rho):
    """Qubit-a coherence relative to t=0 (off-diagonal of Tr_b rho)."""
    ra = rho.reshape(2, 2, 2, 2)
    off = ra[0, 0, 1, 0] + ra[0, 1, 1, 1]              # <+|rho_a|->
    return float(np.abs(off) / 0.5)


def run_arm(label, N, m, d, eps, ts):
    w, U = modes(N, m)
    xa, xb = N // 2 - d // 2, N // 2 + (d + 1) // 2
    rows = []
    for t in ts:
        rho = two_qubit_rho(t, w, U, xa, xb, eps)
        rows.append(dict(t=float(t), EN=logneg(rho), V=visibility(rho)))
    return dict(label=label, N=N, m=m, d=d, rows=rows)


if __name__ == '__main__':
    eps = 0.05
    ts = np.linspace(0, 400, 81)

    print("N6 v3 — the BMV invoice, executed exactly in the lattice.")
    print(f"eps = {eps}; E_N = entanglement transacted, D = -ln V = fee paid\n")

    arms = []
    for label, N, m, d in (
        ("GAPPED       m=0.1,  d=10, N=400", 400, 0.1, 10),
        ("GAPPED       m=0.1,  d=20, N=400", 400, 0.1, 20),
        ("NEAR-MASSLESS m=1e-3, d=10, N=400", 400, 1e-3, 10),
        ("NEAR-MASSLESS m=1e-3, d=20, N=400", 400, 1e-3, 20),
        ("NEAR-MASSLESS m=1e-3, d=10, N=800", 800, 1e-3, 10),
    ):
        a = run_arm(label, N, m, d, eps, ts)
        arms.append(a)
        sel = [a['rows'][k] for k in (10, 20, 40, 80)]
        line = "  ".join(f"t={q['t']:.0f}: E={q['EN']:.4f} D={-np.log(max(q['V'],1e-300)):.4f}"
                         for q in sel)
        print(f"{label}\n    {line}")

    # fee-per-nat at fixed transaction size: first t where E_N >= 0.05
    print("\nfee per nat at E_N = 0.05:")
    for a in arms:
        EN = np.array([q['EN'] for q in a['rows']])
        V = np.array([q['V'] for q in a['rows']])
        k = np.argmax(EN >= 0.05)
        if EN[k] >= 0.05:
            D = -np.log(max(V[k], 1e-300))
            print(f"  {a['label']}:  t={a['rows'][k]['t']:.0f}  "
                  f"D={D:.5f}  F=D/E={D/EN[k]:.4f}")
        else:
            print(f"  {a['label']}:  E_N never reached 0.05 "
                  f"(max {EN.max():.4f})")
    json.dump(arms, open('/Users/antoine/agi/ledger/n6_bmv.json', 'w'), indent=1)
    print("\n-> n6_bmv.json")