#!/usr/bin/env python3
"""Modal app for the MD pass -- self-contained surface-tension MD (GROMACS), pure fluids AND blends.

Each Modal function runs in its own isolated container (no PyTorch shares with GROMACS), and the
toolchain comes from conda packages -- **never built from source on the box**. A job goes
SMILES -> acpype GAFF2/AM1-BCC topology (per component) -> merged multi-molecule topology ->
liquid box -> NPT -> slab -> NVT production -> surface tension via `gmx energy` (#Surf*SurfTen).
GROMACS is the right tool because it exposes the pressure tensor surface tension needs.

Entrypoints:
  modal run scripts/md_pass/modal_md_app.py::toolcheck_main
  modal run scripts/md_pass/modal_md_app.py::sigma_one --smiles CCCCC --name pentane
  modal run scripts/md_pass/modal_md_app.py::components_3seed      # 8 components x 3 seeds (mean +/- SEM)
  modal run scripts/md_pass/modal_md_app.py::blends_sigmaT         # 3 blends x 4 temperatures (Eotvos)

Results land in the 'cf10-md-out' modal.Volume. All MD outputs are computational tier-1 evidence
(sigma), never measured wetlab data. The pentane control calibrates the GAFF2 force-field offset.
"""
from __future__ import annotations

import json

import modal

CONDA = "/opt/conda/envs/md/bin"
IMAGE = (
    modal.Image.from_registry("condaforge/miniforge3:latest", add_python="3.11")
    .run_commands("mamba create -y -n md -c conda-forge python=3.11 gromacs ambertools acpype rdkit numpy",
                  "mamba clean -ay")
)

app = modal.App("cf10-md")
VOL = modal.Volume.from_name("cf10-md-out", create_if_missing=True)


@app.function(image=IMAGE, gpu="L4", timeout=600)
def toolcheck() -> dict:
    import subprocess
    out = {}
    for label, cmd in [("gmx", [f"{CONDA}/gmx", "--version"]),
                       ("acpype", [f"{CONDA}/acpype", "--version"]),
                       ("rdkit", [f"{CONDA}/python", "-c", "import rdkit; print(rdkit.__version__)"])]:
        try:
            r = subprocess.run(cmd, capture_output=True, text=True, timeout=120)
            out[label] = (r.stdout or r.stderr).strip().splitlines()[:1]
        except Exception as e:
            out[label] = [f"ERR {type(e).__name__}: {e}"]
    return out


