massive-cloud/alg/delivery_v1.2.1/underframe_deflection/fitting.py

from __future__ import annotations

from dataclasses import asdict

import numpy as np
from scipy.interpolate import UnivariateSpline

from .structures import CandidateBand, FittedLine, RepresentativePoint, TrackSolution


def _robust_center_2d(points_yz: np.ndarray, max_iter: int) -> tuple[np.ndarray, float]:
    center = np.median(points_yz, axis=0)
    residual = 0.0
    for _ in range(max_iter):
        deltas = points_yz - center
        scales = np.median(np.abs(deltas), axis=0) + 1e-6
        radius = np.sqrt(np.sum((deltas / scales) ** 2, axis=1))
        c = 4.685
        weights = np.square(1.0 - np.square(np.clip(radius / c, 0.0, 1.0)))
        weights[radius >= c] = 0.0
        if not np.any(weights > 0):
            break
        center = np.average(points_yz, axis=0, weights=weights)
        residual = float(np.sqrt(np.average(np.sum(deltas**2, axis=1), weights=weights)))
    return center, residual


def _extract_underside_core_points(
    region_points: np.ndarray,
    candidate: CandidateBand,
    config: dict,
) -> np.ndarray:
    rep_cfg = config["representative"]
    spacing = config["runtime"]["spacing_mm"]
    profile_margin = max(
        rep_cfg["underside_profile_margin_spacing_multiplier"] * spacing,
        0.2 * rep_cfg["band_y_margin_mm"],
    )
    y_lo = max(float(region_points[:, 1].min()), candidate.y_min - profile_margin)
    y_hi = min(float(region_points[:, 1].max()), candidate.y_max + profile_margin)
    if not np.isfinite(y_lo) or not np.isfinite(y_hi) or (y_hi - y_lo) <= 1e-6:
        return np.zeros((0, 3), dtype=np.float64)

    bin_count = max(int(rep_cfg["underside_profile_bin_count"]), 8)
    y_edges = np.linspace(y_lo, y_hi, bin_count + 1)
    valid_bins: list[dict[str, float | int]] = []
    min_bin_points = int(rep_cfg["underside_min_bin_points"])
    underside_quantile = float(rep_cfg["underside_quantile"])

    for index in range(bin_count):
        bin_mask = (region_points[:, 1] >= y_edges[index]) & (region_points[:, 1] < y_edges[index + 1])
        bin_points = region_points[bin_mask]
        if len(bin_points) < min_bin_points:
            continue
        z_floor = float(np.quantile(bin_points[:, 2], underside_quantile))
        valid_bins.append(
            {
                "index": index,
                "y_min": float(y_edges[index]),
                "y_max": float(y_edges[index + 1]),
                "y_center": float(0.5 * (y_edges[index] + y_edges[index + 1])),
                "z_floor": z_floor,
                "count": int(len(bin_points)),
            }
        )

    if len(valid_bins) < 3:
        return np.zeros((0, 3), dtype=np.float64)

    outer_bin_count = max(int(np.ceil(len(valid_bins) * rep_cfg["underside_outer_fraction"])), 2)
    if candidate.side == "left":
        outer_bins = valid_bins[:outer_bin_count]
    else:
        outer_bins = valid_bins[-outer_bin_count:]
    anchor_bin = min(outer_bins, key=lambda item: item["z_floor"])
    anchor_index = int(anchor_bin["index"])
    floor_ref = float(anchor_bin["z_floor"])
    rise_limit = floor_ref + max(
        rep_cfg["underside_rise_spacing_multiplier"] * spacing,
        35.0,
    )

    valid_by_index = {int(item["index"]): item for item in valid_bins}
    selected_bins = [anchor_bin]
    expand_index = anchor_index - 1
    while expand_index in valid_by_index and float(valid_by_index[expand_index]["z_floor"]) <= rise_limit:
        selected_bins.append(valid_by_index[expand_index])
        expand_index -= 1
    expand_index = anchor_index + 1
    while expand_index in valid_by_index and float(valid_by_index[expand_index]["z_floor"]) <= rise_limit:
        selected_bins.append(valid_by_index[expand_index])
        expand_index += 1
    if not selected_bins:
        return np.zeros((0, 3), dtype=np.float64)

    y_min = min(float(item["y_min"]) for item in selected_bins)
    y_max = max(float(item["y_max"]) for item in selected_bins)
    core_ceiling = floor_ref + max(
        rep_cfg["underside_core_spacing_multiplier"] * spacing,
        45.0,
    )
    core_mask = (
        (region_points[:, 1] >= y_min)
        & (region_points[:, 1] <= y_max)
        & (region_points[:, 2] <= core_ceiling)
    )
    core_points = region_points[core_mask]
    if len(core_points) < 6:
        return np.zeros((0, 3), dtype=np.float64)
    return core_points


