perovskite-hpt-test/scripts/41_nature_multround_bo_loop.py

#!/usr/bin/env python3
"""41_nature_multround_bo_loop.py — Multi-round Nature + GP+EI closed loop.

Each round:
  1. Screen batch
  2. Orthogonal dual-gate (UV-Vis ∧ PL)
  3. Hits → reproducibility
  4. Repro-ok → scale-up ledger
  5. Feed all measured Eg into GP+EI → propose next screen (exclude screened)
  6. Repeat for N rounds (default 3)

Single-round backward compat: scripts/40_nature_mobile_robot_loop.py calls
run_campaign(rounds=1) with classic output names.
"""

from __future__ import annotations

import argparse
import json
import sqlite3
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

ROOT = Path(__file__).resolve().parents[1]
_DEC = ROOT / "platform" / "decision"
if str(_DEC) not in sys.path:
    sys.path.insert(0, str(_DEC))

from bo_next_batch import propose_next_screen_batch  # noqa: E402
from heuristic_planner import decide, load_gates, summarize_reproducibility  # noqa: E402
from human_gate import (  # noqa: E402
    apply_override,
    ensure_human_gate_schema,
    get_pending_override,
    mark_override_applied,
)
from measurement_from_spectra import measurements_for_formula  # noqa: E402

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402
import numpy as np  # noqa: E402
import pandas as pd  # noqa: E402

SCHEMA = ROOT / "platform" / "schema" / "schema.sql"
GATES_PATH = ROOT / "platform" / "config" / "gates_perovskite_phase0.yaml"
HALIDE = ROOT / "data" / "pv_window_halides.csv"
AGG = ROOT / "data" / "nomad_perovskite_aggregated.csv"
SPECTRA_DIR = ROOT / "data" / "spectra"
# Naming: data/spectra/{formula}_uvvis.csv + {formula}_pl.csv → real parsers; else synthetic

DEFAULT_ROUNDS = 3
SCREEN_N = 12
RNG_SEED = 42


def utc_now() -> str:
    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


def load_chemistry_pool(rng: np.random.RandomState) -> pd.DataFrame:
    """Full human-chosen chemistry space (GGA pool + halide PV window)."""
    cols = ["formula", "band_gap_eV"]
    frames: list[pd.DataFrame] = []
    if HALIDE.exists():
        frames.append(pd.read_csv(HALIDE).dropna(subset=["band_gap_eV"])[cols])
    if AGG.exists():
        agg = pd.read_csv(AGG)
        pool = agg[
            (agg.functional_class == "gga")
            & (agg.band_gap_eV > 0.05)
            & (agg.band_gap_eV < 5.0)
        ][cols].copy()
        frames.append(pool)
    if not frames:
        raise SystemExit("No chemistry pool data found")
    out = pd.concat(frames, ignore_index=True).drop_duplicates("formula")
    out = out.dropna(subset=["band_gap_eV"]).reset_index(drop=True)
    # Shuffle once for stable but mixed ordering
    return out.sample(frac=1.0, random_state=rng).reset_index(drop=True)


