calibration/tools/rscap_v2/hi13_imu.py

"""Decode Hipnuc / HI13 (HI91/HI92) IMU frames from a V2 .rscap capture.

Matches ``EcarSensorMinimal/RawSerialImu/Hi13Protocol.cs``:
sync ``5A A5``, CRC16 over header[0:4]+payload, tag ``0x91`` / ``0x92``.

HI91 (preferred for calibration):
- accel: float32 in g → m/s² (* 9.80665)
- gyro: float32 in deg/s → rad/s
- device time: uint32 ms at frame offset 14 → ``t_s = ms * 1e-3``
"""

from __future__ import annotations

import struct
from dataclasses import dataclass

import numpy as np

from .capture_format_v2 import CaptureFile
from .n300_imu import ImuSample, samples_to_arrays

G0 = 9.80665
DEG2RAD = np.pi / 180.0


@dataclass(frozen=True)
class Hi13Sample:
    """One CRC-valid native HI91 record.

    t_s/system_time_ms are sensor-owned. Host receive ticks are retained only
    for clock diagnostics and cross-clock fitting.
    """

    t_s: float
    system_time_ms: int
    device_timestamp_us: int
    host_receive_utc_ticks: int
    gyro_rad_s: tuple[float, float, float]
    accel_m_s2: tuple[float, float, float]
    rpy_deg: tuple[float, float, float]
    quaternion_wxyz: tuple[float, float, float, float]
    mag_ut: tuple[float, float, float]
    pps_sync_stamp_ms: int
    temperature_c: int
    air_pressure_pa: float
    frame_tag: int = 0x91


def crc16_hi13(frame: bytes, payload_length: int) -> int:
    crc = 0
    for value in frame[:4]:
        crc = _update_crc16(crc, value)
    for value in frame[6 : 6 + payload_length]:
        crc = _update_crc16(crc, value)
    return crc & 0xFFFF


def _update_crc16(crc: int, value: int) -> int:
    crc ^= (value & 0xFF) << 8
    for _ in range(8):
        if crc & 0x8000:
            crc = ((crc << 1) ^ 0x1021) & 0xFFFF
        else:
            crc = (crc << 1) & 0xFFFF
    return crc


def parse_hi91_sample(raw: bytes, host_receive_utc_ticks: int = 0) -> Hi13Sample | None:
    """Decode every calibration-relevant field from a CRC-valid HI91 frame."""

    if len(raw) < 6 + 76:
        return None
    payload_length = raw[2] | (raw[3] << 8)
    if payload_length < 76 or len(raw) < 6 + payload_length:
        return None
    if raw[6] != 0x91:
        return None
    expected = raw[4] | (raw[5] << 8)
    if crc16_hi13(raw, payload_length) != expected:
        return None
    device_ms = int(struct.unpack_from("<I", raw, 14)[0])
    ax, ay, az = struct.unpack_from("<fff", raw, 18)
    gx, gy, gz = struct.unpack_from("<fff", raw, 30)
    gyro = (gx * DEG2RAD, gy * DEG2RAD, gz * DEG2RAD)
    accel = (ax * G0, ay * G0, az * G0)
    return Hi13Sample(
        t_s=float(device_ms) * 1e-3,
        system_time_ms=device_ms,
        device_timestamp_us=device_ms * 1000,
        host_receive_utc_ticks=int(host_receive_utc_ticks),
        gyro_rad_s=gyro,
        accel_m_s2=accel,
        rpy_deg=struct.unpack_from("<fff", raw, 54),
        quaternion_wxyz=struct.unpack_from("<ffff", raw, 66),
        mag_ut=struct.unpack_from("<fff", raw, 42),
        pps_sync_stamp_ms=int(struct.unpack_from("<H", raw, 7)[0]),
        temperature_c=int(struct.unpack_from("<b", raw, 9)[0]),
        air_pressure_pa=float(struct.unpack_from("<f", raw, 10)[0]),
    )


def parse_hi91_frame(raw: bytes) -> tuple[tuple[float, float, float], tuple[float, float, float], int] | None:
    """Backward-compatible compact HI91 decoder."""

    sample = parse_hi91_sample(raw)
    if sample is None:
        return None
    return sample.gyro_rad_s, sample.accel_m_s2, sample.system_time_ms


def iter_hi13_imu_samples(
    capture: CaptureFile,
    *,
    host_utc_ticks_min: int | None = None,
    host_utc_ticks_max: int | None = None,
) -> list[Hi13Sample]:
    """Return CRC-valid HI91 samples sorted by device timestamp.

    Streams chunk-by-chunk (no giant join) and can skip whole chunks outside the
    host UTC receive window before parsing.
    """

    samples: list[Hi13Sample] = []
    carry = b""
    for chunk in capture.chunks:
        if host_utc_ticks_min is not None and chunk.receive_utc_ticks < host_utc_ticks_min:
            carry = b""
            continue
        if host_utc_ticks_max is not None and chunk.receive_utc_ticks > host_utc_ticks_max:
            # chunks are time-ordered; remaining ones are later
            if chunk.receive_utc_ticks > host_utc_ticks_max:
                break
        stream = carry + chunk.raw
        cursor = 0
        while cursor + 6 < len(stream):
            sync = stream.find(b"\x5A\xA5", cursor)
            if sync < 0:
                carry = b""
                break
            if sync + 6 > len(stream):
                carry = stream[sync:]
                break
            payload_length = stream[sync + 2] | (stream[sync + 3] << 8)
            if payload_length < 1 or payload_length > 512:
                cursor = sync + 1
                continue
            end = sync + 6 + payload_length
            if end > len(stream):
                carry = stream[sync:]
                break
            host_ticks = chunk.receive_utc_ticks
            parsed = parse_hi91_sample(stream[sync:end], host_ticks)
            cursor = end
            if parsed is None:
                continue
            if host_utc_ticks_min is not None and host_ticks < host_utc_ticks_min:
                continue
            if host_utc_ticks_max is not None and host_ticks > host_utc_ticks_max:
                continue
            samples.append(parsed)
        else:
            carry = b""
    samples.sort(key=lambda sample: (sample.t_s, sample.device_timestamp_us))
    return samples


__all__ = [
    "ImuSample",
    "Hi13Sample",
    "crc16_hi13",
    "iter_hi13_imu_samples",
    "parse_hi91_frame",
    "parse_hi91_sample",
    "samples_to_arrays",
]