"""
SPARK (Simulation-based Posterior Amortization for Reaction Kinetics)
Amortized Bayesian inference for outer-sphere electron-transfer cyclic voltammetry.

Gradio web interface for the v4 continuous-manifold model: from one or more cyclic
voltammograms it returns a calibrated posterior over a continuous space of reaction
mechanisms, the alternative mechanisms the data admit, per-parameter posteriors, and a
posterior-predictive reconstruction. Mechanism identity is read post-hoc from which
elementary-step gates the posterior places above kinetic silence -- there is no classifier.
"""

import os
import sys
import json
from pathlib import Path

import numpy as np
import gradio as gr

# spark_app dir first (local inference.py/plotting.py), repo root appended for `manifold`
_HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(_HERE))
sys.path.append(str(_HERE.parent))

from inference import get_predictor, GATE_PRETTY
from preprocessing import parse_cv_csv, estimate_E0
from plotting import (
    plot_mechanism_posterior, plot_alternatives, plot_presence,
    plot_parameter_posteriors, plot_reconstruction, parameter_table,
)
from manifold import reaction_network as RN

DEMO_DIR = _HERE / "demo_data"

# lazy predictor (loaded on first request so the Space boots fast / shows errors cleanly)
_PRED = None
_PRED_ERR = None


def _predictor():
    global _PRED, _PRED_ERR
    if _PRED is None and _PRED_ERR is None:
        try:
            _PRED = get_predictor()
        except Exception as e:  # noqa: BLE001
            _PRED_ERR = str(e)
    return _PRED


def _mech_choice_label(m):
    return f"{m['label']}  ({m['prob']*100:.0f}%)"


def _error(msg):
    # order: mech_plot, presence_plot, summary, dropdown, rec_plot, param_plot, table, state
    return (None, None, f"### Error\n\n{msg}",
            gr.update(choices=[], value=None), None, None, "", None)


# ---------------------------------------------------------------------------
# Core analysis: raw scans -> fast posterior outputs + a mechanism selector.
# Reconstruction + per-mechanism parameter posteriors are computed for the
# selected mechanism (top by default; user can switch via the dropdown).
# ---------------------------------------------------------------------------

def _analyze(scans, E0_V, T_K, A_cm2, C_mM, D_cm2s, n_electrons,
             n_post, n_pred, n_scans):
    pred = _predictor()
    if pred is None:
        return _error(f"Model unavailable: {_PRED_ERR}")
    if not scans:
        return _error("No usable voltammograms were parsed.")

    C_molcm3 = float(C_mM) * 1e-6 if C_mM else 1e-6
    n = int(n_electrons) if n_electrons else 1
    T = float(T_K) if T_K else 298.15
    A = float(A_cm2) if A_cm2 else 0.0707
    D = float(D_cm2s) if D_cm2s else 1e-5
    n_scans = int(n_scans); n_pred = int(n_pred)

    e0 = float(E0_V) if E0_V else float(np.median(
        [estimate_E0(s["E_V"], s["i_A"]) for s in scans]))

    try:
        exp = pred.build_exp(scans, E0_V=e0, T_K=T, A_cm2=A, C_A_molcm3=C_molcm3,
                             D_A_cm2s=D, n=n)
        s = pred.sample_posterior(exp, n_scans=n_scans, n_post=int(n_post))
    except Exception as e:  # noqa: BLE001
        return _error(f"Inference failed: {e}")

    mechs, _ = pred.mechanism_posterior(s, top_k=8)
    fig_mech = plot_mechanism_posterior(mechs)          # fast: probabilities only
    fig_pres = plot_presence(pred.presence(s), GATE_PRETTY)
    summary = _summary_md(mechs)

    choices = [_mech_choice_label(m) for m in mechs]
    labelmap = {c: m["gates"] for c, m in zip(choices, mechs)}
    state = {"s": s, "exp": exp, "labelmap": labelmap,
             "n_scans": n_scans, "n_pred": n_pred}

    # reconstruct + inspect the top mechanism for the initial view
    top = mechs[0]
    info = pred.inspect_mechanism(s, top["gates"], exp, n_scans=n_scans, n_pred=n_pred)
    fig_rec = plot_reconstruction(info["panels"], top["label"], info["nrmse"])
    fig_par = plot_parameter_posteriors(info["subset"], info["params"], RN.SLOT_IDX)
    table = "### Parameter posteriors (selected mechanism)\n\n" + parameter_table(info["params"])
    dd = gr.update(choices=choices, value=choices[0])
    return fig_mech, fig_pres, summary, dd, fig_rec, fig_par, table, state


