"""
_scan_height_filter.py — Pure-Python height filter for 2D LIDAR scans (no ROS2 deps).

Given pitch/roll angles (from IMU accelerometer) and the LIDAR's mount height
above the floor, marks scan rays whose computed world-frame height falls outside
[z_min, z_max] as invalid (range = inf).

Coordinate convention (ROS REP-103):
  x  = forward
  y  = left
  z  = up
  pitch > 0  = nose down (front lower than rear)
  roll  > 0  = left side up (robot leaning right)
"""

from __future__ import annotations

import math
from typing import List, Tuple


def filter_scan_by_height(
    ranges: List[float],
    angle_min: float,
    angle_increment: float,
    lidar_height_m: float,
    pitch_rad: float,
    roll_rad: float,
    z_min_m: float = 0.05,
    z_max_m: float = 2.0,
    range_min: float = 0.0,
    range_max: float = float('inf'),
) -> List[float]:
    """
    Filter a 2D LaserScan by estimated point height in the world frame.

    For each ray at bearing angle θ and range r, the world-frame height is:

        z = lidar_height + r·cos(θ)·sin(pitch) − r·sin(θ)·sin(roll)

    Rays whose estimated z falls outside [z_min_m, z_max_m] are set to inf
    (Nav2 treats inf as "no obstacle" / max range).

    Parameters
    ----------
    ranges          : raw scan ranges; inf / nan = no-return
    angle_min       : bearing of first ray (rad)
    angle_increment : per-step bearing increment (rad)
    lidar_height_m  : height of the LIDAR above the floor (m)
    pitch_rad       : robot pitch angle (rad)
    roll_rad        : robot roll angle (rad)
    z_min_m         : minimum valid height (m); below → ground plane → filtered
    z_max_m         : maximum valid height (m); above → ceiling / out-of-range → filtered
    range_min       : minimum valid range from LaserScan message
    range_max       : maximum valid range from LaserScan message

    Returns
    -------
    List[float]  filtered ranges — invalid points set to inf
    """
    out: List[float] = []
    sin_pitch = math.sin(pitch_rad)
    sin_roll  = math.sin(roll_rad)

    for i, r in enumerate(ranges):
        # Pass through existing invalid readings unchanged
        if math.isinf(r) or math.isnan(r) or r <= range_min or r >= range_max:
            out.append(r)
            continue

        angle   = angle_min + i * angle_increment
        x_laser = r * math.cos(angle)   # forward component in laser frame
        y_laser = r * math.sin(angle)   # lateral component in laser frame

        # Estimated world-frame height of this scan point
        z_world = lidar_height_m + x_laser * sin_pitch - y_laser * sin_roll

        out.append(float('inf') if z_world < z_min_m or z_world > z_max_m else r)

    return out


def pitch_roll_from_accel(ax: float, ay: float, az: float) -> Tuple[float, float]:
    """
    Estimate pitch and roll from raw accelerometer readings (static / low-dynamics).

    Uses the gravity vector projected onto each tilt axis.  The denominator
    sqrt(ay²+az²) / sqrt(ax²+az²) is always non-negative, so the result is
    independent of the sign convention for the Z axis (works for both
    z-up and z-down IMU frames).

    Parameters
    ----------
    ax, ay, az : accelerometer components (m/s²) in body / IMU frame.
        Convention-agnostic: works for both az ≈ +g and az ≈ −g at rest.

    Returns
    -------
    (pitch_rad, roll_rad)
        pitch : rotation around Y axis; positive when nose is lower than tail
        roll  : rotation around X axis; positive when left side is higher
    """
    pitch = math.atan2(-ax, math.sqrt(ay * ay + az * az))
    roll  = math.atan2( ay, math.sqrt(ax * ax + az * az))
    return pitch, roll


class AttitudeEstimator:
    """
    Single-pole low-pass filter applied to pitch and roll derived from
    accelerometer readings.  Smooths out vibration noise while remaining
    responsive to genuine tilt changes.

    Parameters
    ----------
    alpha : float (0 < alpha < 1)
        Smoothing factor.  Higher = faster response, less filtering.
        alpha = 1 − exp(−dt / τ) for desired time constant τ.
    """

    def __init__(self, alpha: float = 0.1):
        self._alpha = alpha
        self._pitch = 0.0
        self._roll  = 0.0
        self._initialised = False

    def update(self, ax: float, ay: float, az: float) -> Tuple[float, float]:
        """Update filter with new accelerometer sample; return (pitch, roll)."""
        pitch_raw, roll_raw = pitch_roll_from_accel(ax, ay, az)
        if not self._initialised:
            self._pitch = pitch_raw
            self._roll  = roll_raw
            self._initialised = True
        else:
            self._pitch += self._alpha * (pitch_raw - self._pitch)
            self._roll  += self._alpha * (roll_raw  - self._roll)
        return self._pitch, self._roll

    @property
    def pitch(self) -> float:
        return self._pitch

    @property
    def roll(self) -> float:
        return self._roll
