#!/usr/bin/env python3
"""Build explicitly weak ME1a scan/CSV diagnostic targets on its own CAD graph."""
from __future__ import annotations

import itertools
import json
from pathlib import Path

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

from match_targets import geodesic_distance, heater_values, interpolate_between_heaters, provisional_heater_projection

PROJECT = Path(__file__).resolve().parents[1]
ROOT = PROJECT / "experiments/ME1a_DT"
RUN = "14h06"


def pca_frame(points: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    center = np.mean(points, axis=0)
    values, vectors = np.linalg.eigh(np.cov((points-center).T))
    axes = vectors[:, np.argsort(values)[::-1]]
    if np.linalg.det(axes) < 0: axes[:, -1] *= -1
    return center, axes


def rigid_fit(source: np.ndarray, target: np.ndarray) -> np.ndarray:
    a, b = source.mean(0), target.mean(0)
    u, _, vt = np.linalg.svd((source-a).T @ (target-b))
    rotation = vt.T @ u.T
    if np.linalg.det(rotation) < 0:
        vt[-1] *= -1; rotation = vt.T @ u.T
    matrix = np.eye(4)
    matrix[:3, :3] = rotation
    matrix[:3, 3] = b - rotation @ a
    return matrix


def transform(points: np.ndarray, matrix: np.ndarray) -> np.ndarray:
    return points @ matrix[:3, :3].T + matrix[:3, 3]


def register(scan: np.ndarray, cad: np.ndarray) -> tuple[np.ndarray, list[dict]]:
    # Fixed subsets keep this reproducible and prevent dense scan regions dominating.
    source = scan[np.linspace(0, len(scan)-1, min(16000, len(scan)), dtype=int)].astype(float)
    target = cad.astype(float)
    cs, vs = pca_frame(source); ct, vt = pca_frame(target)
    tree = cKDTree(target)
    candidates = []
    for perm in itertools.permutations(range(3)):
        for signs in itertools.product((-1, 1), repeat=3):
            basis = np.eye(3)[:, perm] @ np.diag(signs)
            if np.linalg.det(basis) < 0: continue
            rotation = vt @ basis @ vs.T
            matrix = np.eye(4); matrix[:3, :3] = rotation; matrix[:3, 3] = ct-rotation@cs
            dist = tree.query(transform(source, matrix), workers=-1)[0]
            candidates.append((float(np.mean(np.minimum(dist, 80)**2)**.5), perm, signs, matrix))
    candidates.sort(key=lambda row: row[0])
    reports = []
    for initial, perm, signs, matrix in candidates[:8]:
        matrix = matrix.copy()
        for threshold in (80, 35, 15, 8):
            for _ in range(18):
                moved = transform(source, matrix)
                dist, index = tree.query(moved, workers=-1)
                good = dist < min(threshold, float(np.quantile(dist, .90)))
                if good.sum() < 100: break
                change = rigid_fit(moved[good], target[index[good]])
                matrix = change @ matrix
                if np.linalg.norm(change[:3, 3]) < .005 and np.linalg.norm(change[:3,:3]-np.eye(3)) < 1e-5:
                    break
        dist = tree.query(transform(source, matrix), workers=-1)[0]
        reports.append({"initial_capped_rmse80_mm": initial,
                        "capped_rmse80_mm": float(np.mean(np.minimum(dist,80)**2)**.5),
                        "node_nn_median_mm": float(np.median(dist)),
                        "node_nn_p95_mm": float(np.quantile(dist,.95)),
                        "determinant": float(np.linalg.det(matrix[:3,:3])),
                        "pca_permutation": perm, "pca_signs": signs, "transform": matrix.tolist()})
    reports.sort(key=lambda row: row["capped_rmse80_mm"])
    return np.asarray(reports[0]["transform"]), reports


def main() -> None:
    source = ROOT / "analysis/results"
    output = ROOT / "data_processed/paired"
    inspection = ROOT / "results/inspection"
    output.mkdir(parents=True, exist_ok=True); inspection.mkdir(parents=True, exist_ok=True)
    manifest = json.loads((source/"source_manifest.json").read_text())
    defects = json.loads((source/"defects.json").read_text())
    with np.load(ROOT/"data_processed/graphs/graph_12000.npz", allow_pickle=False) as d:
        graph = {k:d[k] for k in d.files}
    with np.load(ROOT/"data_processed/tabular/process_and_geometry.npz", allow_pickle=False) as d:
        tab = {k:d[k] for k in d.files}
    with np.load(source/f"{RUN}_all_merge_pointcloud_full.npz", allow_pickle=False) as d:
        scan = np.asarray(d["points"], np.float32)
    cad = graph["pos_mm"]; normal = graph["normal"]
    matrix, candidates = register(scan, cad)
    scan_cad = transform(scan, matrix).astype(np.float32)

    # Scanner raster convention inferred from the AC28 run and checked against
    # this ME1a fused height image at measured scan locations.
    fused = PROJECT / manifest["fused_ply"]
    raster = fused.with_name("Merged_Height.tif")
    height = np.asarray(Image.open(raster))
    low = np.asarray(manifest["fused_check"]["min_xyz"])
    high = np.asarray(manifest["fused_check"]["max_xyz"])
    pitch_x = (high[0]-low[0]) / (height.shape[1]-1)
    pitch_y = (high[1]-low[1]) / (height.shape[0]-1)
    check = scan[::max(1,len(scan)//30000)]
    ix = np.rint((check[:,0]-low[0])/.1).astype(int)
    iy = np.rint((high[1]-check[:,1])/.1).astype(int)
    in_raster = (ix>=0)&(ix<height.shape[1])&(iy>=0)&(iy<height.shape[0])
    raster_z = np.full(len(check), np.nan)
    raster_z[in_raster] = height[iy[in_raster],ix[in_raster]]
    good = in_raster & (raster_z>-900)
    abs_z = np.abs(raster_z[good]-check[good,2])
    anchors=[]
    for row in defects:
        px,py = int(float(row["X"])),int(float(row["Y"]))
        z = float(height[py,px]) if 0<=px<height.shape[1] and 0<=py<height.shape[0] else np.nan
        raw = np.array([low[0]+.1*px, high[1]-.1*py, z],float)
        mapped = transform(raw[None],matrix)[0]
        distance,node = cKDTree(cad).query(mapped)
        anchors.append({"csv_no":row["No"],"class":row["Error"],"pixel_xy":[px,py],
                        "width_height_px":[row.get("W", ""),row.get("H", "")],
                        "raw_xyz_mm":raw.tolist(),"cad_xyz_mm":mapped.tolist(),
                        "nearest_cad_node":int(node),"nearest_cad_node_mm":float(distance),
                        "valid_height":bool(np.isfinite(z) and z>-900)})

    edge = graph["edge_index"]; length = graph["edge_attr"][:,0]
    radii = np.asarray([5.,10.,20.],np.float32)
    distances=[]
    for row in anchors:
        if row["valid_height"] and row["nearest_cad_node_mm"] < 30:
            distances.append(geodesic_distance(len(cad),edge,length,row["nearest_cad_node"]))
    if not distances:
        raise RuntimeError("No defect anchor is near the CAD; diagnostic training labels cannot be built")
    minimum = np.min(np.stack(distances),axis=0)
    weak_mask = minimum[:,None] <= radii[None]
    weak_probability = np.exp(-.5*(minimum[:,None]/radii[None])**2).astype(np.float32)
    scan_tree = cKDTree(scan_cad)
    cad_distance,cad_index = scan_tree.query(cad,workers=-1)
    cad_delta = scan_cad[cad_index]-cad
    cad_signed = np.einsum("ij,ij->i",cad_delta,normal)
    top_raw,bottom_raw = heater_values(tab["process_feature_names"],tab["process_x"][0],tab["process_feature_valid"][0])
    top,top_imputed=interpolate_between_heaters(top_raw)
    bottom,bottom_imputed=interpolate_between_heaters(bottom_raw)
    zone,node_top,node_bottom,node_delta=provisional_heater_projection(cad,top,bottom)
    paired=output/f"{RUN}_cad_ply_targets.npz"
    np.savez_compressed(paired,run_id=np.asarray(RUN),cad_pos_mm=cad,cad_normal=normal,
        cad_edge_index=edge,process_x=tab["process_x"][0],process_feature_valid=tab["process_feature_valid"][0],
        process_feature_names=tab["process_feature_names"],scan_pos_raw_mm=scan,
        scan_pos_cad_mm=scan_cad,cad_to_scan_distance_mm=cad_distance.astype(np.float32),
        cad_to_scan_normal_residual_mm=cad_signed.astype(np.float32),defect_class=np.asarray("multiple_unvalidated"),
        defect_anchor_raw_mm=np.asarray([a["raw_xyz_mm"] for a in anchors],np.float32),
        defect_anchor_cad_mm=np.asarray([a["cad_xyz_mm"] for a in anchors],np.float32),
        geodesic_distance_to_defect_seed_mm=minimum,weak_area_radii_mm=radii,
        weak_defect_mask=weak_mask,weak_defect_probability=weak_probability,
        heater_top_c=top,heater_bottom_c=bottom,heater_top_c_raw=top_raw,heater_bottom_c_raw=bottom_raw,
        heater_top_imputed=top_imputed,heater_bottom_imputed=bottom_imputed,
        heater_zone_index_zero_based=zone,provisional_heater_top_c_per_cad_node=node_top,
        provisional_heater_bottom_c_per_cad_node=node_bottom,
        provisional_heater_delta_c_per_cad_node=node_delta,heater_projection_valid=np.asarray(False),
        transform_raw_ply_to_cad=matrix,transform_determinant=np.asarray(np.linalg.det(matrix[:3,:3])),
        registration_independently_validated=np.asarray(False),training_valid=np.zeros(len(cad),bool))

    fig,axes=plt.subplots(2,3,figsize=(16,8),constrained_layout=True)
    for ax,(i,j) in zip(axes[0],((0,1),(0,2),(1,2))):
        ax.scatter(cad[::2,i],cad[::2,j],s=.4,c="#234b69",label="ME1a CAD",rasterized=True)
        ax.scatter(scan_cad[::25,i],scan_cad[::25,j],s=.3,c="#e38635",label="fused scan",rasterized=True)
        ax.set_aspect("equal");ax.set_xlabel("XYZ"[i]+" mm");ax.set_ylabel("XYZ"[j]+" mm");ax.legend(markerscale=8)
    axes[1,0].scatter(cad[:,0],cad[:,2],s=2,c=weak_probability[:,1],cmap="magma",vmin=0,vmax=1,rasterized=True)
    axes[1,0].set_title("6 CSV seeds, assumed 10 mm geodesic risk")
    axes[1,1].scatter(cad[:,0],cad[:,2],s=2,c=np.minimum(cad_distance,30),cmap="viridis",vmin=0,vmax=30,rasterized=True)
    axes[1,1].set_title("CAD to fused scan nearest distance, capped 30 mm")
    axes[1,2].bar(np.arange(17)+1,top,label="upper");axes[1,2].plot(np.arange(17)+1,bottom,"o-",label="lower")
    axes[1,2].set_title("XLS heaters; zone position on CAD unverified");axes[1,2].legend()
    for ax in axes[1,:2]:ax.set_aspect("equal");ax.set_xlabel("CAD X (mm)");ax.set_ylabel("CAD Z (mm)")
    fig.suptitle("ME1a visual inspection — proper rigid fit is diagnostic only")
    fig.savefig(inspection/"alignment_targets.png",dpi=160);plt.close(fig)

    review={"run":RUN,"fused_ply_vertices":manifest["fused_check"]["vertex_count"],
      "raster_size_xy":[height.shape[1],height.shape[0]],"inferred_pitch_xy_mm":[float(pitch_x),float(pitch_y)],
      "raster_mapping_check":{"sample_count":len(check),"valid_fraction":float(good.mean()),
                              "height_abs_median_mm":float(np.median(abs_z)),"height_abs_p90_mm":float(np.quantile(abs_z,.9))},
      "registration_candidates":candidates,"selected_pose_is_independently_validated":False,
      "anchors":anchors,"anchors_within_30mm":len(distances),
      "weak_10mm_positive_nodes":int(weak_mask[:,1].sum()),
      "training_valid_nodes":0,"paired_npz":str(paired.relative_to(ROOT)),
      "limitations":["PCA/ICP pose is a best-fit diagnostic, not an independently verified scanner extrinsic",
          "CSV pixel convention inferred from raster geometry but defect identity and extent remain unreviewed",
          "Nearest surface residual is not material-point deformation",
          "All model fits below are training-set weak-label diagnostics only"]}
    (inspection/"target_review.json").write_text(json.dumps(review,ensure_ascii=False,indent=2))
    (output/"index.json").write_text(json.dumps({"entries":[{"run":RUN,"npz":paired.name,
        "training_valid_nodes":0,"status":"diagnostic, unvalidated"}],"review":"../../results/inspection/target_review.json"},indent=2))
    print(json.dumps({"pose_capped_rmse80_mm":candidates[0]["capped_rmse80_mm"],
        "raster_height_p90_mm":review["raster_mapping_check"]["height_abs_p90_mm"],
        "anchors_within_30mm":len(distances),"weak_10mm_positive_nodes":int(weak_mask[:,1].sum())},indent=2))


if __name__=="__main__":main()