def on_mechanism_change(label, state):
    """Reconstruct + show parameter posteriors for the mechanism picked in the dropdown."""
    if not state or not label:
        return None, None, ""
    pred = _predictor()
    gates = state["labelmap"].get(label, [])
    try:
        info = pred.inspect_mechanism(state["s"], gates, state["exp"],
                                      n_scans=state["n_scans"], n_pred=state["n_pred"])
    except Exception as e:  # noqa: BLE001
        return None, None, f"Reconstruction failed: {e}"
    name = label.split("  (")[0]
    fig_rec = plot_reconstruction(info["panels"], name, info["nrmse"])
    fig_par = plot_parameter_posteriors(info["subset"], info["params"], RN.SLOT_IDX)
    table = "### Parameter posteriors (selected mechanism)\n\n" + parameter_table(info["params"])
    return fig_rec, fig_par, table


def _summary_md(mechs):
    lines = ["### Mechanism inference"]
    top = mechs[0]
    lines.append(f"Most probable mechanism: **{top['label']}** "
                 f"(posterior probability {top['prob']*100:.0f}%).")
    strong = [m for m in mechs if m["prob"] >= 0.15]
    if len(strong) >= 2:
        alt = "; ".join(f"{m['label']} ({m['prob']*100:.0f}%)" for m in strong)
        lines.append(
            f"\n**This voltammogram is consistent with more than one mechanism:** {alt}. "
            "Use the selector below to reconstruct each and inspect its parameter posterior. "
            "A single point fit, given one assumed mechanism, would report only one of these.")
    else:
        lines.append("\nUse the selector below to reconstruct a mechanism and inspect its "
                     "parameter posterior.")
    return "\n".join(lines)


# ---------------------------------------------------------------------------
# CSV entry point
# ---------------------------------------------------------------------------

def analyze_cv_csv(files, scan_rates_text, E0_V, T_K, A_cm2, C_mM, D_cm2s,
                   n_electrons, n_post, n_pred, n_scans):
    if not files:
        return _error("Please upload at least one CSV file (one per scan rate).")
    rates_text = (scan_rates_text or "").strip()
    user_rates = None
    if rates_text:
        try:
            user_rates = [float(s.strip()) for s in rates_text.split(",")]
        except ValueError:
            return _error("Invalid scan rates. Enter comma-separated numbers in V/s.")
        if len(user_rates) != len(files):
            return _error(f"Number of files ({len(files)}) must match number of "
                          f"scan rates ({len(user_rates)}).")
    scans = []
    for idx, f in enumerate(files):
        try:
            parsed = parse_cv_csv(Path(f.name).read_text())
        except Exception as e:  # noqa: BLE001
            return _error(f"Could not parse '{Path(f.name).name}': {e}")
        if user_rates is not None:
            v = user_rates[idx]
        elif "scan_rate_Vs" in parsed:
            v = parsed["scan_rate_Vs"]
        else:
            return _error(f"Cannot determine scan rate for '{Path(f.name).name}'. "
                          "Provide scan rates (V/s) or include a Time (s) column.")
        scans.append({"E_V": parsed["E_V"], "i_A": parsed["i_A"], "v_Vs": v})
    return _analyze(scans, E0_V, T_K, A_cm2, C_mM, D_cm2s, n_electrons,
                    n_post, n_pred, n_scans)


# ---------------------------------------------------------------------------
# Image entry point (digitize plot images)
# ---------------------------------------------------------------------------