# ---- the physics, run inside the container via the conda interpreter ----
MD_DRIVER = r'''
import json, subprocess, sys, os
from pathlib import Path
CONDA = "__CONDA__"
NTOMP = int(os.environ.get("OMP_NUM_THREADS") or os.cpu_count() or 4)
os.environ["PATH"] = CONDA + ":" + os.environ.get("PATH", "")
os.environ["AMBERHOME"] = "/opt/conda/envs/md"
from rdkit import Chem
from rdkit.Chem import AllChem, Descriptors

def sh(c): return subprocess.run(c, shell=True, capture_output=True, text=True)

def split_itp(text):
    """Split an acpype _GMX.itp into (atomtypes_block, moleculetype_onward)."""
    i = text.find("[ moleculetype ]")
    return text[:i], text[i:]

def atomtype_lines(block):
    out, inat = [], False
    for ln in block.splitlines():
        s = ln.strip()
        if s.startswith("["):
            inat = s.startswith("[ atomtypes ]"); continue
        if inat and s and not s.startswith(";"):
            out.append(ln)
    return out

def parametrize(smiles, resname):
    """SMILES -> 3D -> acpype GAFF2/AM1-BCC. Returns (gro_path, atomtype_lines, moltype_block, natom, mw)."""
    m = Chem.AddHs(Chem.MolFromSmiles(smiles))
    if AllChem.EmbedMolecule(m, AllChem.ETKDGv3()) != 0:
        return None
    AllChem.MMFFOptimizeMolecule(m)
    pdb = f"{resname}.pdb"; Chem.MolToPDBFile(m, pdb)
    r = sh(f"acpype -i {pdb} -b {resname} -n {Chem.GetFormalCharge(m)} -a gaff2 -c bcc")
    itp = Path(f"{resname}.acpype/{resname}_GMX.itp")
    gro = Path(f"{resname}.acpype/{resname}_GMX.gro")
    if not itp.exists() or not gro.exists():
        return None
    at, mt = split_itp(itp.read_text())
    return {"gro": str(gro), "atomtypes": atomtype_lines(at), "moltype": mt,
            "natom": m.GetNumAtoms(), "mw": Descriptors.MolWt(m), "resname": resname}

def _parse_xvg(path):
    import numpy as np
    t, cols = [], []
    for ln in Path(path).read_text().splitlines():
        if ln.startswith(("#", "@")):
            continue
        parts = ln.split()
        if len(parts) >= 2:
            t.append(float(parts[0])); cols.append([float(x) for x in parts[1:]])
    return np.array(t), np.array(cols)

def _gk_viscosity(edr, T_K):
    """Green-Kubo mu from fine-sampled off-diagonal pressure ACFs (3-component average).
    Retains negative draws (noise floor) per the repo's F2 convention."""
    import numpy as np
    r = sh(f"printf 'Pres-XY\\nPres-XZ\\nPres-YZ\\nVolume\\n\\n' | gmx energy -f {edr} -o pres.xvg")
    if not Path("pres.xvg").exists():
        return None
    t, cols = _parse_xvg("pres.xvg")
    if len(t) < 1000 or cols.shape[1] < 4:
        return None
    dt_ps = float(t[1] - t[0])
    V_m3 = float(np.mean(cols[:, 3])) * 1e-27          # nm3 -> m3
    kT = 1.380649e-23 * T_K
    mus = []
    n = len(t)
    upto = n // 4                                       # integrate ACF over the first quarter
    for k in range(3):
        p = cols[:, k] - np.mean(cols[:, k])            # bar
        f = np.fft.rfft(p, 2 * n)
        acf = np.fft.irfft(f * np.conj(f))[:n] / (n - np.arange(n))
        integ = np.cumsum(acf[:upto]) * dt_ps * 1e-12   # bar^2 * s
        mu_pa_s = (V_m3 / kT) * integ[-1] * 1e10        # (1e5 Pa/bar)^2 = 1e10
        mus.append(mu_pa_s * 1e3)                       # mPa*s
    mu = float(np.mean(mus))
    return {"mu_mPas_gk": round(mu, 4), "mu_components": [round(m, 4) for m in mus],
            "noise_floor": bool(any(m < 0 for m in mus)), "n_frames": int(n)}

def _epsilon(tpr, xtc, T_K):
    r = sh(f"printf '0\\n' | gmx dipoles -f {xtc} -s {tpr} -temp {T_K} -o mtot.xvg")
    for line in (r.stdout + r.stderr).splitlines():
        s = line.strip()
        if s.startswith("Epsilon"):
            try:
                return float(s.split("=")[1].split()[0])
            except Exception:
                pass
    return None

def _demix_index(tpr, xtc, parsed, counts):
    """Subcell mole-fraction variance of species 0 vs binomial expectation (0=mixed, ->1 demixed)."""
    import numpy as np
    if len(parsed) < 2:
        return None
    sh(f"printf '0\\n' | gmx trjconv -f {xtc} -s {tpr} -dump 1e9 -o final.gro 2>/dev/null")
    if not Path("final.gro").exists():
        return None
    lines = Path("final.gro").read_text().splitlines()
    natoms = int(lines[1])
    box = [float(x) for x in lines[2 + natoms].split()[:3]]
    coords = []
    for ln in lines[2:2 + natoms]:
        try:
            coords.append((float(ln[20:28]), float(ln[28:36]), float(ln[36:44])))
        except Exception:
            return None
    coords = np.array(coords)
    # species assignment by atom index: [molecules] lists species sequentially
    n0 = counts[parsed[0]["resname"]] * parsed[0]["natom"]
    mol_pos, species = [], []
    idx = 0
    for si, p in enumerate(parsed):
        for _ in range(counts[p["resname"]]):
            mol_pos.append(coords[idx])                  # first atom of each molecule
            species.append(si)
            idx += p["natom"]
    mol_pos = np.array(mol_pos); species = np.array(species)
    cells = np.minimum((mol_pos / (np.array(box) / 3)).astype(int), 2)
    cell_id = cells[:, 0] * 9 + cells[:, 1] * 3 + cells[:, 2]
    x_global = float(np.mean(species == 0))
    fracs, weights = [], []
    for c in range(27):
        m = cell_id == c
        if m.sum() >= 5:
            fracs.append(float(np.mean(species[m] == 0)))
            weights.append(int(m.sum()))
    if len(fracs) < 8:
        return None
    var_local = float(np.average((np.array(fracs) - x_global) ** 2, weights=weights))
    var_max = x_global * (1 - x_global) or 1.0
    return round(min(1.0, var_local / var_max), 3)

def run_properties(spec, parsed, counts, T_K, prod_ps, seed, base):
    """Cubic-box production with frames + fine pressure sampling -> epsilon, GK viscosity,
    density, and (blends) a demix index. No slab -- bulk properties."""
    Path("prop.mdp").write_text(base + f"nsteps={int(prod_ps/0.002)}\npcoupl=berendsen\n"
                                "pcoupltype=isotropic\nref-p=1.0\ncompressibility=4.5e-5\ntau-p=2.0\n"
                                f"gen-vel=yes\ngen-temp={T_K}\ngen-seed={seed}\n"
                                "nstxout-compressed=1000\nnstcalcenergy=10\nnstenergy=10\n")
    if sh("gmx grompp -f prop.mdp -c npt.gro -p topol.top -o prop -maxwarn 10").returncode:
        return {"status": "grompp_prop_failed", "label": spec["label"]}
    sh(f"gmx mdrun -deffnm prop -ntmpi 1 -ntomp {NTOMP} -nb cpu")
    if not Path("prop.edr").exists():
        return {"status": "prop_failed", "label": spec["label"]}
    eps = _epsilon("prop.tpr", "prop.xtc", T_K)
    gk = _gk_viscosity("prop.edr", T_K)
    d = sh("printf 'Density\\n\\n' | gmx energy -f prop.edr -o dprop.xvg")
    rho = None
    for line in (d.stdout + d.stderr).splitlines():
        if line.strip().startswith("Density"):
            try: rho = float(line.split()[1])
            except Exception: pass
    demix = _demix_index("prop.tpr", "prop.xtc", parsed, counts)
    return {"status": "completed", "mode": "properties", "label": spec["label"], "T_K": T_K,
            "seed": seed, "epsilon_md": round(eps, 3) if eps else None,
            "viscosity": gk, "density_kgm3": round(rho, 1) if rho else None,
            "demix_index": demix, "counts": counts, "meta": spec.get("meta", {})}

def build(spec):
    work = Path(spec["work"]); work.mkdir(parents=True, exist_ok=True); os.chdir(work)
    sh("rm -rf ./* 2>/dev/null")
    T_K, prod_ps, seed = float(spec["T_K"]), float(spec["prod_ps"]), int(spec.get("seed", 1))
    comps = spec["components"]   # [{smiles, resname, count}]
    parsed = []
    for c in comps:
        p = parametrize(c["smiles"], c["resname"])
        if p is None:
            return {"status": "acpype_failed", "label": spec["label"], "failed": c["resname"]}
        p["count"] = int(c["count"]); parsed.append(p)
    # box sized to a loose density so inserts succeed; NPT compresses
    total_mass = sum(p["mw"] * p["count"] for p in parsed)
    rho0 = 550.0
    edge = ((total_mass / 1000.0) / (rho0 * 6.022e23) * 1e27) ** (1/3)
    # insert each species sequentially into the same box; read the ACTUAL count placed from the
    # "Added N molecules" report (gro residue names are unreliable for distinguishing species).
    box, counts = None, {}
    for p in parsed:
        if box is None:
            r = sh(f"gmx insert-molecules -ci {p['gro']} -nmol {p['count']} -box {edge} {edge} {edge} -o box.gro")
        else:
            r = sh(f"gmx insert-molecules -f box.gro -ci {p['gro']} -nmol {p['count']} -o box.gro")
        box = "box.gro"
        added = p["count"]
        for line in (r.stdout + r.stderr).splitlines():
            if "Added" in line and "molecule" in line:
                try: added = int(line.split("Added")[1].split()[0])
                except: pass
        counts[p["resname"]] = added
    if not Path("box.gro").exists() or sum(counts.values()) == 0:
        return {"status": "empty_box", "label": spec["label"], "counts": counts}
    # merged topology: [defaults] + unique [atomtypes] + each [moleculetype] + [molecules]
    seen, at_lines = set(), []
    for p in parsed:
        for ln in p["atomtypes"]:
            key = ln.split()[0]
            if key not in seen:
                seen.add(key); at_lines.append(ln)
    top = ["[ defaults ]", "1 2 yes 0.5 0.8333", "", "[ atomtypes ]", *at_lines, ""]
    for p in parsed:
        top.append(p["moltype"])
    top += ["", "[ system ]", spec["label"], "", "[ molecules ]"]
    for p in parsed:
        top.append(f"{p['resname']} {counts[p['resname']]}")
    Path("topol.top").write_text("\n".join(top) + "\n")
    base = (f"integrator=md\ndt=0.002\nconstraints=h-bonds\ncutoff-scheme=Verlet\nrcoulomb=1.0\n"
            f"rvdw=1.0\ncoulombtype=PME\ntcoupl=berendsen\ntc-grps=System\ntau-t=0.5\nref-t={T_K}\n")
    Path("min.mdp").write_text("integrator=steep\nnsteps=10000\nemtol=100\nemstep=0.01\n"
                               "cutoff-scheme=Verlet\nrcoulomb=1.0\nrvdw=1.0\ncoulombtype=PME\n")
    Path("npt.mdp").write_text(base + f"nsteps=250000\npcoupl=berendsen\npcoupltype=isotropic\nref-p=1.0\n"
                               f"compressibility=4.5e-5\ntau-p=2.0\ngen-vel=yes\ngen-temp={T_K}\n"
                               f"gen-seed={seed}\nnstenergy=500\n")
    Path("nvt.mdp").write_text(base + f"nsteps={int(prod_ps/0.002)}\ngen-vel=yes\ngen-temp={T_K}\n"
                               f"gen-seed={seed}\nnstcalcenergy=10\nnstenergy=200\n")
    def grompp(mdp, c, o): return sh(f"gmx grompp -f {mdp} -c {c} -p topol.top -o {o} -maxwarn 10")
    def mdrun(d): return sh(f"gmx mdrun -deffnm {d} -ntmpi 1 -ntomp {NTOMP} -nb cpu")
    if grompp("min.mdp", "box.gro", "min").returncode or not Path("min.tpr").exists():
        return {"status": "grompp_min_failed", "label": spec["label"]}
    mdrun("min")
    start = "min.gro" if Path("min.gro").exists() else "box.gro"
    if grompp("npt.mdp", start, "npt").returncode:
        return {"status": "grompp_npt_failed", "label": spec["label"]}
    mdrun("npt")
    if not Path("npt.gro").exists():
        return {"status": "npt_failed", "label": spec["label"]}
    if spec.get("mode") == "properties":
        return run_properties(spec, parsed, counts, T_K, prod_ps, seed, base)
    last = [l for l in Path("npt.gro").read_text().splitlines() if l.strip()][-1].split()
    x, y, z = float(last[0]), float(last[1]), float(last[2])
    sh(f"gmx editconf -f npt.gro -o slab.gro -box {x} {y} {z*3} -c")
    if grompp("nvt.mdp", "slab.gro", "nvt").returncode:
        return {"status": "grompp_nvt_failed", "label": spec["label"]}
    mdrun("nvt")
    if not Path("nvt.edr").exists():
        return {"status": "nvt_failed", "label": spec["label"]}
    half = int(prod_ps/2)
    e = sh(f"printf '#Surf*SurfTen\\n\\n' | gmx energy -f nvt.edr -o st.xvg -b {half}")
    st = None
    for line in (e.stdout+e.stderr).splitlines():
        if line.strip().startswith("#Surf*SurfTen"):
            try: st = float(line.split()[1])
            except: pass
    sigma = round(st*0.1/2.0, 2) if st is not None else None
    d = sh("printf 'Density\\n\\n' | gmx energy -f npt.edr -o d.xvg -b 250")
    rho = None
    for line in (d.stdout+d.stderr).splitlines():
        if line.strip().startswith("Density"):
            try: rho = float(line.split()[1])
            except: pass
    return {"status": "completed" if sigma is not None else "no_surften", "label": spec["label"],
            "T_K": T_K, "seed": seed, "sigma_mNm": sigma, "raw_surften_barnm": st,
            "density_kgm3": round(rho, 1) if rho else None, "counts": counts, "meta": spec.get("meta", {})}

print("RESULT:" + json.dumps(build(json.loads(sys.argv[1]))))
'''.replace("__CONDA__", CONDA)


