inference.py

"""SPARK v4 outer-sphere ET predictor for the Hugging Face demo.

Loads the v4 continuous-manifold model (ManifoldModel over the 42-slot reaction-network
vector) and, from one or more nondimensionalized cyclic voltammograms, returns:
  - a categorical posterior over MECHANISMS (distinct active elementary-step sets),
  - per-parameter posterior summaries and elementary-step presence probabilities,
  - a posterior-predictive signal reconstruction for a chosen mechanism.

This is the continuous-manifold replacement for the old discrete classifier+flow app: the
mechanism is read post-hoc from which gates the posterior places above kinetic silence, and
because the posterior is global it surfaces the ALTERNATIVE mechanisms a voltammogram admits.

Design notes for the HF (CPU) Space:
  - Encoder input normalization stats are loaded from a shipped JSON (v4_norm_stats.json),
    NOT recomputed from a training dataset (which is not on the Space).
  - Reconstruction is lightweight: a few posterior draws with a per-solve timeout.
"""
from __future__ import annotations

import json
import os
import signal
import threading
from collections import Counter
from pathlib import Path

import numpy as np
import torch

from manifold.model_manifold import ManifoldModel
from manifold import reaction_network as RN
from manifold.master_sim_v4 import run_master_v4

from preprocessing import nondimensionalize_cv

# ---- decode / degeneracy constants (match manifold/eval_realcv_v4 + alt_mechanisms_v4) ----
ACT_MARGIN = 0.5          # a gate is "present" only if its median rate clears silence by this
ABS_TOL, REL_TOL = 0.05, 1.5   # a mechanism reconstructs "within tolerance" if NRMSE <= max(ABS, REL*best)
_SOLVE_TIMEOUT = 6        # hard cap (s) per forward solve on CPU

# key always-on / commonly-active parameters to summarize for the user
KEY_SLOTS = ["log10_K0_1", "log10_lambda_1", "alpha_1", "E0_offset",
             "log10_dA", "log10_dB", "log10_Cdl"]

PRETTY = {
    "log10_K0_1": "log10 k0 (standard rate)",
    "log10_lambda_1": "log10 lambda (reorganization)",
    "alpha_1": "alpha (transfer coeff.)",
    "E0_offset": "E0 offset",
    "log10_dA": "log10 D_A (reactant diffusivity)",
    "log10_dB": "log10 D_B (product diffusivity)",
    "log10_Cdl": "log10 C_dl (double layer)",
}

GATE_PRETTY = {
    "log10_K0_2": "2nd electron transfer (EE)",
    "log10_kc": "following reaction (EC)",
    "log10_K0_Z": "ET on chemical product (ECE)",
    "log10_k_dispX": "disproportionation (DISP)",
    "log10_kcat": "catalytic regeneration (EC')",
    "log10_kf": "preceding equilibrium (CE)",
    "log10_kdim": "radical-radical dimerization",
    "log10_kradsub": "radical-substrate coupling",
    "log10_k_dc": "disprop/comprop",
    "log10_K0_med": "mediator ET",
    "log10_k_med": "mediated cross-ET",
    "log10_Kads_A": "adsorption",
    "log10_K0_surf": "surface ET (Laviron)",
    "log10_k_surfchem": "surface chemistry",
    "log10_k_PT": "proton transfer (PCET)",
}


class _Timeout:
    """Per-solve timeout so a stiff posterior draw cannot hang the CPU Space.

    Uses SIGALRM, which only works in the main thread. Gradio runs callbacks in worker
    threads, so we activate the alarm ONLY when on the main thread; otherwise this is a
    no-op (the solve runs to completion). Guarding this is essential: setting a signal
    handler off the main thread raises ValueError, which previously made every
    reconstruction draw fail -> "Reconstruction unavailable"."""
    def __init__(self, seconds):
        self.seconds = int(seconds)
        self.active = False

    def _handler(self, signum, frame):
        raise TimeoutError

    def __enter__(self):
        if threading.current_thread() is threading.main_thread():
            self._old = signal.signal(signal.SIGALRM, self._handler)
            signal.alarm(self.seconds)
            self.active = True
        return self

    def __exit__(self, *a):
        if self.active:
            signal.alarm(0)
            signal.signal(signal.SIGALRM, self._old)


