inference.py

# Allow running without installation by adding ./src to sys.path (done lazily in main before model import)

from pathlib import Path
from typing import Any, Dict, List, Sequence, Set, Tuple

import argparse
import json
import ast
import requests
import sys
import time

# Reuse the per-SMILES checker from scripts/check_accessible.py
from scripts.check_accessible import check_accessible  # type: ignore


def _iter_all_smiles(node: Dict[str, Any]) -> List[str]:
    smiles: List[str] = []
    smi = node.get("smiles")
    if isinstance(smi, str) and smi:
        smiles.append(smi)
    for child in node.get("children", []) or []:
        smiles.extend(_iter_all_smiles(child))
    return smiles


def _load_stock(stock_file: Path) -> Set[str]:
    stock: Set[str] = set()
    if stock_file.is_file():
        with stock_file.open("r", encoding="utf-8") as f:
            for line in f:
                t = line.strip()
                if not t or t.startswith("#"):
                    continue
                stock.add(t)
    return stock


def filter_accessible_paths(
    paths: Sequence[Any], stock_file: Path, include_root: bool = True, sleep: float = 0.2
) -> Tuple[List[Dict[str, Any]], Dict[str, Dict[str, str]], List[Dict[str, Any]]]:
    """Return only those path strings where all nodes are accessible.

    Uses local stock first, then PubChem via check_accessible(). Returns the filtered
    list and a map of SMILES -> {status, detail} used for the decision.
    """
    stock = _load_stock(stock_file)
    session = requests.Session()
    status_cache: Dict[str, Dict[str, str]] = {}

    def is_accessible_smiles(smi: str) -> bool:
        if smi in status_cache:
            return status_cache[smi]["status"] == "purchasable"
        if smi in stock:
            status_cache[smi] = {"status": "purchasable", "detail": "found in stock file"}
            return True
        status, detail = check_accessible(smi, session)
        # Be polite to PubChem
        time.sleep(sleep)
        status_cache[smi] = {"status": status, "detail": detail}
        return status == "purchasable"

    filtered: List[Dict[str, Any]] = []
    reports: List[Dict[str, Any]] = []
    for path_item in paths:
        # Accept either a stringified Python/JSON literal or a dict-like object
        node: Dict[str, Any]
        if isinstance(path_item, str):
            try:
                node = ast.literal_eval(path_item)
            except Exception:
                # skip malformed path strings
                continue
            path_string = path_item
        elif isinstance(path_item, dict):
            node = path_item
            path_string = None  # not used for output, keep node structure
        else:
            # Attempt to coerce unknown objects exposing mapping-like API
            try:
                node = dict(path_item)  # type: ignore[arg-type]
                path_string = None
            except Exception:
                continue
        smiles_chain = _iter_all_smiles(node)
        # optionally exclude root
        if not include_root and smiles_chain:
            smiles_chain = smiles_chain[1:]
        ok = True
        bad_nodes: List[str] = []
        for smi in smiles_chain:
            if not is_accessible_smiles(smi):
                ok = False
                bad_nodes.append(smi)
        if ok:
            filtered.append(node)
        # Collect a per-path report
        reports.append({
            "path": node,
            "not_accessible": len(bad_nodes),
            "not_accessible_smiles": bad_nodes,
        })
    return filtered, status_cache, reports


def _load_paths_from_file(p: Path) -> List[Any]:
    text = p.read_text(encoding="utf-8").strip()
    if not text:
        return []
    # Try JSON array first
    try:
        arr = json.loads(text)
        if isinstance(arr, list):
            # Accept lists of strings or dicts
            if all(isinstance(x, (str, dict)) for x in arr):
                return arr
    except Exception:
        pass
    # Fallback: newline-delimited strings
    return [line for line in text.splitlines() if line.strip()]


