"""Constant IMU-to-LiDAR clock-offset estimation via angular-rate correlation."""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np
from scipy import signal

from .contracts import ImuSeries, LidarFrame
from .geometry import rotation_angle_deg, so3_log
from .registration import estimate_frame_rotations


@dataclass(frozen=True)
class TimeOffsetResult:
    delta_t_s: float
    correlation_peak: float
    search_s: float
    notes: tuple[str, ...] = ()
    ok: bool = True


def _magnitude_series(times: np.ndarray, values: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    mag = np.linalg.norm(values, axis=1) if values.ndim == 2 else np.asarray(values, dtype=float)
    return np.asarray(times, dtype=float), np.asarray(mag, dtype=float)


def _correlate_offset(
    imu_t: np.ndarray,
    imu_mag: np.ndarray,
    lidar_t: np.ndarray,
    lidar_mag: np.ndarray,
    *,
    search_s: float,
    sample_hz: float,
) -> tuple[float, float]:
    """Return ``(delta_t, peak)`` for ``t_imu = t_lidar + delta_t``.

    Implementation: resample both on LiDAR-relative grid, shift IMU by candidate
    offsets, maximize normalized correlation. This avoids ambiguous lag signs.
    """

    t_start = float(lidar_t[0])
    t_end = float(lidar_t[-1])
    if t_end - t_start < 0.5:
        return 0.0, 0.0
    dt = 1.0 / sample_hz
    grid = np.arange(t_start, t_end, dt)
    lidar_sig = np.interp(grid, lidar_t, lidar_mag, left=0.0, right=0.0)
    lidar_sig = lidar_sig - np.mean(lidar_sig)
    lidar_norm = float(np.linalg.norm(lidar_sig)) + 1e-12

    best_delta = 0.0
    best_peak = -1.0
    for delta in np.arange(-search_s, search_s + 1e-12, dt):
        imu_sig = np.interp(grid + delta, imu_t, imu_mag, left=0.0, right=0.0)
        imu_sig = imu_sig - np.mean(imu_sig)
        denom = lidar_norm * (float(np.linalg.norm(imu_sig)) + 1e-12)
        peak = float(np.dot(imu_sig, lidar_sig) / denom)
        if peak > best_peak:
            best_peak = peak
            best_delta = float(delta)

    # Local parabolic refinement.
    deltas = np.array([best_delta - dt, best_delta, best_delta + dt], dtype=float)
    peaks = []
    for delta in deltas:
        imu_sig = np.interp(grid + delta, imu_t, imu_mag, left=0.0, right=0.0)
        imu_sig = imu_sig - np.mean(imu_sig)
        denom = lidar_norm * (float(np.linalg.norm(imu_sig)) + 1e-12)
        peaks.append(float(np.dot(imu_sig, lidar_sig) / denom))
    y0, y1, y2 = peaks
    denom = y0 - 2 * y1 + y2
    if abs(denom) > 1e-12:
        refined = float(best_delta + 0.5 * (y0 - y2) / denom * dt)
        # Parabola can jump outside the searched window; keep it clamped.
        if abs(refined) <= search_s + dt:
            best_delta = refined
        best_peak = float(y1)
    return best_delta, best_peak


def estimate_time_offset(
    imu: ImuSeries,
    frames: list[LidarFrame],
    *,
    gyro_bias_rad_s: np.ndarray | None = None,
    search_s: float = 1.0,
    sample_hz: float = 50.0,
) -> TimeOffsetResult:
    """Estimate ``t_imu = t_lidar + delta_t``.

    Positive ``delta_t`` means the IMU clock reading is ahead of the LiDAR clock
    for the same physical instant (IMU timestamps are larger).
    """

    notes: list[str] = []
    if len(frames) < 5:
        return TimeOffsetResult(0.0, 0.0, search_s, ("not enough LiDAR frames",), False)

    bias = np.zeros(3) if gyro_bias_rad_s is None else np.asarray(gyro_bias_rad_s, dtype=float)
    gyro = imu.gyro_rad_s - bias

    # Use short consecutive (or near-consecutive) pairs. A large stride (e.g.
    # len//20) averages over many seconds and destroys |ω| correlation even when
    # host/device clocks are already aligned.
    stride = 1 if len(frames) < 80 else 2
    rotations, pair_times = estimate_frame_rotations(frames, stride=stride)
    if len(rotations) < 8:
        rotations, pair_times = estimate_frame_rotations(frames, stride=1)
    if len(rotations) < 4:
        return TimeOffsetResult(0.0, 0.0, search_s, ("not enough LiDAR relative rotations",), False)

    lidar_t = []
    lidar_w = []
    for (t_a, t_b), rotation in zip(pair_times, rotations):
        dt_pair = max(t_b - t_a, 1e-3)
        omega = so3_log(rotation) / dt_pair
        lidar_t.append(0.5 * (t_a + t_b))
        lidar_w.append(omega)
    lidar_t_arr = np.asarray(lidar_t, dtype=float)
    lidar_w_arr = np.asarray(lidar_w, dtype=float)

    imu_t, imu_mag = _magnitude_series(imu.t_s, gyro)
    lidar_t_mag, lidar_mag = _magnitude_series(lidar_t_arr, lidar_w_arr)
    delta, peak = _correlate_offset(
        imu_t,
        imu_mag,
        lidar_t_mag,
        lidar_mag,
        search_s=search_s,
        sample_hz=sample_hz,
    )

    notes.append(
        f"LiDAR mean pair rotation {np.mean([rotation_angle_deg(r) for r in rotations]):.2f} deg"
    )
    notes.append(f"searched delta_t in ±{search_s:.3f}s by direct correlation")
    ok = peak > 0.15
    if not ok:
        notes.append("correlation peak is weak; check overlapping motion and axis units")
    return TimeOffsetResult(
        delta_t_s=delta,
        correlation_peak=peak,
        search_s=search_s,
        notes=tuple(notes),
        ok=ok,
    )


def lidar_time_to_imu_time(t_lidar_s: float, delta_t_s: float) -> float:
    """Convert a LiDAR timestamp to the IMU clock using ``t_imu = t_lidar + delta_t``."""

    return float(t_lidar_s + delta_t_s)


def _lidar_omega_series(
    frames: list[LidarFrame],
    *,
    stride: int,
) -> tuple[np.ndarray, np.ndarray]:
    rotations, pair_times = estimate_frame_rotations(frames, stride=stride)
    if len(rotations) < 4:
        rotations, pair_times = estimate_frame_rotations(frames, stride=1)
    lidar_t: list[float] = []
    lidar_w: list[np.ndarray] = []
    for (t_a, t_b), rotation in zip(pair_times, rotations):
        dt_pair = max(t_b - t_a, 1e-3)
        omega = so3_log(rotation) / dt_pair
        lidar_t.append(0.5 * (t_a + t_b))
        lidar_w.append(omega)
    return np.asarray(lidar_t, dtype=float), np.asarray(lidar_w, dtype=float)


def refine_time_offset_signed(
    imu: ImuSeries,
    frames: list[LidarFrame],
    *,
    delta_t_s: float,
    R_IMU_lidar: np.ndarray,
    gyro_bias_rad_s: np.ndarray | None = None,
    search_s: float = 0.08,
    sample_hz: float = 50.0,
) -> TimeOffsetResult:
    """Refine ``δt`` with signed 3-axis rates using a known ``R_IMU_lidar``.

    Cost: mean squared error between ``gyro_imu(t_lidar+δt)`` and
    ``R_IMU_lidar @ omega_lidar(t_lidar)`` on a common grid around the coarse ``δt``.
    """

    notes: list[str] = [f"signed refine around coarse delta_t={delta_t_s:.6f}s"]
    if len(frames) < 5:
        return TimeOffsetResult(delta_t_s, 0.0, search_s, ("not enough LiDAR frames",), False)

    bias = np.zeros(3) if gyro_bias_rad_s is None else np.asarray(gyro_bias_rad_s, dtype=float)
    gyro = imu.gyro_rad_s - bias
    r_x = np.asarray(R_IMU_lidar, dtype=float).reshape(3, 3)

    stride = max(1, len(frames) // 20)
    lidar_t, lidar_w = _lidar_omega_series(frames, stride=stride)
    if lidar_t.size < 4:
        return TimeOffsetResult(delta_t_s, 0.0, search_s, ("not enough LiDAR omega samples",), False)

    # Predicted IMU-frame angular rate from LiDAR relative rotations.
    pred = (r_x @ lidar_w.T).T
    t_start = float(lidar_t[0])
    t_end = float(lidar_t[-1])
    if t_end - t_start < 0.5:
        return TimeOffsetResult(delta_t_s, 0.0, search_s, ("LiDAR span too short for signed refine",), False)

    dt = 1.0 / sample_hz
    grid = np.arange(t_start, t_end, dt)
    pred_grid = np.column_stack(
        [np.interp(grid, lidar_t, pred[:, axis], left=np.nan, right=np.nan) for axis in range(3)]
    )

    def _cost_and_corr(delta: float) -> tuple[float, float]:
        meas = np.column_stack(
            [
                np.interp(grid + delta, imu.t_s, gyro[:, axis], left=np.nan, right=np.nan)
                for axis in range(3)
            ]
        )
        mask = np.isfinite(pred_grid).all(axis=1) & np.isfinite(meas).all(axis=1)
        if int(np.count_nonzero(mask)) < 10:
            return float("inf"), -1.0
        err = meas[mask] - pred_grid[mask]
        cost = float(np.mean(np.sum(err * err, axis=1)))
        a = meas[mask].reshape(-1)
        b = pred_grid[mask].reshape(-1)
        a = a - np.mean(a)
        b = b - np.mean(b)
        corr = float(np.dot(a, b) / ((np.linalg.norm(a) + 1e-12) * (np.linalg.norm(b) + 1e-12)))
        return cost, corr

    coarse_cost, coarse_corr = _cost_and_corr(float(delta_t_s))
    best_delta = float(delta_t_s)
    best_cost = coarse_cost
    best_corr = coarse_corr
    half = abs(float(search_s))
    for delta in np.arange(delta_t_s - half, delta_t_s + half + 1e-12, dt):
        cost, corr = _cost_and_corr(float(delta))
        if cost < best_cost:
            best_cost = cost
            best_delta = float(delta)
            best_corr = corr

    # Parabolic refine on cost around the best discrete delta.
    samples = []
    for delta in (best_delta - dt, best_delta, best_delta + dt):
        cost, _ = _cost_and_corr(float(delta))
        samples.append(cost if np.isfinite(cost) else best_cost)
    y0, y1, y2 = samples
    denom = y0 - 2 * y1 + y2
    if abs(denom) > 1e-12 and y1 <= y0 and y1 <= y2:
        candidate = float(best_delta + 0.5 * (y0 - y2) / denom * dt)
        cand_cost, cand_corr = _cost_and_corr(candidate)
        if cand_cost < best_cost:
            best_delta = candidate
            best_cost = cand_cost
            best_corr = cand_corr

    # Guard with magnitude correlation so ICP-biased signed minima cannot wander.
    imu_t, imu_mag = _magnitude_series(imu.t_s, gyro)
    lidar_t_mag, lidar_mag = _magnitude_series(lidar_t, lidar_w)

    def _mag_score(delta: float) -> float:
        t_start_l = float(lidar_t_mag[0])
        t_end_l = float(lidar_t_mag[-1])
        grid_m = np.arange(t_start_l, t_end_l, dt)
        lidar_sig = np.interp(grid_m, lidar_t_mag, lidar_mag, left=0.0, right=0.0)
        lidar_sig = lidar_sig - np.mean(lidar_sig)
        imu_sig = np.interp(grid_m + delta, imu_t, imu_mag, left=0.0, right=0.0)
        imu_sig = imu_sig - np.mean(imu_sig)
        denom = (float(np.linalg.norm(lidar_sig)) + 1e-12) * (float(np.linalg.norm(imu_sig)) + 1e-12)
        return float(np.dot(imu_sig, lidar_sig) / denom)

    mag_at_coarse = _mag_score(float(delta_t_s))
    mag_at_best = _mag_score(best_delta)

    notes.append(
        f"signed 3-axis refine: delta_t={best_delta:.6f}s, "
        f"mse={best_cost:.4g} (coarse_mse={coarse_cost:.4g}), "
        f"corr={best_corr:.3f}, mag_corr={mag_at_best:.3f} (coarse_mag={mag_at_coarse:.3f}), "
        f"search=±{half:.3f}s"
    )
    improved = (
        np.isfinite(best_cost)
        and best_cost < coarse_cost * 0.999
        # Do not sacrifice the more reliable magnitude alignment for a noisy signed MSE gain.
        and mag_at_best + 1e-4 >= mag_at_coarse
    )
    if not improved:
        notes.append("signed refine rejected by MSE/mag-consistency; keeping previous delta_t")
        return TimeOffsetResult(
            delta_t_s=float(delta_t_s),
            correlation_peak=mag_at_coarse if mag_at_coarse > 0 else best_corr,
            search_s=search_s,
            notes=tuple(notes),
            ok=True,
        )
    return TimeOffsetResult(
        delta_t_s=best_delta,
        correlation_peak=mag_at_best,
        search_s=search_s,
        notes=tuple(notes),
        ok=True,
    )