def extract_representative_points(
    points_vehicle: np.ndarray,
    track: TrackSolution,
    config: dict,
) -> list[RepresentativePoint]:
    rep_cfg = config["representative"]
    representatives: list[RepresentativePoint] = []
    for candidate in track.candidates:
        region_mask = (
            (points_vehicle[:, 0] >= candidate.x_start)
            & (points_vehicle[:, 0] < candidate.x_end)
            & (points_vehicle[:, 1] >= candidate.y_min - rep_cfg["band_y_margin_mm"])
            & (points_vehicle[:, 1] <= candidate.y_max + rep_cfg["band_y_margin_mm"])
            & (
                points_vehicle[:, 2]
                <= candidate.z_low + rep_cfg["band_z_margin_mm"]
            )
        )
        region_points = points_vehicle[region_mask]
        if len(region_points) < max(8, candidate.low_point_count):
            continue
        core_points = _extract_underside_core_points(region_points, candidate, config)
        if len(core_points) < 6:
            z_quantile = float(np.quantile(region_points[:, 2], rep_cfg["core_quantile"]))
            core_ceiling = min(z_quantile, candidate.z_low + rep_cfg["core_margin_mm"])
            core_points = region_points[region_points[:, 2] <= core_ceiling]
        if len(core_points) < 6:
            core_points = region_points[np.argsort(region_points[:, 2])[: max(6, len(region_points) // 4)]]
        center_yz, residual = _robust_center_2d(core_points[:, 1:3], rep_cfg["max_iter"])
        confidence = float(
            np.clip(
                0.35
                + 0.3 * min(len(core_points) / max(candidate.low_point_count, 1), 1.5)
                + 0.35 * max(0.0, 1.0 - residual / max(candidate.width, 1.0)),
                0.0,
                1.0,
            )
        )
        representatives.append(
            RepresentativePoint(
                side=track.side,
                slice_index=candidate.slice_index,
                x=float(np.median(core_points[:, 0])),
                y=float(center_yz[0]),
                z=float(center_yz[1]),
                width=float(candidate.width),
                point_count=int(len(core_points)),
                residual=residual,
                confidence=confidence,
                x_start=candidate.x_start,
                x_end=candidate.x_end,
            )
        )
    representatives.sort(key=lambda item: item.x)
    return representatives


def _fit_spline(
    x: np.ndarray,
    values: np.ndarray,
    smooth_scale: float,
    iterations: int = 4,
) -> tuple[UnivariateSpline, float]:
    weights = np.ones_like(values)
    spline = None
    residual_rms = 0.0
    smoothing = max(len(x) * smooth_scale**2, 1e-6)
    for _ in range(iterations):
        spline = UnivariateSpline(x, values, w=weights, s=smoothing)
        fitted = spline(x)
        residual = values - fitted
        scale = np.median(np.abs(residual)) * 1.4826 + 1e-6
        ratio = np.abs(residual) / (4.685 * scale)
        weights = np.square(1.0 - np.square(np.clip(ratio, 0.0, 1.0)))
        weights[ratio >= 1.0] = 0.05
        residual_rms = float(np.sqrt(np.mean(residual**2)))
    assert spline is not None
    return spline, residual_rms


def _reject_contaminated_representatives(
    x: np.ndarray,
    z: np.ndarray,
    window_mm: float,
    sigma: float,
) -> np.ndarray:
    """Flag representative points that deviate from a robust local baseline.

    When the cloud is sampled directly from CAD, gussets and cross-members
    intermittently raise the detected beam underside by tens of mm at
    isolated X positions. A sliding-window median is robust to such localized
    spikes (up to ~50% local contamination) while passing smooth, genuine
    deflection through almost unchanged. Points whose height deviates from the
    median baseline by more than ``sigma`` robust standard deviations are
    rejected before the spline fit.
    """
    n = len(x)
    keep = np.ones(n, dtype=bool)
    if n < 6:
        return keep
    baseline = np.empty(n, dtype=np.float64)
    for index in range(n):
        local = z[(x >= x[index] - window_mm) & (x <= x[index] + window_mm)]
        baseline[index] = float(np.median(local))
    residual = z - baseline
    mad = float(np.median(np.abs(residual - np.median(residual))))
    scale = 1.4826 * mad + 1e-6
    keep = np.abs(residual) <= sigma * scale
    if np.count_nonzero(keep) < 4:
        return np.ones(n, dtype=bool)
    return keep


def fit_representative_line(
    representatives: list[RepresentativePoint],
    config: dict,
    side: str,
) -> FittedLine | None:
    if len(representatives) < 4:
        return None
    x = np.asarray([item.x for item in representatives], dtype=np.float64)
    y = np.asarray([item.y for item in representatives], dtype=np.float64)
    z = np.asarray([item.z for item in representatives], dtype=np.float64)
    order = np.argsort(x)
    x = x[order]
    y = y[order]
    z = z[order]
    reps_sorted = [representatives[index] for index in order]
    if np.any(np.diff(x) <= 0):
        unique_mask = np.ones_like(x, dtype=bool)
        unique_mask[1:] = np.diff(x) > 1e-6
        x = x[unique_mask]
        y = y[unique_mask]
        z = z[unique_mask]
        reps_sorted = [item for item, keep in zip(reps_sorted, unique_mask) if keep]
    if len(x) < 4:
        return None

    rep_cfg = config["representative"]
    spacing = config["runtime"]["spacing_mm"]

    if rep_cfg.get("outlier_rejection_enabled", True) and len(x) >= 6:
        window_mm = rep_cfg.get("outlier_window_slices", 2.5) * config["beam"]["slice_step_mm"]
        sigma = rep_cfg.get("outlier_sigma", 2.5)
        keep = _reject_contaminated_representatives(x, z, window_mm, sigma)
        if np.count_nonzero(keep) >= 4:
            x = x[keep]
            y = y[keep]
            z = z[keep]
            reps_sorted = [item for item, flag in zip(reps_sorted, keep) if flag]
    y_spline, y_residual = _fit_spline(
        x,
        y,
        smooth_scale=rep_cfg["spline_smoothing_y_spacing_multiplier"] * spacing,
    )
    z_spline, z_residual = _fit_spline(
        x,
        z,
        smooth_scale=rep_cfg["spline_smoothing_z_spacing_multiplier"] * spacing,
    )
    max_gap = rep_cfg["max_gap_slices"] * config["beam"]["slice_step_mm"]
    gaps: list[tuple[float, float]] = []
    for x0, x1 in zip(x[:-1], x[1:]):
        if (x1 - x0) > max_gap:
            gaps.append((float(x0), float(x1)))
    low_confidence_segments = list(gaps)
    return FittedLine(
        side=side,
        x_support=x,
        y_support=y,
        z_support=z,
        y_spline=y_spline,
        z_spline=z_spline,
        y_residual_rms=y_residual,
        z_residual_rms=z_residual,
        low_confidence_segments=low_confidence_segments,
        gaps=gaps,
    )


def evaluate_line(line: FittedLine, x_grid: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    y = line.y_spline(x_grid)
    z = line.z_spline(x_grid)
    confidence = np.ones_like(x_grid, dtype=np.float64)
    support_min = float(line.x_support.min())
    support_max = float(line.x_support.max())
    outside = (x_grid < support_min) | (x_grid > support_max)
    confidence[outside] = 0.0
    for start, end in line.low_confidence_segments:
        mask = (x_grid >= start) & (x_grid <= end)
        confidence[mask] = np.minimum(confidence[mask], 0.25)
    return y, z, confidence


def representatives_to_rows(representatives: list[RepresentativePoint]) -> list[dict]:
    return [asdict(item) for item in representatives]