#!/usr/bin/env python3
"""Sample trimmed STEP faces in UV and construct topology-derived face adjacency.

This is intentionally a B-Rep representation, not a tessellated bounding-box UV
patch: every UV sample is accepted only when OpenCascade's face classifier says
it lies in the trimmed face domain.
"""
from __future__ import annotations

import argparse
import json
import platform
import resource
import time
from itertools import islice
from collections import defaultdict
from pathlib import Path

import numpy as np


def _root() -> Path: return Path(__file__).resolve().parents[1]


def _peak_memory_mb() -> float:
    rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
    return float(rss / (1024 * 1024 if platform.system() == 'Darwin' else 1024))


def _ocp():
    """Delayed imports give a useful CLI error while the optional OCP wheel installs."""
    try:
        from OCP.STEPControl import STEPControl_Reader
        from OCP.IFSelect import IFSelect_RetDone
        from OCP.TopAbs import TopAbs_FACE, TopAbs_EDGE, TopAbs_IN, TopAbs_ON
        from OCP.TopExp import TopExp, TopExp_Explorer
        from OCP.TopoDS import TopoDS
        from OCP.BRepTools import BRepTools
        from OCP.BRepAdaptor import BRepAdaptor_Surface
        from OCP.BRepLProp import BRepLProp_SLProps
        from OCP.BRepClass import BRepClass_FaceClassifier
        from OCP.gp import gp_Pnt2d
        from OCP.GProp import GProp_GProps
        from OCP.BRepGProp import BRepGProp
        from OCP.collections import IndexedDataMap_TopoDS_Shape_List_TopoDS_Shape_TopTools_ShapeMapHasher, IndexedMap_TopoDS_Shape_TopTools_ShapeMapHasher
    except ImportError as exc:
        raise RuntimeError("OCP/cadquery-ocp is required for STEP preprocessing. Use the project .venv.") from exc
    return locals()


def _shape_hash(shape) -> int:
    # HashCode is stable for one OpenCascade process.  IsSame resolves a rare hash
    # collision below, so it cannot create false face adjacency.
    # OCP 8 Python wrappers do not expose TopoDS_Shape.HashCode; the underlying
    # TShape handle is stable for an imported model.  IsSame still verifies every
    # bucket candidate, so a Python-hash collision cannot add a false adjacency.
    return int(hash(shape.TShape()))


def _uv_bounds(BRepTools, face):
    try: return tuple(float(x) for x in BRepTools.UVBounds_s(face))
    except AttributeError: return tuple(float(x) for x in BRepTools.UVBounds(face))


def _surface_area(BRepGProp, GProp_GProps, face) -> float:
    props = GProp_GProps()
    try: BRepGProp.SurfaceProperties_s(face, props)
    except AttributeError: BRepGProp.SurfaceProperties(face, props)
    return float(props.Mass())


def _edge_length(BRepGProp, GProp_GProps, edge) -> float:
    props = GProp_GProps()
    try: BRepGProp.LinearProperties_s(edge, props)
    except AttributeError: BRepGProp.LinearProperties(edge, props)
    return float(props.Mass())


def _faces(shape, O):
    out = []; ex = O['TopExp_Explorer'](shape, O['TopAbs_FACE'])
    while ex.More():
        out.append(O['TopoDS'].Face(ex.Current())); ex.Next()
    return out


