#!/usr/bin/env python3
"""Demodulate every burst in a capture: bits, preamble length, frame comparison."""
import numpy as np, json, sys
from scipy.signal import firwin, lfilter

FS = 2.048e6


def load(p):
    raw = np.fromfile(p, dtype=np.uint8).astype(np.float32)
    iq = (raw - 127.5) / 127.5
    return iq[0::2] + 1j * iq[1::2]


def events(x, thr_frac=0.25, min_ms=2.0, k=20):
    ps = np.convolve(np.abs(x) ** 2, np.ones(k) / k, mode="same")
    base, peak = np.percentile(ps, 10), np.percentile(ps, 99.9)
    on = ps > base + thr_frac * (peak - base)
    d = np.diff(on.astype(np.int8))
    st, en = np.where(d == 1)[0], np.where(d == -1)[0]
    if len(en) < len(st):
        en = np.append(en, len(ps) - 1)
    return [(s, e) for s, e in zip(st, en) if (e - s) / FS * 1e3 >= min_ms]


def demod_burst(x, s, e):
    seg = x[s + 100:e - 100].copy()
    ph = np.unwrap(np.angle(seg))
    a = np.polyfit(np.arange(len(ph)), ph, 1)[0]
    f_off = a * FS / (2 * np.pi)
    y = seg * np.exp(-1j * a * np.arange(len(seg)))
    inst = np.diff(np.unwrap(np.angle(y))) * FS / (2 * np.pi)
    b = firwin(101, 20e3 / (FS / 2))
    f = lfilter(b, 1, inst)
    return f, f_off


def symbols(f, fs=FS, n_pre=200000):
    """Clock from the alternating preamble: find the dominant tone, sample at 2x it."""
    d = f - f.mean()
    S = np.abs(np.fft.rfft(d * np.hanning(len(d)))) ** 2
    fr = np.fft.rfftfreq(len(d), 1 / fs)
    m = (fr > 2e3) & (fr < 40e3)
    tone = fr[m][np.argmax(S[m])]
    rate = 2 * tone                     # alternating bits -> half the symbol rate
    period = fs / rate
    # phase-align on the first strong transition
    b = f > 0
    edges = np.where(np.diff(b.astype(np.int8)) != 0)[0]
    if len(edges) < 10:
        return None, rate
    ph0 = edges[0]
    idx = np.arange(ph0 + period / 2, len(b) - 1, period).astype(int)
    return b[idx].astype(int), rate


def main(path, fc):
    x = load(path)
    ev = events(x)
    print(f"{path}: {len(ev)} bursts in {len(x)/FS:.0f} s")
    out = []
    for (s, e) in ev:
        f, f_off = demod_burst(x, s, e)
        bits, rate = symbols(f)
        if bits is None or len(bits) < 40:
            continue
        out.append({"t": s / FS, "f": round(float(fc + f_off / 1e6), 5), "rate": round(float(rate)),
                    "bits": "".join(map(str, bits))})
    print(f"demodulated: {len(out)}")
    if not out:
        return
    rates = np.array([o["rate"] for o in out])
    print(f"symbol rate: median {np.median(rates):.0f} Bd (p10 {np.percentile(rates,10):.0f}, p90 {np.percentile(rates,90):.0f})")
    # how much of each frame is the alternating preamble
    def pre_len(b):
        k = 0
        for i in range(2, min(len(b), 400)):
            if b[i] == b[i - 2] and b[i - 1] != b[i]:
                k = i
            else:
                break
        return k
    pres = [pre_len(o["bits"]) for o in out]
    print(f"alternating preamble length: median {np.median(pres):.0f} symbols "
          f"({np.median(pres)/np.median(rates)*1e3:.2f} ms)")
    seen = {}
    for o in out:
        seen.setdefault(o["bits"][:160], 0)
        seen[o["bits"][:160]] += 1
    print(f"distinct 160-symbol prefixes: {len(seen)} for {len(out)} bursts")
    top = sorted(seen.items(), key=lambda kv: -kv[1])[:3]
    for pre, n in top:
        print(f"   x{n}: {pre[:96]}")
    json.dump(out, open("deep/frames.json", "w"), indent=1)
    print("wrote deep/frames.json")


if __name__ == "__main__":
    main(sys.argv[1], float(sys.argv[2]))