"""
One cut, two consequences -- a computable model for the Entanglement-Continuity Theory.
=======================================================================================

A follow-on to the eight-qubit Heisenberg toy (`toy_model.py`). That toy showed
the SUBJECT half of the theory -- a bounded, integrated region inside a global
pure state -- but contained no spatial structure at all. This model closes that
gap. In a SINGLE many-body state it exhibits both consequences the theory says
follow from one factorization:

  Part A  -- an approximate emergent GEOMETRY, recovered from mutual information
             alone (dimension by MDS, curvature by Ollivier-Ricci), and
  Part B  -- a bounded, integrated SUBJECT (a Markov-blanket region) living
             inside that same state, with the toy's boundary/integration sweep.

  The bridge -- the payoff -- shows the two are the same cut: regions that make
             good subjects (strongly integrated, controllably bounded) are also
             geometrically COMPACT patches in the recovered geometry. Subjecthood
             and spatial localization track each other in one computable model.

Why this is tractable -- free-fermion (Gaussian) states
-------------------------------------------------------
Every entropy of a fermionic Gaussian state, and therefore every mutual
information, follows from the two-point correlation matrix C_ij = <c_i^dag c_j>
by diagonalizing a submatrix. This scales to hundreds of sites on a laptop,
where a generic interacting state dies near 20. That is the step that buys the
system sizes a geometry needs to be definable at all. The lattice we build the
hopping model on is the GROUND-TRUTH geometry; the test is whether we can read
that geometry back out of the entanglement without being told.

Geometries used (ground truth):
  * line   -- 1D chain                  (effective dimension ~1, curvature ~0)
  * grid   -- 2D square lattice         (effective dimension ~2, curvature ~0)
  * tree   -- regular Bethe lattice      (negatively curved; the Bruhat-Tits tree
              is exactly the bulk of p-adic AdS/CFT, so a tree is a legitimate
              discrete hyperbolic / holographic geometry, not a toy stand-in).

A regular {p,q} Poincare-disk tiling would be the textbook hyperbolic case; we
use the tree instead because it is exact and drift-free, whereas naive reflection
tilings accumulate floating-point error past the first ring. The scientific claim
-- recover NEGATIVE curvature from the mutual-information structure and tell it
from flat --
served by any genuinely negatively curved graph; the tree is the cleanest one.

How to read the result (and how NOT to). This is a computational-physics toy that
supports the STRUCTURAL PLAUSIBILITY of the Entanglement-Continuity picture: in
gapped free-fermion states on local graphs, mutual information can recover an
approximate geometry, pick out compact modular regions, and reproduce a boundary/
integration trade-off in one state. It is NOT a proof of the metaphysics. In
particular it does NOT:
  * derive consciousness or produce any felt quality (the hard problem is untouched;
    calling an integrated module a "subject" is the theory's reading, not a result);
  * prove emergent spacetime -- it READS geometry BACK from a hopping Hamiltonian
    that is local by construction; it does NOT derive locality or the factorization
    from a structureless Hilbert space (the theory's actual move). Recovering a
    static metric in a tractable class is also far weaker than Ryu-Takayanagi-style
    emergence; no dynamics, no gravity, no cosmology;
  * establish a literal Markov blanket -- O(S)=I(S:E) is boundary openness (cross-
    boundary information), not a conditional-independence screening-off; the blanket
    reading is an inspiration, not a result, and would need extra machinery.

Requires: numpy, scipy, networkx, matplotlib.
Run:  python3 continuity_geometry.py
"""

import itertools
import numpy as np
import networkx as nx


# ===========================================================================
#  1.  GROUND-TRUTH LATTICES
#      Each returns a networkx graph G with node attribute 'pos' (the true
#      coordinates, used only for the ground-truth overlay) plus metadata.
# ===========================================================================
def line_lattice(n=64):
    G = nx.path_graph(n)
    for i in G:
        G.nodes[i]["pos"] = (float(i), 0.0)
    G.graph.update(name="line", true_dim=1, curv_sign=0)
    return G


def grid_lattice(L=12):
    G = nx.grid_2d_graph(L, L)
    G = nx.convert_node_labels_to_integers(G, label_attribute="cell")
    for i in G:
        x, y = G.nodes[i]["cell"]
        G.nodes[i]["pos"] = (float(x), float(y))
    G.graph.update(name="grid", true_dim=2, curv_sign=0)
    return G


def tree_lattice(branching=2, depth=6):
    """Regular tree: root has `branching+1` children-equivalent; every internal
    node has total degree branching+1. Radial Poincare-style embedding: depth d
    sits at radius 1-2^{-d}, children fan out within the parent's angular wedge."""
    G = nx.Graph()
    G.add_node(0)
    G.nodes[0]["pos"] = (0.0, 0.0)
    G.nodes[0]["depth"] = 0
    nxt = 1
    # root gets (branching+1) subtrees so it is a regular internal node too
    frontier = [(0, -np.pi, np.pi, 0)]
    while frontier:
        new = []
        for parent, a0, a1, d in frontier:
            if d >= depth:
                continue
            k = (branching + 1) if parent == 0 else branching
            for j in range(k):
                child = nxt
                nxt += 1
                lo = a0 + (a1 - a0) * j / k
                hi = a0 + (a1 - a0) * (j + 1) / k
                ang = 0.5 * (lo + hi)
                rad = 1.0 - 2.0 ** (-(d + 1))
                G.add_edge(parent, child)
                G.nodes[child]["pos"] = (rad * np.cos(ang), rad * np.sin(ang))
                G.nodes[child]["depth"] = d + 1
                new.append((child, lo, hi, d + 1))
        frontier = new
    G.graph.update(name="tree", true_dim=None, curv_sign=-1)
    return G


# ===========================================================================
#  2.  FREE-FERMION (GAUSSIAN) MACHINERY
#      H = sum_ij T_ij c_i^dag c_j .  Ground state = fill the lowest N_fill
#      single-particle modes.  Correlation matrix C_ij = <c_i^dag c_j> is the
#      projector onto occupied modes; every entropy follows from a submatrix.
# ===========================================================================
def _bipartite_sign(G):
    """+/-1 two-colouring of a bipartite graph (all lattices here are bipartite)."""
    colour = nx.algorithms.bipartite.color(G)
    return np.array([1.0 if colour[i] == 0 else -1.0 for i in range(G.number_of_nodes())])


