tools/online_screening.py

"""Online tools for evaluating an arbitrary uploaded CIF."""
from __future__ import annotations

import importlib.util
import math
import sys
from functools import lru_cache
from pathlib import Path
from typing import Any

import numpy as np
from pymatgen.core import Structure
from rdkit import Chem, RDConfig

from tools.adsorption import predict_benzene, predict_toluene
from tools.descriptors import compute_scm_eigenvalues
from tools.linker_extraction import extract_linker
from tools.toxicity import predict_toxicity


ROOT = Path(__file__).resolve().parent.parent
COPRINET_ROOT = ROOT / "coprinet"

ACTIVE_TOOLS = [
    "parse_cif",
    "predict_adsorption",
    "identify_linker",
    "check_metals",
    "predict_sa_score",
    "predict_aquatic_toxicity",
    "predict_price",
]

ALLOWED_METALS = {"Mg", "Al", "Ca", "Ti", "Mn", "Fe", "Cu", "Zn", "Zr", "Ag"}
ALL_METALS = {
    "Li", "Be", "Na", "Mg", "Al", "K", "Ca", "Sc", "Ti", "V", "Cr", "Mn", "Fe", "Co", "Ni",
    "Cu", "Zn", "Ga", "Rb", "Sr", "Y", "Zr", "Nb", "Mo", "Tc", "Ru", "Rh", "Pd", "Ag", "Cd",
    "In", "Sn", "Cs", "Ba", "La", "Ce", "Pr", "Nd", "Pm", "Sm", "Eu", "Gd", "Tb", "Dy", "Ho",
    "Er", "Tm", "Yb", "Lu", "Hf", "Ta", "W", "Re", "Os", "Ir", "Pt", "Au", "Hg", "Tl", "Pb",
    "Bi", "Po", "Fr", "Ra", "Ac", "Th", "Pa", "U", "Np", "Pu",
}


def _jsonable(value: Any) -> Any:
    if isinstance(value, dict):
        return {str(k): _jsonable(v) for k, v in value.items()}
    if isinstance(value, list | tuple):
        return [_jsonable(v) for v in value]
    if isinstance(value, np.ndarray):
        return _jsonable(value.tolist())
    if isinstance(value, np.generic):
        return _jsonable(value.item())
    if isinstance(value, float) and (math.isnan(value) or math.isinf(value)):
        return None
    if isinstance(value, Path):
        return str(value)
    return value


def _error(tool: str, message: str) -> dict[str, Any]:
    return {"tool": tool, "status": "error", "summary": message, "error": message}


def _needs_linker(tool: str) -> dict[str, Any]:
    return {"tool": tool, "status": "unavailable", "summary": "A linker SMILES must be identified first.", "pass": None}


def parse_cif_tool(cif_path: str) -> dict[str, Any]:
    try:
        structure = Structure.from_file(cif_path)
        elements = sorted({site.specie.symbol for site in structure})
        metals = [element for element in elements if element in ALL_METALS]
        return {
            "tool": "parse_cif",
            "status": "success",
            "formula": structure.composition.reduced_formula,
            "num_sites": structure.num_sites,
            "elements": elements,
            "metals": metals,
            "summary": f"Parsed {structure.num_sites} atomic sites; metal elements: {', '.join(metals) or 'none'}",
        }
    except Exception as exc:
        return _error("parse_cif", f"CIF parsing failed: {exc}")


def predict_adsorption_tool(cif_path: str) -> dict[str, Any]:
    try:
        b_desc = compute_scm_eigenvalues(cif_path, 520)
        t_desc = compute_scm_eigenvalues(cif_path, 584)
        benzene = predict_benzene(b_desc["eigenvalues"])
        toluene = predict_toluene(t_desc["eigenvalues"])
        warnings = [w for w in [b_desc.get("applicability_warning"), t_desc.get("applicability_warning")] if w]
        return {
            "tool": "predict_adsorption",
            "status": "success",
            "benzene_uptake_mg_g": benzene["uptake_mg_g"],
            "toluene_uptake_mg_g": toluene["uptake_mg_g"],
            "benzene_model": benzene["model_version"],
            "toluene_model": toluene["model_version"],
            "descriptor_raw_dim_benzene": b_desc["raw_dim"],
            "descriptor_raw_dim_toluene": t_desc["raw_dim"],
            "warnings": warnings,
            "summary": f"Benzene {benzene['uptake_mg_g']} mg/g; toluene {toluene['uptake_mg_g']} mg/g",
        }
    except Exception as exc:
        return _error("predict_adsorption", f"Adsorption prediction failed: {exc}")


