app.py

import gradio as gr
import matplotlib.pyplot as plt
import numpy as np
from io import BytesIO
from PIL import Image
from matplotlib.patches import Rectangle, Circle, FancyArrowPatch, PathPatch
from matplotlib.path import Path

# ================== Content (transcript language) ==================
STEPS = [
    dict(
        title="Step 1: Motor neuron → NMJ",
        where=("A motor neuron goes directly from the spinal cord to the skeletal muscle. "
               "The action potential travels down the axon to the axon terminal, calcium channels open, "
               "and acetylcholine is released into the synapse."),
        question="When the motor neuron releases acetylcholine, what happens next?",
        options=[
            "The muscle relaxes immediately.",
            "Acetylcholine moves across the synapse and binds to receptors on the muscle cell membrane.",
            "Calcium leaves the muscle fiber."
        ],
        correct=1,
        visual="neuron"
    ),
    dict(
        title="Step 2: Motor end plate (ligand-gated channels)",
        where=("Acetylcholine moves to the motor end plate and binds nicotinic acetylcholine receptors. "
               "These are ligand-gated channels that open when acetylcholine binds."),
        question="When acetylcholine binds to its receptor at the motor end plate, which ion moves into the muscle cell?",
        options=[
            "Sodium moves into the cell.",
            "Potassium moves into the cell.",
            "Calcium leaves the sarcoplasmic reticulum."
        ],
        correct=0,
        visual="nmj"
    ),
    dict(
        title="Step 3: Threshold → action potential on sarcolemma",
        where=("Sodium entry changes the voltage until threshold potential is reached. "
               "Voltage-gated channels open and an action potential travels along the sarcolemma."),
        question="What allows the electrical signal to travel quickly across the muscle cell membrane?",
        options=[
            "Voltage-gated sodium channels opening along the sarcolemma.",
            "Continuous acetylcholine release.",
            "ATP from mitochondria."
        ],
        correct=0,
        visual="sarcolemma"
    ),
    dict(
        title="Step 4: T-tubule voltage sensing (DHP)",
        where=("The sarcolemma dives into the cell as the T-tubule. "
               "When the action potential reaches this area, the DHP receptor senses the voltage change."),
        question="When the DHP receptor senses the voltage change, what does it do?",
        options=[
            "It moves the ryanidine receptor so calcium can leave the sarcoplasmic reticulum.",
            "It brings more sodium into the cell.",
            "It breaks down ATP."
        ],
        correct=0,
        visual="t_tubule"
    ),
    dict(
        title="Step 5: Calcium leaves the SR",
        where=("The ryanidine receptor opens. Calcium moves out of the sarcoplasmic reticulum into the cytoplasm "
               "following a concentration gradient (high in SR → lower in cytoplasm)."),
        question="Why does calcium move out of the sarcoplasmic reticulum?",
        options=[
            "There is a high concentration of calcium inside the SR and a lower concentration in the cytoplasm.",
            "Calcium is pushed out by sodium.",
            "It is actively pumped out using ATP."
        ],
        correct=0,
        visual="sr_release"
    ),
    dict(
        title="Step 6: Troponin → tropomyosin moves",
        where=("Once calcium is in the sarcoplasm, it binds to troponin, which causes tropomyosin to move away "
               "from the binding sites on actin."),
        question="What is exposed when tropomyosin moves?",
        options=[
            "The myosin binding sites on actin.",
            "The ATP-binding sites on myosin.",
            "The calcium pumps on the sarcoplasmic reticulum."
        ],
        correct=0,
        visual="thin_filament"
    ),
    dict(
        title="Step 7: Cross-bridge cycling (ATP’s role)",
        where=("Myosin binds to actin. ATP causes detachment; ATP hydrolysis re-cocks the myosin head. "
               "Without ATP, myosin stays attached and the muscle is stiff."),
        question="If there is no ATP available, what happens?",
        options=[
            "The myosin remains attached to actin, causing stiffness.",
            "The muscle continues to contract rapidly.",
            "The SR releases more calcium."
        ],
        correct=0,
        visual="crossbridge"
    ),
    dict(
        title="Step 8: Relaxation",
        where=("When the excitatory signal stops, acetylcholine esterase breaks down acetylcholine. "
               "Calcium is pumped back into the sarcoplasmic reticulum, and tropomyosin moves back over the binding sites."),
        question="What two actions cause relaxation?",
        options=[
            "Acetylcholine breakdown and calcium re-uptake into the sarcoplasmic reticulum.",
            "More acetylcholine release and ATP depletion.",
            "Sodium leaving the muscle fiber."
        ],
        correct=0,
        visual="relax"
    ),
]