def hopping_matrix(G, t=1.0, mass=1.0, seed=0):
    """Hermitian hopping matrix H = -t sum_<ij> c_i^dag c_j + m sum_i s_i n_i,
    with s_i a +/-1 staggered (sublattice) sign. The staggered mass m opens a
    spectral gap at half filling, so correlations -- and therefore the mutual
    information -- decay EXPONENTIALLY with lattice distance. That is what makes
    d = -log I track true distance linearly and the recovered metric faithful;
    a gapless (m=0) state decays polynomially and warps the embedding. A tiny
    seeded jitter lifts any residual degeneracy so the ground state is unique."""
    n = G.number_of_nodes()
    T = np.zeros((n, n), dtype=float)
    for i, j in G.edges():
        T[i, j] = T[j, i] = -t
    rng = np.random.default_rng(seed)
    T[np.diag_indices(n)] = mass * _bipartite_sign(G) + 1e-6 * rng.standard_normal(n)
    return T


def correlation_matrix(T, filling=0.5):
    """Ground-state two-point matrix C_ij = <c_i^dag c_j> at the given filling.
    C = sum over occupied modes of v_a v_a^* (a real projector here)."""
    w, v = np.linalg.eigh(T)
    n = T.shape[0]
    n_occ = max(1, int(round(filling * n)))
    Vocc = v[:, :n_occ]                      # lowest-energy modes
    C = Vocc @ Vocc.conj().T
    return C


def _entropy_from_eigs(nu):
    """von Neumann entropy (bits) from correlation-matrix eigenvalues nu in [0,1]."""
    nu = np.clip(nu.real, 1e-12, 1.0 - 1e-12)
    return float(-np.sum(nu * np.log2(nu) + (1.0 - nu) * np.log2(1.0 - nu)))


def region_entropy(C, region):
    region = list(region)
    if not region:
        return 0.0
    sub = C[np.ix_(region, region)]
    nu = np.linalg.eigvalsh(sub)
    return _entropy_from_eigs(nu)


def mutual_information(C, A, B):
    A, B = list(A), list(B)
    return (region_entropy(C, A) + region_entropy(C, B)
            - region_entropy(C, A + B))


def mi_matrix(C):
    """Full pairwise single-site mutual-information matrix I[i,j], i != j.
    Vectorized over the cheap single-site entropies; pair entropy via 2x2."""
    n = C.shape[0]
    d = np.clip(np.real(np.diag(C)), 1e-12, 1 - 1e-12)
    s1 = -(d * np.log2(d) + (1 - d) * np.log2(1 - d))      # single-site entropies
    I = np.zeros((n, n))
    for i in range(n):
        for j in range(i + 1, n):
            sub = C[np.ix_([i, j], [i, j])]
            sij = _entropy_from_eigs(np.linalg.eigvalsh(sub))
            I[i, j] = I[j, i] = max(s1[i] + s1[j] - sij, 0.0)
    return I


# ===========================================================================
#  3.  PART A -- RECOVER THE GEOMETRY FROM MUTUAL INFORMATION
#      MI -> distance (monotone) -> local graph -> geodesic distances ->
#      (i) effective dimension by classical MDS, (ii) Ollivier-Ricci curvature.
# ===========================================================================
def mi_graph(I, tau=0.35, monotone="neglog"):
    """'Proximity' graph from mutual information. Keep links whose MI exceeds a
    fraction `tau` of the global maximum. With a spectral gap the lattice is
    near-translation-invariant and the nearest-neighbour shell sits at ~max while
    the next shell collapses far below it, so a single global threshold recovers
    the local neighbourhood on any lattice without being told its coordination
    number. Edge weights are a distance monotone of I; distant pairs get no edge
    and their separation is recovered as a GEODESIC (shortest path), Isomap-style."""
    n = I.shape[0]
    H = nx.Graph()
    H.add_nodes_from(range(n))
    gmax = I.max()
    thresh = tau * gmax

    def dist(val):
        v = np.clip(val / gmax, 1e-6, 1.0)
        if monotone == "neglog":
            return -np.log(v)
        if monotone == "inv":
            return 1.0 / v - 1.0
        if monotone == "sqrt":
            return np.sqrt(-np.log(v))
        raise ValueError(monotone)

    iu = np.argwhere(np.triu(I, 1) >= thresh)
    for i, j in iu:
        H.add_edge(int(i), int(j), weight=dist(I[i, j]))
    # ensure connectivity: bridge components by their single strongest cross-link
    if H.number_of_nodes() and not nx.is_connected(H):
        comps = list(nx.connected_components(H))
        main = set(max(comps, key=len))
        for comp in comps:
            if comp <= main:
                continue
            a, b = max(((a, b) for a in comp for b in main),
                       key=lambda ab: I[ab[0], ab[1]])
            H.add_edge(a, b, weight=dist(max(I[a, b], 1e-6 * gmax)))
            main |= comp
    return H


def geodesic_distances(H):
    """All-pairs shortest-path distances on the weighted MI graph."""
    n = H.number_of_nodes()
    D = np.full((n, n), np.inf)
    for src, dd in nx.all_pairs_dijkstra_path_length(H, weight="weight"):
        for dst, val in dd.items():
            D[src, dst] = val
    # patch any disconnected pairs with the largest finite distance * 1.5
    finite = D[np.isfinite(D)]
    if finite.size:
        D[~np.isfinite(D)] = finite.max() * 1.5
    np.fill_diagonal(D, 0.0)
    return D


def mds_spectrum(D):
    """Classical MDS: eigenvalues of the double-centred -1/2 D^2 (the Gram
    matrix). Returns sorted positive eigenvalues and the participation-ratio
    effective dimension PR = (sum l)^2 / sum l^2."""
    n = D.shape[0]
    J = np.eye(n) - np.ones((n, n)) / n
    B = -0.5 * J @ (D ** 2) @ J
    ev = np.linalg.eigvalsh(B)
    ev = np.sort(ev[ev > 1e-9])[::-1]
    if ev.size == 0:
        return ev, 0.0
    pr = (ev.sum() ** 2) / np.sum(ev ** 2)
    return ev, float(pr)