def analyze_cv_image(files, scan_rates_text, E0_V, threshold, n_post, n_pred,
                     n_scans, x_min, x_max, y_min, y_max):
    if not files:
        return _error("Please upload at least one plot image.")
    try:
        from digitizer import digitize_plot, auto_detect_axis_bounds
        from PIL import Image as PILImage
    except ImportError:
        return _error("Image digitization requires opencv-python-headless and Pillow.")
    rates_text = (scan_rates_text or "").strip()
    if not rates_text:
        return _error("Please enter scan rate(s) (V/s), comma-separated.")
    try:
        rates = [float(s.strip()) for s in rates_text.split(",")]
    except ValueError:
        return _error("Invalid scan rates.")
    if len(rates) == 1 and len(files) > 1:
        rates = rates * len(files)
    if len(rates) != len(files):
        return _error(f"Number of scan rates ({len(rates)}) must match images ({len(files)}).")
    has_bounds = all(v not in (None, 0) for v in [x_min, x_max, y_min, y_max])
    scans = []
    for idx, f in enumerate(files):
        fpath = f.name if hasattr(f, "name") else str(f)
        try:
            img = np.array(PILImage.open(fpath).convert("RGB"))
        except Exception as e:  # noqa: BLE001
            return _error(f"Could not read image {idx+1}: {e}")
        if has_bounds:
            b = {"x_min": float(x_min), "x_max": float(x_max),
                 "y_min": float(y_min), "y_max": float(y_max)}
        else:
            try:
                b = auto_detect_axis_bounds(img)
            except Exception:  # noqa: BLE001  (e.g. easyocr unavailable)
                b = None
            if b is None:
                return _error(f"Could not auto-detect axes for image {idx+1}. "
                              "Enter E min/max and I min/max under 'Axis overrides' and retry.")
        try:
            E_V, I_raw = digitize_plot(img, b["x_min"], b["x_max"], b["y_min"], b["y_max"],
                                       threshold=int(threshold))
        except Exception as e:  # noqa: BLE001
            return _error(f"Digitization failed for image {idx+1}: {e}")
        i_max = np.max(np.abs(I_raw))
        i_A = I_raw * (1e-6 if i_max > 100 else 1e-3 if i_max > 0.1 else 1.0)
        scans.append({"E_V": E_V, "i_A": i_A, "v_Vs": rates[idx]})
    return _analyze(scans, E0_V, 298.15, 0.0707, 1.0, 1e-5, 1, n_post, n_pred, n_scans)


# ---------------------------------------------------------------------------
# Demo examples
# ---------------------------------------------------------------------------

def _demo_examples():
    """Group bundled demo CSVs by mechanism -> (label, [csv paths], scan_rates, E0)."""
    out = []
    # os_* = real experimental case studies; sim_* = synthetic held-out test examples
    metas = sorted(DEMO_DIR.glob("os_*_metadata.json")) + sorted(DEMO_DIR.glob("sim_*_metadata.json"))
    for meta in metas:
        try:
            m = json.load(open(meta))
        except Exception:  # noqa: BLE001
            continue
        stem = meta.name.replace("_metadata.json", "")
        csvs = sorted(DEMO_DIR.glob(f"{stem}_*mVs.csv"))
        if not csvs:
            continue
        # files are sorted by zero-padded mV (ascending); pair with ascending rates
        out.append({"label": m.get("system", stem), "csvs": [str(c) for c in csvs],
                    "rates": sorted(m.get("scan_rates_Vs", [])),
                    "E0": m.get("physical_params", {}).get("E0_V", 0.0)})
    return out


# ---------------------------------------------------------------------------
# UI: theme, CSS, hero header, About
# ---------------------------------------------------------------------------

HERO_HTML = """
<div id="hero">
  <div class="hero-title">SPARK</div>
  <div class="hero-sub">Simulation-based Posterior Amortization for Reaction Kinetics</div>
  <div class="hero-tag">
    Calibrated Bayesian inference of electrochemical reaction mechanisms from cyclic
    voltammetry in a single forward pass. SPARK returns a full posterior over a
    <b>continuous space of mechanisms</b>, so it reports not just the single most likely
    mechanism but the <b>alternatives your data genuinely support</b>, with honest uncertainty.
  </div>
  <div class="hero-flow">
    upload voltammograms &nbsp;&rarr;&nbsp; posterior over mechanisms &nbsp;&rarr;&nbsp;
    pick a mechanism to reconstruct &amp; inspect its parameters
  </div>
</div>
"""

CSS = """
.gradio-container {max-width: 1200px !important; margin: auto !important;}
#hero {text-align: center; padding: 22px 18px 18px; margin-bottom: 8px;
       border-radius: 16px;
       background: linear-gradient(135deg, rgba(99,102,241,0.12), rgba(16,185,129,0.12));
       border: 1px solid rgba(99,102,241,0.20);}
#hero .hero-title {font-size: 2.6rem; font-weight: 800; letter-spacing: 2px;
       background: linear-gradient(90deg,#4f46e5,#059669); -webkit-background-clip: text;
       -webkit-text-fill-color: transparent; background-clip: text;}
#hero .hero-sub {font-size: 1.05rem; font-weight: 600; color: #475569; margin-top: 2px;}
#hero .hero-tag {font-size: 0.97rem; color: #334155; max-width: 820px; margin: 12px auto 0;
       line-height: 1.5;}
#hero .hero-flow {font-size: 0.9rem; color: #4f46e5; margin-top: 12px; font-weight: 600;}
.card {border: 1px solid rgba(100,116,139,0.18); border-radius: 14px; padding: 14px 16px;
       box-shadow: 0 1px 3px rgba(15,23,42,0.06);}
.verdict {background: rgba(16,185,129,0.06); border-color: rgba(16,185,129,0.28);}
.gr-plot {min-height: 340px;}
"""

