#!/usr/bin/env python3
"""Merge the 5k GNN runs with the existing graph-size experiment."""
from __future__ import annotations

import csv
import json
import shutil
import sys
from pathlib import Path

PROJECT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(PROJECT))
from models import gnn_joint_experiment as experiment

MAIN = PROJECT / "results/ME1a_DT_comparison/gnn_joint"
ADDED = PROJECT / "results/ME1a_DT_comparison/gnn_joint_5000"


def main() -> None:
    report = json.loads((MAIN / "report.json").read_text())
    extra = json.loads((ADDED / "report.json").read_text())
    if report["epochs_per_fit"] != extra["epochs_per_fit"]:
        raise ValueError("Graph-size fits must use the same number of epochs")
    new_rows = extra["rows"]
    for row in new_rows:
        source = PROJECT / row["checkpoint"]
        target = MAIN / source.name
        shutil.copy2(source, target)
        row["checkpoint"] = str(target.relative_to(PROJECT))
    new_budgets = {r["budget"] for r in new_rows}
    rows = [r for r in report["rows"] if r["budget"] not in new_budgets] + new_rows
    for row in rows:
        row.setdefault("loss_mode", "soft")
    rows.sort(key=lambda r: (r["budget"], r["kind"], r["setup"], r["eval_run"]))
    report["rows"] = rows
    (MAIN / "report.json").write_text(json.dumps(report, indent=2), encoding="utf-8")
    with (MAIN / "report.csv").open("w", newline="") as fh:
        writer = csv.DictWriter(fh, fieldnames=list(rows[0]))
        writer.writeheader()
        writer.writerows(rows)
    experiment.OUT = MAIN
    budgets = sorted({r["budget"] for r in rows})
    kinds = ["mesh", "residual_geo", "geo_mlp"]
    experiment.make_figures(rows, budgets, kinds)
    experiment.make_html(rows, budgets, kinds)
    print(f"{len(rows)} rows across {budgets} requested graph nodes")


if __name__ == "__main__":
    main()