"""ICP fusion between Multi-HMR SMPL-X meshes and iPhone LiDAR point clouds.

All operations happen in the **webcam camera frame** (meters, OpenCV
convention: +X right, +Y down, +Z forward). LiDAR points must be
pre-transformed via `Extrinsic.T_arkit_to_cam`.
"""
from __future__ import annotations

import logging
from dataclasses import dataclass

import numpy as np

try:
    import open3d as o3d
except ImportError:  # pragma: no cover - exercised via skipif at import sites
    o3d = None  # type: ignore[assignment]

_LOG = logging.getLogger(__name__)

MIN_LIDAR_POINTS = 200
MIN_FITNESS = 0.30
MAX_RMSE_M = 0.05
CROP_MARGIN_M = 0.30


@dataclass
class IcpConfig:
    voxel_size_m: float = 0.02
    max_correspondence_m: float = 0.05
    max_iterations: int = 30


@dataclass
class IcpResult:
    vertices_registered: np.ndarray
    accepted: bool
    fitness: float
    rmse_m: float
    iterations: int


def register_mesh_to_lidar(
    smplx_verts_cam: np.ndarray,
    lidar_points_cam: np.ndarray,
    config: IcpConfig | None = None,
) -> IcpResult:
    """Register SMPL-X verts onto a cropped LiDAR neighborhood."""
    if o3d is None:
        raise RuntimeError("open3d not installed — install with `uv sync --extra lidar`")

    cfg = config or IcpConfig()
    src = np.ascontiguousarray(smplx_verts_cam, dtype=np.float32)

    if not np.isfinite(src).all():
        _LOG.debug("ICP rejected: NaN/Inf in SMPL-X verts")
        return IcpResult(src, False, 0.0, float("inf"), 0)

    lidar = _crop_to_bbox(lidar_points_cam, src, margin_m=CROP_MARGIN_M)
    if lidar.shape[0] < MIN_LIDAR_POINTS or not np.isfinite(lidar).all():
        _LOG.debug("ICP rejected: insufficient LiDAR points (%d)", lidar.shape[0])
        return IcpResult(src, False, 0.0, float("inf"), 0)

    src_pcd = _to_pcd(src, cfg.voxel_size_m, estimate_normals=True)
    tgt_pcd = _to_pcd(lidar, cfg.voxel_size_m, estimate_normals=True)

    if len(src_pcd.points) < 10 or len(tgt_pcd.points) < 10:
        return IcpResult(src, False, 0.0, float("inf"), 0)

    criteria = o3d.pipelines.registration.ICPConvergenceCriteria(
        max_iteration=cfg.max_iterations,
        relative_fitness=1e-6,
        relative_rmse=1e-6,
    )
    # Coarse-to-fine: a wide first pass handles translations larger than the
    # final correspondence threshold, then the strict pass refines and gates.
    coarse = o3d.pipelines.registration.registration_icp(
        src_pcd, tgt_pcd, max(cfg.max_correspondence_m * 5.0, 0.20),
        np.eye(4),
        o3d.pipelines.registration.TransformationEstimationPointToPlane(),
        criteria,
    )
    result = o3d.pipelines.registration.registration_icp(
        src_pcd, tgt_pcd, cfg.max_correspondence_m,
        coarse.transformation,
        o3d.pipelines.registration.TransformationEstimationPointToPlane(),
        criteria,
    )

    accepted = (result.fitness >= MIN_FITNESS) and (result.inlier_rmse <= MAX_RMSE_M)
    if not accepted:
        _LOG.debug("ICP rejected: fitness=%.3f rmse=%.4f", result.fitness, result.inlier_rmse)
        return IcpResult(src, False, float(result.fitness), float(result.inlier_rmse), 0)

    T = np.asarray(result.transformation, dtype=np.float32)
    homog = np.concatenate([src, np.ones((src.shape[0], 1), dtype=np.float32)], axis=1)
    fused = (homog @ T.T)[:, :3]
    if not np.isfinite(fused).all():
        return IcpResult(src, False, float(result.fitness), float(result.inlier_rmse), 0)

    return IcpResult(
        vertices_registered=np.ascontiguousarray(fused, dtype=np.float32),
        accepted=True,
        fitness=float(result.fitness),
        rmse_m=float(result.inlier_rmse),
        iterations=cfg.max_iterations,
    )