def identify_linker_tool(cif_path: str) -> dict[str, Any]:
    try:
        result = extract_linker(cif_path)
        status = "success" if result.get("linker_smiles") else "unavailable"
        return {
            "tool": "identify_linker",
            "status": status,
            "metals": result.get("metals") or [],
            "linker_smiles": result.get("linker_smiles"),
            "linker_name": result.get("linker_name"),
            "linker_formula": result.get("linker_formula"),
            "extraction_level": result.get("extraction_level"),
            "extraction_note": result.get("extraction_note"),
            "summary": result.get("extraction_note"),
        }
    except Exception as exc:
        return _error("identify_linker", f"Linker identification failed: {exc}")


def check_metals_tool(context: dict[str, Any], cif_path: str) -> dict[str, Any]:
    parse = context.get("parse_cif")
    if not parse or parse.get("status") == "error":
        parse = parse_cif_tool(cif_path)
    metals = parse.get("metals") or []
    illegal = [metal for metal in metals if metal not in ALLOWED_METALS]
    passed = len(illegal) == 0
    return {
        "tool": "check_metals",
        "status": "pass" if passed else "fail",
        "pass": passed,
        "detected_metals": metals,
        "illegal_metals": illegal,
        "allowed_metals": sorted(ALLOWED_METALS),
        "summary": "Metal check passed" if passed else f"Disallowed metals detected: {', '.join(illegal)}",
    }


@lru_cache(maxsize=1)
def _sa_module():
    path = Path(RDConfig.RDContribDir) / "SA_Score" / "sascorer.py"
    spec = importlib.util.spec_from_file_location("mofscreen_sascorer", path)
    if spec is None or spec.loader is None:
        raise RuntimeError(f"cannot load sascorer: {path}")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def predict_sa_score_tool(context: dict[str, Any]) -> dict[str, Any]:
    smiles = _linker_smiles(context)
    if not smiles:
        return _needs_linker("predict_sa_score")
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        return _error("predict_sa_score", f"Invalid linker SMILES: {smiles}")
    try:
        score = float(_sa_module().calculateScore(mol))
        return {
            "tool": "predict_sa_score",
            "status": "success",
            "linker_smiles": smiles,
            "sa_score": round(score, 3),
            "scale": "1 easy, 10 hard",
            "easy_synthesis": score <= 6.0,
            "summary": f"SA score {score:.3f} (1 easy to synthesize, 10 hard to synthesize)",
        }
    except Exception as exc:
        return _error("predict_sa_score", f"SA score calculation failed: {exc}")


def predict_aquatic_toxicity_tool(context: dict[str, Any]) -> dict[str, Any]:
    smiles = _linker_smiles(context)
    if not smiles:
        return _needs_linker("predict_aquatic_toxicity")
    try:
        tox = predict_toxicity(smiles)
        values = [
            tox.get("IBC50_Vibrio"),
            tox.get("IGC50_Tetrahymena"),
            tox.get("LC50_Pimephales"),
            tox.get("LC50_Daphnia"),
        ]
        present = [float(v) for v in values if v is not None]
        status = "success" if present else "unavailable"
        tox.update({
            "tool": "predict_aquatic_toxicity",
            "status": status,
            "linker_smiles": smiles,
            "mean_toxicity": round(sum(present) / len(present), 4) if present else None,
            "worst_toxicity": max(present) if present else None,
            "summary": "Aquatic toxicity prediction completed" if present else "The toxicity model did not return valid endpoints",
        })
        return tox
    except Exception as exc:
        return _error("predict_aquatic_toxicity", f"Aquatic toxicity prediction failed: {exc}")


@lru_cache(maxsize=1)
def _price_predictor():
    if str(COPRINET_ROOT) not in sys.path:
        sys.path.insert(0, str(COPRINET_ROOT))

    import torch

    if not getattr(torch.load, "_mofscreen_weights_patch", False):
        original_load = torch.load

        def load_with_trusted_checkpoint(*args, **kwargs):
            # ponytail: local trusted CoPriNet checkpoint; remove when upstream supports torch>=2.6 defaults.
            kwargs.setdefault("weights_only", False)
            return original_load(*args, **kwargs)

        load_with_trusted_checkpoint._mofscreen_weights_patch = True
        torch.load = load_with_trusted_checkpoint

    from pricePrediction.predict.predict import GraphPricePredictor

    return GraphPricePredictor(n_gpus=0, n_cpus=0, batch_size=1)


