#!/usr/bin/env python3
"""
N5 — THE LOCKED RESERVE, LOCATED
=================================
Monogamy Gravity, Folio VI numerical experiment.

CLAIM UNDER TEST (the Locked Reserve conjecture, Folio I / Folio VI):
    Part of the vacuum's correlation account is LOCKED -- it exists
    (mutual information > 0) but cannot be withdrawn by any local
    operation (no distillable entanglement). The conjecture: the locked
    component dominates at large separation, and at cosmological
    distances the vacuum's accounts are pure reserve -- which is why
    that sector gravitates as a constant (Lambda) rather than as a
    force.

GATE (written before the run):
    If the drainable component (log-negativity) persists wherever the
    account (mutual information) persists -- no sudden death, locked
    fraction never reaching 1 -- then the vacuum has no locked sector
    and Folio VI's reading of Lambda loses its mechanism.

PROTOCOL:
    Two intervals of n sites at separation d in the chain vacuum.
      MI(A:B)   : the total account (mutual information)
      E_N(A:B)  : the drainable part (logarithmic negativity --
                  for Gaussian states, NPT is the distillability
                  criterion; E_N = 0 means nothing can be drawn by
                  any protocol based on nonpositive partial transpose)
      locked fraction = 1 - [accounts with E_N > 0]/[account total]
    Sweep d. Find d* where E_N dies while MI persists.
    Arms: near-critical (m = 1e-3) and gapped (m = 0.1).

Plain numpy.
"""

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


def chain_covs(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)
    w = np.sqrt(np.clip(w2, 1e-30, None))
    return (U * (0.5 / w)) @ U.T, (U * (0.5 * w)) @ U.T


def entropy(X, P, R):
    idx = np.ix_(R, R)
    nu = np.sqrt(np.clip(np.linalg.eigvals(X[idx] @ P[idx]).real, 0.25, None))
    up, dn = nu + .5, nu - .5
    return float(np.sum(up * np.log(up)) -
                 np.sum(np.where(dn > 1e-12,
                                 dn * np.log(np.clip(dn, 1e-300, None)), 0)))


def MI(X, P, A, B):
    return entropy(X, P, A) + entropy(X, P, B) - entropy(X, P, A + B)


def logneg(X, P, A, B):
    R = A + B
    idx = np.ix_(R, R)
    D = np.ones(len(R)); D[len(A):] = -1.0
    Pt = (P[idx] * D).T * D
    nu = np.sqrt(np.clip(np.linalg.eigvals(X[idx] @ Pt).real, 1e-30, None))
    return float(np.sum(np.where(nu < 0.5, -np.log(2 * nu), 0.0)))


def arm(label, N, m, n, seps):
    X, P = chain_covs(N, m)
    c = N // 2
    print(f"\n=== {label}: N={N}, m={m}, intervals of n={n} sites")
    print(f"{'d':>4} {'MI(A:B)':>10} {'E_N(A:B)':>10} {'locked %':>9}")
    rows = []
    for d in seps:
        A = list(range(c - d // 2 - n, c - d // 2))
        B = list(range(c + (d + 1) // 2, c + (d + 1) // 2 + n))
        mi, en = MI(X, P, A, B), logneg(X, P, A, B)
        locked = 100.0 if en < 1e-12 else 100.0 * max(0.0, 1.0 - en / mi)
        rows.append(dict(d=d, MI=mi, EN=en, locked=locked))
        print(f"{d:4d} {mi:10.5f} {en:10.6f} {locked:9.1f}")
    # d* = first separation with E_N = 0 while MI > 0
    dstar = next((r['d'] for r in rows if r['EN'] < 1e-12 and r['MI'] > 1e-6), None)
    print(f"--- sudden death of the drainable part: d* = {dstar} "
          f"(interval size n = {n})" if dstar else
          "--- no sudden death in range")
    return dict(label=label, N=N, m=m, n=n, rows=rows, dstar=dstar)


if __name__ == '__main__':
    print("N5 — the locked reserve, located.")
    arms = [
        arm("NEAR-CRITICAL", 400, 1e-3, 8, (1, 2, 3, 4, 6, 8, 12, 16, 24, 32, 48, 64)),
        arm("GAPPED (m=0.1)", 400, 0.1, 8, (1, 2, 3, 4, 6, 8, 12, 16, 24, 32)),
        arm("NEAR-CRITICAL, larger accounts (n=16)", 400, 1e-3, 16,
            (2, 4, 8, 12, 16, 24, 32, 48, 64, 96)),
    ]
    print("\n" + "=" * 64)
    print("VERDICT")
    for a in arms:
        last = a['rows'][-1]
        print(f"{a['label']}: d* = {a['dstar']}, account at far end "
              f"MI = {last['MI']:.4f} with E_N = {last['EN']:.6f} "
              f"-> locked {last['locked']:.0f}%")
    json.dump(arms, open('/Users/antoine/agi/ledger/n5_results.json', 'w'),
              indent=1)
    print("-> n5_results.json")