#!/usr/bin/env python3
"""Bulk decoder: stream a capture, decode every burst, write one JSON line per frame.

    python3 bulk_demod.py <iq_file> <centre_mhz> <out.jsonl>

Chunked so a multi-GB capture never needs to be in RAM. Symbol period is re-estimated per
burst (cheap, from the crossing spacing) and the sampling phase is chosen by matching the
0x2DD4 sync, so no drift across the frame.
"""
import numpy as np, json, sys, os
from scipy.signal import firwin, lfilter

FS = 2.048e6
SYNC = np.array([int(c) for c in "0010110111010100"])
CUT = 40e3
CHUNK_S = 5.0
MARGIN_S = 0.12
LPF = firwin(101, CUT / (FS / 2))


def discriminator(x):
    ph = np.unwrap(np.angle(x))
    a = np.polyfit(np.arange(len(ph)), ph, 1)[0]
    y = x * np.exp(-1j * a * np.arange(len(x)))
    inst = np.diff(np.unwrap(np.angle(y))) * FS / (2 * np.pi)
    return lfilter(LPF, 1, inst), a * FS / (2 * np.pi)


def crossings(f):
    sg = np.signbit(f)
    idx = np.where(sg[1:] != sg[:-1])[0]
    if not len(idx):
        return np.array([])
    y0, y1 = f[idx], f[idx + 1]
    return idx + np.clip(-y0 / (y1 - y0 + 1e-12), 0, 1)


def period(f):
    c = crossings(f)
    if len(c) < 20:
        return None
    d = np.diff(c)
    d = d[(d > 40) & (d < 130)]
    if len(d) < 10:
        return None
    m = float(np.median(d))
    runs, cur = [], [c[0]]
    for i, ok in enumerate(np.abs(np.diff(c) - m) < 0.12 * m):
        if ok:
            cur.append(c[i + 1])
        else:
            if len(cur) >= 6:
                runs.append(np.array(cur))
            cur = [c[i + 1]]
    if len(cur) >= 6:
        runs.append(np.array(cur))
    if not runs:
        return None
    longest = max(runs, key=len)
    T = np.polyfit(np.arange(len(longest)), longest, 1)[0]
    return T if FS / 26000 < T < FS / 20000 else None


def sample(f, T, phase, nsym):
    n = len(f)
    centres = phase + (np.arange(nsym) + 0.5) * T
    lo = np.clip((centres - 0.30 * T).astype(int), 0, n - 1)
    hi = np.clip((centres + 0.30 * T).astype(int), 1, n)
    cs = np.concatenate(([0.0], np.cumsum(f)))
    return (cs[hi] - cs[lo]) / np.maximum(hi - lo, 1)


def decode_burst(x):
    f, f_off = discriminator(x)
    T = period(f)
    if T is None:
        return None
    nsym = int(len(f) / T) - 1
    if nsym < 100:
        return None
    best = None
    for phase in np.arange(0.0, T, T / 24):
        soft = sample(f, T, phase, nsym)
        bits = (soft > 0).astype(int)
        if len(bits) < len(SYNC) + 8:
            continue
        win = np.lib.stride_tricks.sliding_window_view(bits, len(SYNC))
        eq = (win == SYNC).all(axis=1)
        if not eq.any():
            continue
        idx = int(np.argmax(eq))
        conf = float(np.mean(np.abs(soft[max(0, idx - 16):idx + 1])))
        score = conf - 0.02 * idx
        if best is None or score > best[0]:
            best = (score, idx, bits, soft)
    if best is None:
        return None
    _, idx, bits, soft = best
    b = "".join(map(str, bits))
    start = idx + len(SYNC)
    # payload ends where >=24 bits of alternation resume
    j = start
    while j < len(b) - 24:
        w = b[j:j + 24]
        if all(w[t] != w[t - 1] and w[t] == w[t - 2] for t in range(2, 24)):
            break
        j += 1
    payload = b[start:j]
    return {"sync": idx, "T_us": T / FS * 1e6, "f_off": float(f_off),
            "payload": payload, "plen": len(payload),
            "conf": float(np.mean(np.abs(soft[start:start + len(payload)]))) if payload else 0.0}


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 main(path, fc, out):
    """Walk the file in CHUNK_S steps; each read extends MARGIN_S past the step so bursts
    straddling a boundary are still decoded exactly once."""
    total = os.path.getsize(path) // 2
    pos = 0
    recs = 0
    step = int(CHUNK_S * FS)
    span = int((CHUNK_S + 2 * MARGIN_S) * FS)
    with open(path, "rb") as fh, open(out, "w") as fo:
        while pos < total:
            fh.seek(pos * 2)
            buf = fh.read(span * 2)
            if len(buf) < 4096:
                break
            raw = np.frombuffer(buf, dtype=np.uint8).astype(np.float32)
            iq = (raw - 127.5) / 127.5
            x = iq[0::2] + 1j * iq[1::2]
            base_s = pos / FS
            for s, e in events(x):
                if s < MARGIN_S * FS:              # belonging to the previous chunk
                    continue
                if s > step + MARGIN_S * 0.5 * FS:  # beyond this step; next chunk takes it
                    continue
                r = decode_burst(x[s + 100:e - 100])
                if r is None:
                    continue
                r["t"] = round(base_s + s / FS, 3)
                r["mhz"] = round(fc + r["f_off"] / 1e6, 5)
                fo.write(json.dumps(r) + "\n")
                recs += 1
            pos += step
            print(f"  {pos/FS:7.1f} s  frames {recs}", flush=True)
    print(f"done: {recs} frames -> {out}")


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