"""Frame-to-frame IMU preintegration (Phase-A rotation + Phase-C full factor).

Phase-A: ``ΔR``, 3×3 ``Σ``, ``J_bg``.
Phase-C: ``ΔR/Δv/Δp``, 9×9 ``Σ`` (with bias RW process noise), ``J_bg``/``J_ba``.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np

from .geometry import orthonormalize_rotation, so3_exp, so3_log, skew


@dataclass(frozen=True)
class GyroPreintegration:
    """Rotation-only preintegration on ``[t0, t1]`` (IMU clock)."""

    delta_R: np.ndarray
    duration_s: float
    mean_gyro_norm: float
    sigma_rad: float
    weight: float
    bias_rad_s: np.ndarray
    cov: np.ndarray
    J_bg: np.ndarray


@dataclass(frozen=True)
class ImuPreintegration:
    """Full IMU preintegration on ``[t0, t1]`` (IMU clock).

    ``delta_R`` maps vectors from IMU frame at ``t1`` into IMU frame at ``t0``.
    ``delta_v`` / ``delta_p`` are body-frame increments (no gravity).
    Error-state order in ``cov`` / Jacobians: ``[δθ, δv, δp]`` (9).
    ``J_bg`` / ``J_ba`` are 9×3: first-order correction w.r.t. constant bias deltas.
    """

    delta_R: np.ndarray
    delta_v: np.ndarray
    delta_p: np.ndarray
    duration_s: float
    mean_gyro_norm: float
    sigma_rad: float
    weight: float
    gyro_bias_rad_s: np.ndarray
    acc_bias_m_s2: np.ndarray
    cov: np.ndarray
    J_bg: np.ndarray
    J_ba: np.ndarray


def _right_jacobian(phi: np.ndarray) -> np.ndarray:
    """SO(3) right Jacobian ``Jr(φ)`` with ``Exp(φ+δ)≈Exp(φ)Exp(Jr δ)``."""

    phi = np.asarray(phi, dtype=float).reshape(3)
    angle = float(np.linalg.norm(phi))
    if angle < 1e-8:
        return np.eye(3) - 0.5 * skew(phi)
    axis = phi / angle
    s = skew(axis)
    return (
        np.eye(3)
        - ((1.0 - np.cos(angle)) / angle) * s
        + ((angle - np.sin(angle)) / angle) * (s @ s)
    )


def _interp_vec(times_s: np.ndarray, values: np.ndarray, t: float) -> np.ndarray:
    """Linear interpolate a 3-vector series at an exact time."""

    return np.array(
        [float(np.interp(t, times_s, values[:, axis])) for axis in range(3)],
        dtype=float,
    )


def _interp_gyro(times_s: np.ndarray, gyro_rad_s: np.ndarray, t: float) -> np.ndarray:
    """Linear interpolate gyro at an exact time."""

    return _interp_vec(times_s, gyro_rad_s, t)


def _pair_weight(duration_s: float, mean_gyro_norm: float, cov_trace: float) -> float:
    """Larger weight for short, excited, low-covariance intervals."""

    duration_term = 1.0 / max(duration_s, 0.05)
    excite_term = min(max(mean_gyro_norm, 1e-3), 1.0)
    avg_var = max(cov_trace / 3.0, 1e-8)
    return float(duration_term * excite_term / avg_var)


def preintegrate_gyro(
    times_s: np.ndarray,
    gyro_rad_s: np.ndarray,
    t0: float,
    t1: float,
    bias_rad_s: np.ndarray | None = None,
    *,
    sigma_g_rad_s_sqrt_hz: float = 1.5e-3,
) -> GyroPreintegration:
    """Discrete mid-point gyro preintegration with exact endpoints.

    ``delta_R`` maps vectors from IMU frame at ``t1`` into IMU frame at ``t0``
    via right-invariant updates ``ΔR ← ΔR Exp((ω-b) dt)``.

    Also returns:
    - ``cov``: 3×3 covariance of the right tangent noise on ``ΔR``
    - ``J_bg``: ``ΔR(b+δb) ≈ ΔR Exp(J_bg δb)``
    """

    times_s = np.asarray(times_s, dtype=float).reshape(-1)
    gyro_rad_s = np.asarray(gyro_rad_s, dtype=float).reshape(-1, 3)
    bias = np.zeros(3) if bias_rad_s is None else np.asarray(bias_rad_s, dtype=float).reshape(3)
    duration = float(max(t1 - t0, 0.0))

    empty = GyroPreintegration(
        delta_R=np.eye(3),
        duration_s=0.0,
        mean_gyro_norm=0.0,
        sigma_rad=1e3,
        weight=1e-6,
        bias_rad_s=bias.copy(),
        cov=np.eye(3) * 1e6,
        J_bg=np.zeros((3, 3)),
    )
    if times_s.size < 2 or duration <= 0:
        return empty

    t0 = float(np.clip(t0, times_s[0], times_s[-1]))
    t1 = float(np.clip(t1, times_s[0], times_s[-1]))
    duration = float(max(t1 - t0, 0.0))
    if duration <= 0:
        return empty

    left = int(np.searchsorted(times_s, t0, side="left") - 1)
    right = int(np.searchsorted(times_s, t1, side="right"))
    left = max(left, 0)
    right = min(right, times_s.size - 1)
    if right <= left:
        return empty

    delta_r = np.eye(3)
    j_bg = np.zeros((3, 3))
    cov = np.zeros((3, 3))
    sigma2 = float(sigma_g_rad_s_sqrt_hz) ** 2
    gyro_norms: list[float] = []

    for index in range(left, right):
        t_a = float(times_s[index])
        t_b = float(times_s[index + 1])
        if t_b <= t0 or t_a >= t1:
            continue
        seg0 = max(t_a, t0)
        seg1 = min(t_b, t1)
        dt = seg1 - seg0
        if dt <= 0:
            continue

        # Exact endpoint gyro via linear interpolation inside the sample interval.
        g_a = _interp_gyro(times_s, gyro_rad_s, seg0)
        g_b = _interp_gyro(times_s, gyro_rad_s, seg1)
        omega = 0.5 * (g_a + g_b) - bias
        gyro_norms.append(float(np.linalg.norm(omega)))

        theta = omega * dt
        jr = _right_jacobian(theta)
        a_mat = so3_exp(-theta)
        j_bg = a_mat @ j_bg - jr * dt
        cov = a_mat @ cov @ a_mat.T + jr @ (sigma2 * dt * np.eye(3)) @ jr.T
        delta_r = delta_r @ so3_exp(theta)

    delta_r = orthonormalize_rotation(delta_r)
    mean_gyro_norm = float(np.mean(gyro_norms)) if gyro_norms else 0.0
    cov = 0.5 * (cov + cov.T)
    cov = cov + np.eye(3) * 1e-12
    if mean_gyro_norm < 0.02:
        cov = cov * 4.0
    cov_trace = float(np.trace(cov))
    sigma_rad = float(np.sqrt(max(cov_trace / 3.0, 1e-12)))
    weight = _pair_weight(duration, mean_gyro_norm, cov_trace)
    return GyroPreintegration(
        delta_R=delta_r,
        duration_s=duration,
        mean_gyro_norm=mean_gyro_norm,
        sigma_rad=sigma_rad,
        weight=weight,
        bias_rad_s=bias.copy(),
        cov=cov,
        J_bg=np.asarray(j_bg, dtype=float),
    )


def preintegrate_imu(
    times_s: np.ndarray,
    gyro_rad_s: np.ndarray,
    acc_m_s2: np.ndarray,
    t0: float,
    t1: float,
    gyro_bias_rad_s: np.ndarray | None = None,
    acc_bias_m_s2: np.ndarray | None = None,
    *,
    sigma_g_rad_s_sqrt_hz: float = 1.5e-3,
    sigma_a_m_s2_sqrt_hz: float = 2.0e-2,
    sigma_bg_rw_rad_s_sqrt_hz: float = 1.0e-5,
    sigma_ba_rw_m_s2_sqrt_hz: float = 1.0e-3,
) -> ImuPreintegration:
    """Mid-point IMU preintegration with exact endpoints and bias-RW noise.

    Discrete updates (right-invariant)::

        ΔR ← ΔR Exp((ω-bg) dt)
        Δv ← Δv + ΔR (a-ba) dt
        Δp ← Δp + Δv_old dt + 0.5 ΔR (a-ba) dt²

    Propagates a 15-DoF error state ``[δθ, δv, δp, δbg, δba]`` then returns the
    top-left 9×9 covariance (bias RW already folded in) and 9×3 Jacobians.
    """

    times_s = np.asarray(times_s, dtype=float).reshape(-1)
    gyro_rad_s = np.asarray(gyro_rad_s, dtype=float).reshape(-1, 3)
    acc_m_s2 = np.asarray(acc_m_s2, dtype=float).reshape(-1, 3)
    bg = np.zeros(3) if gyro_bias_rad_s is None else np.asarray(gyro_bias_rad_s, dtype=float).reshape(3)
    ba = np.zeros(3) if acc_bias_m_s2 is None else np.asarray(acc_bias_m_s2, dtype=float).reshape(3)

    empty = ImuPreintegration(
        delta_R=np.eye(3),
        delta_v=np.zeros(3),
        delta_p=np.zeros(3),
        duration_s=0.0,
        mean_gyro_norm=0.0,
        sigma_rad=1e3,
        weight=1e-6,
        gyro_bias_rad_s=bg.copy(),
        acc_bias_m_s2=ba.copy(),
        cov=np.eye(9) * 1e6,
        J_bg=np.zeros((9, 3)),
        J_ba=np.zeros((9, 3)),
    )
    if times_s.size < 2 or acc_m_s2.shape != gyro_rad_s.shape:
        return empty

    t0 = float(np.clip(t0, times_s[0], times_s[-1]))
    t1 = float(np.clip(t1, times_s[0], times_s[-1]))
    duration = float(max(t1 - t0, 0.0))
    if duration <= 0:
        return empty

    left = int(np.searchsorted(times_s, t0, side="left") - 1)
    right = int(np.searchsorted(times_s, t1, side="right"))
    left = max(left, 0)
    right = min(right, times_s.size - 1)
    if right <= left:
        return empty

    delta_r = np.eye(3)
    delta_v = np.zeros(3)
    delta_p = np.zeros(3)
    # Jacobians of [δθ, δv, δp] w.r.t. constant bias (accumulated analytically).
    j_bg = np.zeros((9, 3))
    j_ba = np.zeros((9, 3))
    # 15×15 covariance: [θ, v, p, bg, ba]
    cov15 = np.zeros((15, 15))
    sg2 = float(sigma_g_rad_s_sqrt_hz) ** 2
    sa2 = float(sigma_a_m_s2_sqrt_hz) ** 2
    sbg2 = float(sigma_bg_rw_rad_s_sqrt_hz) ** 2
    sba2 = float(sigma_ba_rw_m_s2_sqrt_hz) ** 2
    gyro_norms: list[float] = []

    for index in range(left, right):
        t_a = float(times_s[index])
        t_b = float(times_s[index + 1])
        if t_b <= t0 or t_a >= t1:
            continue
        seg0 = max(t_a, t0)
        seg1 = min(t_b, t1)
        dt = seg1 - seg0
        if dt <= 0:
            continue

        g_a = _interp_vec(times_s, gyro_rad_s, seg0)
        g_b = _interp_vec(times_s, gyro_rad_s, seg1)
        a_a = _interp_vec(times_s, acc_m_s2, seg0)
        a_b = _interp_vec(times_s, acc_m_s2, seg1)
        omega = 0.5 * (g_a + g_b) - bg
        acc = 0.5 * (a_a + a_b) - ba
        gyro_norms.append(float(np.linalg.norm(omega)))

        theta = omega * dt
        jr = _right_jacobian(theta)
        r_dt = so3_exp(theta)
        r_mid = delta_r  # rotate body accel into i0 frame before update

        # Bias Jacobians (Forster-style first-order recursion).
        j_r_bg = j_bg[0:3]
        j_v_bg = j_bg[3:6]
        j_p_bg = j_bg[6:9]
        j_r_ba = j_ba[0:3]
        j_v_ba = j_ba[3:6]
        j_p_ba = j_ba[6:9]

        acc_skew = skew(acc)
        j_p_bg_new = j_p_bg + j_v_bg * dt - 0.5 * r_mid @ acc_skew @ j_r_bg * (dt**2)
        j_v_bg_new = j_v_bg - r_mid @ acc_skew @ j_r_bg * dt
        j_r_bg_new = r_dt.T @ j_r_bg - jr * dt

        j_p_ba_new = j_p_ba + j_v_ba * dt - 0.5 * r_mid * (dt**2)
        j_v_ba_new = j_v_ba - r_mid * dt
        j_r_ba_new = r_dt.T @ j_r_ba

        j_bg = np.vstack([j_r_bg_new, j_v_bg_new, j_p_bg_new])
        j_ba = np.vstack([j_r_ba_new, j_v_ba_new, j_p_ba_new])

        # Nominal state update (use pre-update Δv in position).
        delta_p = delta_p + delta_v * dt + 0.5 * r_mid @ acc * (dt**2)
        delta_v = delta_v + r_mid @ acc * dt
        delta_r = orthonormalize_rotation(delta_r @ r_dt)

        # Linearized error-state transition (15×15).
        f = np.eye(15)
        a_mat = so3_exp(-theta)
        f[0:3, 0:3] = a_mat
        f[0:3, 9:12] = -jr * dt
        f[3:6, 0:3] = -r_mid @ acc_skew * dt
        f[3:6, 12:15] = -r_mid * dt
        f[6:9, 0:3] = -0.5 * r_mid @ acc_skew * (dt**2)
        f[6:9, 3:6] = np.eye(3) * dt
        f[6:9, 12:15] = -0.5 * r_mid * (dt**2)

        # Noise: continuous densities σ²; Var(∫n dt)=σ² dt. Columns: n_g, n_a, n_bg, n_ba.
        g_mat = np.zeros((15, 12))
        g_mat[0:3, 0:3] = jr
        g_mat[3:6, 3:6] = r_mid
        g_mat[6:9, 3:6] = 0.5 * r_mid * dt
        g_mat[9:12, 6:9] = np.eye(3)
        g_mat[12:15, 9:12] = np.eye(3)
        q = np.zeros((12, 12))
        q[0:3, 0:3] = sg2 * dt * np.eye(3)
        q[3:6, 3:6] = sa2 * dt * np.eye(3)
        q[6:9, 6:9] = sbg2 * dt * np.eye(3)
        q[9:12, 9:12] = sba2 * dt * np.eye(3)
        cov15 = f @ cov15 @ f.T + g_mat @ q @ g_mat.T

    delta_r = orthonormalize_rotation(delta_r)
    mean_gyro_norm = float(np.mean(gyro_norms)) if gyro_norms else 0.0
    cov9 = cov15[0:9, 0:9]
    cov9 = 0.5 * (cov9 + cov9.T) + np.eye(9) * 1e-12
    if mean_gyro_norm < 0.02:
        cov9 = cov9.copy()
        cov9[0:3, 0:3] = cov9[0:3, 0:3] * 4.0
    cov_trace = float(np.trace(cov9[0:3, 0:3]))
    sigma_rad = float(np.sqrt(max(cov_trace / 3.0, 1e-12)))
    weight = _pair_weight(duration, mean_gyro_norm, cov_trace)
    return ImuPreintegration(
        delta_R=delta_r,
        delta_v=np.asarray(delta_v, dtype=float),
        delta_p=np.asarray(delta_p, dtype=float),
        duration_s=duration,
        mean_gyro_norm=mean_gyro_norm,
        sigma_rad=sigma_rad,
        weight=weight,
        gyro_bias_rad_s=bg.copy(),
        acc_bias_m_s2=ba.copy(),
        cov=np.asarray(cov9, dtype=float),
        J_bg=np.asarray(j_bg, dtype=float),
        J_ba=np.asarray(j_ba, dtype=float),
    )


def apply_bias_correction_imu(
    preint: ImuPreintegration,
    delta_gyro_bias: np.ndarray | None = None,
    delta_acc_bias: np.ndarray | None = None,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """First-order bias correction of ``ΔR/Δv/Δp``.

    Returns ``(delta_R, delta_v, delta_p)``.
    """

    dbg = np.zeros(3) if delta_gyro_bias is None else np.asarray(delta_gyro_bias, dtype=float).reshape(3)
    dba = np.zeros(3) if delta_acc_bias is None else np.asarray(delta_acc_bias, dtype=float).reshape(3)
    j_bg = np.asarray(preint.J_bg, dtype=float).reshape(9, 3)
    j_ba = np.asarray(preint.J_ba, dtype=float).reshape(9, 3)
    delta_r = orthonormalize_rotation(preint.delta_R @ so3_exp(j_bg[0:3] @ dbg))
    delta_v = preint.delta_v + j_bg[3:6] @ dbg + j_ba[3:6] @ dba
    delta_p = preint.delta_p + j_bg[6:9] @ dbg + j_ba[6:9] @ dba
    return delta_r, np.asarray(delta_v, dtype=float), np.asarray(delta_p, dtype=float)


def relative_rotation_from_lidar(R_X: np.ndarray, R_B: np.ndarray) -> np.ndarray:
    """Map LiDAR relative rotation into IMU frame: ``R_X R_B R_X^T``."""

    r_x = orthonormalize_rotation(R_X)
    r_b = orthonormalize_rotation(R_B)
    return orthonormalize_rotation(r_x @ r_b @ r_x.T)


def preintegration_rotation_residual(
    delta_R: np.ndarray,
    R_X: np.ndarray,
    R_B: np.ndarray,
) -> np.ndarray:
    """``log( delta_R^T * R_X R_B R_X^T )`` in so(3)."""

    predicted = relative_rotation_from_lidar(R_X, R_B)
    return so3_log(delta_R.T @ predicted)


def apply_bias_jacobian_correction(
    delta_R: np.ndarray,
    J_bg: np.ndarray,
    delta_bias_rad_s: np.ndarray,
) -> np.ndarray:
    """First-order update ``ΔR(b+δb) ≈ ΔR Exp(J_bg δb)``."""

    db = np.asarray(delta_bias_rad_s, dtype=float).reshape(3)
    j_bg = np.asarray(J_bg, dtype=float).reshape(3, 3)
    return orthonormalize_rotation(delta_R @ so3_exp(j_bg @ db))


def apply_constant_bias_correction(
    delta_R: np.ndarray,
    duration_s: float,
    delta_bias_rad_s: np.ndarray,
) -> np.ndarray:
    """Legacy first-order correction when ``J_bg`` is unavailable.

    ``ΔR(b+δb) ≈ ΔR Exp(-δb Δt)`` (identity Jacobian approximation).
    """

    db = np.asarray(delta_bias_rad_s, dtype=float).reshape(3)
    return orthonormalize_rotation(delta_R @ so3_exp(-db * float(duration_s)))


def residual_whiten_matrix(cov: np.ndarray) -> np.ndarray:
    """Return ``W`` such that ``W @ e`` is approximately information-whitened.

    Accepts square ``n×n`` covariances (3×3 rotation or 9×9 full IMU).
    """

    matrix = np.asarray(cov, dtype=float)
    if matrix.ndim != 2 or matrix.shape[0] != matrix.shape[1]:
        raise ValueError("cov must be square")
    n = matrix.shape[0]
    matrix = 0.5 * (matrix + matrix.T) + np.eye(n) * 1e-10
    try:
        info = np.linalg.inv(matrix)
        return np.linalg.cholesky(info).T
    except np.linalg.LinAlgError:
        scale = 1.0 / max(float(np.sqrt(np.trace(matrix) / n)), 1e-6)
        return np.eye(n) * scale