from __future__ import annotations

import argparse

from .simulator import PROFILES, print_result, run_simulation, with_overrides
from .visualization import animate_history


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Roundabout network stochastic shuffler"
    )
    parser.add_argument("--seed", type=int, default=7)
    parser.add_argument("--size", type=int, default=14)
    parser.add_argument("--profile", choices=["all", *PROFILES.keys()], default="all")
    parser.add_argument("--max-steps", type=int, default=None)
    parser.add_argument("--slowdown-prob", type=float, default=None)
    parser.add_argument("--injection-prob", type=float, default=None)
    parser.add_argument(
        "--animate", action="store_true", help="Create GIF (single profile only)"
    )
    parser.add_argument("--output", type=str, default="roundabout.gif")
    parser.add_argument("--fps", type=int, default=6)
    args = parser.parse_args()

    if args.size < 1:
        raise SystemExit("--size must be >= 1")
    if args.slowdown_prob is not None and not (0.0 <= args.slowdown_prob <= 1.0):
        raise SystemExit("--slowdown-prob must be in [0, 1]")
    if args.injection_prob is not None and not (0.0 <= args.injection_prob <= 1.0):
        raise SystemExit("--injection-prob must be in [0, 1]")

    original = list(range(args.size))
    print("original:", original)

    if args.profile == "all":
        if args.animate:
            raise SystemExit("--animate requires a concrete --profile")
        for name in PROFILES:
            profile = with_overrides(
                PROFILES[name], args.slowdown_prob, args.injection_prob
            )
            result = run_simulation(
                original,
                args.seed,
                profile,
                record_history=False,
                max_steps_override=args.max_steps,
            )
            print_result(profile, result)
        return

    profile = with_overrides(
        PROFILES[args.profile], args.slowdown_prob, args.injection_prob
    )
    result = run_simulation(
        original,
        args.seed,
        profile,
        record_history=args.animate,
        max_steps_override=args.max_steps,
    )
    print_result(profile, result)

    if args.animate:
        out = animate_history(result.history, profile.name, args.output, args.fps)
        print("animation:", out)