def predict_price_tool(context: dict[str, Any]) -> dict[str, Any]:
    smiles = _linker_smiles(context)
    if not smiles:
        return _needs_linker("predict_price")
    try:
        pred = float(next(_price_predictor().yieldPredictions([smiles])))
        mol = Chem.MolFromSmiles(smiles)
        mw = Chem.Descriptors.ExactMolWt(mol) if mol is not None else None
        usd_per_mmol = math.exp(pred)
        usd_per_g = usd_per_mmol * 1000 / mw if mw else None
        return {
            "tool": "predict_price",
            "status": "success",
            "linker_smiles": smiles,
            "coprinet_log_usd_per_mmol": round(pred, 6),
            "usd_per_mmol": round(usd_per_mmol, 6),
            "usd_per_g": round(usd_per_g, 6) if usd_per_g is not None else None,
            "exact_mol_wt": round(mw, 6) if mw else None,
            "summary": f"CoPriNet predicted price {usd_per_mmol:.3f} USD/mmol",
        }
    except Exception as exc:
        return {"tool": "predict_price", "status": "unavailable", "summary": f"CoPriNet price prediction unavailable: {exc}", "pass": None}


def run_online_tools(cif_path: str, requested_tools: list[str]) -> dict[str, Any]:
    context: dict[str, Any] = {}
    trace: list[dict[str, Any]] = []
    tools = _with_dependencies(requested_tools)
    for name in tools:
        result = _run_one(name, context, cif_path)
        context[name] = _jsonable(result)
        trace.append({"agent": "Tool Execution Agent", "action": f"ran {name}", "output": context[name]})
    row = _row(context, Path(cif_path).stem)
    return _jsonable({
        "mof_id": Path(cif_path).stem,
        "results": context,
        "tool_results": context,
        "row": row,
        "agent_trace": trace,
        "gate_status": _gate_status(context),
        "recommendation": _recommendation(context),
        "final_score": _score(_gate_status(context)),
        "explanation": _explanation(context),
        "warnings": [],
        "errors": [],
    })


def _run_one(name: str, context: dict[str, Any], cif_path: str) -> dict[str, Any]:
    if name == "parse_cif":
        return parse_cif_tool(cif_path)
    if name == "predict_adsorption":
        return predict_adsorption_tool(cif_path)
    if name == "identify_linker":
        return identify_linker_tool(cif_path)
    if name == "check_metals":
        return check_metals_tool(context, cif_path)
    if name == "predict_sa_score":
        return predict_sa_score_tool(context)
    if name == "predict_aquatic_toxicity":
        return predict_aquatic_toxicity_tool(context)
    if name == "predict_price":
        return predict_price_tool(context)
    return _error(name, f"Unknown tool: {name}")


def _with_dependencies(tools: list[str]) -> list[str]:
    out: list[str] = []
    for tool in tools:
        if tool in {"check_metals"}:
            out.append("parse_cif")
        if tool in {"predict_sa_score", "predict_aquatic_toxicity", "predict_price"}:
            out.append("identify_linker")
        out.append(tool)
    return list(dict.fromkeys(out))


def _linker_smiles(context: dict[str, Any]) -> str | None:
    linker = context.get("identify_linker") or {}
    smiles = linker.get("linker_smiles")
    return str(smiles) if smiles else None


def _gate_status(results: dict[str, Any]) -> str:
    if any(r.get("status") == "fail" for r in results.values()):
        return "failed"
    if any(r.get("status") in {"error", "unavailable"} for r in results.values()):
        return "incomplete"
    return "completed"


def _recommendation(results: dict[str, Any]) -> str:
    if any(r.get("status") == "fail" for r in results.values()):
        return "not_recommended"
    if any(r.get("status") in {"error", "unavailable"} for r in results.values()):
        return "needs_manual_review"
    return "continue_evaluation"


def _score(status: str) -> float:
    return {"completed": 10.0, "incomplete": 5.0, "failed": 0.0}.get(status, 5.0)


def _explanation(results: dict[str, Any]) -> str:
    return "; ".join(f"{name}: {payload.get('status')}" for name, payload in results.items())


def _row(results: dict[str, Any], mof_id: str) -> dict[str, Any]:
    adsorption = results.get("predict_adsorption") or {}
    metals = results.get("check_metals") or {}
    linker = results.get("identify_linker") or {}
    sa = results.get("predict_sa_score") or {}
    tox = results.get("predict_aquatic_toxicity") or {}
    price = results.get("predict_price") or {}
    return {
        "MOF ID": mof_id,
        "Benzene (mg/g)": adsorption.get("benzene_uptake_mg_g"),
        "Toluene (mg/g)": adsorption.get("toluene_uptake_mg_g"),
        "Metal status": metals.get("status"),
        "Metals": ", ".join(metals.get("detected_metals") or []),
        "Linker": linker.get("linker_name") or linker.get("linker_formula"),
        "Linker SMILES": linker.get("linker_smiles"),
        "SA score": sa.get("sa_score"),
        "Mean toxicity": tox.get("mean_toxicity"),
        "Price USD/g": price.get("usd_per_g"),
        "Price USD/mmol": price.get("usd_per_mmol"),
        "Gate status": _gate_status(results),
        "Recommendation": _recommendation(results),
    }