r"""
================================================================================
A METASTABILITY MODEL OF THE ACCESSORY SLEEP STATE
Supplementary methods and reference implementation
================================================================================

PURPOSE
-------
This script formalises and tests a single argument: that quiet wakefulness is
not a cheaper substitute for sleep but a *leaky* one, and that the withdrawal of
responsiveness which defines sleep is the mechanism that makes sustained
inactivity stable, rather than an incidental cost of it. It answers the standard
objection to an "adaptive inactivity" / accessory account of sleep — "if the
point is just to stay inactive, why not stay quietly awake instead?" — and, as a
by-product, recovers the ecological titration of sleep depth from a single
parameter. It generates the four-panel figure used in the manuscript and prints
the key quantities and the closed-form optimum.

--------------------------------------------------------------------------------
1. STATE SPACE AND THE CORE STOCHASTIC MODEL
--------------------------------------------------------------------------------
An animal in its rest window can be in one of two behavioural macrostates:

    REST (inactive, sensory-disengaged)  <-->  ACTIVE WAKEFULNESS (engaged)

We treat departures from REST as a memoryless (Poisson) process with a constant
hazard rate lambda [events/min]. The memorylessness is a deliberate
simplification: it assumes that, conditional on the current arousal threshold,
the chance of being tipped into wakefulness in the next instant does not depend
on how long the animal has already rested. Under this assumption:

    Survival of continuous rest:     P(t) = exp(-lambda * t)
    Prob. of >=1 involuntary arousal
    over a window of length T:        1 - exp(-lambda * T)
    Mean uninterrupted bout length:   1 / lambda

These give Figure panel A and the "quiet wakefulness almost surely fails over a
night" result: for QW-magnitude hazards, 1 - exp(-lambda*T) -> 1 well within an
hour.

--------------------------------------------------------------------------------
2. THE AROUSAL THRESHOLD AS THE CONTROL PARAMETER
--------------------------------------------------------------------------------
Sleep depth is represented by a single scalar, the arousal threshold theta >= 0.
Stimuli impinge on the animal and each one triggers a transition to wakefulness
only if its salience exceeds theta. Two classes of stimulus are distinguished:

  (a) BENIGN / IRRELEVANT stimuli (rustling leaves, conspecific movement, etc.)
      arrive at rate a [events/min]. Their salience is exponentially distributed,
      so the fraction that breach threshold theta is

          Fbar_benign(theta) = exp(-K * theta).

      Every breach is a "false alarm" that ejects the animal from rest. Hence the
      metastability hazard is

          lambda(theta) = a * Fbar_benign(theta) = a * exp(-K * theta).

      Raising theta (deeper sleep) lowers lambda; quiet wakefulness is the
      low-theta / high-lambda limit.

  (b) GENUINE THREATS (approaching predators) arrive at rate beta [events/min].
      They are MORE SALIENT than benign stimuli: a looming predator breaches a
      higher threshold than a rustling leaf. We encode salience with a factor
      s > 1, so the probability that a genuine threat is salient enough to wake
      the animal in time is

          Fbar_threat(theta) = exp(-K * theta / s),

      and the probability of FAILING to respond in time is

          p_miss(theta) = 1 - exp(-K * theta / s).

The assumption s > 1 is the single load-bearing premise of the model: it creates
a band of raised thresholds that suppress false-alarm arousals (killing lambda)
while still admitting genuine threats. That band is sleep. No result below
depends on the exact functional forms (exponential CCDFs are chosen for
analytic transparency); they depend only on threats being more salient than
benign stimuli, i.e. on the threat CCDF decaying more slowly in theta than the
benign CCDF.

--------------------------------------------------------------------------------
3. OCCUPANCY OF THE REST STATE (panel B)
--------------------------------------------------------------------------------
The survival function describes an *uninterrupted* bout. Across a whole window an
animal that is ejected simply resettles and tries again, so the biologically
relevant quantity is the long-run fraction of the window actually spent
inactive. Modelling REST<->ACTIVE as a two-state continuous-time Markov chain
with rest-exit rate lambda and a resettling (active->rest) rate mu = 1/tau, the
stationary occupancy of REST is

    occupancy = (1/lambda) / (1/lambda + tau) = mean_bout / (mean_bout + tau),

where tau [min] is the mean time to calm down and re-enter rest after an
involuntary arousal. Quiet wakefulness (short mean_bout) captures only a small
fraction of the window; sleep (long mean_bout) captures almost all of it. This
is the accessory benefit the animal is trying to bank, and only sleep banks it.

--------------------------------------------------------------------------------
4. THE FITNESS TRADE-OFF AND THE INTERIOR OPTIMUM (panel C)
--------------------------------------------------------------------------------
Deeper sleep is not free: the same threshold that filters benign stimuli also
filters genuine threats. We write the expected cost over the window as the sum
of two terms:

    INSTABILITY cost  (false-alarm arousals; exposure, energy, conspicuousness):
        C_inst(theta)  = c_FA   * a    * Fbar_benign(theta) * T
                       = c_FA   * a    * exp(-K*theta)      * T      (decreases in theta)

    PREDATION cost    (missed genuine threats):
        C_pred(theta)  = c_pred * beta * p_miss(theta)      * T
                       = c_pred * beta * (1-exp(-K*theta/s)) * T     (increases in theta)

    TOTAL:  C(theta) = C_inst(theta) + C_pred(theta)

Because C_inst falls and C_pred rises with theta, and because s > 1 makes them
fall/rise at different rates, C(theta) has an interior minimum. Setting
dC/dtheta = 0 gives the closed form

        theta*  =  ln( s * c_FA * a / (c_pred * beta) )  /  ( K * (1 - 1/s) )

which is positive (an interior optimum exists) iff

        s * c_FA * a  >  c_pred * beta.

The shallow limit (quiet wakefulness) is strictly dominated: it pays the full
instability cost for no compensating gain in safety. The deep limit (bottomless
sleep) is dominated too: it pays the full predation cost. Sleep is the optimum
between them.

--------------------------------------------------------------------------------
5. ECOLOGICAL TITRATION (panel D)
--------------------------------------------------------------------------------
Differentiating the closed form,

        d(theta*)/d(beta)  =  -1 / ( beta * K * (1 - 1/s) )  <  0,

so the optimal threshold — and hence mean bout length and sleep depth — FALLS as
predation pressure beta rises. Safe / refuge conditions (small beta) select deep,
consolidated sleep; dangerous conditions (large beta) select light, arousable
sleep, i.e. the empirical regime of one-eye-open and unihemispheric sleep. This
is the same "buffer" knob invoked elsewhere in the manuscript, here derived
rather than asserted.

--------------------------------------------------------------------------------
6. FALSIFIABLE PREDICTIONS AND MAPPING TO DATA
--------------------------------------------------------------------------------
The model is stated in measurable quantities:
  * lambda_QW (quiet-wake -> active-wake transition hazard) and lambda_S (sleep
    -> wake hazard) for MATCHED stimuli are directly estimable from high-
    throughput behavioural tracking (e.g. ethoscope data in Drosophila). The
    model predicts lambda_QW >> lambda_S.
  * Sleep depth / mean bout length should titrate against ecological risk as in
    panel D.
  * (Within-night extension, not plotted) If the important components are
    discharged first, the optimal threshold should not be constant but should
    start high and relax across the night as the core is completed — i.e. the
    observed front-loading and dissipation of deep (slow-wave) sleep is the
    trajectory of theta(t) predicted by discharging the useful/vital core before
    the accessory buffer.

--------------------------------------------------------------------------------
7. ASSUMPTIONS AND LIMITATIONS
--------------------------------------------------------------------------------
  * Poisson (memoryless) exits from rest: ignores refractoriness and sleep-cycle
    structure; adequate for the qualitative occupancy/survival results.
  * A single scalar arousal threshold stands in for the full neurophysiology of
    sleep depth; the mechanistic counterpart is the mutually-inhibitory sleep-
    wake "flip-flop" switch (Saper et al., 2010), for which theta indexes the
    depth of the resting basin.
  * Exponential stimulus-salience CCDFs are chosen for closed-form transparency;
    the qualitative conclusions require only s > 1 (threats more salient than
    benign stimuli).
  * Parameters below are ILLUSTRATIVE, chosen to place the optimum in a
    biologically sensible range (mean bouts of tens of minutes) and to make QW
    clearly leaky; they are not fitted to a particular species. Absolute cost
    magnitudes are arbitrary units; only their ratio and the salience s matter
    for the location of the optimum.

--------------------------------------------------------------------------------
8. PARAMETERS (see values in the PARAMETERS block below)
--------------------------------------------------------------------------------
  T       [min]        length of the rest window (one night of unsafe hours)
  a       [1/min]      arrival rate of benign/irrelevant stimuli
  K       [1/theta]    sensitivity of the arousal hazard to the threshold
  s       [-]          salience of genuine threats relative to benign stimuli (>1)
  tau     [min]        mean resettling time after an involuntary arousal (mu=1/tau)
  c_FA    [cost]       cost per false-alarm arousal
  c_pred  [cost]       cost per missed genuine threat
  beta0   [1/min]      baseline genuine-threat arrival rate (predation pressure)
  theta_QW/light/deep  representative thresholds for the three plotted regimes

OUTPUTS
-------
  * stdout: per-regime hazards, mean bouts, P(>=1 arousal/night); numeric and
    closed-form optimum; QW-vs-optimum cost ratio.
  * metastability_figure.png / .pdf : the four-panel figure.

RUN
---
  python3 metastability_model.py     (requires numpy, matplotlib)
================================================================================
"""

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import font_manager
from matplotlib.ticker import MaxNLocator