@app.function(image=IMAGE, cpu=16.0, timeout=7200, volumes={"/out": VOL})
def run_md(spec: dict) -> dict:
    import subprocess
    from pathlib import Path
    spec = dict(spec)
    spec["work"] = f"/out/{spec['label']}"
    Path("/driver.py").write_text(MD_DRIVER)
    r = subprocess.run([f"{CONDA}/python", "/driver.py", json.dumps(spec)], capture_output=True, text=True)
    VOL.commit()
    for line in reversed((r.stdout or "").splitlines()):
        if line.startswith("RESULT:"):
            return json.loads(line[len("RESULT:"):])
    return {"status": "driver_error", "label": spec.get("label"), "stderr": (r.stderr or "")[-500:]}


# ----- fluid definitions -----
COMPONENTS = [
    {"name": "pentane", "cas": "109-66-0", "smiles": "CCCCC", "lit_sigma": 15.5},
    {"name": "DMC", "cas": "616-38-6", "smiles": "COC(=O)OC", "lit_sigma": 28.4},
    {"name": "methanol", "cas": "67-56-1", "smiles": "CO", "lit_sigma": 22.1},
    {"name": "TriEG", "cas": "112-27-6", "smiles": "OCCOCCOCCO", "lit_sigma": 45.2},
    {"name": "HFO-1336mzz-Z", "cas": "692-49-9", "smiles": "F/C(=C\\C(F)(F)F)/C(F)(F)F", "lit_sigma": None},
    {"name": "vinylene_carbonate", "cas": "872-36-6", "smiles": "O=C1OC=CO1", "lit_sigma": None},
    {"name": "FK-5-1-12", "cas": "756-13-8", "smiles": "FC(F)(F)C(F)(F)C(=O)C(F)(C(F)(F)F)C(F)(F)F", "lit_sigma": None},
    {"name": "HCFO-1233zd(E)", "cas": "102687-65-0", "smiles": "Cl/C=C/C(F)(F)F", "lit_sigma": None},
    {"name": "TFE-amine", "cas": "753-90-2", "smiles": "NCC(F)(F)F", "lit_sigma": None},
]
# blends as mole-ratio component lists (counts derived from a target total)
BLENDS = [
    {"name": "HFO-VC_70-30", "cas_pair": ["692-49-9", "872-36-6"],
     "components": [("F/C(=C\\C(F)(F)F)/C(F)(F)F", 0.70), ("O=C1OC=CO1", 0.30)]},
    {"name": "MeOH-TriEG_5-95", "cas_pair": ["67-56-1", "112-27-6"],
     "components": [("CO", 0.30), ("OCCOCCOCCO", 0.70)]},   # mole-fraction box (5:95 wt ~ skewed; use 30:70 mol for a stable box)
    {"name": "Pentane-DMC_80-20", "cas_pair": ["109-66-0", "616-38-6"],
     "components": [("CCCCC", 0.80), ("COC(=O)OC", 0.20)]},
]
TEMPS = [288.15, 298.15, 308.15, 318.15]