def growth_dimension(D, interior=None):
    """Ball-growth (Hausdorff-like) dimension: the average number of sites N(r)
    within geodesic radius r grows like r^d for a flat d-dimensional space, so
    d = d log N / d log r. Robust to the Manhattan-vs-Euclidean distinction that
    inflates MDS on a square lattice. For a tree N(r) grows EXPONENTIALLY, so the
    fitted slope keeps climbing with r and no finite dimension exists -- the
    operational signature of negative curvature. Returns (dim, log-log fit pts)."""
    n = D.shape[0]
    nodes = list(interior) if interior is not None else list(range(n))
    finite = D[np.isfinite(D)]
    rmax = np.percentile(finite[finite > 0], 55)        # stay clear of the boundary
    radii = np.linspace(rmax * 0.30, rmax, 9)
    counts = []
    for r in radii:
        counts.append(np.mean([(D[i] <= r).sum() for i in nodes]))
    lr, lc = np.log(radii), np.log(counts)
    slope = np.polyfit(lr, lc, 1)[0]
    return float(slope), (radii, np.array(counts))


def is_exponential_growth(D, interior=None):
    """A tell-tale of hyperbolicity: log N(r) is ~linear in r (exponential growth)
    rather than ~linear in log r. Returns R^2 of a log-N-vs-r straight-line fit."""
    n = D.shape[0]
    nodes = list(interior) if interior is not None else list(range(n))
    finite = D[np.isfinite(D)]
    rmax = np.percentile(finite[finite > 0], 60)
    radii = np.linspace(rmax * 0.25, rmax, 10)
    logN = np.array([np.log(np.mean([(D[i] <= r).sum() for i in nodes]))
                     for r in radii])
    a, b = np.polyfit(radii, logN, 1)
    resid = logN - (a * radii + b)
    ss = 1.0 - np.sum(resid ** 2) / np.sum((logN - logN.mean()) ** 2)
    return float(ss)


def mds_embed(D, dim=2):
    """Classical MDS coordinates (top `dim` axes) for drawing the recovered space."""
    n = D.shape[0]
    J = np.eye(n) - np.ones((n, n)) / n
    B = -0.5 * J @ (D ** 2) @ J
    w, v = np.linalg.eigh(B)
    order = np.argsort(w)[::-1][:dim]
    L = np.sqrt(np.clip(w[order], 0, None))
    return v[:, order] * L


def ollivier_ricci(H, D, alpha=0.0):
    """Ollivier-Ricci curvature on every edge of the weighted MI graph H.
    kappa(x,y) = 1 - W1(m_x, m_y)/d(x,y), with m_x a lazy walk: mass alpha on x,
    (1-alpha) spread over neighbours inversely to edge distance. W1 by exact LP."""
    from scipy.optimize import linprog
    nodes = list(H.nodes())
    idx = {u: i for i, u in enumerate(nodes)}

    def measure(x):
        nbrs = list(H[x])
        w = np.array([1.0 / max(H[x][v]["weight"], 1e-9) for v in nbrs])
        w = w / w.sum() * (1.0 - alpha)
        supp = [x] + nbrs
        mass = np.concatenate([[alpha], w])
        return supp, mass

    curv = {}
    for x, y in H.edges():
        sx, mx = measure(x)
        sy, my = measure(y)
        cost = np.array([[D[idx[a], idx[b]] for b in sy] for a in sx]).ravel()
        nA, nB = len(sx), len(sy)
        # transport LP: sum_j t_ij = mx_i ; sum_i t_ij = my_j ; t>=0
        Aeq, beq = [], []
        for i in range(nA):
            row = np.zeros(nA * nB)
            row[i * nB:(i + 1) * nB] = 1.0
            Aeq.append(row); beq.append(mx[i])
        for j in range(nB):
            row = np.zeros(nA * nB)
            row[j::nB] = 1.0
            Aeq.append(row); beq.append(my[j])
        res = linprog(cost, A_eq=np.array(Aeq), b_eq=np.array(beq),
                      bounds=(0, None), method="highs")
        W1 = res.fun if res.success else np.nan
        dxy = D[idx[x], idx[y]]
        curv[(x, y)] = 1.0 - W1 / dxy if dxy > 0 else 0.0
    return curv


def interior_nodes(H, frac=0.6):
    """Nodes whose graph-eccentricity is small -- i.e. away from the boundary.
    Used so dimension/curvature statistics are not dominated by edge effects."""
    ecc = nx.eccentricity(H)
    cutoff = np.percentile(list(ecc.values()), frac * 100)
    return [u for u in H if ecc[u] <= cutoff]


def recover_geometry(G, C=None, monotone="neglog"):
    """Full Part-A pipeline on one lattice. Returns a dict of recovered
    quantities for figures and gates."""
    if C is None:
        C = correlation_matrix(hopping_matrix(G))
    I = mi_matrix(C)
    H = mi_graph(I, monotone=monotone)
    D = geodesic_distances(H)
    inter = interior_nodes(H)
    ev, pr = mds_spectrum(D)
    gdim, growth = growth_dimension(D, inter)
    exp_r2 = is_exponential_growth(D, inter)
    curv = ollivier_ricci(H, D)
    inter_set = set(inter)
    kvals = np.array([curv[e] for e in curv])
    kint = np.array([v for e, v in curv.items()
                     if e[0] in inter_set and e[1] in inter_set])
    if kint.size == 0:
        kint = kvals
    return dict(I=I, H=H, D=D, mds_ev=ev, pr=pr,
                growth_dim=gdim, growth=growth, exp_r2=exp_r2,
                curv=curv, kappa=kvals, kappa_int=kint,
                kappa_mean=float(np.mean(kint)),
                neg_frac=float(np.mean(kint < 0)))