def _face_adjacency(shape, face_count: int, O):
    """Kernel indexed edge→face map: O(E+F), no geometric or O(F²) matching."""
    edge_map = O['IndexedDataMap_TopoDS_Shape_List_TopoDS_Shape_TopTools_ShapeMapHasher']()
    O['TopExp'].MapShapesAndAncestors_s(shape, O['TopAbs_EDGE'], O['TopAbs_FACE'], edge_map)
    face_map = O['IndexedMap_TopoDS_Shape_TopTools_ShapeMapHasher']()
    O['TopExp'].MapShapes_s(shape, O['TopAbs_FACE'], face_map)
    pairs: list[tuple[int, int, float]] = []; boundary = 0; nonmanifold = 0
    for i in range(1, edge_map.Extent()+1):
        edge = O['TopoDS'].Edge(edge_map.FindKey(i))
        ancestors = edge_map.FindFromIndex(i)
        # OCP's list wrapper has an iterator but (in several 8.x wheels) no
        # Python StopIteration; bound it by the kernel-reported list length.
        owners = sorted({face_map.FindIndex(face)-1 for face in islice(ancestors, ancestors.Length()) if 0 < face_map.FindIndex(face) <= face_count})
        # When a debug max-face subset is requested, links to excluded faces are
        # intentionally absent and do not alter the full-model adjacency logic.
        if len(owners) == 1: boundary += 1; continue
        if len(owners) < 2: continue
        if len(owners) > 2: nonmanifold += 1
        length = _edge_length(O['BRepGProp'], O['GProp_GProps'], edge)
        # A non-manifold edge receives all pairwise face links, explicitly
        # reported in metadata instead of being silently dropped.
        for oi, a in enumerate(owners):
            for b in owners[oi+1:]: pairs.extend(((a, b, length), (b, a, length)))
    arr = np.asarray(pairs, dtype=np.float64) if pairs else np.empty((0, 3), dtype=np.float64)
    return (arr[:, :2].T.astype(np.int32), arr[:, 2:3].astype(np.float32), boundary, nonmanifold)


def _kind(surface) -> str:
    try: return str(surface.GetType()).replace('GeomAbs_', '')
    except Exception: return 'unknown'


def _sample_face(face, grid: int, O):
    U0, U1, V0, V1 = _uv_bounds(O['BRepTools'], face)
    us, vs = np.linspace(U0, U1, grid), np.linspace(V0, V1, grid)
    pos = np.zeros((grid, grid, 3), np.float32); normal = np.zeros_like(pos)
    curv = np.zeros((grid, grid, 3), np.float32)  # |kmax|, |kmin|, signed Gaussian
    valid = np.zeros((grid, grid), bool); normal_valid = np.zeros((grid, grid), bool); curvature_valid = np.zeros((grid, grid), bool)
    surface = O['BRepAdaptor_Surface'](face, True)
    for ui, u in enumerate(us):
        for vi, v in enumerate(vs):
            # Classifier applies wires/holes/trims. Bounds are merely a candidate domain.
            classifier = O['BRepClass_FaceClassifier'](face, O['gp_Pnt2d'](float(u), float(v)), 1e-7)
            if classifier.State() not in (O['TopAbs_IN'], O['TopAbs_ON']): continue
            try:
                props = O['BRepLProp_SLProps'](surface, float(u), float(v), 2, 1e-8)
                p = props.Value(); xyz = np.array([p.X(), p.Y(), p.Z()], dtype=np.float64)
                if not np.isfinite(xyz).all(): continue
                pos[ui, vi] = xyz; valid[ui, vi] = True
                if props.IsNormalDefined():
                    q = props.Normal(); nn = np.array([q.X(), q.Y(), q.Z()], dtype=np.float64); ln = np.linalg.norm(nn)
                    if ln > 1e-12 and np.isfinite(nn).all(): normal[ui, vi] = nn / ln; normal_valid[ui, vi] = True
                if props.IsCurvatureDefined():
                    kmax, kmin = float(props.MaxCurvature()), float(props.MinCurvature())
                    if np.isfinite([kmax, kmin]).all(): curv[ui, vi] = (abs(kmax), abs(kmin), kmax*kmin); curvature_valid[ui, vi] = True
            except Exception:
                # Singularities and individual failed evaluations remain explicit masks.
                continue
    return pos, normal, curv, valid, normal_valid, curvature_valid, (U0, U1, V0, V1), _kind(surface)


def _preview(valid: np.ndarray, kinds: list[str], areas: np.ndarray, path: Path) -> None:
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    fig, ax = plt.subplots(1, 2, figsize=(10, 4))
    ax[0].imshow(valid.mean(axis=0), vmin=0, vmax=1, cmap='viridis'); ax[0].set_title('trim-valid fraction at each UV grid cell'); ax[0].set_xlabel('v index'); ax[0].set_ylabel('u index')
    names, count = np.unique(kinds, return_counts=True); ax[1].barh(names, count); ax[1].set_title('surface kinds'); ax[1].set_xlabel('faces')
    fig.tight_layout(); fig.savefig(path, dpi=160); plt.close(fig)