# ---- palette (matches the manuscript) ----
TAN   = "#C9B79C"   # accessory
TEAL  = "#2E7C7C"   # useful / light sleep
NAVY  = "#1E2233"   # vital / deep sleep
RUST  = "#B5493F"   # quiet wakefulness (the unstable, "bad" state)
GREY  = "#6b6b6b"

plt.rcParams.update({
    # keep text as text in vector output so it can be edited downstream
    "svg.fonttype": "none",
    "pdf.fonttype": 42,
    "font.size": 10.5,
    "font.family": "sans-serif",
    "axes.spines.top": False,
    "axes.spines.right": False,
    "axes.linewidth": 0.9,
    "xtick.direction": "out",
    "ytick.direction": "out",
    "axes.titlesize": 11.5,
    "axes.titleweight": "bold",
    "figure.dpi": 130,
})

# =====================================================================
# MODEL
# =====================================================================
# Rest window (one night of unsafe hours)
T = 480.0                      # minutes (8 h)
t = np.linspace(0, T, 1000)
tA = np.linspace(0, 60, 600)   # short window for the metastability panel

# Benign (irrelevant) stimuli arrive at rate a; a stimulus triggers arousal
# iff its intensity exceeds the arousal threshold theta.
a = 1.0                        # benign stimuli per minute (theta=0 -> mean bout 1 min)
# Complementary CDF of benign-stimulus intensity: P(intensity > theta) = exp(-theta)
K = 0.8                        # threshold sensitivity of the arousal hazard
def Fbar_benign(theta):        # fraction of benign stimuli that breach threshold
    return np.exp(-K*theta)