def main() -> int:
    data_path = Path("./data")
    default_ckpt = data_path / "checkpoints"
    default_config = data_path / "configs" / "dms_dictionary.yaml"
    default_stock = data_path / "compounds" / "buyables-stock.txt"
    default_output = data_path / "accessible_paths_from_inference.json"

    ap = argparse.ArgumentParser(description="Generate routes (optional) and filter paths where all nodes are purchasable.")
    src = ap.add_argument_group("source of paths")
    src.add_argument("--paths-file", type=Path, default=None, help="If provided, read candidate paths from file (JSON array or newline separated)")
    src.add_argument("--target", default=None, help="Target SMILES. Required if not using --paths-file")
    src.add_argument("--n-steps", type=int, default=None, help="Number of steps (None lets the model decide)")
    src.add_argument("--starting-material", default=None, help="Optional starting material SMILES")
    src.add_argument("--model", default="explorer", help="Model name or checkpoint; defaults to explorer")
    src.add_argument("--beam-size", type=int, default=32, help="Beam size (default: 32)")
    src.add_argument("--ckpt-dir", type=Path, default=default_ckpt, help="Checkpoint directory (default: data/checkpoints)")
    src.add_argument("--config", type=Path, default=default_config, help="Config path (default: data/configs/dms_dictionary.yaml)")

    filt = ap.add_argument_group("filter options")
    filt.add_argument("--stock-file", type=Path, default=default_stock, help="Local purchasable SMILES list")
    excl = filt.add_mutually_exclusive_group()
    excl.add_argument("--include-root", action="store_true", help="Require root (product) to be purchasable")
    excl.add_argument("--children-only", action="store_true", help="Only require children (reactants) to be purchasable")
    filt.add_argument("--sleep", type=float, default=0.2, help="Delay between PubChem requests (default: 0.2s)")

    ap.add_argument("--output", type=Path, default=default_output, help="Output JSON file path")

    args = ap.parse_args()

    # Determine include_root flag
    include_root = True
    if args.children_only:
        include_root = False
    elif args.include_root:
        include_root = True

    # Load or generate candidate paths
    if args.paths_file is not None:
        if not args.paths_file.is_file():
            print(f"Error: paths file not found: {args.paths_file}", file=sys.stderr)
            return 2
        paths = _load_paths_from_file(args.paths_file)
    else:
        if not args.target:
            print("Error: --target is required when --paths-file is not provided", file=sys.stderr)
            return 2
        # Lazy import to avoid requiring torch when only checking paths
        # Also add ./src to sys.path if available
        repo_root = Path(__file__).resolve().parent
        src_dir = repo_root / "src"
        if (src_dir / "directmultistep").exists():
            sys.path.insert(0, str(src_dir))
        from directmultistep.generate import generate_routes  # type: ignore
        paths = generate_routes(
            target=args.target,
            n_steps=args.n_steps,
            starting_material=args.starting_material,
            model=args.model,
            beam_size=args.beam_size,
            config_path=args.config,
            ckpt_dir=args.ckpt_dir,
        )

    # Filter by accessibility
    filtered_paths, statuses, reports = filter_accessible_paths(
        paths=paths,
        stock_file=args.stock_file,
        include_root=include_root,
        sleep=args.sleep,
    )

    selection_mode = "all_access"
    min_inacc = None
    selected_paths: List[Dict[str, Any]] = filtered_paths
    # If none fully accessible, pick those with minimal number of not-accessible nodes
    if not selected_paths:
        if reports:
            min_inacc = min(r["not_accessible"] for r in reports)
            selection_mode = "min_inaccessible"
            selected_paths = [r["path"] for r in reports if r["not_accessible"] == min_inacc]

    args.output.parent.mkdir(parents=True, exist_ok=True)
    with args.output.open("w", encoding="utf-8") as f:
        json.dump(
            {
                "total": len(paths),
                "accessible": len(filtered_paths),
                "include_root": include_root,
                "selection": selection_mode,
                "min_not_accessible": min_inacc,
                "paths": selected_paths,
                "statuses": statuses,
                "path_reports": reports,
            },
            f,
            indent=2,
            ensure_ascii=True,
        )
    print(
        f"Candidates: {len(paths)}; fully-accessible ({'including root' if include_root else 'children only'}): {len(filtered_paths)}; "
        f"selected: {len(selected_paths)} (mode={selection_mode}, min_not_accessible={min_inacc}); saved: {args.output}"
    )
    return 0 if selected_paths else 3


if __name__ == "__main__":
    sys.exit(main())