from __future__ import annotations

import argparse
from pathlib import Path

from scripts._experiment_utils import REPO_ROOT, add_model_config_arguments, model_config_from_namespace, run_nested_gbd_and_save
from scripts.exports.export_sensitivity_csv import export_sensitivity_csv


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Replicate the fare sensitivity experiments.")
    parser.add_argument("--output-dir", type=Path, default=REPO_ROOT / "results" / "price_sensitivity")
    parser.add_argument("--csv-output", type=Path, default=None, help="CSV output path. Defaults to <output-dir>/price.csv.")
    parser.add_argument(
        "--export-only",
        action="store_true",
        help="Skip solving and only aggregate existing JSON results into CSV.",
    )
    add_model_config_arguments(parser, exclude={"tau"})
    return parser


def export_results_csv(input_dir: Path, output_path: Path | None = None) -> Path:
    return export_sensitivity_csv(input_dir, "price", output_path)


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

        for value in [round(index / 10, 1) for index in range(0, 40)]:
            config = base_config.with_overrides(tau=value)
            output_name = f"L-{config.route_num}-OD-{config.od_num}-P-{len(config.freq_set)}-TP-{value:.1f}.json"
            run_nested_gbd_and_save(config, args.output_dir / output_name)

    csv_path = export_results_csv(args.output_dir, args.csv_output)
    print(f"Wrote CSV summary to {csv_path}")


if __name__ == "__main__":
    main()
