cg_extract.py

"""Shared extraction: turn a ChemGraph agent `state` into the rich payload the
site renders. Used by both app.py (live) and precompute.py (cached heavy runs)
so the shape never diverges."""
from __future__ import annotations
import os
import re
from typing import Optional


# ---- element data for the client (kept here so formula/labels are consistent) ----
def parse_xyz(path: Optional[str]) -> dict:
    """Read an .xyz file → {atoms:[{el,x,y,z}], natoms}. Empty dict on failure."""
    if not path or not os.path.exists(path):
        return {}
    try:
        with open(path, "r", encoding="utf-8") as f:
            lines = f.read().splitlines()
        n = int(lines[0].strip())
        atoms = []
        for ln in lines[2 : 2 + n]:
            parts = ln.split()
            if len(parts) < 4:
                continue
            el = parts[0]
            x, y, z = (float(parts[1]), float(parts[2]), float(parts[3]))
            atoms.append({"el": el, "x": round(x, 4), "y": round(y, 4), "z": round(z, 4)})
        return {"atoms": atoms, "natoms": len(atoms)}
    except Exception:
        return {}


def _num(x):
    try:
        return float(x)
    except Exception:
        return None


def _clean_freqs(raw) -> list:
    """vib frequencies come as strings like '1592.3' or '-42.1i'. Keep real,
    drop translation/rotation near-zeros, round."""
    out = []
    for v in raw or []:
        s = str(v)
        imag = s.endswith("i")
        val = _num(s[:-1] if imag else s)
        if val is None:
            continue
        if imag:
            out.append({"cm1": round(-abs(val), 1), "imaginary": True})
        elif abs(val) > 1.0:  # drop ~0 modes (translations/rotations)
            out.append({"cm1": round(val, 1), "imaginary": False})
    return out


def _summarize(tool: str, out: dict, task: str) -> str:
    if not isinstance(out, dict):
        return str(out)[:140]
    if tool == "molecule_name_to_smiles":
        return f"{out.get('resolved_name', out.get('name',''))} → SMILES {out.get('smiles','?')} · {out.get('molecular_formula','')} (PubChem)"
    if tool == "smiles_to_coordinate_file":
        return f"built 3D structure · {out.get('natoms','?')} atoms"
    if tool == "run_ase":
        e = out.get("single_point_energy")
        if isinstance(e, (int, float)):
            return f"simulation success · energy {e:.4f} eV"
        if out.get("dipole_moment"):
            return "simulation success · dipole computed"
        if isinstance(out.get("result"), dict):
            r = out["result"]
            if "ir_spectrum" in r:
                return "vibrational + IR analysis success"
            if "vibrational_frequencies" in r:
                return "vibrational analysis success"
            if "thermochemistry" in r:
                return "thermochemistry computed"
        return str(out.get("status", "done"))
    return str(out.get("status", ""))[:140]


def _detail(tool: str, args: dict, out: dict) -> dict:
    """Richer, foldable per-step detail (SMILES, formula, args, coords count)."""
    d = {}
    if tool == "molecule_name_to_smiles" and isinstance(out, dict):
        d = {"name": out.get("name"), "resolved": out.get("resolved_name"),
             "smiles": out.get("smiles"), "formula": out.get("molecular_formula"),
             "cid": out.get("cid"), "source": out.get("source")}
    elif tool == "smiles_to_coordinate_file" and isinstance(out, dict):
        d = {"smiles": out.get("smiles"), "natoms": out.get("natoms")}
    elif tool == "run_ase":
        params = (args or {}).get("params", {}) if isinstance(args, dict) else {}
        calc = params.get("calculator")
        calc_s = method = None
        if isinstance(calc, dict):
            calc_s = calc.get("calculator_type")
            method = calc.get("method")
        elif isinstance(calc, str):
            # live runs pass the full repr, e.g. "calculator_type='TBLite' method='GFN2-xTB' ...";
            # keep just the type + method, not the whole settings dump.
            m = re.search(r"calculator_type=['\"]?([\w-]+)", calc)
            calc_s = m.group(1) if m else calc[:24]
            mm = re.search(r"method=['\"]?([\w-]+)", calc)
            method = mm.group(1) if mm else None
        d = {"driver": params.get("driver"), "calculator": calc_s, "method": method}
    return {k: v for k, v in d.items() if v is not None}


