from typing import Self

import numpy as np
from numba import njit
from numpy.typing import ArrayLike, NDArray

from attipy._quatops import _correct_quat_with_gibbs2, _correct_quat_with_rotvec

from ._attitude import Attitude
from ._kalman_fast import (
    _kalman_update_scalar_fast,
    _kalman_update_sequential_fast,
    _project_cov_ahead_fast,
)
from ._statespace import (
    _dyawda,
    _measurement_matrix,
    _process_noise_cov,
    _state_transition,
    _update_state_transition,
)
from ._transforms import _nz_b_from_quat, _yaw_from_quat
from ._vectorops import _normalize_vec, _skew_symmetric

P0 = (
    (1.0e-6, 0.0, 0.0, 0.0, 0.0, 0.0),
    (0.0, 1.0e-6, 0.0, 0.0, 0.0, 0.0),
    (0.0, 0.0, 1.0e-6, 0.0, 0.0, 0.0),
    (0.0, 0.0, 0.0, 1.0e-6, 0.0, 0.0),
    (0.0, 0.0, 0.0, 0.0, 1.0e-6, 0.0),
    (0.0, 0.0, 0.0, 0.0, 0.0, 1.0e-6),
)


def _gravity_nav(g: float, nav_frame: str) -> NDArray[np.float64]:
    """
    Gravity vector expressed in the navigation frame ('NED' or 'ENU').

    Parameters
    ----------
    g : float
        Gravitational acceleration in m/s^2.
    nav_frame : {'NED', 'ENU'}
        Navigation frame in which the gravity vector is expressed.

    Returns
    -------
    ndarray, shape (3,)
        Gravity vector expressed in the navigation frame.
    """
    if nav_frame.lower() == "ned":
        g_n = np.array([0.0, 0.0, g])
    elif nav_frame.lower() == "enu":
        g_n = np.array([0.0, 0.0, -g])
    else:
        raise ValueError(f"Unknown navigation frame: {nav_frame}.")
    return g_n


def _nz2vg(nav_frame: str) -> float:
    """
    Gravity direction along the navigation frame's z-axis. Transforms the z-axis
    of the navigation frame to a gravity reference vector (unit vector).

    Parameters
    ----------
    nav_frame : {'NED', 'ENU'}
        Navigation frame.

    Returns
    -------
    float
        Gravity direction along the navigation frame's z-axis. +1.0 for 'NED' and
        -1.0 for 'ENU'.
    """
    if nav_frame.lower() == "ned":
        return 1.0
    elif nav_frame.lower() == "enu":
        return -1.0
    else:
        raise ValueError(f"Unknown navigation frame: {nav_frame}.")


def _signed_smallest_angle(angle: float) -> float:
    """
    Convert the given angle to the smallest signed angle between [-pi., pi) radians.

    Parameters
    ----------
    angle : float
        Angle in radians.

    Returns
    -------
    float
        The smallest angle between [-pi, pi] radians.
    """
    return (angle + np.pi) % (2.0 * np.pi) - np.pi


@njit  # type: ignore[misc]
def _reset(q_nb, bg_b, dx) -> None:
    """
    Reset states (regulating error-states to zero).

    Moves information from the error-state vector to the nominal state vectors,
    and then resets the error-state vector to zero.

    Parameters
    ----------
    q_nb : ndarray, shape (4,)
        Unit quaternion to be updated in place.
    bg_b : ndarray, shape (3,)
        Gyroscope bias to be updated in place.
    dx : ndarray, shape (6,)
        Error-state vector (da, dbg) to be reset in place.
    """
    _correct_quat_with_gibbs2(q_nb, dx[0:3])
    bg_b[0] += dx[3]
    bg_b[1] += dx[4]
    bg_b[2] += dx[5]
    for i in range(6):
        dx[i] = 0.0