# False-alarm arousal hazard = metastability hazard of the resting state
def lam(theta):
    return a * Fbar_benign(theta)          # per minute

# Genuine threats (predators) are more SALIENT: their intensity distribution
# is heavier, so they still breach a raised threshold. s>1 = threat salience.
s = 4.0
def Fbar_threat(theta):        # P(a genuine threat wakes the animal in time)
    return np.exp(-K*theta / s)
def p_miss(theta):             # P(fail to respond to a genuine threat)
    return 1.0 - Fbar_threat(theta)

# =====================================================================
# MEASURED PARAMETERS  (Joyce et al. 2024 HMM; hmm_bout_durations.csv)
# =====================================================================
# 950,651 Viterbi-decoded bouts across 8 datasets, with a right-censoring flag.
# The two per-species melanogaster models are degenerate (deep-sleep dwell
# collapses to ~1 min) and are excluded; the canonical get_HMM("M") covers that
# lineage. These bouts are the UNCIRCULAR dataset: they are decoded state
# durations, not the HMM's own geometric dwell assumption.
import pandas as pd
BOUTS = "hmm_bout_durations.csv"
STATES = ["Deep sleep", "Light sleep", "Quiet awake", "Full awake"]
try:
    _b = pd.read_csv(BOUTS)
    _b = _b[~_b.species.isin(["D.mel", "wild_d.mel"])]
    _u = _b[~_b.right_censored]
    meas_lambda = {st: 1.0/_u.loc[_u.canonical_label == st, "duration_min"].mean()
                   for st in STATES}
    meas_mean   = {st: 1.0/meas_lambda[st] for st in STATES}
    tau_meas    = meas_mean["Full awake"]          # mean time to resettle after ejection
    have_bouts  = True
    print("\nMEASURED from %s (%d bouts, %.1f%% censored):" %
          (BOUTS, len(_b), 100*_b.right_censored.mean()))
    for st in STATES:
        print("  %-12s mean bout %6.2f min -> lambda = %.4f/min"
              % (st, meas_mean[st], meas_lambda[st]))
except Exception as _e:
    print("\n(bout data not found: %s -- falling back to illustrative parameters)" % _e)
    have_bouts = False
    tau_meas = 6.0

def theta_of(lmbda):
    """Invert lambda(theta) = a*exp(-K*theta) to place a measured state on the axis."""
    return -np.log(lmbda / a) / K

