src/comparison/continuity.py

"""G0/G1/G2 surface continuity checking at boundary regions."""

from __future__ import annotations

import logging
from dataclasses import dataclass, field

import numpy as np
from OCP.BRep import BRep_Tool
from OCP.BRepAdaptor import BRepAdaptor_Surface
from OCP.GeomLProp import GeomLProp_SLProps
from OCP.ShapeAnalysis import ShapeAnalysis_Surface
from OCP.TopAbs import TopAbs_FACE
from OCP.TopExp import TopExp_Explorer
from OCP.TopoDS import TopoDS, TopoDS_Face, TopoDS_Shape
from OCP.gp import gp_Pnt

logger = logging.getLogger(__name__)


@dataclass
class ContinuityResult:
    """Results of continuity checks."""

    g0_deviations: list[float] = field(default_factory=list)
    g1_deviations_deg: list[float] = field(default_factory=list)
    g2_deviations_pct: list[float] = field(default_factory=list)

    g0_max: float = 0.0
    g1_max_deg: float = 0.0
    g2_max_pct: float = 0.0

    g0_pass: bool = True
    g1_pass: bool = True
    g2_pass: bool = True

    num_samples_checked: int = 0


def _get_closest_face_and_uv(
    faces: list[TopoDS_Face],
    point: np.ndarray,
) -> tuple[TopoDS_Face, float, float, float] | None:
    """Find the closest face to a 3D point and return UV parameters."""
    best_face = None
    best_dist = float("inf")
    best_u, best_v = 0.0, 0.0

    gp_point = gp_Pnt(float(point[0]), float(point[1]), float(point[2]))

    for face in faces:
        try:
            surface = BRep_Tool.Surface_s(face)
            sas = ShapeAnalysis_Surface(surface)
            uv = sas.ValueOfUV(gp_point, 1.0)
            proj_pnt = sas.Value(uv.X(), uv.Y())
            dist = gp_point.Distance(proj_pnt)

            if dist < best_dist:
                best_dist = dist
                best_face = face
                best_u = uv.X()
                best_v = uv.Y()
        except Exception:
            continue

    if best_face is None:
        return None
    return best_face, best_u, best_v, best_dist


def _evaluate_surface_props(
    face: TopoDS_Face,
    u: float,
    v: float,
) -> tuple[np.ndarray | None, np.ndarray | None, float | None]:
    """Evaluate surface normal and curvature at (u, v)."""
    try:
        surface = BRep_Tool.Surface_s(face)
        props = GeomLProp_SLProps(surface, u, v, 2, 1e-6)

        normal = None
        if props.IsNormalDefined():
            n = props.Normal()
            normal = np.array([n.X(), n.Y(), n.Z()])

        curvature = None
        if props.IsCurvatureDefined():
            curvature = props.MeanCurvature()

        point = props.Value()
        position = np.array([point.X(), point.Y(), point.Z()])

        return position, normal, curvature
    except Exception:
        return None, None, None


def _collect_faces(shape: TopoDS_Shape) -> list[TopoDS_Face]:
    """Collect all faces from a shape."""
    faces = []
    exp = TopExp_Explorer(shape, TopAbs_FACE)
    while exp.More():
        faces.append(TopoDS.Face_s(exp.Current()))
        exp.Next()
    return faces


def check_continuity(
    shape_a: TopoDS_Shape,
    shape_b: TopoDS_Shape,
    sample_points: np.ndarray,
    g0_tolerance_mm: float = 0.05,
    g1_tolerance_deg: float = 0.3,
    g2_tolerance_pct: float = 3.0,
    max_check_points: int = 1000,
) -> ContinuityResult:
    """Check G0/G1/G2 continuity between two shapes at sample points.

    Args:
        shape_a: First shape (reference).
        shape_b: Second shape (comparison).
        sample_points: Points at which to evaluate continuity.
        g0_tolerance_mm: G0 positional tolerance.
        g1_tolerance_deg: G1 tangent angle tolerance.
        g2_tolerance_pct: G2 curvature relative tolerance.
        max_check_points: Limit on number of points to check.

    Returns:
        ContinuityResult with pass/fail per grade.
    """
    faces_a = _collect_faces(shape_a)
    faces_b = _collect_faces(shape_b)

    if not faces_a or not faces_b:
        logger.warning("One or both shapes have no faces -- skipping continuity check")
        return ContinuityResult()

    # Subsample if needed
    check_points = sample_points
    if len(sample_points) > max_check_points:
        rng = np.random.default_rng(42)
        indices = rng.choice(len(sample_points), max_check_points, replace=False)
        check_points = sample_points[indices]

    result = ContinuityResult()

    for point in check_points:
        match_a = _get_closest_face_and_uv(faces_a, point)
        match_b = _get_closest_face_and_uv(faces_b, point)

        if match_a is None or match_b is None:
            continue

        face_a, u_a, v_a, dist_a = match_a
        face_b, u_b, v_b, dist_b = match_b

        # Skip if projection is too far (point not on either surface)
        if dist_a > g0_tolerance_mm * 10 or dist_b > g0_tolerance_mm * 10:
            continue

        pos_a, norm_a, curv_a = _evaluate_surface_props(face_a, u_a, v_a)
        pos_b, norm_b, curv_b = _evaluate_surface_props(face_b, u_b, v_b)

        if pos_a is None or pos_b is None:
            continue

        result.num_samples_checked += 1

        # G0: positional
        g0_dev = float(np.linalg.norm(pos_a - pos_b))
        result.g0_deviations.append(g0_dev)

        # G1: tangent (normal angle)
        if norm_a is not None and norm_b is not None:
            cos_angle = np.clip(np.dot(norm_a, norm_b), -1.0, 1.0)
            angle_deg = float(np.degrees(np.arccos(abs(cos_angle))))
            result.g1_deviations_deg.append(angle_deg)

        # G2: curvature
        if curv_a is not None and curv_b is not None:
            denom = max(abs(curv_a), abs(curv_b), 1e-12)
            g2_pct = abs(curv_a - curv_b) / denom * 100.0
            result.g2_deviations_pct.append(float(g2_pct))

    # Compute max and pass/fail
    if result.g0_deviations:
        result.g0_max = max(result.g0_deviations)
        result.g0_pass = result.g0_max <= g0_tolerance_mm

    if result.g1_deviations_deg:
        result.g1_max_deg = max(result.g1_deviations_deg)
        result.g1_pass = result.g1_max_deg <= g1_tolerance_deg

    if result.g2_deviations_pct:
        result.g2_max_pct = max(result.g2_deviations_pct)
        result.g2_pass = result.g2_max_pct <= g2_tolerance_pct

    logger.info(
        "Continuity check (%d points): G0 max=%.4f mm (%s), "
        "G1 max=%.2f deg (%s), G2 max=%.1f%% (%s)",
        result.num_samples_checked,
        result.g0_max, "PASS" if result.g0_pass else "FAIL",
        result.g1_max_deg, "PASS" if result.g1_pass else "FAIL",
        result.g2_max_pct, "PASS" if result.g2_pass else "FAIL",
    )
    return result