from __future__ import annotations

import argparse
from pathlib import Path

from _experiment_utils import REPO_ROOT, add_model_config_arguments, build_base_grid, model_config_from_namespace, run_nested_gbd_and_save


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Replicate the main computational experiments.")
    parser.add_argument("--output-dir", type=Path, default=REPO_ROOT / "results" / "main_experiments")
    add_model_config_arguments(parser, exclude={"route_num", "od_num", "freq_set", "accelerate"})
    return parser


def main() -> None:
    args = build_parser().parse_args()
    args.output_dir.mkdir(parents=True, exist_ok=True)
    base_config = model_config_from_namespace(args)

    experiment_groups = {
        "with_acceleration": True,
        "without_acceleration": False,
    }

    for subdir_name, accelerate in experiment_groups.items():
        output_dir = args.output_dir / subdir_name
        output_dir.mkdir(parents=True, exist_ok=True)

        for od_num, route_num, freq_set in build_base_grid():
            config = base_config.with_overrides(
                od_num=od_num,
                route_num=route_num,
                freq_set=[float(value) for value in freq_set],
                accelerate=accelerate,
            )
            output_name = f"L-{route_num}-OD-{od_num}-P-{len(freq_set)}.json"
            run_nested_gbd_and_save(config, output_dir / output_name)


if __name__ == "__main__":
    main()
