#!/usr/bin/env python3
"""Measure how CAD graph node budgets alter image, point, and voxel inputs."""
from __future__ import annotations

import csv
import json
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from scipy.spatial import cKDTree

from representations import make_multiview, cad_points, surface_voxels

PROJECT = Path(__file__).resolve().parents[1]
PARTS = {"AC28": PROJECT, "ME1a_DT": PROJECT/"experiments/ME1a_DT"}
OUT = PROJECT/"results/ME1a_DT_comparison"
BUDGETS = (12000,6000,5000,3000,1500,750,300,100)


def prepared_cad(path: Path) -> dict[str,np.ndarray]:
    with np.load(path,allow_pickle=False) as d:
        p,n,f,e,a,x,v = (d[k] for k in ("pos_mm","normal","face","edge_index","edge_attr","x","feature_valid"))
    angle=np.rad2deg(np.arccos(np.clip(a[:,4],-1,1)))
    turn=np.bincount(e[0],weights=angle,minlength=len(p))/np.maximum(np.bincount(e[0],minlength=len(p)),1)
    return {"pos_mm":p,"normal":n,"face":f,"normal_turn_deg":turn.astype(np.float32),
            "k1_per_mm":x[:,13],"k2_per_mm":x[:,14],"gaussian_signed_per_mm2":x[:,15],
            "curvature_valid":v[:,13]&v[:,14]&v[:,15]}


def physical_voxels(data: dict[str,np.ndarray], origin: np.ndarray, pitch: float) -> set[tuple[int,int,int]]:
    center=data["grid_origin_mm"]+(data["voxel_coords"]+.5)*float(data["voxel_pitch_mm"])
    coord=np.floor((center-origin)/pitch).astype(int)
    return set(map(tuple,coord.tolist()))


def main() -> None:
    OUT.mkdir(parents=True,exist_ok=True)
    records=[]; panels={}
    for name,root in PARTS.items():
        graph_dir=root/"data_processed/graphs"
        index=json.loads((graph_dir/"index.json").read_text())
        by_budget={r["requested_node_budget"]:r for r in index["resolutions"]}
        canonical=None
        panels[name]={}
        for budget in BUDGETS:
            cad=prepared_cad(graph_dir/f"graph_{budget}.npz")
            image=make_multiview(cad,128)
            points=cad_points(cad,4096)["pos_mm"]
            voxel=surface_voxels(cad,32)
            if canonical is None:
                canonical=(image,points,voxel)
            base_image,base_points,base_voxel=canonical
            a=np.asarray(base_image["silhouette"],bool); b=np.asarray(image["silhouette"],bool)
            intersection=int((a&b).sum()); union=int((a|b).sum())
            overlap=a&b
            depth=np.abs(base_image["depth_mm"][overlap]-image["depth_mm"][overlap])
            d1=cKDTree(base_points).query(points,workers=-1)[0]
            d2=cKDTree(points).query(base_points,workers=-1)[0]
            common_origin=base_voxel["grid_origin_mm"]
            common_pitch=float(base_voxel["voxel_pitch_mm"])
            va=physical_voxels(base_voxel,common_origin,common_pitch)
            vb=physical_voxels(voxel,common_origin,common_pitch)
            graph=by_budget[budget]
            records.append({"part":name,"requested_nodes":budget,"actual_nodes":len(cad["pos_mm"]),
                "source_to_graph_p95_mm":graph["surface_distance_mm"]["source_to_simplified"]["p95_mm"],
                "area_change_pct":100*graph["area_change_fraction"],"components":graph["components"],
                "analytic_curvature_valid_pct":100*float(cad["curvature_valid"].mean()),
                "image_silhouette_iou_128":intersection/max(1,union),
                "image_depth_mae_mm_common_pixels":float(depth.mean()) if len(depth) else None,
                "image_common_visible_pixels":int(overlap.sum()),
                "point_symmetric_chamfer_mean_mm":float((d1.mean()+d2.mean())/2),
                "point_symmetric_p95_mm":float((np.quantile(d1,.95)+np.quantile(d2,.95))/2),
                "surface_voxel_jaccard_32":len(va&vb)/max(1,len(va|vb)),
                "surface_voxels_32":len(vb)})
            if budget in (12000,5000,3000,100):panels[name][budget]=image["depth_mm"][2].copy(),image["silhouette"][2].copy()

    (OUT/"representation_sensitivity.json").write_text(json.dumps({
        "reference":"Each part's own 12,000-node request; actual node counts are reported",
        "image":"six 128×128 orthographic views; silhouette IoU and depth error on common visible pixels",
        "points":"4,096 deterministic area-uniform triangle samples; symmetric nearest-neighbour distances",
        "voxels":"32 cells along the longest axis; occupied surface-cell centers rebinned into each part's finest graph physical grid",
        "records":records},indent=2))
    with (OUT/"representation_sensitivity.csv").open("w",newline="") as fh:
        writer=csv.DictWriter(fh,fieldnames=records[0].keys());writer.writeheader();writer.writerows(records)
    fig,axes=plt.subplots(2,4,figsize=(17,7),constrained_layout=True)
    for row,(name,budgets) in enumerate(panels.items()):
        for col,budget in enumerate((12000,5000,3000,100)):
            depth,valid=budgets[budget]
            axes[row,col].imshow(np.ma.masked_where(~valid,depth),cmap="viridis",origin="lower")
            item=next(x for x in records if x["part"]==name and x["requested_nodes"]==budget)
            axes[row,col].set_title(f"{name} request {budget:,}, actual {item['actual_nodes']:,}\n"
                                    f"silhouette IoU {item['image_silhouette_iou_128']:.3f}")
            axes[row,col].axis("off")
    fig.suptitle("Same +Y CAD view and 128² pixels for each graph resolution")
    fig.savefig(OUT/"resolution_views.png",dpi=160);plt.close(fig)
    fig,axes=plt.subplots(1,3,figsize=(14,4),constrained_layout=True)
    for name in PARTS:
        part=[r for r in records if r["part"]==name]
        x=[r["actual_nodes"] for r in part]
        for ax,key,label in zip(axes,("source_to_graph_p95_mm","image_silhouette_iou_128","surface_voxel_jaccard_32"),
                                ("Source→graph P95 (mm)","Six-view silhouette IoU","Surface voxel Jaccard")):
            ax.plot(x,[r[key] for r in part],"o-",label=name)
            ax.set_xscale("log");ax.invert_xaxis();ax.set_xlabel("Actual graph nodes");ax.set_ylabel(label);ax.grid(alpha=.25)
    axes[0].legend();fig.savefig(OUT/"sensitivity_curves.png",dpi=160);plt.close(fig)
    print(OUT/"representation_sensitivity.json")


if __name__=="__main__":main()