#!/usr/bin/env python3
"""Merge separately generated 5k CAD graphs into both part resolution sets."""
from __future__ import annotations

import json
import shutil
from pathlib import Path

import numpy as np

from cad_graphs import _save_figure, _write_viewer

PROJECT = Path(__file__).resolve().parents[1]
PREP = PROJECT / "results/ME1a_DT_comparison/graph_5000_prep"
PARTS = {"AC28": PROJECT, "ME1a_DT": PROJECT / "experiments/ME1a_DT"}


def integrate(part: str, root: Path) -> None:
    source = PREP / part / "graphs"
    target = root / "data_processed/graphs"
    new = json.loads((source / "index.json").read_text())
    index_path = target / "index.json"
    index = json.loads(index_path.read_text())
    assert new["source_vertices"] == index["source_vertices"]
    with np.load(source / "graph_5000.npz", allow_pickle=False) as d, np.load(
        target / "graph_12000.npz", allow_pickle=False
    ) as baseline:
        assert np.array_equal(d["feature_names"], baseline["feature_names"])
        assert np.array_equal(d["edge_feature_names"], baseline["edge_feature_names"])
    shutil.copy2(source / "graph_5000.npz", target / "graph_5000.npz")
    index["resolutions"] = sorted(
        [r for r in index["resolutions"] if r["requested_node_budget"] != 5000]
        + new["resolutions"], key=lambda r: -r["requested_node_budget"]
    )
    index_path.write_text(json.dumps(index, indent=2), encoding="utf-8")
    graphs = []
    for record in index["resolutions"]:
        budget = record["requested_node_budget"]
        with np.load(target / f"graph_{budget}.npz", allow_pickle=False) as d:
            names = list(d["feature_names"].astype(str))
            curvature = d["x"][:, names.index("mean_curvature_abs_1_per_mm")]
            graphs.append({"requested_nodes": budget,
                           "actual_nodes": record["actual_nodes"],
                           "pos_norm": d["pos_norm"], "faces": d["face"].T,
                           "edge_index": d["edge_index"],
                           "curvature": curvature,
                           "distance": record["surface_distance_mm"]})
    out = root / "results/graphs"
    _save_figure(graphs, out / "comparison.png")
    _write_viewer(graphs, out)
    print(part, [(g["requested_nodes"], g["actual_nodes"]) for g in graphs], flush=True)


if __name__ == "__main__":
    for name, path in PARTS.items():
        integrate(name, path)