# ---- FINAL SPRINT (FS3): gap blends + new component (TFE-amine) + dossier-aligned ternaries ----
# Track-B lead blends that had no blend sigma/miscibility yet; mole-fraction component lists.
GAP_BLENDS = [
    {"name": "FK-VC_80-20", "cas_pair": ["756-13-8", "872-36-6"],
     "components": [("FC(F)(F)C(F)(F)C(=O)C(F)(C(F)(F)F)C(F)(F)F", 0.80), ("O=C1OC=CO1", 0.20)]},
    {"name": "HCFO-VC_80-20", "cas_pair": ["102687-65-0", "872-36-6"],
     "components": [("Cl/C=C/C(F)(F)F", 0.80), ("O=C1OC=CO1", 0.20)]},
    {"name": "HFO-TFEA_90-10", "cas_pair": ["692-49-9", "753-90-2"],
     "components": [("F/C(=C\\C(F)(F)F)/C(F)(F)F", 0.90), ("NCC(F)(F)F", 0.10)]},
]
# dossier-aligned ternaries (C-05 MeOH+TriEG+DMC, C-04 HFO+VC+DMC): 1-seed sigma + properties (demix)
GAP_TERNARIES = [
    {"name": "MeOH-TriEG-DMC", "cas_triple": ["67-56-1", "112-27-6", "616-38-6"],
     "components": [("CO", 0.30), ("OCCOCCOCCO", 0.50), ("COC(=O)OC", 0.20)]},
    {"name": "HFO-VC-DMC", "cas_triple": ["692-49-9", "872-36-6", "616-38-6"],
     "components": [("F/C(=C\\C(F)(F)F)/C(F)(F)F", 0.55), ("O=C1OC=CO1", 0.25), ("COC(=O)OC", 0.20)]},
]


