"""`ImuPreintegrator` — single owner of GTSAM IMU preintegration (AZ-276 / E-CC-HELPERS).

Implements the `imu_preintegrator` contract v1.0.0 at
`_docs/02_document/contracts/shared_helpers/imu_preintegrator.md`.

Single-threaded by design — no internal lock. The composition root
binds one instance per writer thread. Strict-monotonic ``ts_ns`` is the
hard timestamp invariant; non-monotonic samples raise
``ImuPreintegrationError`` without mutating internal state.

Bias drift remains the consumer's responsibility — call
``reset_with_bias`` when the C1/C5 estimator's bias estimate changes.
"""

from __future__ import annotations

from typing import Any, Final

import gtsam
import numpy as np

from gps_denied_onboard._types.calibration import CameraCalibration
from gps_denied_onboard._types.nav import ImuBias, ImuSample, ImuWindow

__all__ = [
    "CombinedImuFactor",
    "ImuPreintegrationError",
    "ImuPreintegrator",
    "make_imu_preintegrator",
]


# Documented defaults pulled from a Bosch BMI088-class IMU running at
# 200 Hz — used when ``CameraCalibration.metadata`` does not carry an
# explicit ``imu_noise_model`` block. The contract owner notes the FC's
# per-deployment IMU noise model lives in ``CameraCalibration``; these
# defaults are only for bring-up + unit tests.
_DEFAULT_ACCEL_NOISE_DENSITY: Final[float] = 1.86e-3  # m/s^2 / sqrt(Hz)
_DEFAULT_GYRO_NOISE_DENSITY: Final[float] = 1.87e-4  # rad/s / sqrt(Hz)
_DEFAULT_ACCEL_BIAS_RW: Final[float] = 4.33e-4  # m/s^3 / sqrt(Hz)
_DEFAULT_GYRO_BIAS_RW: Final[float] = 2.66e-5  # rad/s^2 / sqrt(Hz)
_DEFAULT_INTEGRATION_NOISE: Final[float] = 1e-8
_DEFAULT_GRAVITY_M_S2: Final[float] = 9.80665

# Re-export GTSAM's combined IMU factor so consumers do not import GTSAM.
CombinedImuFactor = gtsam.CombinedImuFactor


class ImuPreintegrationError(RuntimeError):
    """Raised on schema, monotonicity, or empty-window violations.

    Carries the offending timestamp and the last accepted timestamp in
    its message so the consumer's catch-and-log path can record it as
    an FDR ``kind="imu.skew"`` event (per AZ-276 Risk 2 mitigation).
    """


def _bias_to_gtsam(bias: ImuBias) -> gtsam.imuBias.ConstantBias:
    return gtsam.imuBias.ConstantBias(
        np.asarray(bias.accel_bias, dtype=np.float64),
        np.asarray(bias.gyro_bias, dtype=np.float64),
    )


def _zero_bias() -> gtsam.imuBias.ConstantBias:
    return gtsam.imuBias.ConstantBias(np.zeros(3, dtype=np.float64), np.zeros(3, dtype=np.float64))


def _read_imu_noise(metadata: dict[str, Any]) -> dict[str, float]:
    """Pull IMU noise densities from ``CameraCalibration.metadata``.

    Falls back to the documented defaults when the block is absent —
    every key in the noise block is optional and independently
    defaulted, so partial blocks are honoured.
    """
    block = metadata.get("imu_noise_model", {}) if isinstance(metadata, dict) else {}
    return {
        "accel_noise_density": float(
            block.get("accel_noise_density", _DEFAULT_ACCEL_NOISE_DENSITY)
        ),
        "gyro_noise_density": float(block.get("gyro_noise_density", _DEFAULT_GYRO_NOISE_DENSITY)),
        "accel_bias_rw": float(block.get("accel_bias_rw", _DEFAULT_ACCEL_BIAS_RW)),
        "gyro_bias_rw": float(block.get("gyro_bias_rw", _DEFAULT_GYRO_BIAS_RW)),
        "integration_noise": float(block.get("integration_noise", _DEFAULT_INTEGRATION_NOISE)),
        "gravity_m_s2": float(block.get("gravity_m_s2", _DEFAULT_GRAVITY_M_S2)),
    }


