#!/usr/bin/env python3
"""
rtl_433 .cu8 IQ  ->  Flipper-style signed pulse train (µs)
==========================================================

Turns the rtl_433 project's ``.cu8`` test captures (interleaved I/Q uint8) into
the same signed-duration pulse array our RAW feature extractor already consumes
(``+high`` / ``-low`` in microseconds), so genuinely-independent, non-Flipper
captures can augment the statistical category classifier's training set.

Why amplitude (OOK) demod: the target Sub-GHz device classes we care about
(weather sensors, doorbells, PIR/door/smoke security sensors, simple remotes)
are overwhelmingly OOK/ASK — the carrier is keyed on/off, so the signal
*envelope* (magnitude of the complex sample) is the message. FSK captures keep
constant amplitude, so this demod yields no usable transitions on them; that is
a feature, not a bug — ``demod_file`` returns ``None`` and such captures
self-filter out of the training set via ``looks_like_ook``.

The sample rate and centre frequency are read from the rtl_433 filename
convention ``*_<freq>M_<rate>k.cu8`` (e.g. ``g003_433.92M_250k.cu8``), falling
back to 250 kHz / 433.92 MHz.

This is a training-data-prep tool (offline), not part of the serving path.
"""

from __future__ import annotations

import re
from pathlib import Path
from typing import List, Optional, Tuple

import numpy as np

DEFAULT_RATE = 250_000
DEFAULT_FREQ = 433_920_000

# Quality gate — reject captures that did not demodulate into a plausible
# OOK pulse train (FSK, pure noise, or empty).
MIN_TRANSITIONS = 16      # need real structure, not one long burst
MIN_SHORT_US = 40         # ignore sub-symbol glitches
MAX_KEEP_PULSES = 1024    # cap runaway repeats (Flipper captures are bounded)


def parse_rate_freq(name: str) -> Tuple[int, int]:
    """Extract (sample_rate_hz, frequency_hz) from an rtl_433 filename."""
    rate_m = re.search(r"_(\d+)k", name)
    rate = int(rate_m.group(1)) * 1000 if rate_m else DEFAULT_RATE
    freq_m = re.search(r"_(\d+(?:\.\d+)?)M", name)
    freq = int(float(freq_m.group(1)) * 1_000_000) if freq_m else DEFAULT_FREQ
    return rate, freq


def _magnitude(path: Path) -> Optional[np.ndarray]:
    raw = np.fromfile(path, dtype=np.uint8)
    if raw.size < 4:
        return None
    raw = raw.astype(np.float32) - 127.5  # centre uint8 IQ
    i, q = raw[0::2], raw[1::2]
    n = min(i.size, q.size)
    if n == 0:
        return None
    return np.sqrt(i[:n] * i[:n] + q[:n] * q[:n])


def _runs_to_pulses(hyst: np.ndarray, us_per_sample: float) -> List[int]:
    """Run-length encode a boolean high/low mask into signed µs durations."""
    if hyst.size == 0:
        return []
    # Indices where the level changes.
    change = np.flatnonzero(np.diff(hyst.view(np.int8))) + 1
    bounds = np.concatenate(([0], change, [hyst.size]))
    pulses: List[int] = []
    for a, b in zip(bounds[:-1], bounds[1:]):
        dur = (b - a) * us_per_sample
        if dur < MIN_SHORT_US:
            continue
        sign = 1 if hyst[a] else -1
        pulses.append(int(round(dur)) * sign)
    return pulses


def _trim_gaps(pulses: List[int]) -> List[int]:
    """Drop leading/trailing low (gap) runs so the train starts on a pulse."""
    start = 0
    while start < len(pulses) and pulses[start] < 0:
        start += 1
    end = len(pulses)
    while end > start and pulses[end - 1] < 0:
        end -= 1
    return pulses[start:end]


def demod_file(path: Path) -> Tuple[Optional[List[int]], int]:
    """Demodulate one .cu8 file to (pulse_train | None, frequency_hz)."""
    path = Path(path)
    rate, freq = parse_rate_freq(path.name)
    us_per_sample = 1_000_000.0 / rate
    mag = _magnitude(path)
    if mag is None or mag.size == 0:
        return None, freq

    # Threshold between the noise floor and the signal peak. OOK captures have
    # a large high/low separation, so a fixed fraction of the range is robust;
    # hysteresis is unnecessary given the >15 dB SNR of these test files.
    lo = float(np.percentile(mag, 50))
    hi = float(np.percentile(mag, 99))
    if hi - lo < 1.0:  # no amplitude modulation (FSK / dead air)
        return None, freq
    thr = lo + 0.3 * (hi - lo)
    hyst = mag > thr

    pulses = _trim_gaps(_runs_to_pulses(hyst, us_per_sample))
    if len(pulses) > MAX_KEEP_PULSES:
        pulses = pulses[:MAX_KEEP_PULSES]
    if not looks_like_ook(pulses):
        return None, freq
    return pulses, freq


def looks_like_ook(pulses: List[int]) -> bool:
    """True when the train has enough OOK structure to be a real signal."""
    if pulses is None or len(pulses) < MIN_TRANSITIONS:
        return False
    highs = [p for p in pulses if p > 0]
    if len(highs) < MIN_TRANSITIONS // 2:
        return False
    # A single constant width with no variation is not a keyed message.
    return len(set(highs)) > 1


if __name__ == "__main__":
    import argparse

    ap = argparse.ArgumentParser(description="Demodulate an rtl_433 .cu8 to a pulse train")
    ap.add_argument("file", type=Path)
    args = ap.parse_args()
    pulses, freq = demod_file(args.file)
    if pulses is None:
        print(f"{args.file.name}: no OOK pulse train (FSK/degenerate)")
    else:
        print(f"{args.file.name}: freq={freq} pulses={len(pulses)}")
        print("first 24:", pulses[:24])