@app.local_entrypoint()
def sigma_one(smiles: str, name: str, n_mol: int = 400, t_k: float = 298.15, prod_ps: float = 800.0):
    spec = {"label": f"{name}_{int(t_k)}K_s1", "T_K": t_k, "prod_ps": prod_ps, "seed": 1,
            "components": [{"smiles": smiles, "resname": "MOL", "count": n_mol}]}
    print(json.dumps(run_md.remote(spec), indent=2))


@app.local_entrypoint()
def components_3seed(n_mol: int = 400, prod_ps: float = 800.0, seeds: int = 3):
    import os
    specs = []
    for c in COMPONENTS:
        for s in range(1, seeds + 1):
            specs.append({"label": f"{c['name']}_298K_s{s}", "T_K": 298.15, "prod_ps": prod_ps, "seed": s,
                          "components": [{"smiles": c["smiles"], "resname": "MOL", "count": n_mol}],
                          "meta": {"name": c["name"], "cas": c["cas"], "lit_sigma": c.get("lit_sigma")}})
    results = list(run_md.map(specs))
    out = {"results": results, "kind": "components_3seed", "seeds": seeds}
    print(json.dumps(out, indent=2))
    path = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "data", "processed",
                                        "md_components_3seed.json"))
    with open(path, "w") as fh:
        json.dump(out, fh, indent=2)
    print("wrote data/processed/md_components_3seed.json")


