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_with_fixed_upper_solution_and_save,
)
from src.busopt.runner import run_nested_gbd_hpr, save_result


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Replicate (full) Cooperative scheme.")
    parser.add_argument("--output-dir", type=Path, default=REPO_ROOT / "results" / "mode_comparison")
    add_model_config_arguments(parser, exclude={"s_bar"})
    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)
    config = base_config.with_overrides(s_bar=1.5, Psi0=100)
    ideal_output_path = args.output_dir / "FC-Ideal.json"
    implemented_output_path = args.output_dir / "FC-Implemented.json"

    ideal_result = run_nested_gbd_hpr(config)
    save_result(ideal_result, ideal_output_path)
    if ideal_result.get("optimal_up_solution") is None:
        raise RuntimeError(f"FC-Ideal did not produce a feasible upper-level solution: {ideal_output_path}")
    run_nested_gbd_with_fixed_upper_solution_and_save(
        config,
        ideal_result["optimal_up_solution"],
        implemented_output_path,
    )


if __name__ == "__main__":
    main()
