"""Train or replay the roller landing on Pollen's Got a Talent simulation."""

import argparse
import json
from pathlib import Path

import learn as landing
from arena.cast import cast_of
from arena.hub_fetch import check_manifest, rig_of
from arena.policy import load_policy
from arena.runner import run_stage
from arena.timeline import load_troupe
from export_actor import export


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    commands = parser.add_subparsers(dest="command", required=True)
    training = commands.add_parser("train")
    training.add_argument("--output", type=Path, required=True)
    training.add_argument(
        "--start-from",
        type=Path,
        help="checkpoint.json with the 32 learned coefficients",
    )
    training.add_argument("--generations", type=int, default=25)
    training.add_argument("--population", type=int, default=40)
    training.add_argument(
        "--batch",
        type=int,
        default=8,
        help="simulated starts per candidate per generation",
    )
    training.add_argument("--workers", type=int, default=12)
    training.add_argument("--random-seed", type=int, default=11062026)
    training.add_argument("--speed", type=float, default=0.65)
    training.add_argument("--roll-at", type=float, default=1.65)
    build = commands.add_parser("export")
    build.add_argument("checkpoint", type=Path)
    build.add_argument("output", type=Path)
    replay = commands.add_parser("replay")
    replay.add_argument("model_folder", type=Path)
    replay.add_argument("--seeds", nargs="+", type=int, default=[1])
    args = parser.parse_args()
    if args.command == "export":
        ck = json.loads(args.checkpoint.read_text())
        export(ck["theta"], args.output)
        return
    if args.command == "train":
        if (
            args.population % 2
            or args.population < 10
            or min(args.generations, args.batch, args.workers) < 1
            or args.batch > 520
        ):
            parser.error(
                "Use at least one generation, worker and seed (up to 520) and an even population of at least 10."
            )
        if not 0 <= args.speed <= 2 or not 0.7 <= args.roll_at <= 12:
            parser.error(
                "Use a speed from 0 to 2 m/s and a roll time from 0.7 to 12 seconds."
            )
        if args.output.exists() and any(args.output.iterdir()):
            parser.error("Choose a new or empty training directory.")
        args.output.mkdir(parents=True, exist_ok=True)
        warm = (
            json.loads(args.start_from.read_text())["theta"]
            if args.start_from
            else None
        )
        landing.configure(args.speed, args.roll_at)
        landing.save(args.output / "timeline.json", landing.timeline_dict())
        landing.train(
            args.generations,
            args.population,
            args.batch,
            warm,
            args.output,
            random_seed=args.random_seed,
            workers=args.workers,
        )
        candidates = json.loads((args.output / "checkpoints.json").read_text())
        landing.save(args.output / "checkpoint.json", candidates[0])
        return
    folder = args.model_folder
    event = landing.EVENT
    manifest = json.loads((folder / "manifest.json").read_text())
    check_manifest(manifest, event)
    policy = load_policy(folder / "policy.onnx")
    troupe = load_troupe(folder / "timeline.json", event)
    cast = cast_of(event, troupe, policy, rig_of(manifest), {})
    schedule = troupe.ducks[0].timeline
    landing.configure(schedule.commands[1][0], schedule.times[2])
    for seed in args.seeds:
        result = run_stage(event, cast.performers, seed)
        row = {"seed": seed, "status": result.status}
        if result.status == "finish":
            row.update(landing.metrics(result.trace[0]))
        print(json.dumps(row))


if __name__ == "__main__":
    main()