class MEKF:
    """
    Multiplicative extended Kalman filter (MEKF) for attitude estimation.

    Parameters
    ----------
    fs : float
        Sampling rate in Hz.
    q : Attitude or array_like, shape (4,), optional
        Initial attitude estimate given as an Attitude instance or a unit quaternion
        (qw, qx, qy, qz). Defaults to the identity quaternion (1.0, 0.0, 0.0, 0.0)
        (i.e., no rotation).
    bg : array_like, shape (3,), optional
        Initial gyroscope bias estimate (bgx, bgy, bgz) in rad/s. Defaults to zero bias.
    P : array_like, shape (6, 6), optional
        Initial error covariance matrix estimate. Defaults to a small diagonal matrix
        (1e-6 * eye(6)). The order of the (error) states is: dx = (da, dbg), where
        da is the attitude error, and dbg is the gyroscope bias error.
    gyro_noise_density : float, optional
        Gyroscope noise density (angular random walk) in (rad/s)/√Hz. Defaults to
        0.0001 (typical value for low-cost MEMS IMUs).
    gyro_bias_stability : float, optional
        Gyroscope bias stability (1-sigma) in rad/s. Defaults to 0.00005 (typical
        value for low-cost MEMS IMUs).
    gyro_bias_corr_time : float, optional
        Gyroscope bias correlation time in seconds. Defaults to 50.0 s.
    nav_frame : {'NED', 'ENU'}, optional
        Specifies the assumed inertial-like navigation frame. Should be 'NED'
        (North-East-Down) (default) or 'ENU' (East-North-Up).
    """

    def __init__(
        self,
        fs: float,
        q: Attitude | ArrayLike = (1.0, 0.0, 0.0, 0.0),
        bg: ArrayLike = (0.0, 0.0, 0.0),
        P: ArrayLike = P0,
        gyro_noise_density: float = 0.0001,
        gyro_bias_stability: float = 0.00005,
        gyro_bias_corr_time: float = 50.0,
        nav_frame: str = "NED",
    ) -> None:
        self._fs = fs
        self._dt = 1.0 / fs
        self._nav_frame = nav_frame.lower()
        self._nz2vg = _nz2vg(self._nav_frame)
        self._tmp = np.empty((6, 6))  # preallocated workspace

        # IMU noise parameters
        self._arw = gyro_noise_density  # angular random walk
        self._gbs = gyro_bias_stability  # gyro bias stability
        self._gbc = gyro_bias_corr_time  # gyro bias correlation time

        # Initial state and covariance estimates
        self._att_nb = Attitude(q) if not isinstance(q, Attitude) else q
        self._bg_b = np.asarray_chkfinite(bg).reshape(3).copy()
        self._P = np.asarray_chkfinite(P).reshape(6, 6).copy()
        self._dx = np.zeros(6)

        # Discrete state-space model
        vg_b = self._nz2vg * _nz_b_from_quat(self._att_nb._q)
        self._phi = _state_transition(self._dt, np.zeros(3), self._gbc)
        self._Q = _process_noise_cov(self._dt, self._arw, self._gbs, self._gbc)
        self._dhdx = _measurement_matrix(self._att_nb._q, vg_b)
        self._dhdx_gref = self._dhdx[0:3]
        self._dhdx_yaw = self._dhdx[3]

    @property
    def P(self) -> NDArray[np.float64]:
        """
        Copy of the error covariance matrix estimate.
        """
        return self._P.copy()

    @property
    def attitude(self) -> Attitude:
        """Attitude estimate (no copy)."""
        return self._att_nb

    def bias_gyro(self, degrees=False) -> NDArray[np.float64]:
        """
        Gyroscope bias estimate expressed in the body frame.

        Parameters
        ----------
        degrees : bool, optional
            Specifies whether to return the bias estimate in deg/s or rad/s. Defaults
            to rad/s.

        Returns
        -------
        ndarray, shape (3,)
            Copy of the gyroscope bias estimate.
        """
        if degrees:
            return np.degrees(self._bg_b.copy())
        return self._bg_b.copy()

    def _aiding_update_gref(self, vg_meas: ArrayLike, vg_var: ArrayLike | None) -> None:
        """
        Update state and covariance with gravity reference vector aiding measurement.
        """
        if vg_var is None:
            raise ValueError("'vg_var' not provided.")

        vg_b = self._nz2vg * _nz_b_from_quat(self._att_nb._q)
        self._dhdx_gref[0:3, 0:3] = _skew_symmetric(vg_b)

        _kalman_update_sequential_fast(
            vg_meas - vg_b,
            vg_var,
            self._dhdx_gref,
            self._dx,
            self._P,
            self._tmp[0],
            self._tmp[1],
        )

    def _aiding_update_yaw(
        self, yaw_meas: float, yaw_var: float | None, yaw_degrees: bool
    ) -> None:
        """
        Update state and covariance with heading (yaw angle) aiding measurement.
        """
        if yaw_var is None:
            raise ValueError("'yaw_var' not provided.")

        if yaw_degrees:
            yaw_meas = (np.pi / 180.0) * yaw_meas
            yaw_var = (np.pi / 180.0) ** 2 * yaw_var

        self._dhdx_yaw[0:3] = _dyawda(self._att_nb._q)

        _kalman_update_scalar_fast(
            _signed_smallest_angle(yaw_meas - _yaw_from_quat(self._att_nb._q)),
            yaw_var,
            self._dhdx_yaw,
            self._dx,
            self._P,
            self._tmp[0],
            self._tmp[1],
        )

    def update(
        self,
        dv: ArrayLike,
        dtheta: ArrayLike,
        *,
        degrees: bool = False,
        yaw: float | None = None,
        yaw_var: float | None = None,
        yaw_degrees: bool = False,
        gref: bool = True,
        gref_var: ArrayLike = (0.001, 0.001, 0.001),
    ) -> Self:
        """
        Update state estimates with IMU and aiding measurements.

        Parameters
        ----------
        dv : array_like, shape (3,)
            Sculling integral in m/s. I.e., the integral of specific force, f, over
            the sampling interval, dt. The simple approximation dv = f * dt can be
            used if sculling-corrected integrals are not available.
        dtheta : array_like, shape (3,)
            Coning integral (see ``degrees`` parameter for units). I.e., the integral
            of angular velocity, w, over the sampling interval, dt. The simple approximation
            dtheta = w * dt can be used if coning-corrected integrals are not available.
        degrees : bool, optional
            Specifies whether ``dtheta`` is given in degrees or radians. Defaults to radians.
        yaw : float, optional
            Heading (yaw angle) aiding measurement (see ``yaw_degrees`` for units).
            Defaults to ``None`` (no yaw aiding).
        yaw_var : float, optional
            Variance of heading (yaw angle) measurement (see ``yaw_degrees`` for units).
            Required for ``yaw``.
        yaw_degrees : bool, optional
            Specifies whether the units of ``yaw`` and ``yaw_var`` are deg and deg^2
            or rad and rad^2 (default).
        gref : bool, optional
            Specifies whether to use accelerometer measurements (dv) and the known
            direction of gravity as aiding. Defaults to ``True``.
        gref_var : array_like, shape (3,), optional
            Variance of gravity reference vector measurement noise (dimensionless).
            Required for ``gref``. Defaults to (0.001, 0.001, 0.001).

        Returns
        -------
        MEKF
            A reference to the instance itself after the update.
        """
        dv = np.asarray(dv)
        dtheta = np.asarray(dtheta)

        if degrees:
            dtheta = np.radians(dtheta)

        dtheta = dtheta - self._dt * self._bg_b

        # Update state-space model
        _update_state_transition(self._phi, dtheta)

        # Project attitude estimate ahead (strapdown algorithm) (a priori)
        _correct_quat_with_rotvec(self._att_nb._q, dtheta)

        # Project error covariance matrix estimate ahead (a priori)
        _project_cov_ahead_fast(self._P, self._phi, self._Q, self._tmp)

        # Update state and covariance estimates with aiding measurements (a posteriori)
        if gref is True:
            self._aiding_update_gref(-_normalize_vec(dv), gref_var)
        if yaw is not None:
            self._aiding_update_yaw(yaw, yaw_var, yaw_degrees)

        # Reset state (regulating error-state to zero)
        _reset(self._att_nb._q, self._bg_b, self._dx)

        return self