NODES = [
    "ACh released at NMJ",
    "ACh binds nicotinic receptor",
    "Na⁺ entry → threshold → sarcolemma AP",
    "T-tubule DHP senses voltage",
    "RYR opens; Ca²⁺ leaves SR",
    "Ca²⁺ binds troponin; tropomyosin moves",
    "Cross-bridge cycling (ATP present)",
    "Relaxation: AChE + SERCA"
]

EDGES = {
    "ACh released at NMJ": ["ACh binds nicotinic receptor"],
    "ACh binds nicotinic receptor": ["Na⁺ entry → threshold → sarcolemma AP"],
    "Na⁺ entry → threshold → sarcolemma AP": ["T-tubule DHP senses voltage"],
    "T-tubule DHP senses voltage": ["RYR opens; Ca²⁺ leaves SR"],
    "RYR opens; Ca²⁺ leaves SR": ["Ca²⁺ binds troponin; tropomyosin moves"],
    "Ca²⁺ binds troponin; tropomyosin moves": ["Cross-bridge cycling (ATP present)"],
    "Cross-bridge cycling (ATP present)": ["Relaxation: AChE + SERCA"],
    "Relaxation: AChE + SERCA": []
}

# ================== Visual style helpers (HD schematics) ==================
PALETTE = {
    "membrane": "#222831",
    "t_tubule": "#3e8ed0",
    "sr": "#f59e0b",
    "channel": "#6b7280",
    "receptor": "#7c3aed",
    "vesicle": "#22c55e",
    "ach": "#16a34a",
    "na": "#2563eb",
    "ca": "#ef4444",
    "text": "#111827",
}

def _ion(ax, x, y, label, color, r=0.035):
    ax.add_patch(Circle((x, y), r, facecolor=color, edgecolor="white", linewidth=1.2))
    ax.text(x, y, label, ha="center", va="center", color="white", fontsize=9, fontweight="bold")

def _membrane(ax, x0=0.05, x1=0.95, y=0.5, thickness=0.03):
    ax.add_patch(Rectangle((x0, y - thickness/2), x1-x0, thickness,
                           facecolor=PALETTE["membrane"], alpha=0.15, edgecolor=PALETTE["membrane"]))

def _t_tubule(ax, x=0.5, y0=0.12, y1=0.88, w=0.06):
    ax.add_patch(Rectangle((x-w/2, y0), w, y1-y0, facecolor=PALETTE["t_tubule"], alpha=0.12, edgecolor=PALETTE["t_tubule"]))
    ax.text(x, y1+0.05, "T-tubule", ha="center", va="bottom", fontsize=11, color=PALETTE["t_tubule"])

def _sr(ax, x0=0.3, x1=0.7, y=0.18, h=0.06, label=True):
    ax.add_patch(Rectangle((x0, y), x1-x0, h, facecolor=PALETTE["sr"], alpha=0.10, edgecolor=PALETTE["sr"]))
    if label:
        ax.text((x0+x1)/2, y-0.03, "SR", ha="center", va="top", fontsize=11, color=PALETTE["sr"])

def _nicotinic_receptor(ax, x=0.8, y=0.5, w=0.06, h=0.04):
    # dimer-like shapes
    ax.add_patch(Rectangle((x-w, y-h/2), w, h, facecolor=PALETTE["receptor"], alpha=0.25, edgecolor=PALETTE["receptor"]))
    ax.add_patch(Rectangle((x, y-h/2), w, h, facecolor=PALETTE["receptor"], alpha=0.25, edgecolor=PALETTE["receptor"]))
    ax.text(x, y-0.07, "nicotinic AChR", ha="center", va="top", fontsize=9, color=PALETTE["receptor"])

