#!/usr/bin/env python3
"""Prepare the ME1a source contracts in an isolated AC28-style experiment root."""
from __future__ import annotations

import argparse
import csv
import json
from pathlib import Path

import numpy as np
import xlrd

from audit_ply_fusion import ply_info, sample_ply, full_scan_check

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


def sources() -> None:
    from OCP.STEPControl import STEPControl_Reader
    from OCP.IFSelect import IFSelect_RetDone
    from OCP.BRepMesh import BRepMesh_IncrementalMesh
    from OCP.StlAPI import StlAPI_Writer

    step = next(PROJECT.glob("datas_2/4. CATIA File*/ME1a_DT FRT UPR LHD.stp"))
    xls = next(PROJECT.glob("datas_2/2. Excel File*/ME1a_DT*.Xls"))
    defect = next(PROJECT.glob("datas_2/3. CSV File*/ME1a_DT*.csv"))
    scan_dir = next(PROJECT.glob("datas_2/1. PLY File*/ME1a_DT*"))
    fused = scan_dir / "all_merge_pointcloud_full.ply"
    source = OUT / "analysis/results"
    source.mkdir(parents=True, exist_ok=True)

    stl = source / "design_fine.stl"
    if not stl.exists():
        reader = STEPControl_Reader()
        if reader.ReadFile(str(step)) != IFSelect_RetDone:
            raise RuntimeError(f"Cannot read STEP: {step}")
        reader.TransferRoots()
        shape = reader.OneShape()
        BRepMesh_IncrementalMesh(shape, .15, False, .15, True)
        writer = StlAPI_Writer(); writer.ASCIIMode = False
        if not writer.Write(shape, str(stl)):
            raise RuntimeError("STEP tessellation export failed")

    sheet = xlrd.open_workbook(xls).sheet_by_index(0)
    rows = []
    for r in range(1, sheet.nrows):
        value = sheet.cell_value(r, 4)
        rows.append({"row": r + 1, "machine": sheet.cell_value(r, 0),
                     "channel": sheet.cell_value(r, 1), "unit": sheet.cell_value(r, 2),
                     "type": sheet.cell_value(r, 3), "value": None if value == "" else value,
                     "timestamp": str(sheet.cell_value(r, 5)).strip()})
    (source / "process.json").write_text(json.dumps({RUN: {"source": str(xls.relative_to(PROJECT)),
                                         "sheet": sheet.name, "rows": rows}}, ensure_ascii=False, indent=2))
    with defect.open(encoding="utf-8-sig", newline="") as fh:
        defects = list(csv.DictReader(fh))
    (source / "defects.json").write_text(json.dumps(defects, ensure_ascii=False, indent=2))

    info = ply_info(fused)
    scan, _ = sample_ply(info, 350_000)
    np.savez_compressed(source / f"{RUN}_all_merge_pointcloud_full.npz", points=scan.astype(np.float32), n_raw=info.count)
    check = full_scan_check(info)
    (source / "source_manifest.json").write_text(json.dumps({
        "part": "ME1a_DT FRT UPR LHD", "run": RUN,
        "step": str(step.relative_to(PROJECT)), "xls": str(xls.relative_to(PROJECT)),
        "defect_csv": str(defect.relative_to(PROJECT)), "fused_ply": str(fused.relative_to(PROJECT)),
        "component_ply": [str(p.relative_to(PROJECT)) for p in sorted(scan_dir.glob("merged_pointcloud_*.ply"))],
        "fused_check": check, "sampled_scan_points": len(scan),
        "defect_rows": len(defects), "process_rows": len(rows)
    }, ensure_ascii=False, indent=2))
    print(json.dumps({"stl": str(stl), "ply_vertices": info.count,
                      "sampled_scan_points": len(scan), "process_rows": len(rows),
                      "defect_rows": len(defects)}, indent=2))


def curvature() -> None:
    path = OUT / "data_processed/brep/faces_uv.npz"
    with np.load(path, allow_pickle=False) as d:
        valid = d["trim_valid"] & d["normal_valid"] & d["curvature_valid"]
        pos = d["pos_mm"][valid]
        normal = d["normal"][valid]
        values = d["curvature"][valid]
    dest = OUT / "analysis/results/cad_surface_curvature.npz"
    np.savez_compressed(dest, points_mm=pos.astype(np.float32),
                        parameter_normals=normal.astype(np.float32),
                        k1_per_mm=values[:, 0].astype(np.float32),
                        k2_per_mm=values[:, 1].astype(np.float32))
    print(json.dumps({"curvature_samples": len(pos), "output": str(dest)}))


def legacy_graph() -> None:
    graph = OUT / "data_processed/graphs/graph_12000.npz"
    with np.load(graph, allow_pickle=False) as d:
        p = d["pos_mm"].copy(); n = d["normal"].copy()
        face = d["face"].copy(); edge = d["edge_index"].copy()
        attr = d["edge_attr"].copy(); x = d["x"].copy(); valid = d["feature_valid"].copy()
    angle = np.rad2deg(np.arccos(np.clip(attr[:, 4], -1, 1)))
    turn = np.bincount(edge[0], weights=angle, minlength=len(p)) / np.maximum(
        np.bincount(edge[0], minlength=len(p)), 1)
    out = OUT / "analysis/results/cad_graph.npz"
    np.savez_compressed(out, pos_mm=p, normal=n, face=face, edge_index=edge,
                        edge_length_mm=attr[:, 0], 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=valid[:, 13] & valid[:, 14] & valid[:, 15],
                        curvature_mapping_distance_mm=np.full(len(p), np.nan, np.float32))
    (OUT / "analysis/results/cad_graph_provenance.json").write_text(json.dumps({
        "source": str(graph), "analytic_k1_k2": "absolute principal curvature magnitudes",
        "gaussian": "separate signed Gaussian channel", "validity": "all analytic curvature channels valid",
        "normal_turn": "mean angle to topological graph neighbours, degrees"
    }, indent=2))
    print(json.dumps({"graph_nodes": len(p), "directed_edges": edge.shape[1],
                      "analytic_curvature_valid_nodes": int((valid[:, 13] & valid[:, 14] & valid[:, 15]).sum())}))


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--stage", required=True, choices=["sources", "curvature", "legacy-graph"])
    args = parser.parse_args()
    {"sources": sources, "curvature": curvature, "legacy-graph": legacy_graph}[args.stage]()


if __name__ == "__main__":
    main()