@app.local_entrypoint()
def blends_sigmaT(total: int = 400, prod_ps: float = 800.0):
    import os
    specs = []
    for b in BLENDS:
        for T in TEMPS:
            comps = [{"smiles": sm, "resname": f"M{i}", "count": max(1, int(round(frac * total)))}
                     for i, (sm, frac) in enumerate(b["components"])]
            specs.append({"label": f"{b['name']}_{int(T)}K", "T_K": T, "prod_ps": prod_ps, "seed": 1,
                          "components": comps,
                          "meta": {"blend": b["name"], "cas_pair": b["cas_pair"]}})
    results = list(run_md.map(specs))
    out = {"results": results, "kind": "blends_sigmaT", "temps": TEMPS}
    print(json.dumps(out, indent=2))
    path = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "data", "processed",
                                        "md_blend_sigmaT.json"))
    with open(path, "w") as fh:
        json.dump(out, fh, indent=2)
    print("wrote data/processed/md_blend_sigmaT.json")


@app.local_entrypoint()
def properties_pass(n_mol: int = 400, prod_ps: float = 1000.0, high_mu_prod_ps: float = 3000.0,
                    max_usd: float = 250.0):
    """WS4: cubic-box properties (epsilon, GK viscosity, density, demix) for components + blends.

    Labels carry the _props suffix so campaign volume outputs are never clobbered. Budget-guarded:
    refuses to dispatch if the rough CPU-second estimate exceeds --max-usd."""
    import os
    specs = []
    for c in COMPONENTS:
        specs.append({"label": f"{c['name']}_props", "mode": "properties", "T_K": 298.15,
                      "prod_ps": prod_ps, "seed": 1,
                      "components": [{"smiles": c["smiles"], "resname": "MOL", "count": n_mol}],
                      "meta": {"name": c["name"], "cas": c["cas"], "lit_sigma": c.get("lit_sigma")}})
    for b in BLENDS:
        pp = high_mu_prod_ps if "TriEG" in b["name"] else prod_ps   # viscous blend needs longer GK
        comps = [{"smiles": sm, "resname": f"M{i}", "count": max(1, int(round(frac * n_mol)))}
                 for i, (sm, frac) in enumerate(b["components"])]
        specs.append({"label": f"{b['name']}_props", "mode": "properties", "T_K": 298.15,
                      "prod_ps": pp, "seed": 1, "components": comps,
                      "meta": {"blend": b["name"], "cas_pair": b["cas_pair"]}})
    # rough cost: 16-cpu Modal ~ $0.002/cpu-min -> ~$2/job-hour; ~1.2 h per ns of production
    est_hours = sum((s["prod_ps"] / 1000.0) * 1.2 + 0.5 for s in specs)
    est_usd = est_hours * 2.0
    print(f"properties pass: {len(specs)} jobs, est ~{est_hours:.1f} job-hours ~= ${est_usd:.0f} "
          f"(cap ${max_usd})")
    if est_usd > max_usd:
        raise SystemExit(f"REFUSING: estimate ${est_usd:.0f} exceeds --max-usd {max_usd}")
    results = list(run_md.map(specs))
    out = {"results": results, "kind": "properties_pass"}
    print(json.dumps(out, indent=2))
    path = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "data", "processed",
                                        "md_properties_raw.json"))
    with open(path, "w") as fh:
        json.dump(out, fh, indent=2)
    print("wrote data/processed/md_properties_raw.json")