def _mech_label(gate_tuple):
    if not gate_tuple:
        return "E (plain electron transfer)"
    return " + ".join(GATE_PRETTY.get(g, g.replace("log10_", "")) for g in gate_tuple)


def _active_set(theta_row):
    """Active gated steps of a single posterior draw (gates above kinetic silence)."""
    return tuple(g for g in RN.GATE_SLOTS
                 if theta_row[RN.SLOT_IDX[g]] > RN.SILENCE_THRESHOLD[g])


def _load_model(ckpt_path, device):
    """Inline of manifold.corner_v4.load_model (avoids pulling the dataset/flow_model chain)."""
    ck = torch.load(ckpt_path, map_location=device, weights_only=False)
    c = ck["config"]
    m = ManifoldModel(
        theta_dim=RN.N_SLOTS, d_context=c["d_context"], d_model=c["d_model"],
        n_heads=c["n_heads"], flow_family=c["flow_family"],
        n_coupling_layers=c["n_coupling_layers"], hidden_dim=c["hidden_dim"],
        encoder_type=c.get("encoder") or "signal",
    ).to(device)
    m.set_theta_stats(ck["theta_mean"], ck["theta_std"])
    m.load_state_dict(ck["state_dict"])
    m.eval()
    return m


class SPARKPredictor:
    """v4 outer-sphere ET predictor."""

    def __init__(self, checkpoint_path, norm_stats_path, device="cpu"):
        self.device = device
        if not Path(checkpoint_path).exists():
            raise FileNotFoundError(f"SPARK v4 checkpoint not found: {checkpoint_path}")
        if not Path(norm_stats_path).exists():
            raise FileNotFoundError(
                f"v4 normalization stats not found: {norm_stats_path}. "
                "This file (v4_norm_stats.json) is required and must be shipped with the app.")
        self.model = _load_model(checkpoint_path, device)
        ns = json.load(open(norm_stats_path))
        # tuples (mean, std) for potential / flux / time
        self.norm = {k: (float(v[0]), float(v[1])) for k, v in ns.items()}

    # ---- input construction (fixed-scale nondim + encoder batch, mirrors prepare_input) ----
    def build_exp(self, scans, E0_V, T_K=298.15, A_cm2=0.0707,
                  C_A_molcm3=1e-6, D_A_cm2s=1e-5, n=1):
        """scans: list of dicts {E_V, i_A, v_Vs}. Returns an exp_data dict of
        nondimensionalized multi-scan CVs (theta, flux, sigma) + sweep limits."""
        thetas, fluxes, sigmas, raw = [], [], [], []
        for s in scans:
            th, fl, sig = nondimensionalize_cv(
                np.asarray(s["E_V"], float), np.asarray(s["i_A"], float),
                float(s["v_Vs"]), E0_V, T_K=T_K, A_cm2=A_cm2,
                C_A_molcm3=C_A_molcm3, D_A_cm2s=D_A_cm2s, n=n)
            thetas.append(th); fluxes.append(fl); sigmas.append(sig)
            raw.append((np.asarray(s["E_V"], float), np.asarray(s["i_A"], float)))
        order = np.argsort(sigmas)   # ascending scan rate
        thetas = [thetas[i] for i in order]; fluxes = [fluxes[i] for i in order]
        sigmas = [sigmas[i] for i in order]; raw = [raw[i] for i in order]
        return {
            "potentials": thetas, "fluxes": fluxes, "sigmas": np.array(sigmas),
            "theta_i": float(max(t.max() for t in thetas)),
            "theta_v": float(min(t.min() for t in thetas)),
            "raw": raw,
        }

    def _batch(self, exp, n_scans=3):
        n_av = len(exp["potentials"])
        use_n = min(n_scans, n_av)
        idxs = np.linspace(0, n_av - 1, use_n).astype(int)
        n_points = max(len(exp["potentials"][i]) for i in idxs)
        pot = np.zeros((use_n, n_points), np.float32)
        flx = np.zeros((use_n, n_points), np.float32)
        tim = np.zeros((use_n, n_points), np.float32)
        mask = np.zeros((use_n, n_points), bool)
        slog = np.zeros(use_n, np.float32)
        fscale = np.zeros(use_n, np.float32)
        pm, ps = self.norm["potential"]; fm, fs = self.norm["flux"]; tm, ts = self.norm["time"]
        for i, idx in enumerate(idxs):
            p = exp["potentials"][idx].astype(np.float32)
            f = exp["fluxes"][idx].astype(np.float32)
            L = len(p)
            trange = exp["theta_i"] - exp["theta_v"]
            t_cv = np.linspace(0, 2 * trange / exp["sigmas"][idx], L).astype(np.float32)
            peak = np.max(np.abs(f)) + 1e-30
            fscale[i] = np.log10(peak); f = f / peak
            slog[i] = np.log10(exp["sigmas"][idx])
            pot[i, :L] = (p - pm) / ps
            flx[i, :L] = (f - fm) / fs
            tim[i, :L] = (t_cv - tm) / ts
            mask[i, :L] = True
        x = np.stack([pot, flx, tim], axis=1)  # [N,3,T]
        return {
            "input": torch.tensor(x).unsqueeze(0).to(self.device),
            "scan_mask": torch.tensor(mask).unsqueeze(0).to(self.device),
            "sigmas": torch.tensor(slog).unsqueeze(0).to(self.device),
            "flux_scales": torch.tensor(fscale).unsqueeze(0).to(self.device),
        }

    # ---- inference ----
    @torch.no_grad()
    def sample_posterior(self, exp, n_scans=3, n_post=2000):
        batch = self._batch(exp, n_scans=n_scans)
        ctx = self.model.encode(batch)
        s = self.model.flow.sample(ctx, n_samples=n_post)[0].cpu().numpy()
        return s

    def mechanism_posterior(self, s, top_k=6):
        """Categorical posterior over mechanisms (distinct active-gate sets)."""
        sets = [_active_set(s[i]) for i in range(len(s))]
        counts = Counter(sets)
        ranked = sorted(counts.items(), key=lambda kv: -kv[1])[:top_k]
        return [{"gates": list(gt), "label": _mech_label(gt), "prob": c / len(s)}
                for gt, c in ranked], sets

    def presence(self, s):
        return {g: float((s[:, RN.SLOT_IDX[g]] > RN.SILENCE_THRESHOLD[g]).mean())
                for g in RN.GATE_SLOTS}

    def decoded_active(self, s):
        """The single decoded mechanism (margin-based, matches eval_realcv_v4.decode)."""
        return sorted(g for g in RN.GATE_SLOTS
                      if (s[:, RN.SLOT_IDX[g]] > RN.SILENCE_THRESHOLD[g] + ACT_MARGIN).mean() > 0.5)

    def param_summary(self, s, slots=None):
        slots = slots or KEY_SLOTS
        out = {}
        for nm in slots:
            j = RN.SLOT_IDX[nm]
            out[nm] = {
                "label": PRETTY.get(nm, nm.replace("log10_", "")),
                "median": float(np.median(s[:, j])),
                "lo": float(np.percentile(s[:, j], 5)),
                "hi": float(np.percentile(s[:, j], 95)),
            }
        return out

    # ---- reconstruction (lightweight, posterior-predictive) ----
    def reconstruct(self, s, active, exp, n_scans=3, n_pred=10):
        """Posterior-predictive mean reconstruction for one mechanism over the inference
        scans. Returns (panels, median_nrmse). panels: list of (sigma, pot, obs, recon)."""
        n_av = len(exp["potentials"])
        use_n = min(n_scans, n_av)
        idxs = np.linspace(0, n_av - 1, use_n).astype(int)
        rng = np.random.default_rng(0)
        ru_draws = [RN.sample_gen_ru(rng) for _ in range(min(n_pred, len(s)))]
        panels, nrmse = [], []
        for idx in idxs:
            sg = float(exp["sigmas"][idx])
            of = exp["fluxes"][idx].astype(np.float64)
            pt = exp["potentials"][idx].astype(np.float64)
            if of.size < 5:
                continue
            recs = []
            for k in range(min(n_pred, len(s))):
                try:
                    with _Timeout(_SOLVE_TIMEOUT):
                        r = run_master_v4(
                            RN.sim_theta_at_sigma_v4(s[k], sorted(active), sg),
                            sigma=sg, theta_i=exp["theta_i"], theta_v=exp["theta_v"],
                            conditions={"pH": 7.0}, log10_Ru=ru_draws[k])
                except Exception:
                    continue
                fr = np.asarray(r["flux"], np.float64)
                if fr.size >= 5 and np.all(np.isfinite(fr)):
                    recs.append(np.interp(np.linspace(0, 1, of.size),
                                          np.linspace(0, 1, fr.size), fr))
            if not recs:
                continue
            rec = np.mean(recs, 0)
            panels.append((sg, pt, of, rec))
            nrmse.append(np.sqrt(np.mean((of - rec) ** 2)) / (np.ptp(of) + 1e-30))
        return panels, (float(np.median(nrmse)) if nrmse else float("nan"))

    def mechanism_subset(self, s, gates, min_n=60):
        """Posterior samples whose active-gate set equals `gates` (that mechanism's mode).
        Falls back to all samples if too few pure-mode draws (reconstruction still pins the
        active set), so parameter posteriors stay meaningful for rare mechanisms."""
        gs = tuple(gates)
        idx = [i for i in range(len(s)) if _active_set(s[i]) == gs]
        return s[idx] if len(idx) >= min_n else s

    def inspect_mechanism(self, s, gates, exp, n_scans=3, n_pred=8):
        """Per-mechanism view: posterior-predictive reconstruction + parameter posteriors
        conditioned on that mechanism, for the interactive selector."""
        sub = self.mechanism_subset(s, gates)
        panels, nr = self.reconstruct(sub, gates, exp, n_scans=n_scans, n_pred=n_pred)
        slots = KEY_SLOTS + [g for g in gates if g in RN.GATE_SLOTS and g not in KEY_SLOTS]
        params = self.param_summary(sub, slots=slots)
        return {"panels": panels, "nrmse": nr, "params": params, "subset": sub}

    def alternatives(self, s, mechs, exp, n_scans=3, n_pred=8, max_mech=2):
        """Reconstruct the top mechanisms and flag which fit within tolerance."""
        cands = []
        for m in mechs[:max_mech]:
            panels, nr = self.reconstruct(s, m["gates"], exp, n_scans=n_scans, n_pred=n_pred)
            cands.append({**m, "nrmse": nr, "panels": panels})
        finite = [c["nrmse"] for c in cands if np.isfinite(c["nrmse"])]
        best = min(finite) if finite else float("inf")
        tol = max(ABS_TOL, REL_TOL * best) if np.isfinite(best) else float("inf")
        for c in cands:
            c["within_tol"] = bool(np.isfinite(c["nrmse"]) and c["nrmse"] <= tol)
        return cands, tol

    def predict(self, scans, E0_V, T_K=298.15, A_cm2=0.0707, C_A_molcm3=1e-6,
                D_A_cm2s=1e-5, n=1, n_scans=3, n_post=2000, n_pred=8, top_k=6):
        """Full pipeline -> dict for the UI."""
        exp = self.build_exp(scans, E0_V, T_K=T_K, A_cm2=A_cm2,
                             C_A_molcm3=C_A_molcm3, D_A_cm2s=D_A_cm2s, n=n)
        s = self.sample_posterior(exp, n_scans=n_scans, n_post=n_post)
        mechs, _ = self.mechanism_posterior(s, top_k=top_k)
        cands, tol = self.alternatives(s, mechs, exp, n_scans=n_scans, n_pred=n_pred)
        return {
            "exp": exp,
            "samples": s,
            "mechanisms": mechs,
            "alternatives": cands,
            "tol": tol,
            "presence": self.presence(s),
            "params": self.param_summary(s),
            "decoded_active": self.decoded_active(s),
            "n_multimodal": int(sum(c["within_tol"] for c in cands)),
        }


# ---- module-level singleton loader (used by app.py) ----
_PREDICTOR = None


def get_predictor():
    global _PREDICTOR
    if _PREDICTOR is None:
        here = Path(__file__).resolve().parent
        ckpt = os.environ.get(
            "SPARK_CHECKPOINT",
            str(here / "checkpoints" / "manifold_model.pt"))
        norm = os.environ.get(
            "SPARK_NORM_STATS",
            str(here / "checkpoints" / "v4_norm_stats.json"))
        _PREDICTOR = SPARKPredictor(ckpt, norm, device="cpu")
    return _PREDICTOR