#!/usr/bin/env python3
"""
The Supervisor's Trap — reference implementation.

Companion artifact to the working paper "The Supervisor's Trap"
(stokes, 2026). Reproduces the model's results so they can be checked
rather than trusted:

  1. Bistability of the fast (arousal) subsystem under a sweep of the
     maintenance rate u, with hysteresis.                    [panel A]
  2. The closed-form fold condition  g(A) = A(1-A) g'(A)  and the fold
     locations u* = lam*A / [g(A)(1-A)], verified against the sweep.
  3. The cusp: below a critical recruitment steepness the two folds
     merge and the trap disappears (monostable).
  4. The full two-timescale system:
       gated interlock + engagement feedback  -> healthy equilibrium
       gate without engagement feedback       -> fail-safe-but-dead freeze
       ungated (G = C, no arousal gate)       -> boom-bust incident cycle
                                                             [panels B, C]

PROVENANCE NOTE (honesty is load-bearing): the paper's original
parameterization was not preserved. The parameter values below are
RE-DERIVED to exhibit the same structure; where the paper quotes exact
numbers (folds at A* ~ 0.32 / 0.56, window u in [1.54, 7.79], cusp
slope ~ 7.1) this script reports what THIS parameterization yields, and
prints both so the comparison is explicit. Structural claims (bistability,
fold condition, cusp existence, freeze vs. cycle) are reproduced exactly;
the quoted numerical values should be treated as illustrative until the
original parameterization is re-published.

Model (as printed in the paper):
    dA/dt = -lam*(1 + kap*G)*A + u_eff * g(A) * (1 - A)
    dC/dt = eps * [ alf*(1-C)*A - bet*C*R ]
    G     = C * h(A)              (the authority interlock)
    u_eff = u0 + u1*(1 - G)       (the engagement feedback)
with recruitment gain g(A) threshold-shaped and R the risk run
unsupervised, modeled here as G*(1-h(A)): authority exercised while the
supervisor is below the readiness gate. Incidents: latent risk
accumulates at rate rho*R, dissipates at rate del when not compounded,
and surfaces (knocking credence down) when it crosses a threshold.

Requires numpy only. Run:  python3 supervisors-trap-model.py
"""

import numpy as np

# ----------------------------------------------------------------------
# Reference parameterization (re-derived; see provenance note)
# ----------------------------------------------------------------------
LAM   = 1.0     # arousal decay rate
S     = 12.0    # recruitment steepness (the cusp parameter)
THG   = 0.45    # recruitment threshold
KAP   = 0.3     # extra arousal decay under granted authority
EPS   = 0.02    # credence timescale (slow)
ALF   = 1.0     # credence growth under supervised operation
BET   = 0.5     # credence erosion per unit unsupervised risk
SH    = 10.0    # interlock gate steepness
THH   = 0.30    # interlock gate threshold
U0    = 0.5     # baseline maintenance drip
U1    = 9.0     # engagement feedback strength (ceiling must clear the upper fold)
RHO   = 0.15    # latent-risk accumulation rate
RC    = 1.0     # incident threshold
DEL   = 0.05    # latent-risk dissipation when not compounded
DC    = 0.35    # credence knocked off per incident


def sig(x):
    return 1.0 / (1.0 + np.exp(-x))


def g(A, s=S, th=THG):
    """Recruitment gain: threshold-shaped, g(0) ~ 0."""
    return sig(s * (A - th))


def gp(A, s=S, th=THG):
    v = g(A, s, th)
    return s * v * (1.0 - v)


def h(A, s=SH, th=THH):
    """Interlock gate: closes toward zero when the supervisor is cold."""
    return sig(s * (A - th))


# ----------------------------------------------------------------------
# 1. Fast subsystem: equilibria and bistability sweep
#    dA/dt = -lam*A + u*g(A)*(1-A)   (C frozen; reduced form from paper)
# ----------------------------------------------------------------------
def fast_equilibria(u, lam=LAM, s=S, th=THG, n=200001):
    A = np.linspace(1e-9, 1 - 1e-9, n)
    f = -lam * A + u * g(A, s, th) * (1.0 - A)
    idx = np.where(np.sign(f[:-1]) != np.sign(f[1:]))[0]
    roots = []
    for i in idx:
        a, b = A[i], A[i + 1]
        for _ in range(60):  # bisection
            m = 0.5 * (a + b)
            fm = -lam * m + u * g(m, s, th) * (1.0 - m)
            fa = -lam * a + u * g(a, s, th) * (1.0 - a)
            if np.sign(fm) == np.sign(fa):
                a = m
            else:
                b = m
        roots.append(0.5 * (a + b))
    return roots


def bistable_window(lam=LAM, s=S, th=THG):
    """Numerical sweep: range of u with three equilibria."""
    us = np.linspace(0.05, 30.0, 6000)
    lo, hi = None, None
    for u in us:
        if len(fast_equilibria(u, lam, s, th)) >= 3:
            if lo is None:
                lo = u
            hi = u
    return lo, hi


def folds_closed_form(lam=LAM, s=S, th=THG):
    """Roots of g(A) = A(1-A) g'(A), then u* = lam*A/[g(A)(1-A)]."""
    A = np.linspace(1e-6, 1 - 1e-6, 400001)
    q = g(A, s, th) - A * (1.0 - A) * gp(A, s, th)
    idx = np.where(np.sign(q[:-1]) != np.sign(q[1:]))[0]
    out = []
    for i in idx:
        Astar = A[i]
        ustar = lam * Astar / (g(Astar, s, th) * (1.0 - Astar))
        out.append((Astar, ustar))
    return out