def extract(state: dict, query: str, molecule: str, calc_label: str, task: str,
            cached: bool = False) -> dict:
    intent = state.get("intent", {}) or {}
    trace = state.get("tool_trace", []) or []

    smiles = formula = resolved = None
    xyz_path = None
    steps = []
    run_out = {}

    for s in trace:
        tool = s.get("tool")
        args = s.get("args", {}) or {}
        out = s.get("output", {}) or {}
        steps.append({
            "tool": tool,
            "status": s.get("status"),
            "summary": _summarize(tool, out, task),
            "detail": _detail(tool, args, out),
        })
        if tool == "molecule_name_to_smiles" and isinstance(out, dict):
            smiles = out.get("smiles") or smiles
            formula = out.get("molecular_formula") or formula
            resolved = out.get("resolved_name") or resolved
        if tool == "smiles_to_coordinate_file" and isinstance(out, dict):
            xyz_path = out.get("path") or xyz_path
        if tool == "run_ase":
            run_out = out if isinstance(out, dict) else {}
            p = (args.get("params", {}) or {})
            xyz_path = xyz_path or p.get("input_structure_file")

    structure = parse_xyz(xyz_path)

    # ---- task-specific result ----
    result = {"kind": task}
    e = run_out.get("single_point_energy")
    if isinstance(e, (int, float)):
        result["energy"] = {"value": e, "unit": run_out.get("unit", "eV")}
    if run_out.get("dipole_moment"):
        vec = [round(_num(x) or 0.0, 4) for x in run_out["dipole_moment"]]
        mag = round(sum(v * v for v in vec) ** 0.5, 4)
        result["dipole"] = {"vector": vec, "magnitude": mag,
                            "unit": run_out.get("dipole_unit", "e·Å")}
    res = run_out.get("result")
    if isinstance(res, dict):
        vf = res.get("vibrational_frequencies")
        if isinstance(vf, dict):
            result["frequencies"] = _clean_freqs(vf.get("frequencies"))
        irs = res.get("ir_spectrum")
        if isinstance(irs, dict):
            freqs = [(_num(x) or 0.0) for x in irs.get("frequency_cm1", [])]
            ints = [(_num(x) or 0.0) for x in irs.get("intensity", [])]
            modes = [{"cm1": round(f, 1), "intensity": round(i, 4)}
                     for f, i in zip(freqs, ints) if f > 1.0]
            result["ir"] = {"modes": modes}
        th = res.get("thermochemistry")
        if isinstance(th, dict):
            result["thermo"] = {
                "enthalpy": _num(th.get("enthalpy")),
                "entropy": _num(th.get("entropy")),
                "gibbs": _num(th.get("gibbs_free_energy")),
                "unit": th.get("unit", "eV"),
            }

    validation = state.get("validation", {}) or {}
    v_ok = validation.get("passed", validation.get("ok")) if isinstance(validation, dict) else None
    has_result = any(k in result for k in ("energy", "dipole", "frequencies", "ir", "thermo"))

    return {
        "ok": True,
        "cached": cached,
        "query": query,
        "molecule": molecule,
        "formula": formula,
        "smiles": smiles,
        "resolved": resolved,
        "calculator": calc_label,
        "task": task,
        "intent": {
            "task_type": intent.get("task_type"),
            "required_output": intent.get("required_output"),
            "calculator": intent.get("calculator"),
            "target_molecules": intent.get("target_molecules"),
        },
        "structure": structure,
        "steps": steps,
        "result": result,
        "verified": bool(v_ok) if v_ok is not None else has_result,
    }