"""Command-line entry points for data, training, generation, and verification."""

import argparse, json
from pathlib import Path
import numpy as np
from .graph import Graph, toy_graph, grid_graph
from .exact import (
    solve,
    sample,
    endpoint_distribution,
    expected_cost,
    uniform_policy,
    tilted_reference,
)


def main():
    p = argparse.ArgumentParser(prog="dooable")
    sub = p.add_subparsers(dest="command", required=True)
    b = sub.add_parser("build")
    b.add_argument("--parents", required=True)
    b.add_argument("--reagents", required=True)
    b.add_argument("--reactions", required=True)
    b.add_argument("--budget", type=int, default=2)
    b.add_argument("--max-nodes", type=int, default=100000)
    b.add_argument("--parent-limit", type=int)
    b.add_argument("--output", required=True)
    t = sub.add_parser("toy")
    t.add_argument("--kind", choices=["multiplicity", "grid"], default="multiplicity")
    t.add_argument("--output", required=True)
    t.add_argument("--budget", type=int, default=6)
    for name in ["exact", "train"]:
        a = sub.add_parser(name)
        a.add_argument("--graph", required=True)
        a.add_argument("--rewards")
        a.add_argument("--temperature", type=float, default=1.0)
        a.add_argument("--output", required=True)
        a.add_argument("--seed", type=int, default=0)
        if name == "train":
            a.add_argument("--steps", type=int, default=2000)
            a.add_argument("--batch-size", type=int, default=64)
            a.add_argument(
                "--backward",
                choices=["learned", "uniform", "exact", "unnormalized"],
                default="learned",
            )
            a.add_argument("--resume")
            a.add_argument("--initialize")
    a = sub.add_parser("sample")
    a.add_argument("--model", required=True)
    a.add_argument("--n", type=int, default=1000)
    a.add_argument("--seed", type=int, default=0)
    a.add_argument("--output", required=True)
    a = sub.add_parser("fit-properties")
    a.add_argument("--data", default="data/downloads")
    a.add_argument("--output", required=True)
    a.add_argument("--seed", type=int, default=0)
    a = sub.add_parser("score")
    a.add_argument("--graph", required=True)
    a.add_argument("--models")
    a.add_argument("--output", required=True)
    a = sub.add_parser("replay")
    a.add_argument("--samples", required=True)
    a.add_argument("--budget", type=int, required=True)
    args = p.parse_args()
    if args.command == "build":
        from .chemistry import read_catalog, reactions, build_graph

        g = build_graph(
            read_catalog(args.parents, args.parent_limit),
            read_catalog(args.reagents),
            reactions(args.reactions),
            args.budget,
            args.max_nodes,
        )
        g.save(args.output)
        print(
            json.dumps(
                {
                    "nodes": len(g.nodes),
                    "edges": len(g.edges),
                    "outcomes": len(g.terminals),
                }
            )
        )
        return
    if args.command == "toy":
        g = toy_graph() if args.kind == "multiplicity" else grid_graph(6, args.budget)
        g.save(args.output)
        return
    if args.command in ["exact", "train", "score"]:
        g = Graph.load(args.graph)
        if args.command == "score":
            Path(args.output).parent.mkdir(parents=True, exist_ok=True)
            if args.models:
                from .properties import property_rewards

                rewards, scores = property_rewards(g, args.models)
                scores.to_csv(Path(args.output).with_suffix(".csv"), index=False)
            else:
                from .chemistry import descriptor_rewards

                rewards = descriptor_rewards(g)
            Path(args.output).write_text(json.dumps(rewards, indent=2))
            return
        rewards = (
            json.loads(Path(args.rewards).read_text())
            if args.rewards
            else {y: 0.0 for y in g.terminals}
        )
        out = Path(args.output)
        out.mkdir(parents=True, exist_ok=True)
        if args.command == "exact":
            sol = solve(g, rewards, args.temperature)
            g.save(out / "graph.json")
            np.save(out / "forward.npy", sol.forward)
            (out / "solution.json").write_text(
                json.dumps(
                    {
                        "temperature": args.temperature,
                        "target": sol.target,
                        "log_z": sol.log_z,
                        "expected_cost": expected_cost(g, sol.forward),
                    },
                    indent=2,
                )
            )
        else:
            from .learning import train

            _, h = train(
                g,
                rewards,
                args.temperature,
                args.steps,
                args.batch_size,
                seed=args.seed,
                output=out,
                backward=args.backward,
                resume=args.resume,
                initialize=args.initialize,
            )
            print(json.dumps(h[-1]))
        return
    if args.command == "sample":
        d = Path(args.model)
        if (d / "forward.npy").exists():
            g = Graph.load(d / "graph.json")
            forward = np.load(d / "forward.npy")
        else:
            from .learning import load_model

            model, _ = load_model(d)
            g = model.graph
            forward = model.probabilities()
        rows = sample(g, forward, args.n, args.seed)
        Path(args.output).parent.mkdir(parents=True, exist_ok=True)
        Path(args.output).write_text("".join(json.dumps(r) + "\n" for r in rows))
        return
    if args.command == "fit-properties":
        from .properties import fit_property

        print(
            json.dumps(
                [
                    fit_property(n, args.data, args.output, args.seed)
                    for n in ["caco2", "bace"]
                ],
                indent=2,
            )
        )
        return
    if args.command == "replay":
        from .chemistry import replay

        rows = [json.loads(l) for l in Path(args.samples).read_text().splitlines() if l]
        passed = sum(replay(r, args.budget) for r in rows)
        print(
            json.dumps(
                {
                    "paths": len(rows),
                    "replayed": passed,
                    "fraction": passed / max(len(rows), 1),
                }
            )
        )
        if passed != len(rows):
            raise SystemExit(1)


if __name__ == "__main__":
    main()