class ImuPreintegrator:
    """Single owner of GTSAM `PreintegratedCombinedMeasurements`.

    Single-threaded by contract. Strict-monotonic timestamps enforced.
    """

    def __init__(self, params: gtsam.PreintegrationCombinedParams) -> None:
        self._params = params
        self._bias: gtsam.imuBias.ConstantBias = _zero_bias()
        self._pim = gtsam.PreintegratedCombinedMeasurements(self._params, self._bias)
        # ``_last_ts_ns`` is None until the first sample is integrated.
        self._last_ts_ns: int | None = None
        self._sample_count: int = 0

    def reset_with_bias(self, bias: ImuBias) -> None:
        """Replace the active bias.

        Discards the partial integration accumulator — the contract
        specifies that re-bias affects "subsequent samples only", which
        we honour by re-initialising the GTSAM PIM with the new bias
        and a clean monotonic baseline. Consumers MUST close the prior
        window via ``reset_for_new_keyframe`` before changing bias if
        they want to retain its contribution.
        """
        self._bias = _bias_to_gtsam(bias)
        self._pim = gtsam.PreintegratedCombinedMeasurements(self._params, self._bias)
        self._sample_count = 0
        self._last_ts_ns = None

    def integrate_sample(self, sample: ImuSample) -> None:
        """Integrate one IMU sample.

        Strict-monotonic guard runs BEFORE state mutation so a rejected
        sample leaves the accumulator unchanged.
        """
        if self._last_ts_ns is not None and sample.ts_ns <= self._last_ts_ns:
            raise ImuPreintegrationError(
                f"non-monotonic IMU sample: ts_ns={sample.ts_ns} <= last_ts_ns={self._last_ts_ns}"
            )

        if self._last_ts_ns is None:
            dt_seconds = 0.0
        else:
            dt_seconds = (sample.ts_ns - self._last_ts_ns) * 1e-9

        accel = np.asarray(sample.accel_xyz, dtype=np.float64)
        gyro = np.asarray(sample.gyro_xyz, dtype=np.float64)

        # GTSAM rejects dt==0; for the first sample we still record the
        # timestamp without integrating so the next sample sees a real dt.
        if dt_seconds > 0.0:
            try:
                self._pim.integrateMeasurement(accel, gyro, dt_seconds)
            except RuntimeError as exc:
                raise ImuPreintegrationError(
                    f"GTSAM PIM rejected sample at ts_ns={sample.ts_ns}: {exc}"
                ) from exc

        self._last_ts_ns = sample.ts_ns
        self._sample_count += 1

    def integrate_window(self, window: ImuWindow) -> None:
        """Integrate every sample in ``window`` in order."""
        for sample in window.samples:
            self.integrate_sample(sample)

    def current_preintegration(self) -> gtsam.PreintegratedCombinedMeasurements:
        """Return the live PIM without resetting state.

        Raises ``ImuPreintegrationError`` if no integration has run
        since the last reset (per AC-3).
        """
        if self._sample_count == 0:
            raise ImuPreintegrationError("no samples since reset: cannot return preintegration")
        return self._pim

    def reset_for_new_keyframe(self) -> gtsam.PreintegratedCombinedMeasurements:
        """Return the closed PIM and clear internal accumulators.

        Caller MUST capture the return value — the helper does not
        retain a reference past the call.
        """
        if self._sample_count == 0:
            raise ImuPreintegrationError("no samples since reset: cannot close keyframe")
        closed = self._pim
        self._pim = gtsam.PreintegratedCombinedMeasurements(self._params, self._bias)
        self._sample_count = 0
        self._last_ts_ns = None
        return closed


def make_imu_preintegrator(calibration: CameraCalibration) -> ImuPreintegrator:
    """Construct an `ImuPreintegrator` from the per-deployment calibration.

    Noise densities are pulled from
    ``calibration.metadata["imu_noise_model"]``; missing keys fall back
    to the documented BMI088-class defaults.
    """
    noise = _read_imu_noise(calibration.metadata or {})

    params = gtsam.PreintegrationCombinedParams.MakeSharedU(noise["gravity_m_s2"])
    params.setAccelerometerCovariance(np.eye(3) * (noise["accel_noise_density"] ** 2))
    params.setGyroscopeCovariance(np.eye(3) * (noise["gyro_noise_density"] ** 2))
    params.setBiasAccCovariance(np.eye(3) * (noise["accel_bias_rw"] ** 2))
    params.setBiasOmegaCovariance(np.eye(3) * (noise["gyro_bias_rw"] ** 2))
    params.setIntegrationCovariance(np.eye(3) * noise["integration_noise"])

    return ImuPreintegrator(params)
