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]
DEFAULT_INPUT_DIR = REPO_ROOT / "results" / "main_experiments"
DEFAULT_OUTPUT_DIR = REPO_ROOT / "results" / "main_experiments" / "convergence"
DEFAULT_INSTANCES = [
    "L-6-OD-30-P-4",
    "L-8-OD-30-P-3",
]
ACCELERATION_DIRS = (
    "with_acceleration",
    "without_acceleration",
)
CSV_COLUMNS = [
    "iter_step",
    "ub",
    "lb",
    "cumulative_elapsed_time",
    "new_incumbent",
    "relax_bound",
    "relax_activate",
]


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="Export iteration-level convergence information from main experiment JSON files."
    )
    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-dir",
        type=Path,
        default=DEFAULT_OUTPUT_DIR,
        help="Directory where one CSV per result will be written.",
    )
    parser.add_argument(
        "--instances",
        nargs="+",
        default=DEFAULT_INSTANCES,
        help="Instance basenames to export, without the .json suffix.",
    )
    return parser


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


def normalize_iteration_row(row: dict[str, Any]) -> dict[str, Any]:
    return {column: row.get(column) for column in CSV_COLUMNS}


def write_csv(rows: list[dict[str, Any]], output_path: Path) -> None:
    output_path.parent.mkdir(parents=True, exist_ok=True)
    with output_path.open("w", newline="", encoding="utf-8-sig") as csv_file:
        writer = csv.DictWriter(csv_file, fieldnames=CSV_COLUMNS)
        writer.writeheader()
        writer.writerows(rows)


def export_one_result(input_path: Path, output_path: Path) -> int:
    data = load_json(input_path)
    iteration_details = data.get("iteration_details")
    if not isinstance(iteration_details, list):
        raise ValueError(f"Missing or invalid 'iteration_details' in {input_path}")

    rows = [normalize_iteration_row(row) for row in iteration_details if isinstance(row, dict)]
    write_csv(rows, output_path)
    return len(rows)


def export_convergence(input_dir: Path, output_dir: Path, instances: list[str]) -> list[Path]:
    written_files: list[Path] = []
    missing_files: list[Path] = []

    for acceleration_dir in ACCELERATION_DIRS:
        source_dir = input_dir / acceleration_dir
        if not source_dir.exists():
            raise FileNotFoundError(f"Input subdirectory does not exist: {source_dir}")

        for instance in instances:
            input_path = source_dir / f"{instance}.json"
            if not input_path.exists():
                missing_files.append(input_path)
                continue

            output_name = f"{instance}-{acceleration_dir}.csv"
            output_path = output_dir / output_name
            export_one_result(input_path, output_path)
            written_files.append(output_path)

    if missing_files:
        missing_text = "\n".join(str(path) for path in missing_files)
        raise FileNotFoundError(f"Missing result files:\n{missing_text}")

    return written_files


def main() -> None:
    args = build_parser().parse_args()
    input_dir = resolve_repo_path(args.input_dir)
    output_dir = resolve_repo_path(args.output_dir)
    written_files = export_convergence(input_dir, output_dir, list(args.instances))
    print(f"Wrote {len(written_files)} CSV files to {output_dir}")
    for path in written_files:
        print(path)


if __name__ == "__main__":
    main()