def cusp_steepness(lam=LAM, th=THG):
    """Critical recruitment steepness below which the trap disappears."""
    lo, hi = 1.0, 30.0
    for _ in range(40):
        mid = 0.5 * (lo + hi)
        if len(folds_closed_form(lam, mid, th)) >= 2:
            hi = mid
        else:
            lo = mid
    return 0.5 * (lo + hi)


# ----------------------------------------------------------------------
# 2. Full two-timescale system with incidents
# ----------------------------------------------------------------------
def run_full(T=4000.0, dt=0.01, gated=True, engagement=True,
             A0=0.8, C0=0.2, seed_u=None):
    n = int(T / dt)
    A, C, r = A0, C0, 0.0
    incidents = 0
    traj = np.empty((n, 4))
    for i in range(n):
        G = C * (h(A) if gated else 1.0)
        u_eff = (U0 + U1 * (1.0 - G)) if engagement else (
            U0 if seed_u is None else seed_u)
        R = G * (1.0 - h(A))
        dA = -LAM * (1.0 + KAP * G) * A + u_eff * g(A) * (1.0 - A)
        dC = EPS * (ALF * (1.0 - C) * A - BET * C * R)
        dr = RHO * R - DEL * r
        A = min(1.0, max(0.0, A + dt * dA))
        C = min(1.0, max(0.0, C + dt * dC))
        r += dt * dr
        if r >= RC:
            incidents += 1
            C = max(0.0, C - DC)
            r = 0.0
        traj[i] = (A, C, G, r)
    return traj, incidents


def main():
    print(__doc__.split("Model (as printed")[0])
    print("=" * 70)
    print("1. FAST SUBSYSTEM: bistability sweep (numerical)")
    lo, hi = bistable_window()
    print(f"   bistable window (sweep):      u in [{lo:.2f}, {hi:.2f}]")
    print(f"   paper (its parameterization): u in [1.57, 7.77]")
    mid_u = 0.5 * (lo + hi)
    eq = fast_equilibria(mid_u)
    print(f"   at u = {mid_u:.2f}: equilibria at A = "
          + ", ".join(f"{a:.3f}" for a in eq)
          + "   (complacent / unstable threshold / engaged)")

    print()
    print("2. CLOSED-FORM FOLDS:  g(A) = A(1-A) g'(A),  u* = lam*A/[g(1-A)]")
    folds = folds_closed_form()
    for Astar, ustar in folds:
        print(f"   fold at A* = {Astar:.3f},  u* = {ustar:.2f}")
    print(f"   paper (its parameterization): A* ~ 0.32 (u* ~ 7.79), "
          f"A* ~ 0.56 (u* ~ 1.54)")
    if len(folds) >= 2:
        w = sorted(u for _, u in folds)
        agree = abs(w[0] - lo) < 0.05 and abs(w[-1] - hi) < 0.05
        print(f"   closed form vs sweep: [{w[0]:.2f}, {w[-1]:.2f}] vs "
              f"[{lo:.2f}, {hi:.2f}]  ->  {'AGREE' if agree else 'DISAGREE'}")

    print()
    print("3. CUSP: critical recruitment steepness (folds annihilate below)")
    sc = cusp_steepness()
    print(f"   s_crit = {sc:.2f}   (paper, its parameterization: ~7.1)")
    print(f"   check: s = {sc - 1:.1f} -> {len(folds_closed_form(s=sc - 1))} "
          f"folds;  s = {sc + 1:.1f} -> {len(folds_closed_form(s=sc + 1))} folds")

    print()
    print("4. FULL SYSTEM (two timescales, incidents)")
    tr, inc = run_full(gated=True, engagement=True)
    A, C, G = tr[-1, 0], tr[-1, 1], tr[-1, 2]
    print(f"   gated + engagement feedback : A={A:.2f} C={C:.2f} G={G:.2f} "
          f"incidents={inc}   <- healthy interior equilibrium")
    tr, inc = run_full(gated=True, engagement=False, seed_u=U0)
    A, C, G = tr[-1, 0], tr[-1, 1], tr[-1, 2]
    print(f"   gated, no engagement fb     : A={A:.2f} C={C:.2f} G={G:.2f} "
          f"incidents={inc}   <- fail-safe but dead (gate shut)")
    tr, inc = run_full(gated=False, engagement=True)
    A, C, G = tr[-1, 0], tr[-1, 1], tr[-1, 2]
    # detect cycling: spread of C over the last quarter of the run
    tail = tr[-len(tr) // 2:, 1]
    print(f"   UNGATED (G = C)             : A={A:.2f} C={C:.2f} G={G:.2f} "
          f"incidents={inc}   <- boom-bust cycle "
          f"(C oscillates {tail.min():.2f}..{tail.max():.2f})")
    print()
    print("Claims reproduced structurally: bistability, fold condition,")
    print("cusp existence, healthy attractor, freeze, incident cycle.")
    print("Exact numbers differ where the original parameterization was")
    print("not preserved — see PROVENANCE NOTE in the header.")


if __name__ == "__main__":
    main()
