"""
A toy model for the Entanglement-Continuity Theory.
================================================================

A minimal, exactly-solvable "mini-universe" that demonstrates the STRUCTURAL
core of the theory: how a single entangled quantum state parses into a bounded
"subject" and an "environment", and how dialing a single boundary parameter
moves the subject between a sealed-off ego and a state continuous with the whole
-- while a unified subject is still present.

What it does NOT do: produce any *felt* quality. It tests the physics of the
parsing (boundary, integration, merging), not the identification of an interior
with experience. That gap is the hard problem and is left untouched, by design.

Model
-----
8 qubits as two 4-site Heisenberg rings (each a closed "shell": a unique singlet
when isolated, so a fully sealed subject carries ZERO cross-boundary information).

  Subject      S = {0,1,2,3}   ring bonds (0,1),(1,2),(2,3),(3,0)   coupling J_S
  Environment  E = {4,5,6,7}   ring bonds (4,5),(5,6),(6,7),(7,4)   coupling J_E
  Boundary     bonds (3,4) and (0,7)                                coupling g

H = J_S * [S-ring] + J_E * [E-ring] + g * [boundary bonds],
with each bond h(i,j) = X_iX_j + Y_iY_j + Z_iZ_j  (antiferromagnetic Heisenberg).

From the unique ground state |Psi> (a timeless, global state) we read off, for the
subject S:

  O(S)   = I(S:E)        boundary OPENNESS = mutual information across the cut.
                         (It RISES as the boundary opens; it is openness/leakage,
                          not a "strength"/closure. Its inverse is segregation.)
                         O = 0  -> sealed-off, bounded subject.
                         O large-> boundary thinned, continuous with the whole.

  Phi(S) = min over internal bipartitions (A|B) of S of I(A:B)
                         interior integration: is S one irreducible subject?
                         Phi small -> S falls apart into independent pieces.

Result
------
* Boundary sweep (J_S=J_E=1, dial g up): O(S) rises from ~0 to several bits
  (the boundary opens) while Phi(S) stays high through the unity regime and
  only erodes as the boundary approaches full dissolution -- the theory's
  recombination / death limit. The window where O is large AND Phi is still high
  is the structural signature it calls "felt continuity". Mutual information is
  not an entanglement measure (it counts classical correlation too), so as a
  check the trade-off is not an artifact of that choice, O and Phi are recomputed
  with a genuine entanglement monotone (the logarithmic negativity) and the same
  window reappears (dashed curves, Figure 1). The real driver is the bounded
  correlation budget of a finite, globally pure, closed-shell state.
* Phase diagram (dial g AND J_S independently): O(S) is governed mainly by the
  boundary coupling g; Phi(S) mainly by the interior coupling J_S. This is the
  theory's two-axis phase space, computed. The "lucid unity" corner -- boundary
  dissolved AND subject present -- is high-g / high-J_S; weaken J_S and
  integration collapses (the deep-sleep corner) however dissolved the boundary.

Requires: numpy, matplotlib.   Run:  python3 toy_model.py
"""

import string, itertools
import numpy as np

# ---- single-qubit operators -------------------------------------------------
I2 = np.eye(2, dtype=complex)
X = np.array([[0, 1], [1, 0]], dtype=complex)
Y = np.array([[0, -1j], [1j, 0]], dtype=complex)
Z = np.array([[1, 0], [0, -1]], dtype=complex)

N = 8
SUBJECT = (0, 1, 2, 3)
ENVIRON = (4, 5, 6, 7)
S_BONDS = [(0, 1), (1, 2), (2, 3), (3, 0)]
E_BONDS = [(4, 5), (5, 6), (6, 7), (7, 4)]
BOUNDARY_BONDS = [(3, 4), (0, 7)]


def embed(op, i, n=N):
    m = np.array([[1]], dtype=complex)
    for k in range(n):
        m = np.kron(m, op if k == i else I2)
    return m