# ---- representative states (by arousal threshold) ----
if have_bouts:
    theta_QW    = theta_of(meas_lambda["Quiet awake"])   # ~0.98
    theta_light = theta_of(meas_lambda["Light sleep"])   # ~1.75
    theta_deep  = theta_of(meas_lambda["Deep sleep"])    # ~4.12
else:
    theta_QW, theta_light, theta_deep = 1.0, 3.0, 5.0

for name, th in [("quiet wakefulness", theta_QW),
                 ("light sleep", theta_light),
                 ("deep sleep", theta_deep)]:
    l = lam(th)
    print(f"{name:18s}: theta={th:.1f}  lambda={l:.4f}/min  "
          f"mean bout={1/l:6.1f} min  "
          f"P(>=1 arousal over 8h)={1-np.exp(-l*T):.4f}")

# =====================================================================
# COST OPTIMISATION
# =====================================================================
# Expected cost over the window as a function of threshold:
#   instability/false-alarm cost  (~ number of spurious arousals)
# + predation cost               (~ missed genuine threats)
c_FA   = 1.0                   # cost per spurious arousal (exposure, energy, conspicuousness)
c_pred = 1000.0                # cost per missed threat (predation is very costly)
beta0  = 0.0004                # baseline genuine-threat arrival rate (per min)

theta = np.linspace(0.0, 8.0, 800)

def cost_instability(theta, a=a):
    return c_FA * a * Fbar_benign(theta) * T
def cost_predation(theta, beta):
    return c_pred * beta * p_miss(theta) * T
def cost_total(theta, beta):
    return cost_instability(theta) + cost_predation(theta, beta)

C_inst = cost_instability(theta)
C_pred = cost_predation(theta, beta0)
C_tot  = C_inst + C_pred
theta_star = theta[np.argmin(C_tot)]                       # numeric optimum
# closed form: theta* = ln( s*c_FA*a / (c_pred*beta) ) / ( K*(1 - 1/s) )
interior = (s * c_FA * a) > (c_pred * beta0)
theta_star_cf = (np.log(s * c_FA * a / (c_pred * beta0)) / (K * (1 - 1 / s))
                 if interior else 0.0)
print(f"\nInterior optimum exists (s*c_FA*a > c_pred*beta0)? {interior}")
print(f"Optimal threshold theta* = {theta_star:.2f} (numeric) / "
      f"{theta_star_cf:.2f} (closed form)  "
      f"-> lambda* = {lam(theta_star):.4f}/min, mean bout = {1/lam(theta_star):.1f} min")
print(f"Cost at QW (theta={theta_QW}) = {cost_total(theta_QW, beta0):.1f}; "
      f"cost at optimum = {cost_total(theta_star, beta0):.1f}; "
      f"ratio = {cost_total(theta_QW, beta0)/cost_total(theta_star, beta0):.1f}x")

# =====================================================================
# TITRATION: optimum vs predation pressure
# =====================================================================
betas = np.linspace(0.0001, 0.004, 400)
theta_opt = np.array([theta[np.argmin(cost_total(theta, b))] for b in betas])
bout_opt  = 1.0 / lam(theta_opt)

# =====================================================================
# WITHIN-NIGHT TRAJECTORY (panel E)  +  COMPARISON TO OBSERVED DATA
# =====================================================================
# If the useful/vital core is discharged first, the optimal threshold is not
# constant across the night: it starts high (protecting the core, deepest sleep)
# and relaxes toward an accessory-buffer floor as the core is spent. This is the
# two-process dissipation of sleep pressure (Process S) restated as theta(t).
theta0_night = 6.0             # initial (deep) threshold at sleep onset
theta_floor  = 2.5             # asymptotic light-sleep / accessory-buffer level
tau_night    = 120.0           # min; decay constant of the core discharge
t_night      = np.linspace(0, T, 1000)
theta_traj   = theta_floor + (theta0_night - theta_floor) * np.exp(-t_night / tau_night)

# ---- Observed deep-sleep timing (Joyce et al. 2024 HMM data) ------------------
# Data file: sleep_state_24h_summary.csv, columns species,hour(=ZT),n_flies and
# HMM state means/SEMs s0..s3 (which sum to 1). Validation against a robotically
# measured arousal threshold established the ordering by depth/arousability:
#   s0 = DEEP sleep, s1 = light sleep, s2 = quiet wake, s3 = active wake
# (asleep = s0 + s1; s0 is the smaller sleep state and concentrates early night,
#  the arousal-threshold-validated "deep sleep" of Joyce et al. 2024). The night
# runs ZT12 (lights-off) -> ZT24; deep sleep should be front-loaded within it.
import os, csv
NIGHT_ZT = list(range(12, 24))          # ZT12..ZT23 = the 12 h night
def load_joyce_deep_sleep(path="sleep_state_24h_summary.csv"):
    if not os.path.exists(path):
        return None
    per_sp, sem_sp = {}, {}
    for r in csv.DictReader(open(path)):
        sp = r["species"]
        per_sp.setdefault(sp, {})[int(float(r["hour"]))] = float(r["s0_mean"])
        sem_sp.setdefault(sp, {})[int(float(r["hour"]))] = float(r["s0_sem"])
    species = sorted(per_sp)
    deep = {sp: np.array([per_sp[sp][h] for h in NIGHT_ZT]) for sp in species}
    sem  = {sp: np.array([sem_sp[sp][h] for h in NIGHT_ZT]) for sp in species}
    return species, deep, sem