def _channel(ax, x, y, open_state=True, label=None):
    h = 0.06
    ax.add_patch(Rectangle((x-0.01, y-h/2), 0.02, h,
                           facecolor=PALETTE["channel"], alpha=0.20, edgecolor=PALETTE["channel"]))
    if open_state:
        ax.add_line(plt.Line2D([x-0.007, x+0.007], [y-h/2+0.01, y+h/2-0.01], color=PALETTE["channel"], linewidth=2))
    else:
        ax.add_line(plt.Line2D([x-0.01, x+0.01], [y+0.03, y-0.03], color=PALETTE["channel"], linewidth=2))
    if label:
        ax.text(x, y+0.055, label, ha="center", va="bottom", fontsize=9, color=PALETTE["channel"])

def _arrow(ax, x0, y0, x1, y1, color="#111", width=2.4, curve=0.0, label=None, label_pos=0.5):
    if curve == 0:
        arrow = FancyArrowPatch((x0, y0), (x1, y1),
                                arrowstyle="-|>", mutation_scale=12, linewidth=width, color=color)
    else:
        verts = [(x0, y0), ((x0+x1)/2, (y0+y1)/2 + curve), (x1, y1)]
        codes = [Path.MOVETO, Path.CURVE3, Path.CURVE3]
        path = Path(verts, codes)
        arrow = FancyArrowPatch(path=path, arrowstyle="-|>", mutation_scale=12, linewidth=width, color=color)
    ax.add_patch(arrow)
    if label:
        mx = x0 + (x1-x0)*label_pos
        my = y0 + (y1-y0)*label_pos + (curve if curve else 0)
        ax.text(mx, my, label, fontsize=10, color=color, fontweight="bold",
                ha="center", va="bottom")

def _vesicle(ax, x, y, r=0.035, n_ach=3):
    ax.add_patch(Circle((x, y), r, facecolor=PALETTE["vesicle"], edgecolor="#136f45", linewidth=1))
    for k in range(n_ach):
        _ion(ax, x + (k-1)*(r*0.45), y, "ACh", PALETTE["ach"], r=0.017)

def _legend(ax, items):
    # items: list of (color, label)
    x, y = 0.05, 0.05
    for i, (c, t) in enumerate(items):
        ax.add_patch(Rectangle((x, y+i*0.04), 0.02, 0.02, facecolor=c, edgecolor=c))
        ax.text(x+0.025, y+i*0.04+0.01, t, va="center", ha="left", fontsize=9, color=PALETTE["text"])