def prepare(step_path: Path, out_dir: Path, grid: int = 8, max_faces: int | None = None, results_dir: Path | None = None) -> dict:
    O = _ocp(); started = time.perf_counter(); out_dir.mkdir(parents=True, exist_ok=True)
    reader = O['STEPControl_Reader']()
    if reader.ReadFile(str(step_path)) != O['IFSelect_RetDone']: raise ValueError(f'Failed to read STEP: {step_path}')
    reader.TransferRoots(); shape = reader.OneShape(); faces = _faces(shape, O)
    if max_faces is not None: faces = faces[:max_faces]
    F = len(faces)
    pos = np.zeros((F, grid, grid, 3), np.float32); normal = np.zeros_like(pos); curvature = np.zeros((F, grid, grid, 3), np.float32)
    valid = np.zeros((F, grid, grid), bool); normal_valid = valid.copy(); curvature_valid = valid.copy(); uv_bounds = np.zeros((F, 4), np.float64); areas = np.zeros(F, np.float64); kinds: list[str] = []
    for i, face in enumerate(faces):
        pp, nn, cc, vv, nv, cv, bounds, kind = _sample_face(face, grid, O)
        pos[i], normal[i], curvature[i], valid[i], normal_valid[i], curvature_valid[i], uv_bounds[i] = pp, nn, cc, vv, nv, cv, bounds
        areas[i] = _surface_area(O['BRepGProp'], O['GProp_GProps'], face); kinds.append(kind)
    edge_index, edge_attr, boundary, nonmanifold = _face_adjacency(shape, F, O)
    names = np.asarray(['kmax_abs_1_per_mm','kmin_abs_1_per_mm','gaussian_signed_1_per_mm2'], dtype='<U36')
    np.savez_compressed(out_dir/'faces_uv.npz', pos_mm=pos, normal=normal, curvature=curvature, trim_valid=valid, normal_valid=normal_valid, curvature_valid=curvature_valid,
                        uv_bounds=uv_bounds, face_area_mm2=areas.astype(np.float32), surface_kind=np.asarray(kinds, dtype='<U32'), curvature_feature_names=names,
                        edge_index=edge_index, edge_attr=edge_attr, edge_feature_names=np.asarray(['shared_edge_length_mm'], dtype='<U32'),
                        training_valid=np.zeros(F, dtype=bool))
    meta = {'source': str(step_path), 'face_count': F, 'grid': [grid, grid], 'trim_valid_fraction': float(valid.mean()), 'normal_valid_fraction': float(normal_valid.mean()),
            'curvature_valid_fraction': float(curvature_valid.mean()), 'face_adjacency_directed_edges': int(edge_index.shape[1]), 'boundary_topological_edges': boundary,
            'nonmanifold_topological_edges': nonmanifold, 'surface_kind_counts': {k:int(kinds.count(k)) for k in sorted(set(kinds))}, 'elapsed_s': time.perf_counter()-started, 'peak_process_memory_mb': _peak_memory_mb(),
            'provenance': 'OpenCascade STEP B-Rep; UV bounds followed by BRepClass_FaceClassifier trim test. Normals follow face/surface orientation.',
            'curvature_caveat': 'Principal curvature signs are stored as magnitudes because face orientation can flip sign; signed Gaussian curvature is retained and orientation invariant.',
            'training_valid': 'all false: no physical labels were supplied'}
    (out_dir/'metadata.json').write_text(json.dumps(meta, indent=2), encoding='utf8')
    results_dir = results_dir or _root()/'results/brep'; results_dir.mkdir(parents=True, exist_ok=True)
    (results_dir/'index.json').write_text(json.dumps(meta, indent=2), encoding='utf8'); _preview(valid, kinds, areas, results_dir/'preview.png'); return meta


def main() -> None:
    p = argparse.ArgumentParser(description=__doc__); p.add_argument('--step', type=Path, default=_root()/'datas_2/AC28_DT RR UPR STD.stp'); p.add_argument('--out', type=Path, default=_root()/'data_processed/brep'); p.add_argument('--results', type=Path, default=_root()/'results/brep'); p.add_argument('--grid', type=int, default=8); p.add_argument('--max-faces', type=int, default=None, help='debug only; default processes every face')
    a=p.parse_args()
    if a.grid < 2: p.error('--grid must be at least 2')
    print(json.dumps(prepare(a.step, a.out, a.grid, a.max_faces, a.results), indent=2))


if __name__ == '__main__': main()