def initial_screen_batch(pool: pd.DataFrame, n: int, rng: np.random.RandomState) -> pd.DataFrame:
    """Round-0 screen: mix near-1.4 eV hits with out-of-window distractors."""
    cols = ["formula", "band_gap_eV"]
    near = pool.assign(dist=(pool["band_gap_eV"] - 1.4).abs()).sort_values("dist")
    n_near = max(n // 3, 3)
    parts = [near.head(n_near)[cols]]
    far = pool[(pool.band_gap_eV < 1.05) | (pool.band_gap_eV > 1.85)]
    n_far = n - sum(len(p) for p in parts)
    if n_far > 0 and len(far) > 0:
        parts.append(far.sample(n=min(n_far, len(far)), random_state=rng)[cols])
    out = pd.concat(parts, ignore_index=True).drop_duplicates("formula").head(n)
    if len(out) < n:
        extra = pool[~pool.formula.isin(out.formula)]
        need = n - len(out)
        if need > 0 and len(extra):
            out = pd.concat(
                [out, extra.sample(n=min(need, len(extra)), random_state=rng)[cols]],
                ignore_index=True,
            )
    return out.reset_index(drop=True)


def formulas_to_frame(formulas: list[str], pool: pd.DataFrame) -> pd.DataFrame:
    """Map proposed formulas back to pool rows (true Eg for offline sim)."""
    sub = pool[pool.formula.isin(formulas)].drop_duplicates("formula")
    # Preserve proposal order
    order = {f: i for i, f in enumerate(formulas)}
    sub = sub.copy()
    sub["_ord"] = sub.formula.map(order)
    return sub.sort_values("_ord").drop(columns="_ord").reset_index(drop=True)


def synth_measurements(eg: float, rng: np.random.RandomState, noise: float = 1.0) -> tuple[dict, dict]:
    eg_m = float(eg) + rng.normal(0, 0.025 * noise)
    target = 1.40
    shape = float(np.clip(1.0 - abs(eg_m - target) / 0.55 + rng.normal(0, 0.05 * noise), 0.0, 1.0))
    ms = {"band_gap_eV": eg_m, "absorbance_shape_score": shape}
    peak = 1240.0 / max(eg_m, 0.3) + rng.normal(0, 15.0 * noise)
    fwhm = float(np.clip(35.0 + 80.0 * abs(eg_m - target) + rng.normal(0, 5.0 * noise), 15.0, 200.0))
    intensity = float(np.clip(0.85 - 0.9 * abs(eg_m - target) + rng.normal(0, 0.06 * noise), 0.0, 1.0))
    nmr = {"peak_nm": float(peak), "fwhm_nm": fwhm, "intensity": intensity}
    return ms, nmr


def get_measurements(formula: str, eg: float, rng: np.random.RandomState, noise: float = 1.0):
    """Prefer on-disk spectra for formula; else synthetic from band_gap."""
    pair = measurements_for_formula(formula, SPECTRA_DIR)
    if pair is not None:
        return pair[0], pair[1], "spectra"
    ms, nmr = synth_measurements(eg, rng, noise=noise)
    return ms, nmr, "synthetic"


def init_db(path: Path) -> sqlite3.Connection:
    path.parent.mkdir(parents=True, exist_ok=True)
    if path.exists():
        path.unlink()
    conn = sqlite3.connect(path)
    conn.executescript(SCHEMA.read_text(encoding="utf-8"))
    ensure_human_gate_schema(conn)
    # Extra tables for multi-round BO audit (additive; 40 single-round still works)
    conn.executescript(
        """
        CREATE TABLE IF NOT EXISTS loop_rounds (
            round_idx INTEGER NOT NULL,
            campaign_id TEXT NOT NULL,
            n_screened INTEGER,
            n_ms_pass INTEGER,
            n_nmr_pass INTEGER,
            n_hits INTEGER,
            n_repro_ok INTEGER,
            n_scaleup INTEGER,
            bo_method TEXT,
            bo_proposed_json TEXT,
            created_at TEXT NOT NULL,
            PRIMARY KEY (campaign_id, round_idx)
        );
        """
    )
    stations = [
        ("SYNTH-01", "synth", "ISynth-like synthesizer", "bay-A", "Chemspeed-analog"),
        ("MS-01", "analysis_ms", "UV-Vis (MS-role)", "bay-B", "benchtop-UVVis"),
        ("NMR-01", "analysis_nmr", "PL (NMR-role)", "bay-B", "benchtop-PL"),
        ("SCALE-01", "scaleup", "Scale-up reactor", "bay-C", "scaleup-reactor"),
    ]
    conn.executemany(
        "INSERT INTO stations(station_id,kind,label,location,instrument) VALUES (?,?,?,?,?)",
        stations,
    )
    return conn


def log_jsonl(fp, event: dict) -> None:
    fp.write(json.dumps(event, ensure_ascii=False) + "\n")


def insert_run(conn, run_id, campaign_id, condition_id, batch_role, parent_hit_id, meta):
    conn.execute(
        """INSERT INTO runs(run_id,campaign_id,condition_id,batch_role,parent_hit_id,
           station_synth,status,created_at,meta_json)
           VALUES (?,?,?,?,?,?,?,?,?)""",
        (
            run_id,
            campaign_id,
            condition_id,
            batch_role,
            parent_hit_id,
            "SYNTH-01",
            "running",
            utc_now(),
            json.dumps(meta),
        ),
    )


def move(conn, run_id, frm, to, robot):
    conn.execute(
        """INSERT INTO sample_moves(run_id,from_station,to_station,robot_id,t)
           VALUES (?,?,?,?,?)""",
        (run_id, frm, to, robot, utc_now()),
    )


def store_meas(conn, run_id, modality, station, metrics, pass_fail, version):
    conn.execute(
        """INSERT INTO measurements(run_id,modality,station_id,raw_path,metrics_json,
           pass_fail,threshold_version,created_at)
           VALUES (?,?,?,?,?,?,?,?)""",
        (
            run_id,
            modality,
            station,
            f"sim/{modality}/{run_id}.json",
            json.dumps(metrics),
            int(pass_fail) if pass_fail is not None else None,
            version,
            utc_now(),
        ),
    )


def store_decision(conn, run_id, result, human_meta=None):
    human_meta = human_meta or {}
    rationale = {
        "ms_reasons": list(result.ms_reasons),
        "nmr_reasons": list(result.nmr_reasons),
        "modality_ms": result.modality_ms,
        "modality_nmr": result.modality_nmr,
    }
    if human_meta.get("human_override"):
        rationale["human_override"] = human_meta["human_override"]
        rationale["heuristic_next_action"] = human_meta.get("heuristic_next_action")
    conn.execute(
        """INSERT INTO decisions(run_id,ms_pass,nmr_pass,hit,next_action,
           threshold_version,rationale_json,created_at,
           human_override,human_note,human_at,human_by)
           VALUES (?,?,?,?,?,?,?,?,?,?,?,?)""",
        (
            run_id,
            int(result.ms_pass),
            int(result.nmr_pass),
            int(result.hit),
            result.next_action,
            result.threshold_version,
            json.dumps(rationale),
            utc_now(),
            human_meta.get("human_override"),
            human_meta.get("human_note"),
            human_meta.get("human_at"),
            human_meta.get("human_by"),
        ),
    )


def analyze_run(conn, run_id, ms, nmr, gates, context="screen", repro_ok=None):
    version = gates["threshold_version"]
    move(conn, run_id, "SYNTH-01", "MS-01", "ROBOT-A")
    move(conn, run_id, "MS-01", "NMR-01", "ROBOT-B")
    heur = decide(ms, nmr, gates, context=context, repro_ok=repro_ok)
    human_meta = {}
    pend = get_pending_override(conn, run_id)
    if pend:
        result = apply_override(
            heur,
            pend["human_override"],
            note=pend.get("human_note"),
            by=pend.get("human_by"),
            at=pend.get("human_at"),
        )
        mark_override_applied(
            conn,
            pend["gate_id"],
            heuristic_next_action=heur.next_action,
            final_next_action=result.next_action,
        )
        human_meta = {
            "human_override": pend["human_override"],
            "human_note": pend.get("human_note"),
            "human_at": pend.get("human_at"),
            "human_by": pend.get("human_by"),
            "heuristic_next_action": heur.next_action,
        }
    else:
        result = heur
    store_meas(conn, run_id, "uvvis", "MS-01", ms, result.ms_pass, version)
    store_meas(conn, run_id, "pl", "NMR-01", nmr, result.nmr_pass, version)
    store_decision(conn, run_id, result, human_meta=human_meta)
    conn.execute("UPDATE runs SET status=? WHERE run_id=?", ("decided", run_id))
    return result


def plot_funnel(stats: dict, path: Path, title: str = "Nature-style decision funnel") -> None:
    labels = ["screened", "ms_pass", "nmr_pass", "hits", "repro_ok", "scaleup"]
    vals = [
        stats["n_screened"],
        stats["n_ms_pass"],
        stats["n_nmr_pass"],
        stats["n_hits"],
        stats["n_repro_ok"],
        len(stats.get("scaleup_queue", [])),
    ]
    path.parent.mkdir(parents=True, exist_ok=True)
    fig, ax = plt.subplots(figsize=(7.5, 4.2))
    colors = ["#4C72B0", "#55A868", "#C44E52", "#8172B2", "#CCB974", "#64B5CD"]
    bars = ax.bar(labels, vals, color=colors)
    ax.set_ylabel("count")
    ax.set_title(title)
    for b, v in zip(bars, vals):
        ax.text(b.get_x() + b.get_width() / 2, v + 0.1, str(v), ha="center", va="bottom", fontsize=10)
    ax.set_ylim(0, max(vals + [1]) * 1.25)
    fig.tight_layout()
    fig.savefig(path, dpi=140)
    plt.close(fig)


def plot_funnel_per_round(round_stats: list[dict], path: Path) -> None:
    """Stacked / grouped funnel counts across rounds."""
    path.parent.mkdir(parents=True, exist_ok=True)
    metrics = ["n_screened", "n_ms_pass", "n_nmr_pass", "n_hits", "n_repro_ok", "n_scaleup"]
    labels = ["screened", "ms_pass", "nmr_pass", "hits", "repro_ok", "scaleup"]
    n_r = len(round_stats)
    x = np.arange(len(labels))
    width = 0.8 / max(n_r, 1)
    fig, ax = plt.subplots(figsize=(9.0, 4.6))
    cmap = plt.cm.viridis(np.linspace(0.2, 0.85, n_r))
    for i, rs in enumerate(round_stats):
        vals = [rs[m] for m in metrics]
        ax.bar(x + i * width, vals, width=width, label=f"R{rs['round']}", color=cmap[i])
    ax.set_xticks(x + width * (n_r - 1) / 2)
    ax.set_xticklabels(labels)
    ax.set_ylabel("count")
    ax.set_title("Nature multi-round funnel (per round)")
    ax.legend(fontsize=8, ncol=min(n_r, 5))
    ax.grid(axis="y", alpha=0.3)
    fig.tight_layout()
    fig.savefig(path, dpi=140)
    plt.close(fig)


def run_one_screen_round(
    conn,
    ledger,
    campaign_id: str,
    round_idx: int,
    batch_df: pd.DataFrame,
    gates: dict,
    rng: np.random.RandomState,
    n_reps: int,
    measured: list[dict],
    screened: set[str],
    cumulative_scaleups: list[dict],
) -> dict[str, Any]:
    """Screen → dual-gate → repro → scaleup for one batch; mutate measured/screened."""
    n_ms_pass = n_nmr_pass = n_hits = 0
    hits: list[dict] = []

    for i, row in batch_df.iterrows():
        formula = str(row["formula"])
        eg_true = float(row["band_gap_eV"])
        run_id = f"R{round_idx:02d}-SCR-{i:03d}"
        insert_run(
            conn,
            run_id,
            campaign_id,
            formula,
            "screen",
            None,
            {"true_band_gap_eV": eg_true, "round": round_idx},
        )
        ms, nmr, src = get_measurements(formula, eg_true, rng, noise=1.0)
        result = analyze_run(conn, run_id, ms, nmr, gates, context="screen")
        screened.add(formula)
        measured.append(
            {
                "formula": formula,
                "eg": ms["band_gap_eV"],
                "true_eg": eg_true,
                "round": round_idx,
                "metrics_source": src,
            }
        )
        n_ms_pass += int(result.ms_pass)
        n_nmr_pass += int(result.nmr_pass)
        if result.hit:
            n_hits += 1
            hits.append(
                {
                    "run_id": run_id,
                    "condition_id": formula,
                    "eg": ms["band_gap_eV"],
                    "ms": ms,
                    "nmr": nmr,
                }
            )
        log_jsonl(
            ledger,
            {
                "event": "screen_decision",
                "round": round_idx,
                "run_id": run_id,
                "condition_id": formula,
                "ms_pass": result.ms_pass,
                "nmr_pass": result.nmr_pass,
                "hit": result.hit,
                "next_action": result.next_action,
                "ms": ms,
                "nmr": nmr,
                "metrics_source": src,
                "human_override": None if not getattr(result, "ms_reasons", None) else (
                    next((x for x in result.ms_reasons if x.startswith("human_override=")), None)
                ),
                "t": utc_now(),
            },
        )

    repro_ok_hits: list[dict] = []
    for h in hits:
        parent = h["run_id"]
        repro_id = f"R{round_idx:02d}-REPRO-{parent}"
        rep_decisions = []
        eg_vals: list[float] = []
        inten_vals: list[float] = []
        for k in range(n_reps):
            run_id = f"{repro_id}-R{k}"
            insert_run(
                conn,
                run_id,
                campaign_id,
                h["condition_id"],
                "reproduce",
                parent,
                {"rep": k, "round": round_idx},
            )
            ms, nmr, _src = get_measurements(h["condition_id"], h["eg"], rng, noise=0.55)
            result = analyze_run(conn, run_id, ms, nmr, gates, context="reproduce")
            rep_decisions.append(result)
            eg_vals.append(ms["band_gap_eV"])
            inten_vals.append(nmr["intensity"])
            measured.append(
                {
                    "formula": h["condition_id"],
                    "eg": ms["band_gap_eV"],
                    "round": round_idx,
                    "role": "repro",
                }
            )
            log_jsonl(
                ledger,
                {
                    "event": "repro_rep",
                    "round": round_idx,
                    "parent_hit_id": parent,
                    "run_id": run_id,
                    "hit": result.hit,
                    "ms": ms,
                    "nmr": nmr,
                    "t": utc_now(),
                },
            )

        stats = summarize_reproducibility(rep_decisions, eg_vals, inten_vals, gates)
        conn.execute(
            """INSERT INTO reproducibility_batches(
                repro_id,parent_hit_id,campaign_id,n_reps,n_pass,pass_rate,
                mean_metrics_json,std_metrics_json,reproducible,created_at)
               VALUES (?,?,?,?,?,?,?,?,?,?)""",
            (
                repro_id,
                parent,
                campaign_id,
                stats["n_reps"],
                stats["n_pass"],
                stats["pass_rate"],
                json.dumps({"eg_mean": float(np.mean(eg_vals)), "pl_inten_mean": float(np.mean(inten_vals))}),
                json.dumps({"eg_std": stats["eg_std"], "pl_intensity_cv": stats["pl_intensity_cv"]}),
                int(stats["reproducible"]),
                utc_now(),
            ),
        )
        batch_action = "scaleup" if stats["reproducible"] else "fail"
        log_jsonl(
            ledger,
            {
                "event": "repro_batch",
                "round": round_idx,
                "repro_id": repro_id,
                "parent_hit_id": parent,
                "stats": stats,
                "next_action": batch_action,
                "t": utc_now(),
            },
        )
        if stats["reproducible"]:
            repro_ok_hits.append(h)

    scaleup_queue: list[dict] = []
    for h in repro_ok_hits:
        run_id = f"R{round_idx:02d}-SCALE-{h['run_id']}"
        insert_run(conn, run_id, campaign_id, h["condition_id"], "scaleup", h["run_id"], {"round": round_idx})
        move(conn, run_id, "NMR-01", "SCALE-01", "ROBOT-A")
        ms, nmr, _src = get_measurements(h["condition_id"], h["eg"], rng, noise=0.4)
        result = analyze_run(conn, run_id, ms, nmr, gates, context="scaleup")
        conn.execute("UPDATE runs SET status=? WHERE run_id=?", ("done", run_id))
        entry = {
            "run_id": run_id,
            "condition_id": h["condition_id"],
            "parent_hit_id": h["run_id"],
            "ms_pass": result.ms_pass,
            "nmr_pass": result.nmr_pass,
            "round": round_idx,
        }
        scaleup_queue.append(entry)
        cumulative_scaleups.append(entry)
        log_jsonl(
            ledger,
            {
                "event": "scaleup",
                "round": round_idx,
                "run_id": run_id,
                "condition_id": h["condition_id"],
                "hit": result.hit,
                "t": utc_now(),
            },
        )

    return {
        "round": round_idx,
        "n_screened": int(len(batch_df)),
        "n_ms_pass": int(n_ms_pass),
        "n_nmr_pass": int(n_nmr_pass),
        "n_hits": int(n_hits),
        "n_repro_ok": int(len(repro_ok_hits)),
        "n_scaleup": int(len(scaleup_queue)),
        "scaleup_queue": scaleup_queue,
        "hits": hits,
        "repro_ok_hits": repro_ok_hits,
    }


def run_campaign(
    *,
    rounds: int = DEFAULT_ROUNDS,
    batch_size: int = SCREEN_N,
    seed: int = RNG_SEED,
    out_db: Path | None = None,
    out_jsonl: Path | None = None,
    out_summary: Path | None = None,
    out_fig: Path | None = None,
    out_fig_rounds: Path | None = None,
    campaign_id: str | None = None,
) -> dict[str, Any]:
    """Run N rounds of screen→gate→repro→scaleup→GP+EI. Returns summary dict."""
    rng = np.random.RandomState(seed)
    gates = load_gates(GATES_PATH)
    version = gates["threshold_version"]
    n_reps = int(gates.get("reproducibility", {}).get("n_reps_sim", 4))

    multi = rounds > 1
    if out_db is None:
        out_db = ROOT / "output" / ("nature_multround_bo_loop.db" if multi else "nature_robot_loop.db")
    if out_jsonl is None:
        out_jsonl = ROOT / "output" / ("nature_multround_bo_loop.jsonl" if multi else "nature_robot_loop.jsonl")
    if out_summary is None:
        out_summary = ROOT / "output" / (
            "nature_multround_bo_loop_summary.json" if multi else "nature_robot_loop_summary.json"
        )
    if out_fig is None:
        out_fig = ROOT / "figures" / (
            "fig_nature_multround_cumulative_funnel.png" if multi else "fig_nature_robot_decision_funnel.png"
        )
    if out_fig_rounds is None and multi:
        out_fig_rounds = ROOT / "figures" / "fig_nature_multround_funnel_per_round.png"
    if campaign_id is None:
        campaign_id = "CAMP-NATURE-MULTI-001" if multi else "CAMP-NATURE-SIM-001"

    pool = load_chemistry_pool(rng)
    conn = init_db(out_db)
    conn.execute(
        """INSERT INTO campaigns(campaign_id,name,chemistry_space,threshold_version,created_at,notes)
           VALUES (?,?,?,?,?,?)""",
        (
            campaign_id,
            f"Nature mobile-robot {'multi-round BO' if multi else 'offline sim'}",
            "halide PV window / GGA pool",
            version,
            utc_now(),
            f"rounds={rounds}; chemistry space human-chosen; navigation GP+EI autonomous",
        ),
    )

    out_jsonl.parent.mkdir(parents=True, exist_ok=True)
    ledger = out_jsonl.open("w", encoding="utf-8")
    log_jsonl(
        ledger,
        {
            "event": "campaign_start",
            "campaign_id": campaign_id,
            "threshold_version": version,
            "rounds": rounds,
            "batch_size": batch_size,
            "pool_size": int(len(pool)),
            "t": utc_now(),
        },
    )

    measured: list[dict] = []
    screened: set[str] = set()
    cumulative_scaleups: list[dict] = []
    round_stats: list[dict] = []
    bo_proposals: list[dict] = []

    # Round 1 batch: heuristic mix; later rounds: GP+EI
    batch_df = initial_screen_batch(pool, batch_size, rng)

    for r in range(1, rounds + 1):
        log_jsonl(ledger, {"event": "round_start", "round": r, "n_batch": len(batch_df), "t": utc_now()})
        rs = run_one_screen_round(
            conn,
            ledger,
            campaign_id,
            r,
            batch_df,
            gates,
            rng,
            n_reps,
            measured,
            screened,
            cumulative_scaleups,
        )

        # GP+EI next batch from ALL measured Eg so far (exclude screened)
        remaining = pool[~pool.formula.isin(screened)]["formula"].tolist()
        # Deduplicate measured by formula keeping last measured eg
        by_f: dict[str, float] = {}
        for m in measured:
            by_f[m["formula"]] = float(m["eg"])
        known = [{"formula": f, "eg": eg} for f, eg in by_f.items()]
        bo = propose_next_screen_batch(known, remaining, batch_size=batch_size, random_state=seed + r)
        bo_proposals.append({"round": r, **bo})
        log_jsonl(ledger, {"event": "bo_propose", "round": r, **bo, "t": utc_now()})

        conn.execute(
            """INSERT INTO loop_rounds(round_idx,campaign_id,n_screened,n_ms_pass,n_nmr_pass,
               n_hits,n_repro_ok,n_scaleup,bo_method,bo_proposed_json,created_at)
               VALUES (?,?,?,?,?,?,?,?,?,?,?)""",
            (
                r,
                campaign_id,
                rs["n_screened"],
                rs["n_ms_pass"],
                rs["n_nmr_pass"],
                rs["n_hits"],
                rs["n_repro_ok"],
                rs["n_scaleup"],
                bo["method"],
                json.dumps(bo["proposed"]),
                utc_now(),
            ),
        )
        rs_out = {k: rs[k] for k in ("round", "n_screened", "n_ms_pass", "n_nmr_pass", "n_hits", "n_repro_ok", "n_scaleup")}
        rs_out["bo_n_proposed"] = len(bo["proposed"])
        rs_out["bo_top"] = bo["formulas"][:5]
        round_stats.append(rs_out)

        if r < rounds:
            if not bo["formulas"]:
                log_jsonl(ledger, {"event": "bo_exhausted", "round": r, "t": utc_now()})
                break
            batch_df = formulas_to_frame(bo["formulas"], pool)
            # Distinctness check vs previous screened
            assert not set(batch_df.formula) & screened, "BO proposed already-screened formulas"

    conn.commit()

    # Cumulative funnel (all rounds)
    cum = {
        "n_screened": int(sum(s["n_screened"] for s in round_stats)),
        "n_ms_pass": int(sum(s["n_ms_pass"] for s in round_stats)),
        "n_nmr_pass": int(sum(s["n_nmr_pass"] for s in round_stats)),
        "n_hits": int(sum(s["n_hits"] for s in round_stats)),
        "n_repro_ok": int(sum(s["n_repro_ok"] for s in round_stats)),
        "scaleup_queue": cumulative_scaleups,
    }
    plot_funnel(
        cum,
        out_fig,
        title=f"Nature-style decision funnel ({'multi-round cumulative' if multi else 'Phase-0'})",
    )
    if multi and out_fig_rounds is not None:
        plot_funnel_per_round(round_stats, out_fig_rounds)

    # Distinctness of successive BO batches
    proposed_sets = [set(p["formulas"]) for p in bo_proposals]
    pairwise_overlap = []
    for i in range(len(proposed_sets) - 1):
        a, b = proposed_sets[i], proposed_sets[i + 1]
        pairwise_overlap.append(
            {
                "rounds": [i + 1, i + 2],
                "overlap": sorted(a & b),
                "n_overlap": len(a & b),
            }
        )

    summary = {
        "rounds_requested": rounds,
        "rounds_completed": len(round_stats),
        "threshold_version": version,
        "n_reps_per_hit": n_reps,
        "batch_size": batch_size,
        "pool_size": int(len(pool)),
        "n_unique_screened": len(screened),
        "cumulative": {
            "n_screened": cum["n_screened"],
            "n_ms_pass": cum["n_ms_pass"],
            "n_nmr_pass": cum["n_nmr_pass"],
            "n_hits": cum["n_hits"],
            "n_repro_ok": cum["n_repro_ok"],
            "n_scaleup": len(cumulative_scaleups),
            "scaleup_queue": cumulative_scaleups,
        },
        "per_round": round_stats,
        "bo_proposals": [
            {
                "round": p["round"],
                "method": p["method"],
                "formulas": p["formulas"],
                "acquisition_scores": p["acquisition_scores"],
            }
            for p in bo_proposals
        ],
        "bo_pairwise_overlap": pairwise_overlap,
        "db": str(out_db.relative_to(ROOT)),
        "jsonl": str(out_jsonl.relative_to(ROOT)),
        "figure": str(out_fig.relative_to(ROOT)),
        "figure_per_round": str(out_fig_rounds.relative_to(ROOT)) if (multi and out_fig_rounds) else None,
        # Backward-compat flat keys (script 40 consumers)
        "n_screened": cum["n_screened"],
        "n_ms_pass": cum["n_ms_pass"],
        "n_nmr_pass": cum["n_nmr_pass"],
        "n_hits": cum["n_hits"],
        "n_repro_ok": cum["n_repro_ok"],
        "scaleup_queue": cumulative_scaleups,
        "bo_stub": bo_proposals[-1] if bo_proposals else {"method": "gp_ei", "proposed": []},
    }
    # Prefer real key name in multi-round summaries
    if multi:
        summary["bo_last"] = summary.pop("bo_stub")
    else:
        # keep bo_stub alias but also expose bo_last with real GP+EI payload
        summary["bo_last"] = summary["bo_stub"]

    out_summary.write_text(json.dumps(summary, indent=2, ensure_ascii=False), encoding="utf-8")
    ledger.close()
    conn.close()

    print("=== Nature multi-round BO loop summary ===" if multi else "=== Nature mobile-robot loop summary ===")
    print(f"rounds       : {summary['rounds_completed']} / {rounds}")
    print(f"n_screened   : {summary['n_screened']}")
    print(f"n_ms_pass    : {summary['n_ms_pass']}  (UV-Vis / MS-role)")
    print(f"n_nmr_pass   : {summary['n_nmr_pass']}  (PL / NMR-role)")
    print(f"n_hits       : {summary['n_hits']}  (dual-pass)")
    print(f"n_repro_ok   : {summary['n_repro_ok']}")
    print(f"scaleup_queue: {[s['condition_id'] for s in cumulative_scaleups]}")
    for p in bo_proposals:
        print(f"BO R{p['round']} [{p['method']}]: {p['formulas'][:4]}{'…' if len(p['formulas'])>4 else ''}")
    print(f"db           : {out_db}")
    print(f"jsonl        : {out_jsonl}")
    print(f"figure       : {out_fig}")
    if multi and out_fig_rounds:
        print(f"figure_rounds: {out_fig_rounds}")
    return summary


def main(argv: list[str] | None = None) -> int:
    p = argparse.ArgumentParser(description="Nature multi-round GP+EI closed loop")
    p.add_argument("--rounds", type=int, default=DEFAULT_ROUNDS, help="Number of screen→BO rounds (default 3)")
    p.add_argument("--batch-size", type=int, default=SCREEN_N)
    p.add_argument("--seed", type=int, default=RNG_SEED)
    args = p.parse_args(argv)
    run_campaign(rounds=max(1, args.rounds), batch_size=args.batch_size, seed=args.seed)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())