# ===========================================================================
#  4.  PART B -- FIND THE SUBJECT IN THE SAME STATE
#      A contiguous region S on the lattice. A boundary parameter g scales the
#      hopping across dS; an interior parameter J_S scales the hopping inside S.
#      Two read-offs (exactly as the toy):
#        O(S)   = I(S:E)   boundary OPENNESS / leakage -- the information crossing
#                          the cut. It RISES as the boundary opens (it is not a
#                          "strength"/closure). A Markov-blanket-INSPIRED proxy:
#                          true blankethood is conditional independence, not just
#                          low I(S:E) (see the writeup's limits).
#        Phi(S) = interior integration, by TWO measures that should agree:
#                 - MIP: min over interior bipartitions of I(A:B)  (toy's proxy)
#                 - Fiedler value: algebraic connectivity of the interior MI graph
#      Robustness: the boundary is also read in a genuine entanglement monotone
#      (Renyi-1/2 entropy across the cut = logarithmic negativity for the pure
#      global state), so the trade-off is not an artifact of mutual information.
# ===========================================================================
def geodesic_ball(G, center, radius):
    """Contiguous region: all sites within graph distance `radius` of center."""
    d = nx.single_source_shortest_path_length(G, center, cutoff=radius)
    return sorted(d.keys())


def subject_hopping(G, S, g=1.0, J_S=1.0, t=1.0, mass=1.0, seed=0):
    """Hopping matrix with intra-S bonds scaled by J_S and S<->E boundary bonds
    scaled by g (the direct analogue of the toy's interior/boundary couplings)."""
    n = G.number_of_nodes()
    Sset = set(S)
    T = np.zeros((n, n))
    for i, j in G.edges():
        ii, jj = (i in Sset), (j in Sset)
        if ii and jj:
            w = -t * J_S
        elif ii != jj:
            w = -t * g
        else:
            w = -t
        T[i, j] = T[j, i] = w
    rng = np.random.default_rng(seed)
    T[np.diag_indices(n)] = mass * _bipartite_sign(G) + 1e-6 * rng.standard_normal(n)
    return T


def renyi_half_entropy(C, region):
    """Renyi-1/2 entropy (bits) of a fermionic Gaussian reduced state from the
    eigenvalues nu of C_region: S_{1/2} = 2 sum log2(sqrt(nu)+sqrt(1-nu)). For a
    PURE global state this equals the logarithmic negativity across the cut -- a
    bona-fide entanglement monotone, used here as the non-MI robustness currency."""
    sub = C[np.ix_(list(region), list(region))]
    nu = np.clip(np.linalg.eigvalsh(sub).real, 1e-12, 1 - 1e-12)
    return float(2.0 * np.sum(np.log2(np.sqrt(nu) + np.sqrt(1.0 - nu))))


