from __future__ import annotations

import argparse
import csv
import json
from dataclasses import dataclass
from pathlib import Path


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


@dataclass(frozen=True)
class InstanceKey:
    route_num: int
    od_num: int
    freq_count: int


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 main_experiments results and export a CSV aligned with paper Table 1."
    )
    parser.add_argument(
        "--input-dir",
        type=Path,
        default=DEFAULT_INPUT_DIR,
        help="Directory containing with_acceleration/ and without_acceleration/ JSON outputs.",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=DEFAULT_OUTPUT_PATH,
        help="CSV output path.",
    )
    return parser


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


def parse_key(data: dict) -> InstanceKey:
    config = data["config"]
    return InstanceKey(
        route_num=int(config["route_num"]),
        od_num=int(config["od_num"]),
        freq_count=len(config["freq_set"]),
    )


def format_decimal(value: float | int | None, digits: int = 2) -> str:
    if value is None:
        return ""
    return f"{float(value):.{digits}f}"


def format_percent(value: float | int | None, digits: int = 2) -> str:
    if value is None:
        return ""
    return f"{float(value):.{digits}f}"


def format_time_seconds(value: float | int | None) -> str:
    if value is None:
        return ""
    return f"{float(value):.0f}"


def format_cardinality(route_num: int, freq_count: int) -> str:
    return f"{2 ** route_num} / {freq_count ** route_num}"


def build_row(no_acc: dict, acc: dict) -> dict[str, str | int]:
    config = acc["config"]
    no_acc_cpu = float(no_acc["CPU (s)"])
    acc_cpu = float(acc["CPU (s)"])
    no_acc_best = float(no_acc.get("time_find_best", no_acc_cpu))
    acc_best = float(acc.get("time_find_best", acc_cpu))
    acc_rate = 0.0 if no_acc_cpu == 0 else (no_acc_cpu - acc_cpu) / no_acc_cpu * 100.0

    return {
        "num_lines": int(config["route_num"]),
        "freq_set": "{" + ",".join(str(int(v)) if float(v).is_integer() else str(v) for v in config["freq_set"]) + "}",
        "cardinality": format_cardinality(int(config["route_num"]), len(config["freq_set"])),
        "num_od": int(config["od_num"]),
        "nested_gbd_obj": format_decimal(no_acc.get("upper_bound")),
        "nested_gbd_gap_pct": format_percent(no_acc.get("Gap", 0.0), 2),
        "nested_gbd_cpu_s": format_time_seconds(no_acc_cpu),
        "nested_gbd_best_s": format_time_seconds(no_acc_best),
        "nested_gbd_a_obj": format_decimal(acc.get("upper_bound")),
        "nested_gbd_a_gap_pct": format_percent(acc.get("Gap", 0.0), 2),
        "nested_gbd_a_cpu_s": format_time_seconds(acc_cpu),
        "nested_gbd_a_best_s": format_time_seconds(acc_best),
        "acc_pct": format_percent(acc_rate, 2),
        "relaxed_activation": int(acc.get("relaxed_activation", 0)),
        "nested_gbd_xf_iterations": f"{int(no_acc['x-f_iterations'])}/{int(no_acc['max_f_iterations'])}",
        "nested_gbd_a_xf_iterations": f"{int(acc['x-f_iterations'])}/{int(acc['max_f_iterations'])}",
    }


def collect_results(input_dir: Path) -> list[dict[str, str | int]]:
    with_dir = input_dir / "with_acceleration"
    without_dir = input_dir / "without_acceleration"
    if not with_dir.exists() or not without_dir.exists():
        raise FileNotFoundError(
            f"Expected both '{with_dir}' and '{without_dir}' to exist."
        )

    with_files = {parse_key(load_json(path)): load_json(path) for path in sorted(with_dir.glob("*.json"))}
    without_files = {parse_key(load_json(path)): load_json(path) for path in sorted(without_dir.glob("*.json"))}

    missing = sorted(set(with_files) ^ set(without_files), key=lambda item: (item.route_num, item.od_num, item.freq_count))
    if missing:
        formatted = ", ".join(
            f"(L={item.route_num}, OD={item.od_num}, |P|={item.freq_count})" for item in missing
        )
        raise ValueError(f"Mismatched result files between acceleration settings: {formatted}")

    rows: list[dict[str, str | int]] = []
    for key in sorted(with_files, key=lambda item: (item.route_num, item.freq_count, item.od_num)):
        rows.append(build_row(without_files[key], with_files[key]))
    return rows


def write_csv(rows: list[dict[str, str | int]], output_path: Path) -> None:
    output_path.parent.mkdir(parents=True, exist_ok=True)
    fieldnames = [
        "num_lines",
        "freq_set",
        "cardinality",
        "num_od",
        "nested_gbd_obj",
        "nested_gbd_gap_pct",
        "nested_gbd_cpu_s",
        "nested_gbd_best_s",
        "nested_gbd_xf_iterations",
        "nested_gbd_a_obj",
        "nested_gbd_a_gap_pct",
        "nested_gbd_a_cpu_s",
        "nested_gbd_a_best_s",
        "nested_gbd_a_xf_iterations",
        "acc_pct",
        "relaxed_activation",
    ]
    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_results(input_dir)
    write_csv(rows, output_path)
    print(f"Wrote {len(rows)} rows to {output_path}")


if __name__ == "__main__":
    main()