def draw_visual(kind: str, detail: str = "HD") -> np.ndarray:
    """
    detail: 'Basic' or 'HD'
    """
    # Larger canvas for HD; use antialiasing and facecolors
    figsize = (7.2, 4.0) if detail == "HD" else (6, 3)
    fig, ax = plt.subplots(figsize=figsize)
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    ax.axis("off")

    # Common elements for several views
    if kind in {"neuron", "nmj", "sarcolemma"}:
        _membrane(ax, y=0.5, thickness=0.035)

    if kind == "neuron":
        ax.text(0.06, 0.86, "Motor neuron → axon terminal (NMJ)", color=PALETTE["text"], fontsize=12, fontweight="bold")
        # Axon
        ax.add_line(plt.Line2D([0.06, 0.38], [0.5, 0.5], color=PALETTE["membrane"], linewidth=5))
        # Vesicles at terminal
        for dx in [0.44, 0.50, 0.56]:
            _vesicle(ax, dx, 0.62, r=0.035 if detail=="HD" else 0.03)
        # ACh diffusion to membrane
        for dx in [0.48, 0.54]:
            _arrow(ax, dx, 0.60, dx+0.10, 0.52, color=PALETTE["ach"], width=2, curve=-0.05)
        ax.text(0.72, 0.44, "ACh in synapse", color=PALETTE["ach"], fontsize=10)
        _legend(ax, [(PALETTE["vesicle"], "Vesicle"), (PALETTE["ach"], "Acetylcholine")])

    elif kind == "nmj":
        ax.text(0.06, 0.9, "Motor end plate (nicotinic receptors open with ACh)", fontsize=12, fontweight="bold", color=PALETTE["text"])
        _nicotinic_receptor(ax, x=0.78, y=0.5)
        # ACh arrows from cleft to receptor
        for dy in [-0.03, 0.0, 0.03]:
            _arrow(ax, 0.62, 0.55+dy, 0.75, 0.50+dy, color=PALETTE["ach"], width=2, curve=-0.02)
        # Na+ entry (if open)
        _channel(ax, 0.78, 0.5, open_state=True, label="Na⁺ channel")
        for k in range(4):
            _ion(ax, 0.80 + 0.03*k, 0.58 + 0.02*np.sin(k), "Na⁺", PALETTE["na"])
            _arrow(ax, 0.80 + 0.03*k, 0.58 + 0.02*np.sin(k), 0.78, 0.52, color=PALETTE["na"], width=1.8, curve=-0.01, label=None)

    elif kind == "sarcolemma":
        ax.text(0.5, 0.9, "Sarcolemma AP via voltage-gated Na⁺", ha="center", fontsize=12, fontweight="bold", color=PALETTE["text"])
        # A wave of open channels (dots) and Na influx
        for xi in np.linspace(0.15, 0.85, 7):
            _channel(ax, xi, 0.5, open_state=True)
            _ion(ax, xi+0.03, 0.62, "Na⁺", PALETTE["na"])
            _arrow(ax, xi+0.03, 0.62, xi, 0.52, color=PALETTE["na"], width=1.6, curve=-0.015)

    elif kind == "t_tubule":
        _t_tubule(ax, x=0.5, y0=0.12, y1=0.88)
        _sr(ax, x0=0.30, x1=0.70, y=0.12, h=0.08, label=True)
        ax.text(0.5, 0.94, "DHP senses voltage → moves RYR", ha="center", fontsize=12, fontweight="bold", color=PALETTE["text"])
        # DHP (on T-tubule wall)
        ax.add_patch(Rectangle((0.48, 0.50-0.04), 0.04, 0.08, facecolor=PALETTE["receptor"], alpha=0.25, edgecolor=PALETTE["receptor"]))
        ax.text(0.50, 0.46, "DHP", ha="center", va="top", fontsize=9, color=PALETTE["receptor"])
        # RYR (on SR)
        ax.add_patch(Rectangle((0.42, 0.12), 0.16, 0.06, facecolor=PALETTE["sr"], alpha=0.18, edgecolor=PALETTE["sr"]))
        ax.text(0.50, 0.19, "RYR", ha="center", va="center", fontsize=10, color=PALETTE["sr"])
        # Link arrow
        _arrow(ax, 0.50, 0.50, 0.50, 0.18, color="#444", width=2, label="Coupling", label_pos=0.55)

    elif kind == "sr_release":
        _t_tubule(ax, x=0.5, y0=0.60, y1=0.95)
        _sr(ax, x0=0.20, x1=0.80, y=0.18, h=0.10, label=True)
        ax.text(0.5, 0.56, "RYR opens; Ca²⁺ leaves SR (high → lower)", ha="center", fontsize=12, fontweight="bold", color=PALETTE["text"])
        # Ca arrows SR -> cytosol
        for x in np.linspace(0.28, 0.72, 5):
            _ion(ax, x, 0.25, "Ca²⁺", PALETTE["ca"])
            _arrow(ax, x, 0.25, x, 0.45, color=PALETTE["ca"], width=2, curve=0.0)

    elif kind == "thin_filament":
        ax.text(0.5, 0.9, "Ca²⁺ binds troponin → Tropomyosin moves", ha="center", fontsize=12, fontweight="bold", color=PALETTE["text"])
        # Actin cable
        ax.add_patch(Rectangle((0.15, 0.45), 0.70, 0.05, facecolor="#9ca3af", edgecolor="#6b7280"))
        # Binding sites reveal (dots)
        for x in np.linspace(0.18, 0.80, 7):
            ax.add_patch(Circle((x, 0.475), 0.01, facecolor="#374151"))
        # Ca icons near troponin
        for x in [0.30, 0.50, 0.70]:
            _ion(ax, x, 0.58, "Ca²⁺", PALETTE["ca"])
            _arrow(ax, x, 0.58, x, 0.48, color=PALETTE["ca"], width=2)

    elif kind == "crossbridge":
        ax.text(0.5, 0.90, "Cross-bridge cycling (ATP detaches; hydrolysis re-cocks)", ha="center",
                fontsize=12, fontweight="bold", color=PALETTE["text"])
        # Actin (top) and myosin (bottom)
        ax.add_patch(Rectangle((0.12, 0.62), 0.76, 0.04, facecolor="#9ca3af", edgecolor="#6b7280"))
        ax.add_patch(Rectangle((0.12, 0.36), 0.76, 0.04, facecolor="#6b7280", edgecolor="#374151"))
        # Heads and binding
        for x in np.linspace(0.18, 0.82, 5):
            _arrow(ax, x, 0.40, x, 0.60, color="#374151", width=2)
        ax.text(0.50, 0.50, "ATP binds → detachment\nATP hydrolysis → re-cock", ha="center", va="center",
                fontsize=10, color=PALETTE["text"])

    elif kind == "relax":
        ax.text(0.5, 0.90, "Relaxation: ACh broken down; Ca²⁺ pumped back to SR", ha="center", fontsize=12, fontweight="bold", color=PALETTE["text"])
        _sr(ax, x0=0.20, x1=0.80, y=0.70, h=0.08, label=True)
        # Ca back to SR
        for x in np.linspace(0.28, 0.72, 5):
            _ion(ax, x, 0.42, "Ca²⁺", PALETTE["ca"])
            _arrow(ax, x, 0.45, x, 0.74, color=PALETTE["ca"], width=2)
        ax.text(0.20, 0.30, "AChE breaks down ACh", fontsize=10, color=PALETTE["ach"])
        _legend(ax, [(PALETTE["ca"], "Calcium"), (PALETTE["ach"], "Acetylcholine")])

    # Render to array
    buf = BytesIO()
    fig.tight_layout()
    fig.savefig(buf, format="png", dpi=(180 if detail == "HD" else 110), bbox_inches="tight")
    plt.close(fig)
    buf.seek(0)
    img = Image.open(buf).convert("RGB")
    return np.array(img)