def phi_mip(C, S):
    """Interior integration as the minimum-information-bipartition MI (the toy's
    proxy): min over bipartitions (A|B) of S of I(A:B). Exact enumeration -- keep
    |S| <= ~16 for the sweep so this stays cheap."""
    S = list(S)
    m = len(S)
    S_whole = region_entropy(C, S)
    best = np.inf
    for r in range(1, m // 2 + 1):
        for A in itertools.combinations(S, r):
            B = [q for q in S if q not in A]
            val = region_entropy(C, A) + region_entropy(C, B) - S_whole
            best = min(best, val)
    return float(best)


def phi_fiedler(C, S):
    """Interior integration as the algebraic connectivity (Fiedler value) of the
    interior mutual-information graph: the 2nd-smallest eigenvalue of its weighted
    Laplacian. Large -> the interior is hard to cut in two -> well integrated. A
    graph-native measure, independent of the contested integrated-information
    metrics, so agreement with the MIP hedges the 'metric quarrel'."""
    S = list(S)
    m = len(S)
    if m < 2:
        return 0.0
    W = np.zeros((m, m))
    for a in range(m):
        for b in range(a + 1, m):
            W[a, b] = W[b, a] = mutual_information(C, [S[a]], [S[b]])
    L = np.diag(W.sum(1)) - W
    ev = np.linalg.eigvalsh(L)
    return float(ev[1])            # smallest is 0; second-smallest = connectivity


# ---- exact fermionic interior logarithmic negativity -----------------------
# The boundary B uses the pure-state identity (Renyi-1/2 across the cut). The
# INTERIOR is a mixed state, so its entanglement needs a mixed-state monotone:
# the fermionic logarithmic negativity. We compute it exactly -- build the
# many-body Gaussian density matrix of rho_S in the occupation basis from its
# correlation matrix (via Jordan-Wigner), apply the fermionic partial transpose
# (partial time-reversal: ordinary transpose on A times the JW phase i^{tau+tau'},
# Shapourian-Shiozaki-Ryu), and take the trace norm. Validated against the
# pure-state Renyi-1/2 identity and the product-state (zero) limit; keep |S|<=~9.
import functools


@functools.lru_cache(maxsize=16)
def _jw_ops(m):
    I2 = np.eye(2); Zp = np.array([[1, 0], [0, -1]], float)
    sm = np.array([[0, 1], [0, 0]], float)            # annihilation |0><1|
    ops = []
    for i in range(m):
        op = np.array([[1.0]])
        for k in range(m):
            op = np.kron(op, Zp if k < i else (sm if k == i else I2))
        ops.append(op.astype(complex))
    return tuple(ops)


def gaussian_density(C_S):
    """Many-body density matrix (2^m x 2^m) of the fermionic Gaussian state with
    correlation matrix C_S, in the site occupation basis."""
    from scipy.linalg import expm
    m = C_S.shape[0]
    w, U = np.linalg.eigh(C_S)
    w = np.clip(w, 1e-12, 1 - 1e-12)
    h = (U * np.log((1 - w) / w)) @ U.conj().T        # entanglement Hamiltonian
    c = _jw_ops(m)
    M = np.zeros((2 ** m, 2 ** m), complex)
    for i in range(m):
        for j in range(m):
            if abs(h[i, j]) > 1e-15:
                M += h[i, j] * (c[i].conj().T @ c[j])
    rho = expm(-M)
    return rho / np.trace(rho)


def _countA(m, A):
    Aset = set(A)
    return np.array([sum(1 for a in Aset if (x >> (m - 1 - a)) & 1)
                     for x in range(2 ** m)], int)


def log_negativity_exact(rho, m, A):
    """E_N = log2 || rho^{T_A}_fermionic ||_1, a genuine entanglement monotone."""
    t = rho.reshape([2] * m + [2] * m)
    perm = list(range(2 * m))
    for a in A:
        perm[a], perm[a + m] = perm[a + m], perm[a]
    pt = np.transpose(t, perm).reshape(2 ** m, 2 ** m)
    ph = (1j) ** _countA(m, A)                        # JW phase i^{tau+tau'}
    pt = pt * np.outer(ph, ph)
    sv = np.linalg.svd(pt, compute_uv=False)
    return float(np.log2(sv.sum()))


def phi_negativity(C, S):
    """Interior integration as the minimum across-bipartition fermionic log-
    negativity of the mixed reduced state rho_S -- the entanglement-monotone
    analogue of the MIP. Exact; |S| <= ~9."""
    S = list(S); m = len(S)
    rho = gaussian_density(C[np.ix_(S, S)])
    best = np.inf
    for r in range(1, m // 2 + 1):
        for A in itertools.combinations(range(m), r):
            best = min(best, log_negativity_exact(rho, m, list(A)))
    return float(max(best, 0.0))


def subject_metrics(G, S, g, J_S, with_neg=False, **kw):
    """All read-offs for region S at couplings (g, J_S), from one ground state."""
    C = correlation_matrix(subject_hopping(G, S, g=g, J_S=J_S, **kw))
    E = [i for i in G if i not in set(S)]
    out = dict(B_mi=mutual_information(C, S, E),
               B_neg=renyi_half_entropy(C, S),          # = log-negativity (pure)
               phi_mip=phi_mip(C, S),
               phi_fiedler=phi_fiedler(C, S))
    if with_neg:
        out["phi_neg"] = phi_negativity(C, S)
    return out


def boundary_sweep(G, S, gs, J_S=1.0, with_neg=False, **kw):
    rows = [subject_metrics(G, S, g, J_S, with_neg=with_neg, **kw) for g in gs]
    return {k: np.array([r[k] for r in rows]) for k in rows[0]}


def n_closest_region(G, center, n):
    """A contiguous region: the n sites closest (graph distance) to center."""
    d = nx.single_source_shortest_path_length(G, center)
    return sorted(sorted(d, key=lambda u: (d[u], u))[:n])


# ===========================================================================
#  5.  THE BRIDGE -- subjecthood and geometric compactness, in ONE state
#      All from the SAME uniform ground state used to recover the geometry in
#      Part A. For many candidate regions of equal size we score each on:
#        subject-quality  Q(S) = (internal binding) / (boundary leakage)
#                              = sum_{i<j in S} I(i:j)  /  I(S:E)
#          -- the Markov-blanket criterion: a good subject talks to itself far
#             more than to the outside.
#        compactness  K(S)     = - mean pairwise RECOVERED geodesic distance
#          -- purely geometric, read from Part A's emergent metric D.
#      The thesis is the JOIN: Q and K track each other -- the regions that make
#      good subjects are exactly the compact geodesic balls. One cut, two
#      consequences, in one computable model.
# ===========================================================================
def internal_mi_total(C, S):
    S = list(S)
    tot = 0.0
    for a in range(len(S)):
        for b in range(a + 1, len(S)):
            tot += mutual_information(C, [S[a]], [S[b]])
    return tot


def subject_quality(C, S, E=None, n=None):
    """Markov-blanket subject score: internal binding / boundary leakage."""
    S = list(S)
    if E is None:
        E = [i for i in range(n) if i not in set(S)]
    B = mutual_information(C, S, E)
    return internal_mi_total(C, S) / (B + 1e-9)


def compactness(D, S):
    """Geometric compactness in the recovered metric: minus the mean pairwise
    recovered geodesic distance (higher = tighter patch). Compared at fixed |S|."""
    S = list(S)
    sub = D[np.ix_(S, S)]
    m = len(S)
    return -float(sub[np.triu_indices(m, 1)].mean())


def candidate_regions(G, size=13, n_balls=18, n_scatter=30, n_blob=14, seed=1):
    """A pool of equal-size regions: compact geodesic balls, scattered random
    sets, and two-blob sets (split subjects) as an intermediate class."""
    rng = np.random.default_rng(seed)
    n = G.number_of_nodes()
    nodes = list(G.nodes())
    ecc = nx.eccentricity(G)
    cut = np.percentile(list(ecc.values()), 70)
    interior = [u for u in nodes if ecc[u] <= cut]
    pool = []

    def trim(region):                       # force exact size by BFS-order trim
        return sorted(region)[:size] if len(region) >= size else None

    # compact geodesic balls
    centers = rng.choice(interior, size=min(n_balls, len(interior)), replace=False)
    for c in centers:
        ball = geodesic_ball(G, int(c), 2)
        reg = trim(ball)
        if reg and len(reg) == size:
            pool.append(("ball", reg))
    # scattered random sets
    for _ in range(n_scatter):
        reg = sorted(rng.choice(n, size=size, replace=False).tolist())
        pool.append(("scatter", reg))
    # two-blob (split) sets: two small balls far apart
    for _ in range(n_blob):
        a, b = rng.choice(interior, size=2, replace=False)
        ba = geodesic_ball(G, int(a), 1)
        bb = geodesic_ball(G, int(b), 1)
        reg = sorted(set(ba) | set(bb))
        # pad/trim to exact size with nearest free sites of blob a
        if len(reg) < size:
            extra = [x for x in geodesic_ball(G, int(a), 3) if x not in reg]
            reg = sorted(set(reg) | set(extra[:size - len(reg)]))
        reg = trim(reg)
        if reg and len(reg) == size:
            pool.append(("blob", reg))
    return pool


def run_bridge(G=None, size=13, **kw):
    """Compute (subject-quality, compactness) for the region pool on the uniform
    state, plus the Pearson correlation that is the paper's payoff number."""
    if G is None:
        G = grid_lattice(15)
    C = correlation_matrix(hopping_matrix(G))
    I = mi_matrix(C)
    H = mi_graph(I)
    D = geodesic_distances(H)
    n = G.number_of_nodes()
    pool = candidate_regions(G, size=size, **kw)
    kinds, Q, K = [], [], []
    for kind, S in pool:
        kinds.append(kind)
        Q.append(subject_quality(C, S, n=n))
        K.append(compactness(D, S))
    Q, K = np.array(Q), np.array(K)
    r = float(np.corrcoef(Q, K)[0, 1])
    return dict(G=G, kinds=np.array(kinds), Q=Q, K=K, pearson=r)


# ===========================================================================
#  6.  FIGURES  (styled to match toy_model.py / the article)
# ===========================================================================
PAPER, INK, INK2, INK3 = "#F1F2F4", "#16181D", "#565B66", "#878C97"
INDIGO, CORAL = "#3A2E7A", "#BE4F2C"


def _style():
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    plt.rcParams.update({
        "font.family": "serif", "font.size": 11,
        "figure.facecolor": PAPER, "axes.facecolor": PAPER, "savefig.facecolor": PAPER,
        "text.color": INK, "axes.labelcolor": INK, "axes.titlecolor": INK,
        "xtick.color": INK2, "ytick.color": INK2, "axes.edgecolor": INK3,
    })
    return plt


def figure_geometry(results):
    """Figure 1 -- Part A: the geometry recovered from mutual information."""
    plt = _style()
    import matplotlib.cm as cm
    from matplotlib.colors import TwoSlopeNorm
    names = ["line", "grid", "tree"]
    fig, ax = plt.subplots(2, 3, figsize=(12.5, 8.0))
    norm = TwoSlopeNorm(vmin=-0.4, vcenter=0.0, vmax=0.4)
    for c, nm in enumerate(names):
        r = results[nm]
        H, curv, G = r["H"], r["curv"], r["G_ref"]
        n = G.number_of_nodes()
        pos = np.array([G.nodes[i]["pos"] for i in range(n)])
        a = ax[0, c]
        # recovered MI-graph edges, drawn on the TRUE coordinates: where they
        # coincide with the lattice, adjacency has been recovered from MI alone.
        true_edges = set(map(lambda e: tuple(sorted(e)), G.edges()))
        for u, v in H.edges():
            ok = tuple(sorted((u, v))) in true_edges
            a.plot([pos[u, 0], pos[v, 0]], [pos[u, 1], pos[v, 1]],
                   color=(INK3 if ok else CORAL), lw=(0.5 if ok else 1.2),
                   alpha=(0.5 if ok else 0.9), zorder=1)
        node_k = np.zeros(n); cnt = np.zeros(n)
        for (u, v), k in curv.items():
            node_k[u] += k; node_k[v] += k; cnt[u] += 1; cnt[v] += 1
        node_k = node_k / np.maximum(cnt, 1)
        a.scatter(pos[:, 0], pos[:, 1], c=node_k, cmap="coolwarm_r", norm=norm,
                  s=20, zorder=2, edgecolors="white", linewidths=0.3)
        frac_ok = np.mean([tuple(sorted(e)) in true_edges for e in H.edges()])
        a.set_title(f"{nm}: recovered graph on true layout", fontsize=12)
        a.set_xticks([]); a.set_yticks([]); a.set_aspect("equal")
        for s in a.spines.values():
            s.set_visible(False)
        dimtxt = ("exponential growth → no finite dim"
                  if nm == "tree" else f"growth dim ≈ {r['growth_dim']:.2f}")
        a.text(0.5, -0.03,
               dimtxt + f"\nκ = {r['kappa_mean']:+.3f}   "
               f"adjacency recovered: {frac_ok*100:.0f}%",
               transform=a.transAxes, ha="center", va="top", fontsize=9, color=INK2)

    # row 2a: curvature histograms
    a = ax[1, 0]
    for nm, col in zip(names, (INK3, INDIGO, CORAL)):
        a.hist(results[nm]["kappa_int"], bins=18, histtype="step", lw=2,
               color=col, label=nm)
    a.axvline(0, color=INK2, lw=0.8, ls=":")
    a.set_title("Ollivier–Ricci curvature", fontsize=12)
    a.set_xlabel("κ  (interior edges)"); a.set_ylabel("count")
    a.legend(frameon=False, fontsize=9)
    for s in ("top", "right"):
        a.spines[s].set_visible(False)

    # row 2b: ball growth -- power law (flat) vs exponential (tree)
    a = ax[1, 1]
    for nm, col in zip(names, (INK3, INDIGO, CORAL)):
        rad, cnt = results[nm]["growth"]
        a.plot(rad, cnt, "o-", color=col, ms=4, lw=1.6, label=nm)
    a.set_xscale("log"); a.set_yscale("log")
    a.set_title("ball growth  N(r)", fontsize=12)
    a.set_xlabel("recovered radius r  (log)")
    a.set_ylabel("sites within r  (log)")
    a.text(0.04, 0.96, "straight on log–log → power law (dim = slope)\n"
           "tree curves up → exponential → hyperbolic",
           transform=a.transAxes, va="top", ha="left", fontsize=8, color=INK2)
    a.legend(frameon=False, fontsize=9, loc="lower right")
    for s in ("top", "right"):
        a.spines[s].set_visible(False)

    # row 2c: MDS eigenvalue spectra
    a = ax[1, 2]
    for nm, col in zip(names, (INK3, INDIGO, CORAL)):
        ev = results[nm]["mds_ev"][:10]
        ev = ev / ev.sum()
        a.plot(range(1, len(ev) + 1), ev, "o-", color=col, ms=4, lw=1.6, label=nm)
    a.set_title("MDS eigenvalue spectrum", fontsize=12)
    a.set_xlabel("component"); a.set_ylabel("variance fraction")
    a.legend(frameon=False, fontsize=9)
    for s in ("top", "right"):
        a.spines[s].set_visible(False)

    fig.suptitle("Part A — geometry recovered from mutual information alone:  "
                 "dimension by ball-growth, curvature by Ollivier–Ricci",
                 fontsize=13, y=1.0)
    fig.tight_layout(rect=[0, 0, 1, 0.98])
    fig.savefig("continuity-geometry-recovery.png", dpi=190, bbox_inches="tight")
    print("wrote continuity-geometry-recovery.png")


def figure_subject(G, S, gs):
    """Figure 2 -- Part B: the bending window, in two currencies, widened by J_S."""
    plt = _style()
    fig, (axL, axR) = plt.subplots(1, 2, figsize=(12.5, 5.0))

    sw = boundary_sweep(G, S, gs, J_S=1.0, with_neg=True)
    # normalise the three integration measures to their g->0 value so the SHAPE
    # (the shared window) is comparable on one axis despite different units
    def nrm(a):
        return a / max(a[0], 1e-9)
    axL.plot(gs, nrm(sw["phi_mip"]), color=CORAL, lw=2.5,
             label=r"$\Phi$  integration (MIP, mutual info)")
    axL.plot(gs, nrm(sw["phi_neg"]), color="#7A1F8B", lw=2.0, ls=(0, (1, 1)),
             label=r"$\Phi$  integration (fermionic log-negativity)")
    axL.plot(gs, nrm(sw["phi_fiedler"]), color=CORAL, lw=1.8, ls=(0, (4, 2)),
             label=r"$\Phi$  integration (Fiedler)")
    axL.set_xlabel(r"boundary coupling  $g$")
    axL.set_ylabel("interior integration  (normalised to sealed value)")
    axL.set_ylim(0, 1.18)
    axB = axL.twinx()
    axB.plot(gs, sw["B_mi"], color=INDIGO, lw=2.5, label=r"$O=I(S{:}E)$  boundary openness (MI)")
    axB.plot(gs, sw["B_neg"], color=INDIGO, lw=1.8, ls=(0, (4, 2)),
             label=r"$O$  boundary openness (log-negativity)")
    axB.set_ylabel(r"boundary openness  $O(S)=I(S{:}E)$  (bits)", color=INDIGO)
    axB.tick_params(axis="y", colors=INDIGO)
    gmin, gmax = gs[0], gs[-1]
    axL.axvspan(gmin, 0.25, color=INK3, alpha=0.10)
    axL.axvspan(0.6, 1.8, color=CORAL, alpha=0.07)
    top = 1.18
    axL.text(0.13, top * 0.92, "sealed\nego", ha="center", color=INK2, fontsize=9)
    axL.text(1.15, top * 0.92, "boundary opening\n+ subject present\n= continuity",
             ha="center", color=CORAL, fontsize=9)
    axL.text(gmax * 0.86, top * 0.92, "recombination", ha="center", color=INK3, fontsize=9)
    axL.set_title("Dialing the boundary: it opens, the subject persists, then recombines",
                  fontsize=11.5, pad=10)
    h1, l1 = axL.get_legend_handles_labels()
    h2, l2 = axB.get_legend_handles_labels()
    axL.legend(h1 + h2, l1 + l2, frameon=False, fontsize=8.5, loc="upper center")
    for s in ("top",):
        axL.spines[s].set_visible(False); axB.spines[s].set_visible(False)

    # right: window widens with interior coupling J_S
    JSs = [0.6, 1.0, 1.6, 2.4]
    cmap = plt.cm.plasma(np.linspace(0.15, 0.8, len(JSs)))
    thresh = None
    widths = []
    for JS, col in zip(JSs, cmap):
        s = boundary_sweep(G, S, gs, J_S=JS)
        phi = s["phi_mip"]
        if thresh is None:
            thresh = 0.35 * phi.max()
        axR.plot(gs, phi, color=col, lw=2.2, label=f"$J_S$ = {JS:.1f}")
        above = gs[phi >= thresh]
        widths.append((JS, above.max() if above.size else gs[0]))
    axR.axhline(thresh, color=INK2, lw=0.8, ls=":")
    axR.text(gs[-1], thresh, "  Φ threshold", va="center", ha="left",
             color=INK2, fontsize=8)
    axR.set_xlabel(r"boundary coupling  $g$")
    axR.set_ylabel(r"interior integration  $\Phi$ (MIP)")
    axR.set_title("Signature prediction: a stronger interior widens the window",
                  fontsize=11.5, pad=10)
    axR.legend(frameon=False, fontsize=9, title="interior coupling")
    for s in ("top", "right"):
        axR.spines[s].set_visible(False)
    note = "  ".join(f"$J_S$={js:.1f}→g*={g:.2f}" for js, g in widths)
    axR.text(0.5, -0.16, "window edge g* (Φ above threshold):  " + note,
             transform=axR.transAxes, ha="center", va="top", fontsize=8, color=INK2)

    fig.suptitle("Part B — the subject in the same state: boundary / integration "
                 "trade-off and its moderation by interior coupling", fontsize=12.5, y=1.02)
    fig.tight_layout()
    fig.savefig("continuity-bending-window.png", dpi=190, bbox_inches="tight")
    print("wrote continuity-bending-window.png")
    return widths, sw


def figure_bridge(br):
    """Figure 3 -- the bridge: subject-quality vs geometric compactness."""
    plt = _style()
    fig, ax = plt.subplots(figsize=(8.2, 6.0))
    styles = {"ball": (CORAL, "o", "geodesic ball (compact subject)"),
              "blob": (INDIGO, "s", "two-blob (split)"),
              "scatter": (INK3, "^", "scattered set")}
    for kind, (col, mk, lab) in styles.items():
        m = br["kinds"] == kind
        ax.scatter(br["K"][m], br["Q"][m], c=col, marker=mk, s=46,
                   alpha=0.8, edgecolors="white", linewidths=0.5, label=lab)
    # regression line
    K, Q = br["K"], br["Q"]
    a, b = np.polyfit(K, Q, 1)
    xs = np.linspace(K.min(), K.max(), 50)
    ax.plot(xs, a * xs + b, color=INK2, lw=1.4, ls="--")
    ax.set_xlabel("geometric compactness  K  (− mean recovered distance) →  tighter")
    ax.set_ylabel("subject quality  Q  =  internal binding / boundary leakage")
    ax.set_title("The bridge — good subjects are compact patches of the emergent space",
                 fontsize=12.5, pad=10)
    ax.text(0.03, 0.97, f"Pearson r = {br['pearson']:.2f}\n"
            "one cut, two consequences:\nthe same MI structure fixes\nboth where things are\n"
            "and what is a subject",
            transform=ax.transAxes, va="top", ha="left", fontsize=9.5, color=INK)
    ax.legend(frameon=False, fontsize=9.5, loc="lower right")
    for s in ("top", "right"):
        ax.spines[s].set_visible(False)
    ax.grid(True, color=INK3, alpha=0.15, lw=0.6)
    fig.tight_layout()
    fig.savefig("continuity-bridge.png", dpi=190, bbox_inches="tight")
    print("wrote continuity-bridge.png")


# ===========================================================================
#  7.  INTERNAL MODEL-VALIDATION GATES (sanity / robustness checks for the toy,
#      NOT an empirical falsification of the theory)
# ===========================================================================
def validate(results, br, widths, sweep):
    """Internal pass/fail gates -- sanity/robustness checks for the model, in the
    series' spirit of saying what would sink it (they validate the toy, not the
    theory)."""
    print("\n=== validation gates ===")
    gates = []

    def gate(name, ok, detail):
        gates.append(ok)
        print(f"  [{'PASS' if ok else 'FAIL'}] {name}: {detail}")

    line, grid, tree = results["line"], results["grid"], results["tree"]
    gate("dimension separates line from grid",
         line["growth_dim"] < 1.4 < grid["growth_dim"],
         f"line≈{line['growth_dim']:.2f}, grid≈{grid['growth_dim']:.2f}")
    gate("curvature separates flat from hyperbolic",
         tree["kappa_mean"] < -0.05 < grid["kappa_mean"] + 0.05
         and tree["kappa_mean"] < grid["kappa_mean"] - 0.08,
         f"grid κ={grid['kappa_mean']:+.3f}, tree κ={tree['kappa_mean']:+.3f}")
    gate("tree shows exponential (not power-law) growth",
         tree["exp_r2"] > 0.985,
         f"exp-fit R²={tree['exp_r2']:.3f}")
    # monotone robustness: grid dimension stable, AND the tree stays clearly
    # negatively curved under a different distance monotone d = sqrt(-log I).
    # (Curvature SIGN is only meaningful where |kappa| is appreciable -- i.e. the
    #  tree -- so the flat grid's near-zero sign is not used as a test.)
    alt_grid = recover_geometry(grid["G_ref"], C=grid["C_ref"], monotone="sqrt")
    alt_tree = recover_geometry(tree["G_ref"], C=tree["C_ref"], monotone="sqrt")
    gate("geometry robust to the distance monotone",
         abs(alt_grid["growth_dim"] - grid["growth_dim"]) < 0.6
         and alt_tree["kappa_mean"] < -0.05,
         f"grid dim {grid['growth_dim']:.2f}→{alt_grid['growth_dim']:.2f}, "
         f"tree κ {tree['kappa_mean']:+.3f}→{alt_tree['kappa_mean']:+.3f} "
         f"under d=√(−logI)")
    # Part B: window widens with interior coupling
    js = [w[0] for w in widths]; gw = [w[1] for w in widths]
    gate("window widens with interior integration (signature prediction)",
         all(x <= y + 1e-9 for x, y in zip(gw, gw[1:])) and gw[-1] > gw[0],
         f"g*({js[0]})={gw[0]:.2f} → g*({js[-1]})={gw[-1]:.2f}")
    # interior window survives under a genuine entanglement monotone
    pn = sweep["phi_neg"]
    gate("interior window survives under exact fermionic log-negativity",
         pn[0] > 0.05 and pn[-1] < 0.5 * pn[0]
         and np.corrcoef(pn, sweep["phi_mip"])[0, 1] > 0.9,
         f"Φ_neg {pn[0]:.3f}→{pn[-1]:.3f}, "
         f"corr with MIP = {np.corrcoef(pn, sweep['phi_mip'])[0,1]:.2f}")
    # the bridge
    gate("good subjects are geometrically compact (the payoff)",
         br["pearson"] > 0.6,
         f"Pearson r(quality, compactness) = {br['pearson']:.2f}")
    print(f"  ----  {sum(gates)}/{len(gates)} gates passed")
    return all(gates)


# ===========================================================================
#  8.  MAIN
# ===========================================================================
def main():
    print("One cut, two consequences -- building the model.\n")

    # --- Part A: recover geometry on three ground-truth lattices ------------
    print("Part A: recovering geometry from mutual information")
    print(f"  {'lattice':6s} {'N':>4s} {'growth_dim':>10s} {'exp_R2':>7s} "
          f"{'kappa_int':>10s} {'neg_frac':>9s}")
    results = {}
    lattices = dict(line=line_lattice(60), grid=grid_lattice(15), tree=tree_lattice(2, 7))
    for nm, G in lattices.items():
        C = correlation_matrix(hopping_matrix(G))
        r = recover_geometry(G, C=C)
        r["G_ref"], r["C_ref"] = G, C
        results[nm] = r
        print(f"  {nm:6s} {G.number_of_nodes():4d} {r['growth_dim']:10.2f} "
              f"{r['exp_r2']:7.3f} {r['kappa_mean']:10.3f} {r['neg_frac']:9.2f}")
    figure_geometry(results)

    # --- Part B: the subject sweep on the grid state ------------------------
    print("\nPart B: the bounded subject and its bending window")
    G = grid_lattice(15)
    center = [n for n in G if G.nodes[n]["pos"] == (7, 7)][0]
    S = n_closest_region(G, center, 7)      # |S|=7: exact fermionic negativity tractable
    gs = np.concatenate([[0.0], np.linspace(0.08, 5.0, 22)])
    widths, sweep = figure_subject(G, S, gs)
    print(f"  region |S|={len(S)}; window edge g* by interior coupling:")
    for js, g in widths:
        print(f"    J_S={js:.1f}  ->  g* = {g:.2f}")
    print("  interior integration in three currencies (g=0 -> g=max):")
    for key, lab in (("phi_mip", "MIP (mutual info)"),
                     ("phi_neg", "fermionic log-negativity"),
                     ("phi_fiedler", "Fiedler value")):
        a = sweep[key]
        print(f"    {lab:28s}: {a[0]:.3f} -> {a[-1]:.3f}")

    # --- The bridge ---------------------------------------------------------
    print("\nThe bridge: subject-quality vs geometric compactness")
    br = run_bridge(grid_lattice(15))
    for k in ("ball", "blob", "scatter"):
        m = br["kinds"] == k
        print(f"  {k:8s} n={int(m.sum()):2d}  meanQ={br['Q'][m].mean():7.2f}  "
              f"meanK={br['K'][m].mean():7.3f}")
    print(f"  Pearson r(quality, compactness) = {br['pearson']:.3f}")
    figure_bridge(br)

    # --- gates --------------------------------------------------------------
    ok = validate(results, br, widths, sweep)
    print("\nALL GATES PASSED" if ok else "\nSOME GATES FAILED -- see above")


if __name__ == "__main__":
    main()
