from __future__ import annotations

import csv
import json
from pathlib import Path
from typing import Mapping


ROOT = Path(__file__).resolve().parents[1]
METRIC_COLUMNS = [
    "Model",
    "Dataset",
    "F1",
    "IoU",
    "mIoU",
    "Precision",
    "Recall",
    "BF1",
    "OA",
    "Kappa",
    "ParamsM",
    "FLOPsG",
    "FPS",
    "GPUUtilMean",
    "GPUMemPeakGB",
    "BestEpoch",
]


def save_metrics(
    model_name: str,
    dataset_name: str,
    split: str,
    metrics_dict: Mapping,
    output_dir: str | Path = "results",
) -> Path:
    base = Path(output_dir)
    if not base.is_absolute():
        base = ROOT / base
    out_dir = base / model_name / dataset_name
    out_dir.mkdir(parents=True, exist_ok=True)
    path = out_dir / f"metrics_{split}.json"
    payload = dict(metrics_dict)
    payload.setdefault("model", model_name)
    payload.setdefault("dataset", dataset_name)
    payload.setdefault("split", split)
    with path.open("w", encoding="utf-8") as f:
        json.dump(payload, f, indent=2, sort_keys=True)
    return path


def append_to_comparison_table(results_dir: str | Path = "results") -> Path:
    base = Path(results_dir)
    if not base.is_absolute():
        base = ROOT / base
    rows = []
    for path in sorted(base.glob("*/*/metrics_test.json")):
        with path.open("r", encoding="utf-8") as f:
            metrics = json.load(f)
        rows.append({
            "Model": metrics.get("model", path.parents[1].name),
            "Dataset": metrics.get("dataset", path.parent.name),
            "F1": metrics.get("f1", metrics.get("F1", "")),
            "IoU": metrics.get("iou", metrics.get("IoU", "")),
            "mIoU": metrics.get("miou", metrics.get("mIoU", "")),
            "Precision": metrics.get("precision", metrics.get("Precision", "")),
            "Recall": metrics.get("recall", metrics.get("Recall", "")),
            "BF1": metrics.get("bf1", metrics.get("BF1", "")),
            "OA": metrics.get("oa", metrics.get("OA", "")),
            "Kappa": metrics.get("kappa", metrics.get("Kappa", "")),
            "ParamsM": metrics.get("params_m", ""),
            "FLOPsG": metrics.get("flops_g", ""),
            "FPS": metrics.get("fps", ""),
            "GPUUtilMean": metrics.get("gpu_util_mean", ""),
            "GPUMemPeakGB": metrics.get("gpu_mem_reserved_peak_gb", metrics.get("gpu_mem_allocated_peak_gb", "")),
            "BestEpoch": metrics.get("best_epoch", metrics.get("epoch", "")),
        })
    out = base / "comparison_table.csv"
    base.mkdir(parents=True, exist_ok=True)
    with out.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=METRIC_COLUMNS)
        writer.writeheader()
        writer.writerows(rows)
    return out