# ================== Step Trainer logic (now passes detail level) ==================
def render_step(i:int, detail:str):
    i = int(i)
    s = STEPS[i]
    img = draw_visual(s["visual"], detail=detail)
    return (f"### {s['title']}",
            s["where"],
            img,
            f"**{s['question']}**",
            gr.update(choices=s["options"], value=s["options"][0]),
            "",  # feedback
            i,   # state
            i)

def submit_step(i:int, picked:str, detail:str):
    i = int(i)
    s = STEPS[i]
    idx = s["options"].index(picked) if picked in s["options"] else -1
    if idx == s["correct"]:
        if i < len(STEPS)-1:
            i += 1
            fb = "✅ **Correct. Advancing…**"
        else:
            fb = "✅ **Done. Relaxation complete.**"
    else:
        i = 0
        fb = "❌ **Incorrect. Returning to Step 1.**"
    title, where, img, q, choices, _, _, _ = render_step(i, detail)
    return title, where, img, q, choices, fb, i, i

def restart_step(_i:int, detail:str):
    title, where, img, q, choices, _, i, p = render_step(0, detail)
    return title, where, img, q, choices, "Restarted.", i, p

# ================== Failure-Point ==================
def failure_check(fails, guess):
    failed_idx = sorted([NODES.index(f) for f in (fails or [])]) if fails else []
    lines = []
    if failed_idx:
        stop = failed_idx[0]
        for j, lab in enumerate(NODES):
            ok = j < stop
            lines.append(("✅ " if ok else "⛔ ") + lab)
            if not ok:
                break
        fb = "✅ **Correct: first failed step located.**" if (guess and NODES.index(guess)==stop) else "❌ **Not quite—identify the FIRST failed step.**"
    else:
        lines = ["✅ " + l for l in NODES]
        lines.append("(No failures set — full propagation to relaxation.)")
        fb = "No failures toggled."
    return "```\n" + "\n".join(lines) + "\n```", fb