@app.local_entrypoint()
def gap_campaign(n_mol: int = 400, prod_ps: float = 800.0, max_usd: float = 250.0):
    """FS3: close the Track-B blend gaps + the new TFE-amine component + dossier-aligned ternaries.

    Per gap blend: 3 sigma seeds @298 + 4 sigma(T) (Eotvos) + 1 properties (epsilon/demix/density).
    TFE-amine component: 3 sigma seeds @298 + 1 properties. Each ternary: 1 sigma @298 + 1 properties
    (demix). Labels carry _props for the properties jobs so campaign outputs are never clobbered.
    Budget-guarded (--max-usd). Writes md_gap_sigmaT.json (sigma seeds + sigma(T)) and
    md_gap_properties_raw.json (cubic-box properties)."""
    import os
    sigma_specs, prop_specs = [], []
    # --- gap blends ---
    for b in GAP_BLENDS:
        comps = [{"smiles": sm, "resname": f"M{i}", "count": max(1, int(round(frac * n_mol)))}
                 for i, (sm, frac) in enumerate(b["components"])]
        for s in (1, 2, 3):     # 3 seeds at 298 K for a SEM on sigma_298
            sigma_specs.append({"label": f"{b['name']}_298K_s{s}", "T_K": 298.15, "prod_ps": prod_ps,
                                "seed": s, "components": comps,
                                "meta": {"blend": b["name"], "cas_pair": b["cas_pair"]}})
        for T in TEMPS:         # sigma(T) Eotvos sweep (seed 1)
            sigma_specs.append({"label": f"{b['name']}_{int(T)}K", "T_K": T, "prod_ps": prod_ps,
                                "seed": 1, "components": comps,
                                "meta": {"blend": b["name"], "cas_pair": b["cas_pair"]}})
        prop_specs.append({"label": f"{b['name']}_props", "mode": "properties", "T_K": 298.15,
                           "prod_ps": 1000.0, "seed": 1, "components": comps,
                           "meta": {"blend": b["name"], "cas_pair": b["cas_pair"]}})
    # --- new component: TFE-amine ---
    tfe = next(c for c in COMPONENTS if c["cas"] == "753-90-2")
    for s in (1, 2, 3):
        sigma_specs.append({"label": f"{tfe['name']}_298K_s{s}", "T_K": 298.15, "prod_ps": prod_ps,
                            "seed": s, "components": [{"smiles": tfe["smiles"], "resname": "MOL", "count": n_mol}],
                            "meta": {"name": tfe["name"], "cas": tfe["cas"], "lit_sigma": None}})
    prop_specs.append({"label": f"{tfe['name']}_props", "mode": "properties", "T_K": 298.15,
                       "prod_ps": 1000.0, "seed": 1,
                       "components": [{"smiles": tfe["smiles"], "resname": "MOL", "count": n_mol}],
                       "meta": {"name": tfe["name"], "cas": tfe["cas"]}})
    # --- ternaries ---
    for t in GAP_TERNARIES:
        comps = [{"smiles": sm, "resname": f"M{i}", "count": max(1, int(round(frac * n_mol)))}
                 for i, (sm, frac) in enumerate(t["components"])]
        sigma_specs.append({"label": f"{t['name']}_298K_s1", "T_K": 298.15, "prod_ps": prod_ps, "seed": 1,
                            "components": comps,
                            "meta": {"blend": t["name"], "cas_pair": t["cas_triple"]}})
        prop_specs.append({"label": f"{t['name']}_props", "mode": "properties", "T_K": 298.15,
                           "prod_ps": 1000.0, "seed": 1, "components": comps,
                           "meta": {"blend": t["name"], "cas_pair": t["cas_triple"]}})
    specs = sigma_specs + prop_specs
    est_hours = sum((s["prod_ps"] / 1000.0) * 1.2 + 0.5 for s in specs)
    est_usd = est_hours * 2.0
    print(f"gap campaign: {len(specs)} jobs ({len(sigma_specs)} sigma + {len(prop_specs)} props), "
          f"est ~{est_hours:.1f} job-hours ~= ${est_usd:.0f} (cap ${max_usd})")
    if est_usd > max_usd:
        raise SystemExit(f"REFUSING: estimate ${est_usd:.0f} exceeds --max-usd {max_usd}")
    sig_results = list(run_md.map(sigma_specs))
    prop_results = list(run_md.map(prop_specs))
    base = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "data", "processed"))
    with open(os.path.join(base, "md_gap_sigmaT.json"), "w") as fh:
        json.dump({"results": sig_results, "kind": "gap_sigmaT", "temps": TEMPS}, fh, indent=2)
    with open(os.path.join(base, "md_gap_properties_raw.json"), "w") as fh:
        json.dump({"results": prop_results, "kind": "gap_properties"}, fh, indent=2)
    n_ok = sum(1 for r in sig_results + prop_results if r.get("status") == "completed")
    print(f"gap campaign done: {n_ok}/{len(specs)} completed; wrote md_gap_sigmaT.json, "
          f"md_gap_properties_raw.json")


