#!/usr/bin/env python3
"""Create a supplementary loss figure from audited 100-epoch logs."""

from __future__ import annotations

import csv
import hashlib
import json
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np


ROOT = Path(__file__).resolve().parent.parent
WORKSPACE = ROOT.parents[2]
RESULTS = WORKSPACE / "论文实践" / "05_远端结果待回传"
FIGURES = ROOT / "figures"

FALLBACK_LOGS = {
    "COCO-full, 1\u00d7": RESULTS / "revision_closure_2026-07-20" / "tipcell_revision_training_closure" / "coco_full_e100_s42" / "loss_log.csv",
    "SimCLR+COCO, 1\u00d7": RESULTS / "second_review_closure" / "simclr_coco_e100_s42_bblr1" / "loss_log.csv",
    "SimCLR+COCO, 5\u00d7": RESULTS / "revision_closure_2026-07-20" / "tipcell_revision_training_closure" / "simclr_coco_e100_s42" / "loss_log.csv",
}

NEW_ROOT = RESULTS / "tipcell_mva_final_revision"
SEED44_LOGS = {
    "COCO-full, 1\u00d7": NEW_ROOT / "coco_full_e100_s44_bblr1" / "loss_log.csv",
    "COCO-full, 5\u00d7": NEW_ROOT / "coco_full_e100_s44_bblr5" / "loss_log.csv",
    "SimCLR+COCO, 1\u00d7": NEW_ROOT / "simclr_coco_e100_s44_bblr1" / "loss_log.csv",
    "SimCLR+COCO, 5\u00d7": NEW_ROOT / "simclr_coco_e100_s44_bblr5" / "loss_log.csv",
}

COLORS = {
    "COCO-full, 1\u00d7": "#0072B2",
    "COCO-full, 5\u00d7": "#56B4E9",
    "SimCLR+COCO, 1\u00d7": "#009E73",
    "SimCLR+COCO, 5\u00d7": "#D55E00",
}


def epoch_means(path: Path, field: str) -> tuple[np.ndarray, np.ndarray]:
    if not path.is_file():
        raise FileNotFoundError(path)
    grouped: dict[int, list[float]] = {}
    with path.open(encoding="utf-8-sig", newline="") as handle:
        for row in csv.DictReader(handle):
            epoch = int(row["epoch"])
            value = float(row[field])
            if not np.isfinite(value):
                raise ValueError(f"Non-finite {field} in {path}")
            grouped.setdefault(epoch, []).append(value)
    if set(grouped) != set(range(100)):
        raise ValueError(f"Expected epochs 0--99 in {path}, got {min(grouped)}--{max(grouped)}")
    epochs = np.array(sorted(grouped)) + 1
    means = np.array([np.mean(grouped[epoch]) for epoch in sorted(grouped)])
    return epochs, means


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def main() -> None:
    using_seed44 = all(path.is_file() for path in SEED44_LOGS.values())
    logs = SEED44_LOGS if using_seed44 else FALLBACK_LOGS
    output_stem = "FigS_training_losses_seed44" if using_seed44 else "FigS_training_losses_seed42"
    plt.rcParams.update(
        {
            "font.family": "DejaVu Sans",
            "font.size": 8,
            "axes.labelsize": 8,
            "axes.titlesize": 9,
            "legend.fontsize": 7,
            "axes.linewidth": 0.7,
            "savefig.bbox": "tight",
            "savefig.pad_inches": 0.03,
        }
    )
    fields = (
        ("total_loss", "Total loss"),
        ("loss_objectness", "RPN objectness"),
        ("loss_rpn_box_reg", "RPN box regression"),
        ("loss_classifier", "RoI classification"),
        ("loss_box_reg", "RoI box regression"),
        ("loss_mask", "Mask loss"),
    )
    fig, axes = plt.subplots(2, 3, figsize=(7.0, 5.1), sharex=True)
    for ax, (field, title) in zip(axes.flat, fields):
        for label, path in logs.items():
            epochs, values = epoch_means(path, field)
            ax.plot(epochs, values, color=COLORS[label], label=label)
        ax.set_title(title, loc="left")
        ax.grid(axis="y", color="#D0D0D0", lw=0.5)
        ax.set_xlabel("Epoch")
        ax.set_ylabel("Mean training loss")
    axes[0, 0].legend(frameon=False)
    fig.subplots_adjust(hspace=0.58, wspace=0.42)
    FIGURES.mkdir(parents=True, exist_ok=True)
    for suffix, kwargs in (("pdf", {}), ("eps", {}), ("png", {"dpi": 600})):
        fig.savefig(FIGURES / f"{output_stem}.{suffix}", **kwargs)
    plt.close(fig)
    audit = {
        "status": "ok",
        "seed": 44 if using_seed44 else 42,
        "complete_factorial": using_seed44,
        "output_stem": output_stem,
        "logs": {
            label: {"path": str(path.resolve()), "sha256": sha256(path)}
            for label, path in logs.items()
        },
    }
    (FIGURES / f"{output_stem}_audit.json").write_text(
        json.dumps(audit, indent=2), encoding="utf-8"
    )
    print(f"TRAINING_CURVES_OK output={FIGURES / output_stem}")


if __name__ == "__main__":
    main()