# ================== Sandbox ==================
def sandbox_metrics(Na_out, Na_in, Ca_sr, Ca_cyto, ATP):
    na_drive = max(0.0, (Na_out - Na_in) / max(1.0, Na_out))
    ap_prob  = min(1.0, na_drive * 1.5)
    dhp_ok   = ap_prob
    ca_drive = max(0.0, (Ca_sr - Ca_cyto) / max(1.0, Ca_sr))
    ca_rel   = dhp_ok * ca_drive
    xbridges = min(1.0, ca_rel * (0.5 + 0.5*min(1.0, ATP)))
    relax_ok = min(1.0, ATP) * 0.7 + (1.0 - ca_rel) * 0.3
    return dict(ap_prob=ap_prob, ca_release=ca_rel, crossbridge=xbridges, relax_ok=relax_ok)

def sandbox_plot(Na_out, Na_in, Ca_sr, Ca_cyto, ATP):
    vals = sandbox_metrics(Na_out, Na_in, Ca_sr, Ca_cyto, ATP)
    fig, ax = plt.subplots(figsize=(6,3))
    keys = ["ap_prob","ca_release","crossbridge","relax_ok"]
    ax.bar(keys, [vals[k] for k in keys], color=[PALETTE["na"], PALETTE["ca"], "#374151", "#10b981"])
    ax.set_ylim(0,1); ax.set_title("Predicted behaviors (0–1)")
    for i, k in enumerate(keys):
        ax.text(i, vals[k]+0.03, f"{vals[k]:.2f}", ha="center", va="bottom", fontsize=9)
    buf = BytesIO(); fig.tight_layout(); fig.savefig(buf, format="png", bbox_inches="tight", dpi=150); plt.close(fig)
    buf.seek(0); img = Image.open(buf).convert("RGB")
    return np.array(img)

# ================== Causality Builder ==================
def chain_add(chain, pick):
    import json
    chain = json.loads(chain)
    remaining = [n for n in NODES if n not in chain]
    if not remaining:
        return (gr.update(value=chain, interactive=False),
                "Chain complete.",
                gr.update(choices=[], interactive=False),
                " → ".join(chain))
    if pick is None:
        return (gr.update(value=chain),
                "Pick a next event.",
                gr.update(),
                " → ".join(chain))
    ok = pick in EDGES.get(chain[-1], [])
    if ok:
        chain.append(pick)
        remaining = [n for n in NODES if n not in chain]
        msg = "✅ **Link accepted.**"
    else:
        msg = "❌ **That event doesn’t logically follow. Try a different next step.**"
    return (gr.update(value=chain),
            msg,
            gr.update(choices=remaining, value=(remaining[0] if remaining else None), interactive=bool(remaining)),
            " → ".join(chain))

def chain_reset():
    import json
    chain = [NODES[0]]
    remaining = [n for n in NODES if n not in chain]
    return (gr.update(value=chain),
            "Reset.",
            gr.update(choices=remaining, value=remaining[0] if remaining else None, interactive=bool(remaining)),
            " → ".join(chain))

