"""Point cloud utilities: crop to bbox in image space + outlier filtering.

Nuwa /api/pointcloud returns points in camera frame; we need to project them
back to pixel space to keep only those inside a YOLO bbox.

Frame note:
  Nuwa publishes points in X-right / Y-up / Z-back (Z<0 in front of camera).
  `cv2.solvePnP` (used by marker-finder to produce T_cam_marker, hence
  T_base_camera after hand-eye) uses X-right / Y-down / Z-forward (Z>0 in front).
  Before applying T_base_camera we flip Y and Z via `nuwa_to_cv_frame`.
"""

from __future__ import annotations

import numpy as np


def nuwa_to_cv_frame(points_xyz: np.ndarray) -> np.ndarray:
    """Flip Y and Z to convert Nuwa camera points to OpenCV convention."""
    pts = np.asarray(points_xyz, dtype=np.float64).copy()
    pts[:, 1] *= -1.0
    pts[:, 2] *= -1.0
    return pts


def project(points_cam: np.ndarray, K: np.ndarray) -> np.ndarray:
    """Project Nx3 camera-frame points to Nx2 pixels via pinhole."""
    pts = np.asarray(points_cam, dtype=np.float64)
    z = pts[:, 2]
    z = np.where(z == 0, 1e-9, z)
    u = K[0, 0] * pts[:, 0] / z + K[0, 2]
    v = K[1, 1] * pts[:, 1] / z + K[1, 2]
    return np.stack([u, v], axis=1)


def crop_bbox(points_cam: np.ndarray, bbox_xyxy: tuple[float, float, float, float],
              K: np.ndarray) -> np.ndarray:
    """Keep only 3D points whose projection falls inside bbox (in pixel coords).

    Works regardless of depth-sign convention (OpenCV: Z>0 in front;
    Nuwa/OpenGL-like: Z<0 in front) — the server-provided point cloud already
    obeys u = fx·x/z + cx with its own convention.
    """
    if points_cam.size == 0:
        return points_cam[:0, :]
    xyz = points_cam[:, :3]
    uv = project(xyz, K)
    x1, y1, x2, y2 = bbox_xyxy
    mask = (
        (uv[:, 0] >= x1) & (uv[:, 0] <= x2) &
        (uv[:, 1] >= y1) & (uv[:, 1] <= y2) &
        (np.abs(xyz[:, 2]) > 1e-6)
    )
    return points_cam[mask]


def filter_outliers(points: np.ndarray, n_std: float = 2.0) -> np.ndarray:
    """Remove points whose distance from centroid exceeds n_std * std."""
    if len(points) < 10:
        return points
    xyz = points[:, :3]
    c = xyz.mean(axis=0)
    d = np.linalg.norm(xyz - c, axis=1)
    keep = d < (d.mean() + n_std * d.std())
    return points[keep]