"""Build, score, train, and evaluate a configured local reaction inventory."""

import argparse
import json
import subprocess
import sys
from pathlib import Path
import yaml
from dooable.chemistry import read_catalog, reactions, build_graph, descriptor_rewards
from dooable.properties import property_rewards


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", required=True)
    args = parser.parse_args()
    cfg = yaml.safe_load(Path(args.config).read_text())
    out = Path(cfg["output"])
    out.mkdir(parents=True, exist_ok=True)
    parents = read_catalog(cfg["parents"], cfg.get("parent_limit"))
    reagents = read_catalog(cfg["reagents"], cfg.get("reagent_limit"))
    templates = reactions(cfg["reactions"])
    for budget in cfg.get("budgets", [0, 1, 2]):
        directory = out / f"budget{budget}"
        directory.mkdir(parents=True, exist_ok=True)
        graph = build_graph(
            parents, reagents, templates, budget, cfg.get("max_nodes", 100000)
        )
        graph.save(directory / "graph.json")
        if cfg.get("property_models"):
            rewards, scores = property_rewards(graph, cfg["property_models"])
            scores.to_csv(directory / "scores.csv", index=False)
        else:
            rewards = descriptor_rewards(graph)
        (directory / "rewards.json").write_text(json.dumps(rewards, indent=2))
        run = dict(
            graph=str(directory / "graph.json"),
            rewards=str(directory / "rewards.json"),
            output=str(directory / "benchmark"),
            methods=cfg.get("methods", ["dooable", "tb_uniform", "exact"]),
            temperature=cfg.get("temperature", 0.7),
            steps=cfg.get("steps", 2000),
            samples=cfg.get("samples", 1000),
            seeds=cfg.get("seeds", [0]),
            batch_size=cfg.get("batch_size", 64),
        )
        run_path = directory / "run.yaml"
        run_path.write_text(yaml.safe_dump(run))
        subprocess.run(
            [
                sys.executable,
                str(Path(__file__).with_name("run_benchmarks.py")),
                "--config",
                str(run_path),
            ],
            check=True,
        )


if __name__ == "__main__":
    main()