def heis(i, j, n=N):
    return (embed(X, i, n) @ embed(X, j, n)
            + embed(Y, i, n) @ embed(Y, j, n)
            + embed(Z, i, n) @ embed(Z, j, n))


# Precompute the bond operators once (they don't depend on the couplings).
_S = sum(heis(i, j) for i, j in S_BONDS)
_E = sum(heis(i, j) for i, j in E_BONDS)
_B = sum(heis(i, j) for i, j in BOUNDARY_BONDS)


def build_H(J_S, J_E, g):
    return J_S * _S + J_E * _E + g * _B


def ground_state(H):
    w, v = np.linalg.eigh(H)
    return v[:, 0], float(w[1] - w[0])          # vector, gap (uniqueness check)


# ---- reduced states, entropy, mutual information ----------------------------
def partial_trace(rho, keep, n=N):
    keep = sorted(keep)
    t = rho.reshape([2] * n + [2] * n)
    L = string.ascii_lowercase
    row, col = list(L[:n]), list(L[n:2 * n])
    for q in range(n):
        if q not in keep:
            col[q] = row[q]                      # contract traced-out qubits
    sub = "".join(row) + "".join(col) + "->" \
          + "".join(row[q] for q in keep) + "".join(col[q] for q in keep)
    d = 2 ** len(keep)
    return np.einsum(sub, t).reshape(d, d)


def vn_entropy(rho):
    ev = np.linalg.eigvalsh(rho)
    ev = ev[ev > 1e-12]
    return float(-np.sum(ev * np.log2(ev)))      # bits


def S_of(rho, subset, n=N):
    if len(subset) in (0, n):
        return 0.0                               # global state is pure
    return vn_entropy(partial_trace(rho, subset, n))


def boundary_openness(rho, S=SUBJECT, E=ENVIRON):
    # O(S) = I(S:E): the information crossing the cut. Rises as the boundary opens.
    return S_of(rho, S) + S_of(rho, E) - S_of(rho, tuple(sorted(S + E)))