joyce = load_joyce_deep_sleep()
obs_is_real = joyce is not None
if obs_is_real:
    species_list, deep_by_sp, sem_by_sp = joyce
    deep_stack = np.vstack([deep_by_sp[sp] for sp in species_list])
    deep_mean  = deep_stack.mean(axis=0)
    deep_msem  = deep_stack.std(axis=0, ddof=1) / np.sqrt(deep_stack.shape[0])
    zt_night   = np.array(NIGHT_ZT)
    # Species that do NOT front-load (documented exceptions: peak in the 2nd half
    # of the night) are detected automatically and highlighted.
    half = len(NIGHT_ZT) // 2
    exceptions = [sp for sp in species_list
                  if np.argmax(deep_by_sp[sp]) >= half]
    print(f"\nPanel E: using REAL Joyce et al. data ({len(species_list)} species).")
    print(f"  cross-species deep-sleep peaks at ZT{zt_night[np.argmax(deep_mean)]}, "
          f"declining thereafter; non-front-loading species: {exceptions}")
else:
    print("\nPanel E: data file not found; skipping real overlay.")

# =====================================================================
# FIGURE
# =====================================================================
fig = plt.figure(figsize=(9.8, 13.2))
gs = fig.add_gridspec(4, 2, height_ratios=[1.0, 1.0, 0.95, 0.72],
                      hspace=0.52, wspace=0.30,
                      left=0.085, right=0.975, top=0.972, bottom=0.045)
axA = fig.add_subplot(gs[0, 0])
axB = fig.add_subplot(gs[0, 1])
axC = fig.add_subplot(gs[1, 0])
axD = fig.add_subplot(gs[1, 1])
axE = fig.add_subplot(gs[2, :])
axF = fig.add_subplot(gs[3, :])

# ---- Panel A: measured bout survival vs the model ----
def kaplan_meier(dur, censored):
    """S(t) with right-censoring. dur in min, censored a boolean array."""
    dur = np.asarray(dur, float); censored = np.asarray(censored, bool)
    ts = np.unique(dur); out_t, out_S, S = [0.0], [1.0], 1.0
    for ti in ts:
        at_risk = (dur >= ti).sum()
        deaths  = ((dur == ti) & (~censored)).sum()
        if at_risk > 0 and deaths > 0:
            S *= 1.0 - deaths / at_risk
        out_t.append(ti); out_S.append(S)
    return np.array(out_t), np.array(out_S)

_PANEL_A = [("Quiet awake", RUST, "quiet wakefulness"),
            ("Light sleep", TEAL, "light sleep"),
            ("Deep sleep",  NAVY, "deep sleep")]
if have_bouts:
    for st, col, lab in _PANEL_A:
        g = _b[_b.canonical_label == st]
        kt, kS = kaplan_meier(g.duration_min.values, g.right_censored.values)
        axA.step(kt, kS, where="post", color=col, lw=2.4,
                 label=f"{lab}  ({meas_mean[st]:.1f} min)")
        axA.plot(tA, np.exp(-meas_lambda[st]*tA), color=col, lw=1.1, ls=":", alpha=0.85)
    axA.plot([], [], color=GREY, lw=1.1, ls=":", label="exponential, same mean")
    axA.set_title("Rest is metastable: measured bout survival", loc="left")
else:
    for th, col, lab in [(theta_QW, RUST, "quiet wakefulness"),
                         (theta_light, TEAL, "light sleep"),
                         (theta_deep, NAVY, "deep sleep")]:
        axA.plot(tA, np.exp(-lam(th)*tA), color=col, lw=2.4, label=lab)
    axA.set_title("Rest is metastable: quiet wakefulness leaks", loc="left")