@app.local_entrypoint()
def prop_one(smiles: str = "CCCCC", name: str = "pentane", n_mol: int = 250, prod_ps: float = 200.0):
    spec = {"label": f"{name}_props_test", "mode": "properties", "T_K": 298.15, "prod_ps": prod_ps,
            "seed": 1, "components": [{"smiles": smiles, "resname": "MOL", "count": n_mol}],
            "meta": {"name": name}}
    print(json.dumps(run_md.remote(spec), indent=2))


@app.local_entrypoint()
def blend_one(which: str = "HFO-VC_70-30", total: int = 300, t_k: float = 298.15, prod_ps: float = 400.0):
    b = next(x for x in BLENDS if x["name"] == which)
    comps = [{"smiles": sm, "resname": f"M{i}", "count": max(1, int(round(frac * total)))}
             for i, (sm, frac) in enumerate(b["components"])]
    spec = {"label": f"{b['name']}_{int(t_k)}K_val", "T_K": t_k, "prod_ps": prod_ps, "seed": 1,
            "components": comps, "meta": {"blend": b["name"], "cas_pair": b["cas_pair"]}}
    print(json.dumps(run_md.remote(spec), indent=2))


@app.local_entrypoint()
def toolcheck_main():
    print(json.dumps(toolcheck.remote(), indent=2))