"""Coarse LiDAR deskew using a constant body rate over the sweep."""

from __future__ import annotations

import numpy as np

from .contracts import ImuSeries, LidarFrame
from .geometry import so3_exp
from .time_offset import lidar_time_to_imu_time


def deskew_lidar_frames(
    frames: list[LidarFrame],
    imu: ImuSeries,
    *,
    delta_t_s: float,
    R_IMU_lidar: np.ndarray | None = None,
    gyro_bias_rad_s: np.ndarray | None = None,
) -> list[LidarFrame]:
    """Return deskewed copies when extrinsic is known; otherwise return originals."""

    if R_IMU_lidar is None:
        return frames

    bias = np.zeros(3) if gyro_bias_rad_s is None else np.asarray(gyro_bias_rad_s, dtype=float)
    r_li = np.asarray(R_IMU_lidar, dtype=float).reshape(3, 3).T
    output: list[LidarFrame] = []

    for frame in frames:
        n = frame.points_xyz.shape[0]
        if n < 10:
            output.append(frame)
            continue
        t_mid_imu = lidar_time_to_imu_time(frame.t_mid_s, delta_t_s)
        index = int(np.clip(np.searchsorted(imu.t_s, t_mid_imu), 1, imu.t_s.size - 1))
        omega_lidar = r_li @ (imu.gyro_rad_s[index] - bias)
        duration = max(frame.t_end_s - frame.t_start_s, 1e-3)
        rel = np.linspace(-0.5, 0.5, n) * duration
        deskewed = np.empty_like(frame.points_xyz)
        # Piecewise-constant rotation over a few time bins.
        bins = 12
        edges = np.linspace(-0.5 * duration, 0.5 * duration, bins + 1)
        for b in range(bins):
            mask = (rel >= edges[b]) & (rel <= edges[b + 1] if b == bins - 1 else rel < edges[b + 1])
            if not np.any(mask):
                continue
            tau = 0.5 * (edges[b] + edges[b + 1])
            rot = so3_exp(omega_lidar * float(tau))
            deskewed[mask] = frame.points_xyz[mask] @ rot.T
        output.append(
            LidarFrame(
                frame_id=frame.frame_id,
                t_start_s=frame.t_start_s,
                t_end_s=frame.t_end_s,
                points_xyz=deskewed,
                path=frame.path,
            )
        )
    return output