axA.set_xlabel("time holding the state continuously (min)")
axA.set_ylabel("P(still continuously in state)")
axA.set_xlim(0, 60); axA.set_ylim(-0.02, 1.02)
axA.legend(frameon=False, fontsize=8.6, loc="upper right")
axA.annotate("quiet wakefulness never exceeded\n21 min in 237,386 bouts",
             xy=(4.5, 0.07), xytext=(21, 0.30), color=RUST, fontsize=8.4,
             arrowprops=dict(arrowstyle="->", color=RUST, lw=1.2))
axA.text(31, 0.62, "deep sleep is not memoryless:\nno bout shorter than ~12 min",
         color=NAVY, fontsize=8.0, va="center")

# ---- Panel B: fraction of the window actually spent at rest ----
# (axB defined above)
lam_grid = a*np.exp(-K*np.linspace(0.2, 7.5, 500))    # hazards from thresholds
mean_bout = 1.0/lam_grid
# expected fraction of window at rest before first ejection is not the point;
# use long-run occupancy with a fixed re-settling time tau after each ejection:
tau = tau_meas   # MEASURED mean full-awake bout: the true cost of one ejection
occupancy = (1.0/lam_grid) / (1.0/lam_grid + tau)      # mu/(lam+mu), mu=1/tau
axB.plot(mean_bout, occupancy, color=NAVY, lw=2.4)
for th, col, lab, dx in [(theta_QW, RUST, "quiet\nwakefulness", 0),
                         (theta_light, TEAL, "light sleep", 0),
                         (theta_deep, NAVY, "deep sleep", 0)]:
    mb = 1/lam(th)
    oc = mb/(mb+tau)
    axB.scatter([mb], [oc], color=col, s=55, zorder=5, edgecolor="white", lw=1)
    dy = 0.075 if lab.startswith("quiet") else (0.055 if lab == "light sleep" else -0.09)
    axB.annotate(lab, xy=(mb, oc), xytext=(mb*1.18, oc + dy),
                 fontsize=8.6, color=col, ha="left")
axB.set_xscale("log")
axB.set_xlabel("mean uninterrupted bout length (min, log)")
axB.set_ylabel("fraction of window spent inactive")
axB.set_title("Only sleep captures the accessory benefit", loc="left")
axB.set_ylim(0, 1.02)
axB.fill_between(mean_bout, occupancy, 0, color=TAN, alpha=0.18)

# ---- Panel C: cost optimisation (why not QW, why not infinitely deep) ----
# (axC defined above)
axC.plot(theta, C_inst, color=RUST, lw=2.0, ls="--", label="instability cost")
axC.plot(theta, C_pred, color=GREY, lw=2.0, ls=":", label="predation cost")
axC.plot(theta, C_tot, color=NAVY, lw=2.6, label="total expected cost")
axC.axvline(theta_star, color=TEAL, lw=1.6)
axC.scatter([theta_star], [C_tot.min()], color=TEAL, s=60, zorder=6,
            edgecolor="white", lw=1)
axC.annotate("optimum\n(= sleep)", xy=(theta_star, C_tot.min()),
             xytext=(theta_star+0.5, C_tot.min()+ C_tot.max()*0.28),
             color=TEAL, fontsize=9,
             arrowprops=dict(arrowstyle="->", color=TEAL, lw=1.2))
axC.axvspan(0, 2.0, color=RUST, alpha=0.08)
axC.text(0.95, C_tot.max()*0.9, "quiet\nwakefulness\n(dominated)",
         color=RUST, fontsize=8.4, ha="center", va="top")
axC.set_xlabel(r"arousal threshold $\theta$  (shallow $\rightarrow$ deep)")
axC.set_ylabel("expected cost over window")
axC.set_title("The optimal arousal threshold is interior", loc="left")
axC.set_xlim(0, 8)
axC.set_ylim(0, C_tot.max()*1.05)
axC.legend(frameon=False, fontsize=8.8, loc="upper center")

# ---- Panel D: titration with predation pressure (comparative prediction) ----
# (axD defined above)
axD_x = betas * 1e3                      # plot on a x10^-3 scale: ticks 1,2,3,4
axD.plot(axD_x, bout_opt, color=NAVY, lw=2.6)
axD.set_xlabel("predation pressure  (genuine-threat rate $\\beta$, $\\times 10^{-3}$/min)")
axD.set_ylabel("optimal mean bout length (min)")
axD.set_title("Sleep depth titrates to ecological risk", loc="left")
axD.annotate("safe / refuge:\ndeep, consolidated sleep",
             xy=(0.4, np.interp(0.0004, betas, bout_opt)),
             xytext=(1.1, bout_opt.max()*0.62), fontsize=8.6, color=NAVY,
             va="center",
             arrowprops=dict(arrowstyle="->", color=NAVY, lw=1.1))
