"""Rigid registration from scan point cloud to CAD reference.

Two-stage approach:

1. **Coarse**: centroid alignment + PCA axis matching with sign
   disambiguation via per-axis skew (third central moment). Resolves
   global orientation flips without requiring a good initial guess.
2. **Fine**: trimmed iterative closest point (ICP) on a downsampled CAD
   target using ``scipy.spatial.cKDTree`` for nearest-neighbour
   queries. Reflection-free rotation via SVD with the standard
   determinant guard.

The result transforms scan points into the CAD coordinate frame:
``cad_point ≈ R @ scan_point + t``.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np
from scipy.spatial import cKDTree

from .preprocess import voxel_downsample


@dataclass(slots=True)
class RegistrationResult:
    R: np.ndarray
    t: np.ndarray
    rms_mm: float
    iterations: int
    inlier_ratio: float
    converged: bool
    coarse_R: np.ndarray
    coarse_t: np.ndarray
    initial_rms_mm: float


def register_to_cad(
    scan: np.ndarray,
    cad: np.ndarray,
    config: dict,
) -> RegistrationResult:
    """Rigidly align ``scan`` to ``cad`` and return the transformation.

    ``config`` is a fully resolved config dict (must contain the
    ``registration`` section with ``icp_max_corr_mm`` already filled by
    :func:`underframe_deflection.config.resolve_numeric_parameters`).
    """
    if len(scan) < 3 or len(cad) < 3:
        raise ValueError("register_to_cad requires at least 3 points in scan and cad.")

    reg_cfg = config["registration"]
    cad_cfg = config.get("cad", {})

    # CAD-side downsampling — keeps NN queries cheap when the CAD has
    # millions of points.
    cad_voxel = float(cad_cfg.get("downsample_voxel_mm", 0.0) or 0.0)
    target_n = int(reg_cfg.get("icp_downsample_target", 30000))
    cad_sampled = cad
    if cad_voxel > 0.0:
        cad_sampled = voxel_downsample(cad_sampled, cad_voxel)
    if target_n > 0 and len(cad_sampled) > target_n:
        idx = np.linspace(0, len(cad_sampled) - 1, target_n, dtype=int)
        cad_sampled = cad_sampled[idx]

    max_corr = float(reg_cfg.get("icp_max_corr_mm", 25.0))
    max_iter = int(reg_cfg.get("icp_max_iter", 60))
    tol_mm = float(reg_cfg.get("icp_tol_mm", 0.05))
    tol_rel = float(reg_cfg.get("icp_tol_rel", 1.0e-4))
    trim = float(reg_cfg.get("trim_ratio", 0.2))

    tree = cKDTree(cad_sampled)

    R0, t0 = _coarse_align_pca(
        scan,
        cad_sampled,
        disambiguate_sign=bool(reg_cfg.get("pca_sign_disambiguation", True)),
    )
    scan_coarse = scan @ R0.T + t0

    # Sanity guard: PCA-based coarse alignment can flip an axis when the scan
    # is already close to the CAD frame (near-symmetric inertia). If the
    # coarse step does not actually improve correspondence over the identity
    # initialisation, fall back to identity so a bad flip cannot wreck ICP.
    guard_max_corr = max(max_corr * 4.0, 0.05 * float(np.ptp(cad_sampled[:, 0]) + 1.0))
    rms_identity = _correspondence_rms(scan, tree, guard_max_corr)
    rms_coarse = _correspondence_rms(scan_coarse, tree, guard_max_corr)
    if rms_identity <= rms_coarse:
        R0 = np.eye(3, dtype=np.float64)
        t0 = np.zeros(3, dtype=np.float64)
        scan_coarse = scan

    # ICP source downsampling — keeps query time bounded.
    src_target = target_n if target_n > 0 else len(scan_coarse)
    if len(scan_coarse) > src_target:
        idx = np.linspace(0, len(scan_coarse) - 1, src_target, dtype=int)
        scan_for_icp = scan_coarse[idx]
    else:
        scan_for_icp = scan_coarse
    R = np.eye(3, dtype=np.float64)
    t = np.zeros(3, dtype=np.float64)
    prev_rms = float("inf")
    current = scan_for_icp.copy()
    iterations_run = 0
    inlier_ratio = 0.0
    rms = float("inf")
    converged = False

    initial_rms = _correspondence_rms(current, tree, max_corr)

    for iteration in range(1, max_iter + 1):
        iterations_run = iteration
        distances, neighbours = tree.query(current, k=1)
        keep = distances <= max_corr
        if trim > 0.0 and np.any(keep):
            kept_d = distances[keep]
            if len(kept_d) > 8:
                cutoff = float(np.quantile(kept_d, 1.0 - trim))
                keep = keep & (distances <= cutoff)
        if np.count_nonzero(keep) < 3:
            break
        src = current[keep]
        dst = cad_sampled[neighbours[keep]]
        delta_R, delta_t = _kabsch(src, dst)
        current = current @ delta_R.T + delta_t
        R = delta_R @ R
        t = delta_R @ t + delta_t

        residuals = current - cad_sampled[neighbours]
        residual_norms = np.linalg.norm(residuals[keep], axis=1)
        rms = float(np.sqrt(np.mean(residual_norms ** 2))) if len(residual_norms) else float("inf")
        inlier_ratio = float(np.count_nonzero(keep) / max(len(current), 1))

        if abs(prev_rms - rms) <= tol_mm or (
            prev_rms > 0.0 and abs(prev_rms - rms) / max(prev_rms, 1e-9) <= tol_rel
        ):
            converged = True
            break
        prev_rms = rms

    R_total = R @ R0
    t_total = R @ t0 + t

    # Near-identity snap: when the scan is already aligned to the CAD (e.g.
    # scan == cad, or a pre-registered scan), ICP can still drift by a few mm
    # of translation / a fraction of a degree because it minimises mean
    # nearest-neighbour distance on a *sparse* cloud (whose floor is ~spacing/2,
    # not 0). That residual wobble shifts a thin ROI band enough to crop the
    # beams. If the recovered transform is within tight identity tolerances and
    # snapping does not worsen correspondence, return exact identity so an
    # already-aligned input is left untouched.
    identity_rot_tol_rad = float(reg_cfg.get("identity_rot_tol_deg", 0.5)) * np.pi / 180.0
    identity_trans_tol_mm = float(reg_cfg.get("identity_trans_tol_mm", 5.0))
    rot_angle = float(np.arccos(np.clip((np.trace(R_total) - 1.0) / 2.0, -1.0, 1.0)))
    trans_norm = float(np.linalg.norm(t_total))
    if rot_angle <= identity_rot_tol_rad and trans_norm <= identity_trans_tol_mm:
        rms_aligned = _correspondence_rms(scan, tree, max_corr)
        rms_transformed = _correspondence_rms(scan @ R_total.T + t_total, tree, max_corr)
        if rms_aligned <= rms_transformed + tol_mm:
            R_total = np.eye(3, dtype=np.float64)
            t_total = np.zeros(3, dtype=np.float64)
            rms = rms_aligned

    return RegistrationResult(
        R=R_total,
        t=t_total,
        rms_mm=rms,
        iterations=iterations_run,
        inlier_ratio=inlier_ratio,
        converged=converged,
        coarse_R=R0,
        coarse_t=t0,
        initial_rms_mm=initial_rms,
    )


def apply_transform(points: np.ndarray, R: np.ndarray, t: np.ndarray) -> np.ndarray:
    """Apply ``R @ p + t`` to a point cloud (broadcasted)."""
    return points @ R.T + t


# ---------------------------------------------------------------------------
# Internals
# ---------------------------------------------------------------------------


def _coarse_align_pca(
    scan: np.ndarray,
    cad: np.ndarray,
    *,
    disambiguate_sign: bool,
) -> tuple[np.ndarray, np.ndarray]:
    """Return (R, t) that map ``scan`` into the CAD PCA frame."""
    scan_axes, scan_center = _pca_basis(scan)
    cad_axes, cad_center = _pca_basis(cad)

    if disambiguate_sign:
        scan_axes = _resolve_sign_by_skew(scan, scan_center, scan_axes, cad, cad_center, cad_axes)

    # R maps scan → CAD: scan_axes^T rotates scan into its own PCA basis,
    # cad_axes rotates back into the CAD frame.
    R = cad_axes @ scan_axes.T
    if np.linalg.det(R) < 0:
        scan_axes[:, -1] *= -1.0
        R = cad_axes @ scan_axes.T
    t = cad_center - R @ scan_center
    return R, t


def _pca_basis(points: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    center = points.mean(axis=0)
    centered = points - center
    covariance = centered.T @ centered / max(len(points), 1)
    eigenvalues, eigenvectors = np.linalg.eigh(covariance)
    order = np.argsort(eigenvalues)[::-1]
    axes = eigenvectors[:, order]
    # ``axes`` is (3, 3) where columns are unit eigenvectors sorted by
    # decreasing variance. Treat columns as the new basis directions.
    return axes, center


def _resolve_sign_by_skew(
    scan: np.ndarray,
    scan_center: np.ndarray,
    scan_axes: np.ndarray,
    cad: np.ndarray,
    cad_center: np.ndarray,
    cad_axes: np.ndarray,
) -> np.ndarray:
    """Pick eigenvector signs that match the chirality of ``cad``."""
    scan_proj = (scan - scan_center) @ scan_axes
    cad_proj = (cad - cad_center) @ cad_axes
    cad_skew = np.array([_third_moment(cad_proj[:, i]) for i in range(3)])
    best_signs = np.array([1.0, 1.0, 1.0])
    best_score = np.inf
    for s0 in (-1.0, 1.0):
        for s1 in (-1.0, 1.0):
            for s2 in (-1.0, 1.0):
                signs = np.array([s0, s1, s2])
                if np.prod(signs) < 0:
                    # Disallow reflections; the kabsch det guard will
                    # cover the residual cases.
                    continue
                scan_skew = np.array(
                    [signs[i] ** 3 * _third_moment(scan_proj[:, i]) for i in range(3)]
                )
                score = float(np.sum(np.abs(scan_skew - cad_skew)))
                if score < best_score:
                    best_score = score
                    best_signs = signs
    return scan_axes * best_signs[None, :]


def _third_moment(values: np.ndarray) -> float:
    if len(values) == 0:
        return 0.0
    std = float(np.std(values))
    if std < 1e-9:
        return 0.0
    centered = values - values.mean()
    return float(np.mean(centered ** 3) / (std ** 3))


def _kabsch(src: np.ndarray, dst: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """Closed-form rigid-body fit (no reflections)."""
    src_center = src.mean(axis=0)
    dst_center = dst.mean(axis=0)
    src_centered = src - src_center
    dst_centered = dst - dst_center
    H = src_centered.T @ dst_centered
    U, _, Vt = np.linalg.svd(H)
    D = np.eye(3, dtype=np.float64)
    if np.linalg.det(Vt.T @ U.T) < 0:
        D[-1, -1] = -1.0
    R = Vt.T @ D @ U.T
    t = dst_center - R @ src_center
    return R, t


def _correspondence_rms(points: np.ndarray, tree: cKDTree, max_corr: float) -> float:
    if len(points) == 0:
        return float("inf")
    distances, _ = tree.query(points, k=1)
    keep = distances <= max_corr
    if not np.any(keep):
        return float(np.sqrt(np.mean(distances ** 2)))
    return float(np.sqrt(np.mean(distances[keep] ** 2)))


__all__ = ["RegistrationResult", "register_to_cad", "apply_transform"]