from __future__ import annotations

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


REPO_ROOT = Path(__file__).resolve().parents[2]

KIND_CONFIG: dict[str, dict[str, str]] = {
    "subsidy": {
        "config_key": "s_bar",
        "column_name": "subsidy",
        "default_output_name": "subsidy.csv",
    },
    "cost": {
        "config_key": "c_bar",
        "column_name": "cost",
        "default_output_name": "cost.csv",
    },
    "price": {
        "config_key": "tau",
        "column_name": "price",
        "default_output_name": "price.csv",
    },
    "min_freq": {
        "config_key": "freq_set",
        "column_name": "min_freq",
        "default_output_name": "min_freq.csv",
    }
}


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Aggregate sensitivity-analysis JSON results into a CSV table."
    )
    parser.add_argument(
        "--kind",
        choices=sorted(KIND_CONFIG),
        required=True,
        help="Sensitivity experiment type.",
    )
    parser.add_argument(
        "--input-dir",
        type=Path,
        required=True,
        help="Directory containing JSON result files.",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=None,
        help="CSV output path. Defaults to <input-dir>/<kind>.csv.",
    )
    return parser


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


def line_sort_key(line_id: str) -> tuple[int, str]:
    suffix = line_id[1:] if line_id.startswith("L") else line_id
    return (int(suffix), line_id) if suffix.isdigit() else (10**9, line_id)


def ordered_solution_values(solution: dict[str, Any] | None) -> list[float]:
    if not solution:
        return []
    return [float(solution[line_id]) for line_id in sorted(solution, key=line_sort_key)]


def format_parameter_value(value: Any) -> str:
    return f"{float(value):.1f}"


def build_row(data: dict[str, Any], kind: str) -> dict[str, Any]:
    cfg = KIND_CONFIG[kind]
    config = data["config"]
    optimal_x = ordered_solution_values(data.get("optimal_up_solution"))
    optimal_f = ordered_solution_values(data.get("optimal_mid_solution"))
    pax = float(data.get("optimal_obj_user_cost", 0.0) or 0.0)
    com = float(data.get("optimal_obj_mid_oper", 0.0) or 0.0)
    funding = float(data.get("optimal_obj_funding", 0.0) or 0.0)
    mu = float(config.get("mu", 0.0) or 0.0)
    funding_cost = mu * funding
    social_cost = float(data.get("upper_bound", pax + com + funding_cost) or 0.0)
    revenue = float(data.get("optimal_mid_revenue", 0.0) or 0.0)
    if kind == "min_freq":
        column_name = format_parameter_value(config[cfg["config_key"]][1])
    else:
        column_name = format_parameter_value(config[cfg["config_key"]])

    return {
        cfg["column_name"]: column_name,
        "optimal_x": str(optimal_x),
        "optimal_f": str(optimal_f),
        "social_cost": social_cost,
        "pax": pax,
        "com": com,
        "funding": funding,
        "revenue": revenue,
        "funding_cost": funding_cost,
        "freq_sum": float(sum(optimal_f)),
    }


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

    rows = [
        build_row(load_json(result_path), kind)
        for result_path in sorted(input_dir.glob("*.json"))
    ]
    column_name = KIND_CONFIG[kind]["column_name"]
    rows.sort(key=lambda row: float(str(row[column_name])))
    return rows


def write_csv(rows: list[dict[str, Any]], output_path: Path, kind: str) -> None:
    output_path.parent.mkdir(parents=True, exist_ok=True)
    fieldnames = [
        KIND_CONFIG[kind]["column_name"],
        "optimal_x",
        "optimal_f",
        "social_cost",
        "pax",
        "com",
        "funding",
        "revenue",
        "funding_cost",
        "freq_sum",
    ]
    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 export_sensitivity_csv(input_dir: Path, kind: str, output_path: Path | None = None) -> Path:
    resolved_output = output_path or input_dir / KIND_CONFIG[kind]["default_output_name"]
    rows = collect_rows(input_dir, kind)
    write_csv(rows, resolved_output, kind)
    return resolved_output


def main() -> None:
    args = build_parser().parse_args()
    rows = collect_rows(args.input_dir, args.kind)
    output_path = args.output or args.input_dir / KIND_CONFIG[args.kind]["default_output_name"]
    write_csv(rows, output_path, args.kind)
    print(f"Wrote {len(rows)} rows to {output_path}")


if __name__ == "__main__":
    main()