def integration(rho, S=SUBJECT):
    """Phi(S): the subject's minimum-information-partition mutual information."""
    S_whole = S_of(rho, S)
    best = np.inf
    for r in range(1, len(S) // 2 + 1):
        for A in itertools.combinations(S, r):
            B = tuple(q for q in S if q not in A)
            best = min(best, S_of(rho, A) + S_of(rho, B) - S_whole)
    return float(best)


# ---- entanglement (negativity) versions of the same two quantities ----------
# Mutual information is not an entanglement measure: it counts classical
# correlation too. So, to check the boundary/interior trade-off is not an artifact
# of that choice, the same two quantities are recomputed with a genuine
# entanglement monotone -- the (logarithmic) negativity. Same window => the result
# is robust to the measure; the driver is the bounded correlation budget of a
# finite, globally pure state, not a theorem about entanglement measures.
def partial_transpose(rho_sub, sites, A):
    """Partial transpose of a density matrix defined on `sites` over subset A."""
    k = len(sites)
    t = rho_sub.reshape([2] * k + [2] * k)
    perm = list(range(2 * k))
    for p, s in enumerate(sites):
        if s in A:
            perm[p], perm[p + k] = perm[p + k], perm[p]
    return np.transpose(t, perm).reshape(2 ** k, 2 ** k)


def log_negativity(rho_sub, sites, A):
    """E_N = log2 ||rho^{T_A}||_1, in bits. Zero on separable states; 1 bit for
       a Bell pair. A bona-fide entanglement monotone (also for mixed states)."""
    pt = partial_transpose(rho_sub, sites, A)
    pt = (pt + pt.conj().T) / 2.0                # PT of Hermitian is Hermitian
    ev = np.linalg.eigvalsh(pt)
    neg = float(np.sum(np.abs(ev[ev < 0])))      # N = (||.||_1 - 1)/2
    return float(np.log2(2.0 * neg + 1.0))


def boundary_negativity(rho, S=SUBJECT):
    """Entanglement across the S:E cut. The global state is pure, so this is a
       true entanglement measure of the boundary (log-negativity, bits)."""
    return log_negativity(rho, tuple(range(N)), S)


def integration_negativity(rho, S=SUBJECT):
    """Interior integration as the minimum across-partition ENTANGLEMENT of the
       (mixed) reduced state rho_S, by log-negativity."""
    rho_S = partial_trace(rho, S)
    best = np.inf
    for r in range(1, len(S) // 2 + 1):
        for A in itertools.combinations(S, r):
            best = min(best, log_negativity(rho_S, S, A))
    return float(best)


def metrics(J_S, J_E, g):
    psi, gap = ground_state(build_H(J_S, J_E, g))
    rho = np.outer(psi, psi.conj())
    return boundary_openness(rho), integration(rho), gap


def metrics_both(J_S, J_E, g):
    """Both currencies at once: (B_mi, Phi_mi, B_neg, Phi_neg)."""
    psi, _ = ground_state(build_H(J_S, J_E, g))
    rho = np.outer(psi, psi.conj())
    return (boundary_openness(rho), integration(rho),
            boundary_negativity(rho), integration_negativity(rho))


# ---- experiments ------------------------------------------------------------
def boundary_sweep(gs):
    out = np.array([metrics(1.0, 1.0, g)[:2] for g in gs])
    return out[:, 0], out[:, 1]


def boundary_sweep_both(gs):
    """B and Phi in both currencies: (B_mi, Phi_mi, B_neg, Phi_neg)."""
    out = np.array([metrics_both(1.0, 1.0, g) for g in gs])
    return out[:, 0], out[:, 1], out[:, 2], out[:, 3]


def phase_diagram(gs, Js):
    Bm = np.zeros((len(Js), len(gs)))
    Pm = np.zeros((len(Js), len(gs)))
    for a, js in enumerate(Js):
        for b, g in enumerate(gs):
            Bm[a, b], Pm[a, b], _ = metrics(js, 1.0, g)
    return Bm, Pm


# ---- figures (styled to match the article) ----------------------------------
def make_figures():
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    PAPER, INK, INK2, INK3 = "#F1F2F4", "#16181D", "#565B66", "#878C97"
    INDIGO, CORAL = "#3A2E7A", "#BE4F2C"
    plt.rcParams.update({
        "font.family": "serif", "font.size": 12,
        "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,
    })

    # Figure 1 -- boundary sweep, in two currencies (robustness overlay)
    gs = np.linspace(0.02, 2.6, 70)
    B, Phi, Bn, Phin = boundary_sweep_both(gs)
    fig, ax = plt.subplots(figsize=(8, 4.7))
    # solid: mutual information (the published measure)
    ax.plot(gs, B, color=INDIGO, lw=2.5, label=r"$O(S)=I(S{:}E)$   boundary openness (mutual info)")
    ax.plot(gs, Phi, color=CORAL, lw=2.5, label=r"$\Phi(S)$   integration (mutual info)")
    # dashed: a genuine entanglement monotone
    ax.plot(gs, Bn, color=INDIGO, lw=1.8, ls=(0, (4, 2)), alpha=0.85,
            label=r"$O(S)$   boundary openness (log-negativity)")
    ax.plot(gs, Phin, color=CORAL, lw=1.8, ls=(0, (4, 2)), alpha=0.85,
            label=r"$\Phi(S)$   integration (log-negativity)")
    ax.set_xlabel(r"boundary parameter  $g$   (coupling across the cut)")
    ax.set_ylabel("bits")
    ax.set_title("Dialing the boundary: the edge dissolves, the subject persists",
                 fontsize=13, pad=12)
    top = max(B.max(), Phi.max()) * 1.12
    ax.set_ylim(-0.1, top)
    ax.axvspan(0.02, 0.35, color=INK3, alpha=0.10)
    ax.axvspan(1.1, 1.9, color=CORAL, alpha=0.07)
    ax.text(0.16, top * 0.9, "sealed\nego", ha="center", color=INK2, fontsize=9)
    ax.text(1.5, top * 0.9, "boundary dissolving\n+ subject present\n= continuity",
            ha="center", color=CORAL, fontsize=9)
    ax.text(2.4, top * 0.9, "recombination", ha="center", color=INK3, fontsize=9)
    for s in ("top", "right"):
        ax.spines[s].set_visible(False)
    ax.legend(frameon=False, fontsize=8.5, loc="center left")
    ax.text(0.99, 0.02,
            "same window under a genuine entanglement monotone (dashed)\n"
            "— the trade-off is not an artifact of mutual information",
            transform=ax.transAxes, ha="right", va="bottom",
            color=INK2, fontsize=8)
    ax.grid(True, color=INK3, alpha=0.15, lw=0.6)
    fig.tight_layout()
    fig.savefig("toy-model-boundary-sweep.png", dpi=200)
    print("wrote toy-model-boundary-sweep.png")

    # Figure 2 -- phase diagram (with honest trade-off contours)
    gg = np.linspace(0.1, 2.5, 34)
    JJ = np.linspace(0.1, 2.5, 34)
    Bm, Pm = phase_diagram(gg, JJ)
    ext = [gg[0], gg[-1], JJ[0], JJ[-1]]
    B_THRESH, P_THRESH = 1.0, 1.0          # "boundary open" / "subject present"
    fig2, axes = plt.subplots(1, 2, figsize=(10.6, 4.8))
    for ax, M, name, cmap in ((axes[0], Bm, r"$O(S)$  boundary openness", "Purples"),
                              (axes[1], Pm, r"$\Phi(S)$  interior integration", "Oranges")):
        im = ax.imshow(M, origin="lower", extent=ext, aspect="auto", cmap=cmap)
        ax.set_xlabel(r"boundary coupling  $g$")
        ax.set_ylabel(r"interior coupling  $J_S$")
        ax.set_title(name, fontsize=12, pad=8)
        fig2.colorbar(im, ax=ax, fraction=0.046, pad=0.04, label="bits")
        # overlay BOTH frontiers on each panel: their overlap is the unity window
        ax.contour(gg, JJ, Bm, levels=[B_THRESH], colors=[INDIGO], linewidths=1.6)
        ax.contour(gg, JJ, Pm, levels=[P_THRESH], colors=[CORAL], linewidths=1.6,
                   linestyles="--")
    axes[1].annotate(r"$O>1$ (open)", xy=(0.62, 0.18), color=INDIGO, fontsize=8)
    axes[1].annotate(r"$\Phi>1$ (subject)", xy=(0.18, 1.7), color=CORAL, fontsize=8)
    axes[1].annotate("continuity\nwindow", xy=(1.15, 2.05), color=INK, fontsize=9,
                     ha="center")
    fig2.suptitle("The phase space, computed:  boundary openness and integration "
                  "trade off (a bounded correlation budget)", fontsize=12, y=1.03)
    fig2.tight_layout()
    fig2.savefig("toy-model-phase-diagram.png", dpi=200, bbox_inches="tight")
    print("wrote toy-model-phase-diagram.png")


def main():
    print("Boundary sweep  (J_S = J_E = 1):")
    print(f"{'g':>6} {'O(S)=I(S:E)':>14} {'Phi(S)':>10} {'gap':>8}")
    for g in [0.05, 0.3, 0.6, 1.0, 1.5, 2.0, 2.5]:
        b, p, gap = metrics(1.0, 1.0, g)
        print(f"{g:6.2f} {b:14.4f} {p:10.4f} {gap:8.4f}")
    print("\nO(S) climbs from ~0 (sealed) as the boundary opens, while Phi(S)")
    print("stays high through the unity regime -- a subject persists.\n")
    make_figures()


if __name__ == "__main__":
    main()