ABOUT_MD = """
### What SPARK does
SPARK turns one or more cyclic voltammograms into a **calibrated posterior over a continuous
space of reaction mechanisms**. A single amortized neural model returns a joint posterior over
a 42-dimensional reaction-network parameter vector in milliseconds; mechanism identity is read
post-hoc from which elementary-step "gates" the posterior places above kinetic silence. There
is no discrete classifier and no per-mechanism refitting.

### Why a posterior over mechanisms
A voltammogram often does not determine a unique mechanism. Because SPARK's posterior spans
the whole manifold, it surfaces the **alternative mechanisms the data genuinely admit** and how
strongly &mdash; something a single point fit (DigiElch / gradient fitting), handed one assumed
mechanism, cannot discover. Use the mechanism selector to reconstruct each candidate and compare
its fit and parameter posterior.

### Elementary steps covered
Built on a unified Marcus-Hush-Chidsey / Butler-Volmer / Nernst electron-transfer law: plain
electron transfer (E), following chemical reaction (EC), second electron transfer (ECE),
disproportionation (DISP), catalytic regeneration (EC'), preceding equilibrium (CE),
two-electron (EE), radical-radical dimerization, radical-substrate coupling, proton-coupled
electron transfer (PCET), surface adsorption (Laviron), plus the double-layer capacitance.

### How to use
Upload potentiostat CSVs (one per scan rate; columns: potential in V, current in A/mA/µA, and
optionally time in s to auto-detect the scan rate) or plot images. Enter scan rates and, if
known, the formal potential and cell parameters for accurate nondimensionalization. Then pick a
mechanism to reconstruct and inspect its parameter posterior.

### Model
Continuous-manifold simulation-based inference (calibrated RQ-NSF density estimator) trained on a
coupled reaction-diffusion-adsorption voltammetry simulator. Outer-sphere electron transfer.
"""


def _output_components():
    """Returns (click_outputs, mech_dropdown, state, per_mech_outputs).
    click_outputs order matches _analyze's return:
      [mech_plot, pres_plot, summary_md, mech_dropdown, rec_plot, par_plot, table_md, state]
    per_mech_outputs (for the dropdown change): [rec_plot, par_plot, table_md]."""
    state = gr.State(None)
    # 1) verdict card
    with gr.Group(elem_classes="card verdict"):
        summary_md = gr.Markdown("Run an analysis to see the mechanism posterior.")
    # 2) overview: posterior over mechanisms + elementary-step presence
    with gr.Group(elem_classes="card"):
        gr.Markdown("#### Posterior over mechanisms")
        with gr.Row():
            mech_plot = gr.Plot(show_label=False)
            pres_plot = gr.Plot(show_label=False)
    # 3) prominent mechanism selector
    with gr.Group(elem_classes="card"):
        gr.Markdown("#### Reconstruct & inspect a mechanism")
        mech_dropdown = gr.Dropdown(
            label="Mechanism (posterior probability shown)", choices=[], interactive=True)
        with gr.Row():
            rec_plot = gr.Plot(show_label=False)
            par_plot = gr.Plot(show_label=False)
        with gr.Accordion("Parameter table (selected mechanism)", open=False):
            table_md = gr.Markdown()
    click_outputs = [mech_plot, pres_plot, summary_md, mech_dropdown,
                     rec_plot, par_plot, table_md, state]
    return click_outputs, mech_dropdown, state, [rec_plot, par_plot, table_md]


def _example_blocks(inputs):
    """Two labeled example blocks: experimental case studies + simulated test examples."""
    ex = _demo_examples()
    real = [e for e in ex if not e["label"].lower().startswith("simulated")]
    sim = [e for e in ex if e["label"].lower().startswith("simulated")]

    def rows(items):
        return [[e["csvs"], ", ".join(f"{r:.3g}" for r in e["rates"]), e["E0"]] for e in items]
    if real:
        gr.Examples(examples=rows(real), inputs=inputs, label="Experimental case studies",
                    example_labels=[e["label"] for e in real])
    if sim:
        gr.Examples(examples=rows(sim), inputs=inputs, label="Simulated test examples (known ground truth)",
                    example_labels=[e["label"] for e in sim])