axD.annotate("dangerous:\nlight, arousable sleep\n(one-eye-open, unihemispheric)",
             xy=(3.2, np.interp(0.0032, betas, bout_opt)),
             xytext=(1.8, bout_opt.max()*0.40), fontsize=8.6, color=RUST,
             va="center",
             arrowprops=dict(arrowstyle="->", color=RUST, lw=1.1))
axD.set_xlim(axD_x.min(), axD_x.max())
axD.set_ylim(0, bout_opt.max()*1.08)
axD.xaxis.set_major_locator(MaxNLocator(nbins=5, integer=True))

# ---- Panel E: model prediction vs REAL deep-sleep timing (Joyce et al. 2024) ----
# NOTE ON WHAT IS PLOTTED. The model predicts a DECLINING theta*(t): the core is
# discharged first, so deep sleep is front-loaded. Five of the eight datasets do
# this; three (D.vir, D.wil, D.yak) place deep sleep later and are the
# lineage-specific exception the framework expects of a labile useful layer.
# The bold line is therefore the mean of the FRONT-LOADING datasets -- the curve
# the model is actually fitted to -- and the exceptions are shown separately as
# their own mean. (Plotting a pooled 8-species mean, as an earlier version did,
# put the dashed fit next to a curve it had never been fitted to, and the two
# visibly diverged after ZT18.)
if obs_is_real:
    # classify by front-loading (early-night mean > late-night mean)
    early = slice(0, 4)      # ZT12-15
    late  = slice(-4, None)  # ZT20-23
    front, excpt = [], []
    for sp in species_list:
        (front if deep_by_sp[sp][early].mean() > deep_by_sp[sp][late].mean()
               else excpt).append(sp)

    front_stack = np.vstack([deep_by_sp[sp] for sp in front])
    fmean = front_stack.mean(axis=0)
    fsem  = front_stack.std(axis=0, ddof=1) / np.sqrt(len(front))
    exc_stack = np.vstack([deep_by_sp[sp] for sp in excpt]) if excpt else None
    emean = exc_stack.mean(axis=0) if excpt else None

    # thin per-dataset lines
    for sp in species_list:
        col = RUST if sp in excpt else TEAL
        axE.plot(zt_night, deep_by_sp[sp], color=col, lw=1.0, alpha=0.38)
    for sp in excpt:
        axE.annotate(sp, xy=(zt_night[-1], deep_by_sp[sp][-1]),
                     xytext=(zt_night[-1] + 0.15, deep_by_sp[sp][-1]),
                     fontsize=7.6, color=RUST, va="center")

    # bold line: mean +/- SEM of the FRONT-LOADING datasets (what the model fits)
    axE.fill_between(zt_night, fmean - fsem, fmean + fsem,
                     color=NAVY, alpha=0.14, lw=0)
    axE.plot(zt_night, fmean, color=NAVY, lw=2.8, marker="o", ms=4,
             label=f"front-loading datasets, mean of {len(front)} $\\pm$ SEM")

    # the exceptions, as their own mean
    if excpt:
        axE.plot(zt_night, emean, color=RUST, lw=2.0, ls=(0, (5, 2)),
                 label=f"exceptions, mean of {len(excpt)} (deep sleep later)")

    # model overlay: bounded exponential decline fitted to the front-loader mean
    pk = int(np.argmax(fmean)); tpk = zt_night[pk]
    try:
        from scipy.optimize import curve_fit
        def mdl(t, A, tau, C): return C + A*np.exp(-(t - tpk)/tau)
        (A_, tau_, C_), _ = curve_fit(
            mdl, zt_night[pk:], fmean[pk:], p0=[0.2, 3, 0.1],
            bounds=([0, 0.5, 0], [1, 12, 1]), maxfev=20000)
        yhat   = mdl(zt_night[pk:], A_, tau_, C_)
        ss_res = np.sum((fmean[pk:] - yhat)**2)
        ss_tot = np.sum((fmean[pk:] - fmean[pk:].mean())**2)
        r2 = 1 - ss_res/ss_tot
        tt = np.linspace(tpk, zt_night[-1], 100)
        axE.plot(tt, mdl(tt, A_, tau_, C_), color="#111111", lw=2.4, ls="--",
                 label=fr"model $\theta^*(t)$ decline (exp. fit, $R^2$={r2:.2f})")
        print(f"  front-loading datasets ({len(front)}): {front}")
        print(f"  exceptions ({len(excpt)}): {excpt}")
        print(f"  model fit to the front-loader mean: R^2={r2:.3f} (tau={tau_:.1f} h)")
    except Exception as e:
        print("  (model fit skipped:", e, ")")

    axE.axvspan(12, 15, color=TEAL, alpha=0.10, lw=0)
    axE.annotate("deep sleep peaks ZT13-15,\nthen declines \u2014 as $\\theta^*(t)$ predicts",
                 xy=(tpk, fmean[pk]), xytext=(15.5, fmean.max()*1.28),
                 fontsize=8.8, color=NAVY, va="center",
                 arrowprops=dict(arrowstyle="->", color=NAVY, lw=1.1))
    axE.text(13.5, 0.02, "core /\ndeep-sleep window", color=TEAL,
             fontsize=7.8, ha="center", va="bottom")

    axE.set_xlabel("Zeitgeber time (h);  ZT12 = lights-off, the night")
    axE.set_ylabel("deep-sleep fraction (HMM state $s_0$)")
    axE.set_title("Model vs observation: deep sleep is front-loaded within the night "
                  "(Drosophila; Joyce et al. 2024)", loc="left")
    axE.set_xlim(12, 23.6)
    top = max((fmean + fsem).max(), emean.max() if excpt else 0,
              max(deep_by_sp[sp].max() for sp in species_list))
    axE.set_ylim(0, top*1.12)

    from matplotlib.lines import Line2D
    handles, labels = axE.get_legend_handles_labels()
    handles += [Line2D([0], [0], color=TEAL, lw=1.4, alpha=0.6),
                Line2D([0], [0], color=RUST, lw=1.4, alpha=0.6)]
    labels  += ["individual front-loading datasets", "individual exceptions"]
    axE.legend(handles, labels, frameon=False, fontsize=7.8,
               loc="upper right", ncol=1, handlelength=1.9)