# ================== Fatigue Lab ==================
def simulate_fatigue(tmax=60, dt=0.5, atp_init=1.0, aerobic=0.3, anaerobic=0.2, load=0.5, serca_load=0.3):
    n=int(tmax/dt)+1; t=np.linspace(0,tmax,n); atp=np.zeros(n); atp[0]=atp_init; rigor=np.zeros(n,dtype=bool)
    for i in range(1,n):
        cons = load*0.4 + serca_load*0.25
        prod = aerobic*0.2 + anaerobic*0.15
        atp[i] = np.clip(atp[i-1] + dt*(prod - cons), 0, 1.2)
        rigor[i] = atp[i] < 0.1
    fig, ax = plt.subplots(1,2, figsize=(8,3))
    ax[0].plot(t, atp, color="#111827"); ax[0].set_xlabel("time (s)"); ax[0].set_ylabel("ATP (a.u.)"); ax[0].set_title("ATP dynamics")
    ax[1].bar(["rigor fraction"], [rigor.mean()], color="#ef4444"); ax[1].set_ylim(0,1); ax[1].set_title("Rigor")
    buf = BytesIO(); fig.tight_layout(); fig.savefig(buf, format="png", bbox_inches="tight", dpi=140); plt.close(fig)
    buf.seek(0); img = Image.open(buf).convert("RGB")
    msg = "Low ATP periods → stiffness (myosin remains attached)." if rigor.mean()>0 else "No stiffness expected (ATP maintained)."
    return np.array(img), f"Estimated rigor fraction: {rigor.mean():.2f}\n{msg}"

# ================== Passport Analyzer ==================
CORE_CONCEPTS = {
  "Flow down gradients": ["gradient","concentration","moves from high to low","diffuse","diffusion"],
  "Cell-to-cell communication": ["neurotransmitter","acetylcholine","receptor","synapse","binds"],
  "Structure–function": ["structure","function","troponin","tropomyosin","binding site","receptor opens"],
  "Energy flow": ["ATP","ADP","Pi","hydrolysis","energy"],
  "Interdependence": ["depends","linked","together","if/then","cascade","pathway"]
}
def passport_analyze(text):
    t = (text or "").lower()
    counts = {k: sum(t.count(w) for w in ws) for k,ws in CORE_CONCEPTS.items()}
    keys = list(counts.keys()); vals = [counts[k] for k in keys]
    fig, ax = plt.subplots(figsize=(6,3)); ax.bar(keys, vals, color="#3e8ed0"); ax.set_title("Core Concept mentions"); ax.tick_params(axis='x', rotation=30)
    for i,k in enumerate(keys):
        ax.text(i, vals[i]+0.05, str(vals[i]), ha="center")
    buf = BytesIO(); fig.tight_layout(); fig.savefig(buf, format="png", bbox_inches="tight", dpi=140); plt.close(fig)
    buf.seek(0); img = Image.open(buf).convert("RGB")
    weak = [k for k in keys if counts[k]==0]
    note = ("Consider adding explicit references to: " + ", ".join(weak)) if weak else "Balanced coverage detected."
    import json
    return np.array(img), json.dumps(counts, indent=2) + "\n\n" + note