def build_app():
    theme = gr.themes.Soft(primary_hue="indigo", secondary_hue="emerald", neutral_hue="slate")
    with gr.Blocks(title="SPARK", theme=theme, css=CSS) as demo:
        gr.HTML(HERO_HTML)
        with gr.Tabs():
            # ---- CSV tab ----
            with gr.Tab("Analyze CV (CSV)"):
                with gr.Row():
                    with gr.Column(scale=4, min_width=320):
                        with gr.Group(elem_classes="card"):
                            gr.Markdown("#### Input")
                            csv_files = gr.File(file_count="multiple",
                                                label="CV CSV files (one per scan rate)")
                            csv_rates = gr.Textbox(label="Scan rates (V/s, comma-separated)",
                                                   placeholder="e.g. 0.05, 0.1, 0.2")
                            with gr.Accordion("Cell parameters (for nondimensionalization)", open=False):
                                csv_e0 = gr.Number(label="Formal potential E0 (V) - blank = auto", value=None)
                                csv_T = gr.Number(label="Temperature (K)", value=298.15)
                                csv_A = gr.Number(label="Electrode area (cm^2)", value=0.0707)
                                csv_C = gr.Number(label="Concentration (mM)", value=1.0)
                                csv_D = gr.Number(label="Diffusion coeff. (cm^2/s)", value=1e-5)
                                csv_n = gr.Number(label="Electrons n", value=1, precision=0)
                            with gr.Accordion("Inference settings", open=False):
                                csv_npost = gr.Slider(500, 4000, value=2000, step=500, label="Posterior samples")
                                csv_npred = gr.Slider(4, 24, value=8, step=2, label="Reconstruction draws")
                                csv_nscans = gr.Slider(1, 3, value=3, step=1, label="Scans used")
                            csv_btn = gr.Button("Analyze", variant="primary", size="lg")
                        _example_blocks([csv_files, csv_rates, csv_e0])
                    with gr.Column(scale=7, min_width=480):
                        csv_out, csv_dd, csv_state, csv_permech = _output_components()
                csv_btn.click(
                    analyze_cv_csv,
                    inputs=[csv_files, csv_rates, csv_e0, csv_T, csv_A, csv_C, csv_D,
                            csv_n, csv_npost, csv_npred, csv_nscans],
                    outputs=csv_out)
                csv_dd.change(on_mechanism_change, inputs=[csv_dd, csv_state],
                              outputs=csv_permech)

            # ---- Image tab ----
            with gr.Tab("Analyze CV (image)"):
                with gr.Row():
                    with gr.Column(scale=4, min_width=320):
                        with gr.Group(elem_classes="card"):
                            gr.Markdown("#### Input")
                            img_files = gr.File(file_count="multiple",
                                                label="CV plot images (one per scan rate)")
                            img_rates = gr.Textbox(label="Scan rates (V/s, comma-separated)",
                                                   placeholder="e.g. 0.05, 0.1, 0.2")
                            img_e0 = gr.Number(label="Formal potential E0 (V) - blank = auto", value=None)
                            img_thr = gr.Slider(20, 200, value=90, step=5, label="Digitization threshold")
                            with gr.Accordion("Axis overrides (if auto-detect fails)", open=False):
                                img_xmin = gr.Number(label="E min (V)", value=None)
                                img_xmax = gr.Number(label="E max (V)", value=None)
                                img_ymin = gr.Number(label="I min", value=None)
                                img_ymax = gr.Number(label="I max", value=None)
                            with gr.Accordion("Inference settings", open=False):
                                img_npost = gr.Slider(500, 4000, value=2000, step=500, label="Posterior samples")
                                img_npred = gr.Slider(4, 24, value=8, step=2, label="Reconstruction draws")
                                img_nscans = gr.Slider(1, 3, value=3, step=1, label="Scans used")
                            img_btn = gr.Button("Analyze", variant="primary", size="lg")
                    with gr.Column(scale=7, min_width=480):
                        img_out, img_dd, img_state, img_permech = _output_components()
                img_btn.click(
                    analyze_cv_image,
                    inputs=[img_files, img_rates, img_e0, img_thr, img_npost, img_npred,
                            img_nscans, img_xmin, img_xmax, img_ymin, img_ymax],
                    outputs=img_out)
                img_dd.change(on_mechanism_change, inputs=[img_dd, img_state],
                              outputs=img_permech)

            with gr.Tab("About"):
                gr.Markdown(ABOUT_MD)
    return demo


if __name__ == "__main__":
    build_app().launch(server_name="0.0.0.0", server_port=7860)