from __future__ import annotations

import argparse
import csv
import json
import math
from pathlib import Path
from typing import Any


REPO_ROOT = Path(__file__).resolve().parents[2]
DEFAULT_INPUT_DIR = REPO_ROOT / "results" / "mode_comparison"
DEFAULT_OUTPUT_PATH = REPO_ROOT / "results" / "tables" / "table2_mode_comparison.csv"

SCENARIO_METADATA: dict[str, dict[str, str | int]] = {
    "public_operation": {"group": "regime", "label": "Public Operation", "order": 10},
    "open_access": {"group": "regime", "label": "Open Access", "order": 20},
    "FC-Ideal": {"group": "contract", "label": "FC-Ideal", "order": 30},
    "FC-Implemented": {"group": "contract", "label": "FC-Implemented", "order": 40},
    "SC-Ideal": {"group": "contract", "label": "SC-Ideal", "order": 50},
    "SC-Implemented": {"group": "contract", "label": "SC-Implemented", "order": 60},
    "no_subsidy": {"group": "ablation", "label": "No Subsidy", "order": 70},
    "no_subsidy_cost_org": {"group": "ablation", "label": "No Subsidy Cost-Ideal", "order": 80},
    "no_subsidy_cost": {"group": "ablation", "label": "No Subsidy Cost-Implemented", "order": 90},
    "no_operating_cost_org": {"group": "ablation", "label": "No Operating Cost-Ideal", "order": 100},
    "no_operating_cost": {"group": "ablation", "label": "No Operating Cost-Implemented", "order": 110},
}


def resolve_repo_path(path: Path) -> Path:
    return path if path.is_absolute() else REPO_ROOT / path


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Collect mode_comparison results and export a CSV aligned with paper Table 2."
    )
    parser.add_argument(
        "--input-dir",
        type=Path,
        default=DEFAULT_INPUT_DIR,
        help="Directory containing mode comparison JSON outputs.",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=DEFAULT_OUTPUT_PATH,
        help="CSV output path.",
    )
    return parser


def load_json(path: Path) -> dict[str, Any]:
    return json.loads(path.read_text(encoding="utf-8-sig"))


def is_finite_number(value: Any) -> bool:
    return isinstance(value, (int, float)) and math.isfinite(float(value))


def format_decimal(value: Any, digits: int = 2) -> str:
    if not is_finite_number(value):
        return ""
    return f"{float(value):.{digits}f}"


def format_integerish(value: Any) -> str:
    if not is_finite_number(value):
        return ""
    return f"{float(value):.0f}"


def classify_status(data: dict[str, Any]) -> tuple[str, str]:
    upper_bound = data.get("upper_bound")
    lower_bound = data.get("lower_bound")
    has_solution = data.get("optimal_up_solution") is not None and data.get("optimal_mid_solution") is not None

    if has_solution and is_finite_number(upper_bound):
        return "feasible", ""
    if upper_bound == float("inf") or lower_bound == float("inf"):
        return "infeasible", "upper/lower bound is Infinity"
    if data.get("optimal_up_solution") is None or data.get("optimal_mid_solution") is None:
        return "infeasible", "missing upper or middle solution"
    return "unknown", "result does not contain a finite incumbent"


def ordered_solution_items(solution: dict[str, Any] | None) -> list[tuple[str, float]]:
    if not solution:
        return []
    return sorted(((str(key), float(value)) for key, value in solution.items()), key=lambda item: item[0])


def format_solution(solution: dict[str, Any] | None) -> str:
    items = ordered_solution_items(solution)
    return "; ".join(f"{key}={value:.0f}" if float(value).is_integer() else f"{key}={value:g}" for key, value in items)


def count_active_lines(solution: dict[str, Any] | None) -> str:
    items = ordered_solution_items(solution)
    if not items:
        return ""
    return str(sum(1 for _, value in items if abs(value) > 1e-9))


def sum_frequency(solution: dict[str, Any] | None) -> str:
    items = ordered_solution_items(solution)
    if not items:
        return ""
    return format_decimal(sum(value for _, value in items), 2)


def infer_scenario_key(path: Path, data: dict[str, Any]) -> str:
    stem = path.stem
    if stem in SCENARIO_METADATA:
        return stem
    for key in ("regime", "ablation"):
        value = data.get(key)
        if isinstance(value, str) and value in SCENARIO_METADATA:
            return value
    return stem