else:
    axE.text(0.5, 0.5, "sleep_state_24h_summary.csv not found",
             ha="center", va="center", transform=axE.transAxes, color=GREY)

# ---- Panel F: the ladder holds in every dataset ----
if have_bouts:
    rest = ["Deep sleep", "Light sleep", "Quiet awake"]
    sp_list = sorted(_u.species.unique())
    xs = np.arange(len(rest))
    ok = 0
    for i, sp in enumerate(sp_list):
        m = [_u.loc[(_u.species == sp) & (_u.canonical_label == st),
                    "duration_min"].mean() for st in rest]
        if m[0] > m[1] > m[2]:
            ok += 1
        axF.plot(xs, m, marker="o", ms=5, lw=1.6, alpha=0.85,
                 color=plt.cm.viridis(i/max(1, len(sp_list)-1)), label=sp)
    pooled = [meas_mean[st] for st in rest]
    axF.plot(xs, pooled, marker="s", ms=8, lw=3.0, color=NAVY, zorder=6,
             label="pooled")
    axF.set_yscale("log")
    from matplotlib.ticker import FixedLocator, NullLocator, FuncFormatter
    axF.yaxis.set_major_locator(FixedLocator([2, 3, 5, 10, 20, 30, 50]))
    axF.yaxis.set_minor_locator(NullLocator())
    axF.yaxis.set_major_formatter(FuncFormatter(lambda v, p: f"{v:g}"))
    axF.set_ylim(1.6, 60)
    axF.set_xticks(xs); axF.set_xticklabels(["deep sleep", "light sleep", "quiet wakefulness"])
    axF.set_ylabel("mean bout length (min, log)")
    axF.set_title(f"The ladder holds in every dataset ({ok}/{len(sp_list)}): "
                  f"deeper states hold rest longer", loc="left")
    axF.legend(frameon=False, fontsize=7.6, ncol=4, loc="lower left")
    axF.set_xlim(-0.25, 2.45)
    print(f"\nPanel F: monotone ladder (deep > light > quiet) in {ok}/{len(sp_list)} datasets")

# Panel letters as large lowercase labels, set outside the axes (matches the
# amended figure used in the manuscript). No suptitle.
for ax, letter in [(axA, "a"), (axB, "b"), (axC, "c"), (axD, "d"),
                   (axE, "e"), (axF, "f")]:
    ax.text(-0.115, 1.16, letter, transform=ax.transAxes,
            fontsize=27, fontweight="normal", va="top", ha="left")

OUT = "Figure_2"
fig.savefig(OUT + ".svg", bbox_inches="tight")            # editable vector (text as text)
fig.savefig(OUT + ".png", dpi=200, bbox_inches="tight")   # raster preview
#fig.savefig(OUT + ".pdf", bbox_inches="tight")            # vector, for direct submission
#fig.savefig(OUT + ".png", dpi=200, bbox_inches="tight")   # raster preview
print(f"\nSaved figure: {OUT}.svg / .pdf / .png")
