"""
ground_plane.py — RANSAC ground plane fitting for D435i depth images (Issue #611).

No ROS2 dependencies — pure Python + numpy, fully unit-testable.

Algorithm
─────────
RANSAC over 3D point cloud in camera optical frame:
  +X = right, +Y = down, +Z = forward (depth)

1. Sample 3 random non-collinear points.
2. Fit plane: normal n, offset d  (n·p = d, ‖n‖=1).
3. Count inliers: |n·p - d| < inlier_thresh_m.
4. Repeat n_iter times; keep plane with most inliers.
5. Optionally refine: refit normal via SVD on inlier subset.

Normal orientation:
  After RANSAC the normal is flipped so it points in the −Y direction
  (upward in world / toward the camera). This makes "height above ground"
  unambiguous: h = d − n·p > 0 means the point is above the plane.

  Equivalently: normal.y < 0 after orientation.

Height above ground:
  h(p) = d − n·p

  h > 0 → above plane (obstacle candidate)
  h ≈ 0 → on the ground
  h < 0 → below ground (artefact / noise)

Obstacle filter:
  min_obstacle_h < h < max_obstacle_h   (configurable, metres)

  Typical D435i camera mount ~0.5 m above ground:
    min_obstacle_h = 0.05 m  (ignore ground noise)
    max_obstacle_h = 0.80 m  (ignore ceiling / upper structure)
"""

from __future__ import annotations

import math
from typing import Optional, Tuple

import numpy as np


PlaneParams = Tuple[np.ndarray, float]   # (normal_3, d_scalar)


# ── RANSAC ─────────────────────────────────────────────────────────────────────

def fit_ground_plane(
    points:          np.ndarray,
    n_iter:          int   = 50,
    inlier_thresh_m: float = 0.06,
    min_inlier_frac: float = 0.30,
    refine:          bool  = True,
) -> Optional[PlaneParams]:
    """
    RANSAC ground-plane estimation from an (N, 3) point array.

    Parameters
    ----------
    points          : (N, 3) float32/64 point cloud in camera optical frame
    n_iter          : RANSAC iterations
    inlier_thresh_m : inlier distance threshold (m)
    min_inlier_frac : minimum fraction of points that must be inliers
    refine          : if True, refit plane via SVD on best inlier set

    Returns
    -------
    (normal, d) if a plane was found, else None.
    Normal is oriented so that normal.y < 0  (points upward in world).
    """
    n = len(points)
    if n < 10:
        return None

    rng          = np.random.default_rng(seed=42)
    best_n_in    = -1
    best_params: Optional[PlaneParams] = None
    pts          = points.astype(np.float64)

    for _ in range(n_iter):
        idx = rng.integers(0, n, size=3)
        p0, p1, p2 = pts[idx[0]], pts[idx[1]], pts[idx[2]]

        edge1 = p1 - p0
        edge2 = p2 - p0
        normal = np.cross(edge1, edge2)
        norm_len = float(np.linalg.norm(normal))
        if norm_len < 1e-8:
            continue

        normal /= norm_len
        d = float(np.dot(normal, p0))

        dists   = np.abs(pts @ normal - d)
        inliers = dists < inlier_thresh_m
        n_in    = int(inliers.sum())

        if n_in > best_n_in:
            best_n_in = n_in
            best_params = (normal.copy(), d)

    if best_params is None:
        return None
    if best_n_in < int(n * min_inlier_frac):
        return None

    normal, d = best_params

    # ── SVD refinement ──────────────────────────────────────────────────────
    if refine:
        dists   = np.abs(pts @ normal - d)
        in_mask = dists < inlier_thresh_m
        if in_mask.sum() >= 4:
            in_pts = pts[in_mask]
            centroid = in_pts.mean(axis=0)
            _, _, Vt = np.linalg.svd(in_pts - centroid)
            normal = Vt[-1]               # smallest singular value = plane normal
            normal /= float(np.linalg.norm(normal))
            d = float(np.dot(normal, centroid))

    # Orient normal upward: ensure n.y < 0 (−Y = up in camera frame)
    if normal[1] > 0:
        normal = -normal
        d      = -d

    return normal, d


# ── Point classification ───────────────────────────────────────────────────────

def height_above_plane(
    points: np.ndarray,
    plane:  PlaneParams,
) -> np.ndarray:
    """
    Signed height of each point above the ground plane.

    h > 0  → above ground (potential obstacle)
    h ≈ 0  → on ground
    h < 0  → subsurface artefact / sensor noise

    Parameters
    ----------
    points : (N, 3) point cloud
    plane  : (normal, d) from fit_ground_plane()

    Returns
    -------
    heights : (N,) float64
    """
    normal, d = plane
    # h = d - n·p  (positive when point is "above" the upward-facing normal)
    return d - (points.astype(np.float64) @ normal)


def obstacle_mask(
    points:             np.ndarray,
    plane:              PlaneParams,
    min_height_m:       float = 0.05,
    max_height_m:       float = 0.80,
) -> np.ndarray:
    """
    Boolean mask: True where a point is an obstacle (above ground, in height window).

    Parameters
    ----------
    points       : (N, 3) point cloud
    plane        : ground plane params
    min_height_m : minimum height above ground to count as obstacle
    max_height_m : maximum height above ground (ceiling / upper structure cut-off)

    Returns
    -------
    mask : (N,) bool
    """
    h = height_above_plane(points, plane)
    return (h > min_height_m) & (h < max_height_m)


def ground_mask(
    points:         np.ndarray,
    plane:          PlaneParams,
    inlier_thresh:  float = 0.06,
) -> np.ndarray:
    """Boolean mask: True for ground-plane inlier points."""
    normal, d = plane
    dists = np.abs(points.astype(np.float64) @ normal - d)
    return dists < inlier_thresh