def build_row(path: Path, data: dict[str, Any]) -> dict[str, str]:
    scenario_key = infer_scenario_key(path, data)
    metadata = SCENARIO_METADATA.get(
        scenario_key,
        {"group": "other", "label": scenario_key, "order": 999},
    )
    config = data.get("config", {})
    status, note = classify_status(data)
    funding = data.get("optimal_obj_funding")
    mu = config.get("mu")
    funding_cost = float(funding) * float(mu) if is_finite_number(funding) and is_finite_number(mu) else None

    return {
        "scenario_group": str(metadata["group"]),
        "scenario": str(metadata["label"]),
        "source_file": path.name,
        "status": status,
        "status_note": note,
        "regime": str(data.get("regime", "")),
        "ablation": str(data.get("ablation", "")),
        "num_lines": format_integerish(config.get("route_num")),
        "num_od": format_integerish(config.get("od_num")),
        "freq_set": "{" + ",".join(str(int(v)) if float(v).is_integer() else str(v) for v in config.get("freq_set", [])) + "}",
        "fare": format_decimal(config.get("tau")),
        "subsidy_cap": format_decimal(config.get("s_bar")),
        "operating_cost_cap": format_decimal(config.get("c_bar")),
        "funding_cost_coeff": format_decimal(config.get("mu")),
        "objective": format_decimal(data.get("upper_bound")),
        "lower_bound": format_decimal(data.get("lower_bound")),
        "gap_pct": format_decimal(data.get("Gap")),
        "passenger_cost": format_decimal(data.get("optimal_obj_user_cost")),
        "operator_cost": format_decimal(data.get("optimal_obj_mid_oper")),
        "funding": format_decimal(funding),
        "funding_cost": format_decimal(funding_cost),
        "revenue": format_decimal(data.get("optimal_mid_revenue")),
        "cpu_s": format_integerish(data.get("CPU (s)")),
        "best_time_s": format_integerish(data.get("time_find_best")),
        "outer_iterations": format_integerish(data.get("iterations")),
        "xf_iterations": format_integerish(data.get("x-f_iterations")),
        "max_xf_iterations": format_integerish(data.get("max_f_iterations")),
        "active_lines": count_active_lines(data.get("optimal_up_solution")),
        "total_frequency": sum_frequency(data.get("optimal_mid_solution")),
        "line_selection": format_solution(data.get("optimal_up_solution")),
        "frequency_plan": format_solution(data.get("optimal_mid_solution")),
        "_sort_order": str(metadata["order"]),
    }


def collect_rows(input_dir: Path) -> list[dict[str, str]]:
    if not input_dir.exists():
        raise FileNotFoundError(f"Input directory does not exist: {input_dir}")

    rows = [build_row(path, load_json(path)) for path in sorted(input_dir.glob("*.json"))]
    rows.sort(key=lambda row: (int(row["_sort_order"]), row["source_file"]))
    for row in rows:
        row.pop("_sort_order", None)
    return rows


def write_csv(rows: list[dict[str, str]], output_path: Path) -> None:
    output_path.parent.mkdir(parents=True, exist_ok=True)
    fieldnames = [
        "scenario_group",
        "scenario",
        "source_file",
        "status",
        "status_note",
        "regime",
        "ablation",
        "num_lines",
        "num_od",
        "freq_set",
        "fare",
        "subsidy_cap",
        "operating_cost_cap",
        "funding_cost_coeff",
        "objective",
        "lower_bound",
        "gap_pct",
        "passenger_cost",
        "operator_cost",
        "funding",
        "funding_cost",
        "revenue",
        "cpu_s",
        "best_time_s",
        "outer_iterations",
        "xf_iterations",
        "max_xf_iterations",
        "active_lines",
        "total_frequency",
        "line_selection",
        "frequency_plan",
    ]
    with output_path.open("w", newline="", encoding="utf-8-sig") as csv_file:
        writer = csv.DictWriter(csv_file, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)


def main() -> None:
    args = build_parser().parse_args()
    input_dir = resolve_repo_path(args.input_dir)
    output_path = resolve_repo_path(args.output)
    rows = collect_rows(input_dir)
    write_csv(rows, output_path)
    print(f"Wrote {len(rows)} rows to {output_path}")


if __name__ == "__main__":
    main()
