r"""Decode the Thor / Micromate (series-4) sensor self-check waveforms from an
IDFW event binary.

Reverse-engineered 2026-09-15 against 4 UM (Thor) oracle events.  The IDFW
binary carries the sensor self-check in its fixed-header region (before the
waveform body), as up to four records tagged ``01 0e 3c/3d/3e/3f`` — the SAME
channel ids as the series-3 MiniMate Plus (Tran / Vert / Long / MicL), which is
the physical self-test:

  * 3c / 3d / 3e = Tran / Vert / Long geophone ring-downs (a damped impulse
    response — resonant frequency + damping).
  * 3f          = MicL pulse train (the mic's known-signal gain check).  Absent
    on three-channel (mic-disabled) units.

Record framing (per record)::

    01 0e [id:1] [flags:3] [count:2 BE] [pad:10] [int16-BE samples × count]
    \___ 18-byte header ___/

Unlike series-3's delta-coded trailing block, series-4 stores each trace as a
raw int16 big-endian array.  ``count`` (the 2-byte field at header offset +8)
is the sample count; the record is padded to a fixed stride after that.
"""
from __future__ import annotations

import struct
from typing import Dict, List

# Record id → channel.  Same ids/order as series-3 (minimateplus.sensor_check).
_ID_TO_CHANNEL = {0x3C: "Tran", 0x3D: "Vert", 0x3E: "Long", 0x3F: "MicL"}
_CHAIN_IDS = (0x3C, 0x3D, 0x3E, 0x3F)

_MARKER = b"\x01\x0e"       # precedes the 1-byte channel id
_HEADER_LEN = 18            # bytes from the marker start to the first sample
_COUNT_OFF = 8              # 2-byte BE sample count, from the marker start
_MAX_COUNT = 4000           # sanity cap (traces are ~70-200 samples)


def _find_chain(raw: bytes):
    """Locate the sensor-check record chain.  Returns a list of
    ``(offset, id, count)`` for the first run of markers whose ids run
    3c, 3d, 3e[, 3f] in order, or ``[]``.

    Records are padded to a fixed stride, so the next marker is not at
    ``header + count*2``; instead collect every ``01 0e [id]`` marker with a
    sane count and take the first id-ordered run.  Validating the id sequence
    (not a lone ``01 0e 3c``) keeps a stray marker in the waveform body from
    matching — the real chain sits in the fixed header, ahead of the body.
    """
    n = len(raw)
    markers = []
    for p in range(n - _HEADER_LEN):
        if raw[p:p + 2] == _MARKER and raw[p + 2] in _ID_TO_CHANNEL:
            count = int.from_bytes(raw[p + _COUNT_OFF:p + _COUNT_OFF + 2], "big")
            if 0 < count <= _MAX_COUNT:
                markers.append((p, raw[p + 2], count))

    for i, (off, rid, _c) in enumerate(markers):
        if rid != 0x3C:
            continue
        run = [markers[i]]
        for m in markers[i + 1:]:
            if len(run) < len(_CHAIN_IDS) and m[1] == _CHAIN_IDS[len(run)]:
                run.append(m)
            else:
                break
        if len(run) >= 3:       # 3-channel (mic-disabled) units are valid
            return run
    return []


def decode_idf_sensor_check(raw: bytes) -> Dict[str, List[int]]:
    """Decode the sensor self-check traces from a Thor/Micromate IDFW binary.

    Returns ``{"Tran": [...], "Vert": [...], "Long": [...], "MicL": [...]}`` in
    raw int16 ADC counts (MicL omitted on 3-channel units), or ``{}`` if the
    binary carries no sensor-check chain (a non-IDF file, or an IDFH histogram).
    """
    chain = _find_chain(raw)
    if not chain:
        return {}
    out: Dict[str, List[int]] = {}
    for off, rid, count in chain:
        start = off + _HEADER_LEN
        blob = raw[start:start + count * 2]
        if len(blob) < count * 2:
            continue
        out[_ID_TO_CHANNEL[rid]] = list(struct.unpack(">%dh" % count, blob))
    return out