# ================== UI ==================
with gr.Blocks(title="EC Coupling Suite") as demo:
    gr.Markdown("# EC Coupling Learning Suite (Transcript Language)")

    with gr.Row():
        detail_picker = gr.Radio(choices=["Basic", "HD"], value="HD", label="Visual detail")

    with gr.Tabs():
        # Step Trainer
        with gr.Tab("Step Trainer"):
            step_state = gr.State(0)
            title = gr.Markdown()
            where = gr.Markdown()
            img = gr.Image(label="Where are we?")
            q = gr.Markdown()
            choice = gr.Radio(choices=[], label="Predict what happens next")
            submit = gr.Button("Submit", variant="primary")
            restart = gr.Button("Restart (Step 1)")
            fb = gr.Markdown()
            prog = gr.Slider(0, len(STEPS)-1, value=0, step=1, interactive=False, label="Progress")

            demo.load(lambda d: render_step(0, d), inputs=[detail_picker],
                      outputs=[title, where, img, q, choice, fb, step_state, prog])

            submit.click(submit_step, [step_state, choice, detail_picker],
                         [title, where, img, q, choice, fb, step_state, prog])
            restart.click(restart_step, [step_state, detail_picker],
                          [title, where, img, q, choice, fb, step_state, prog])

        # Failure-Point
        with gr.Tab("Failure-Point"):
            gr.Markdown("Toggle failures and diagnose the **first** failed step.")
            fails = gr.CheckboxGroup(choices=NODES, label="Failures")
            guess = gr.Dropdown(choices=NODES, label="Your diagnosis")
            check = gr.Button("Test & Check", variant="primary")
            log = gr.Markdown()
            verdict = gr.Markdown()
            check.click(failure_check, [fails, guess], [log, verdict])

        # Sandbox
        with gr.Tab("Sandbox"):
            gr.Markdown("Adjust gradients and ATP; observe predicted behaviors (heuristic).")
            Na_out = gr.Slider(10, 160, value=140, step=1, label='[Na⁺] outside')
            Na_in  = gr.Slider(0, 50, value=15, step=1, label='[Na⁺] inside')
            Ca_sr  = gr.Slider(0.1, 10.0, value=3.0, step=0.1, label='[Ca²⁺] SR')
            Ca_c   = gr.Slider(0.0, 1.0, value=0.1, step=0.01, label='[Ca²⁺] cytoplasm')
            ATP    = gr.Slider(0.0, 1.0, value=0.8, step=0.01, label='ATP (0–1)')
            sb_img = gr.Image(label="Predicted bars")
            for w in [Na_out, Na_in, Ca_sr, Ca_c, ATP]:
                w.change(sandbox_plot, [Na_out, Na_in, Ca_sr, Ca_c, ATP], [sb_img])
            sb_img.value = sandbox_plot(Na_out.value, Na_in.value, Ca_sr.value, Ca_c.value, ATP.value)

        # Causality
        with gr.Tab("Causality"):
            gr.Markdown("Build a valid chain; logic is checked at each link.")
            import json
            chain_state = gr.State(json.dumps([NODES[0]]))
            chain_text = gr.Markdown(" → ".join([NODES[0]]))
            next_pick = gr.Dropdown(choices=[n for n in NODES if n != NODES[0]], label="Next event")
            add = gr.Button("Add link", variant="primary")
            reset = gr.Button("Reset")
            fb2 = gr.Markdown()
            add.click(chain_add, [chain_state, next_pick], [chain_state, fb2, next_pick, chain_text])
            reset.click(chain_reset, [], [chain_state, fb2, next_pick, chain_text])

        # Fatigue
        with gr.Tab("Fatigue"):
            gr.Markdown("Adjust ATP supply/demand; see ATP curve and rigor fraction.")
            aerobic = gr.Slider(0,1,value=0.4,step=0.01,label="Aerobic supply")
            anaer   = gr.Slider(0,1,value=0.3,step=0.01,label="Anaerobic supply")
            work    = gr.Slider(0,1,value=0.6,step=0.01,label="Mechanical load")
            serca   = gr.Slider(0,1,value=0.4,step=0.01,label="SERCA load")
            dur     = gr.Slider(10,180,value=90,step=1,label="Duration (s)")
            ftg_img = gr.Image(label="ATP & Rigor")
            ftg_txt = gr.Markdown()

            def ftg_update(dur_val, aer, anr, load, sl):
                return simulate_fatigue(dur_val, 0.5, 1.0, aer, anr, load, sl)

            for w in [aerobic, anaer, work, serca, dur]:
                w.change(ftg_update, [dur, aerobic, anaer, work, serca], [ftg_img, ftg_txt])

            img0, txt0 = simulate_fatigue(90, 0.5, 1.0, 0.4, 0.3, 0.6, 0.4)
            ftg_img.value, ftg_txt.value = img0, txt0

        # Passport
        with gr.Tab("Passport"):
            gr.Markdown("Paste your notes; see Core Concept emphasis.")
            ta = gr.Textbox(lines=8, label="Notes / reflection")
            pass_img = gr.Image(label="Concept counts")
            pass_txt = gr.Markdown()
            run = gr.Button("Analyze", variant="primary")
            run.click(passport_analyze, [ta], [pass_img, pass_txt])

demo.launch()