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
from src.busopt.runner import (
    run_nested_gbd,
    run_nested_gbd_no_operating_cost,
    run_nested_gbd_with_fixed_solution,
    save_result,
)


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Run ablation studies based on nested GBD.")
    parser.add_argument("--output-dir", type=Path, default=REPO_ROOT / "results" / "mode_comparison")
    add_model_config_arguments(parser)
    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)
    original_config = base_config.with_overrides(s_bar=1.5, Psi0=100)

    no_subsidy_config = original_config.with_overrides(s_bar=0.0)
    no_subsidy_result = run_nested_gbd(no_subsidy_config)
    no_subsidy_result["ablation"] = "no_subsidy"
    save_result(no_subsidy_result, args.output_dir / "no_subsidy.json")

    no_operating_cost_ideal_result = run_nested_gbd_no_operating_cost(original_config)
    if no_operating_cost_ideal_result.get("optimal_up_solution") is None or no_operating_cost_ideal_result.get("optimal_mid_solution") is None:
        raise RuntimeError("no_operating_cost did not produce a feasible upper/middle-level solution.")
    no_operating_cost_ideal_result["ablation"] = "no_operating_cost"
    save_result(no_operating_cost_ideal_result, args.output_dir / "no_operating_cost_org.json")
    no_operating_cost_result = run_nested_gbd_with_fixed_solution(
        original_config,
        no_operating_cost_ideal_result["optimal_up_solution"],
        no_operating_cost_ideal_result["optimal_mid_solution"],
    )
    no_operating_cost_result["ablation"] = "no_operating_cost"
    save_result(no_operating_cost_result, args.output_dir / "no_operating_cost.json")

    no_subsidy_cost_config = original_config.with_overrides(mu=0.0)
    no_subsidy_cost_ideal_result = run_nested_gbd(no_subsidy_cost_config)
    if no_subsidy_cost_ideal_result.get("optimal_up_solution") is None or no_subsidy_cost_ideal_result.get("optimal_mid_solution") is None:
        raise RuntimeError("no_subsidy_cost did not produce a feasible upper/middle-level solution.")
    no_subsidy_cost_ideal_result["ablation"] = "no_subsidy_cost"
    save_result(no_subsidy_cost_ideal_result, args.output_dir / "no_subsidy_cost_org.json")
    no_subsidy_cost_result = run_nested_gbd_with_fixed_solution(
        original_config,
        no_subsidy_cost_ideal_result["optimal_up_solution"],
        no_subsidy_cost_ideal_result["optimal_mid_solution"],
    )
    no_subsidy_cost_result["ablation"] = "no_subsidy_cost"
    save_result(no_subsidy_cost_result, args.output_dir / "no_subsidy_cost.json")


if __name__ == "__main__":
    main()