def _crop_to_bbox(points: np.ndarray, anchor: np.ndarray, margin_m: float) -> np.ndarray:
    if points.size == 0:
        return points.astype(np.float32, copy=False)
    lo = anchor.min(axis=0) - margin_m
    hi = anchor.max(axis=0) + margin_m
    mask = np.all((points >= lo) & (points <= hi), axis=1)
    return points[mask].astype(np.float32, copy=False)


def _to_pcd(points: np.ndarray, voxel_size_m: float, estimate_normals: bool):
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(points.astype(np.float64, copy=False))
    if voxel_size_m > 0:
        pcd = pcd.voxel_down_sample(voxel_size_m)
    if estimate_normals:
        pcd.estimate_normals(
            search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=voxel_size_m * 2, max_nn=30),
        )
    return pcd


def partition_lidar_by_pid(
    lidar_points_cam: np.ndarray,
    pelvises: dict[int, np.ndarray],
    max_dist_m: float = 1.0,
) -> dict[int, np.ndarray]:
    """Assign each LiDAR point to the closest pelvis within ``max_dist_m``.

    Points beyond ``max_dist_m`` from every pelvis (background, furniture)
    are dropped. Returns ``{pid: (M, 3) float32}`` — pids with zero assigned
    points are omitted.
    """
    if not pelvises or lidar_points_cam.size == 0:
        return {}
    pids = list(pelvises.keys())
    centers = np.stack([pelvises[p] for p in pids]).astype(np.float32)
    pts = np.ascontiguousarray(lidar_points_cam, dtype=np.float32)

    diff = pts[:, None, :] - centers[None, :, :]
    d2 = np.einsum("npk,npk->np", diff, diff)
    nearest = d2.argmin(axis=1)
    nearest_d = np.sqrt(d2[np.arange(d2.shape[0]), nearest])

    mask = nearest_d <= max_dist_m
    out: dict[int, np.ndarray] = {}
    for idx, pid in enumerate(pids):
        sel = mask & (nearest == idx)
        if not sel.any():
            continue
        out[pid] = pts[sel]
    return out


PELVIS_VERT_INDEX = 5559  # SMPL-X canonical pelvis vertex


@dataclass
class FusionMetadata:
    applied: set[int]
    fitness: dict[int, float]
    rmse_m: dict[int, float]
    n_lidar_points_used: int


class FusionWorker:
    """Per-frame ICP fusion orchestrator (caller-driven, no internal thread)."""

    def __init__(self, extrinsic, config: IcpConfig | None = None) -> None:
        self._extrinsic = extrinsic
        self._config = config or IcpConfig()

    def set_extrinsic(self, extrinsic) -> None:
        self._extrinsic = extrinsic

    def run_once(self, state) -> FusionMetadata:
        applied: set[int] = set()
        fitness: dict[int, float] = {}
        rmse: dict[int, float] = {}

        lidar = getattr(state, "lidar_points", None)
        if lidar is None or getattr(lidar, "size", 0) == 0 or not state.persons_smplx:
            return FusionMetadata(applied, fitness, rmse, 0)

        T = np.asarray(self._extrinsic.T_arkit_to_cam, dtype=np.float32)
        homog = np.concatenate([lidar, np.ones((lidar.shape[0], 1), dtype=np.float32)], axis=1)
        lidar_cam = (homog @ T.T)[:, :3]

        pelvises = {
            p.pid: p.vertices_3d[PELVIS_VERT_INDEX]
            for p in state.persons_smplx
            if p.vertices_3d is not None
        }
        parts = partition_lidar_by_pid(lidar_cam, pelvises, max_dist_m=1.0)

        for person in state.persons_smplx:
            pts = parts.get(person.pid)
            if pts is None:
                continue
            result = register_mesh_to_lidar(person.vertices_3d, pts, self._config)
            fitness[person.pid] = result.fitness
            rmse[person.pid] = result.rmse_m
            if result.accepted:
                person.vertices_3d = result.vertices_registered
                applied.add(person.pid)

        return FusionMetadata(applied, fitness, rmse, lidar_cam.shape[0])
