#!/usr/bin/env python3
"""Global-grid demodulator: fit one symbol period from the alternating runs, then sample.

The preamble (63 symbols) and the idle tail (~450 symbols) are both pure alternation, so their
zero crossings sit on an exact k*T grid. Fitting T over those runs gives a period accurate to
~1e-5, which is what a 594-symbol frame needs. Bits are sampled at the midpoint of each symbol.

Prints, per burst: preamble length, sync index, payload bits and per-bit confidence.
"""
import numpy as np, json, sys
from scipy.signal import firwin, lfilter

FS = 2.048e6
SYNC = "0010110111010100"
CUT = 20e3


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 discriminator(x, s, e):
    seg = x[s + 100:e - 100].astype(np.complex64)
    ph = np.unwrap(np.angle(seg))
    a = np.polyfit(np.arange(len(ph)), ph, 1)[0]
    y = seg * np.exp(-1j * a * np.arange(len(seg)))
    inst = np.diff(np.unwrap(np.angle(y))) * FS / (2 * np.pi)
    f = lfilter(firwin(101, CUT / (FS / 2)), 1, inst)
    return f, a * FS / (2 * np.pi)


def crossings(f):
    s = np.signbit(f)
    idx = np.where(s[1:] != s[:-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 alternating_runs(c, m, tol=0.12):
    """Split crossings into runs whose spacing is ~m (i.e. alternating bit regions)."""
    if len(c) < 4:
        return []
    d = np.diff(c)
    good = np.abs(d - m) < tol * m
    runs, cur = [], [c[0]]
    for i, ok in enumerate(good):
        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))
    return runs


def fit_run(run):
    """Least-squares slope (period) and phase of a crossing run."""
    k = np.arange(len(run))
    T, p0 = np.polyfit(k, run, 1)
    return T, p0


def decode_burst(x, s, e, verbose=False):
    f, f_off = discriminator(x, s, e)
    n = len(f)
    c = crossings(f)
    if len(c) < 20:
        return None
    d = np.diff(c)
    d = d[(d > 40) & (d < 130)]
    if not len(d):
        return None
    m = float(np.median(d))
    runs = alternating_runs(c, m)
    if not runs:
        return None
    # period from the longest run, phase from the earliest run (the preamble)
    longest = max(runs, key=len)
    T, _ = fit_run(longest)
    first = min(runs, key=lambda r: r[0])
    _, phase = fit_run(first)
    # sanity: period must give 20-26 kBd
    if not (FS / 26000 < T < FS / 20000):
        return None
    # residual check on the longest run using the global T
    k = np.arange(len(longest))
    resid = longest - (longest[0] + k * T)
    if np.abs(resid).max() > 0.35 * T:
        return None

    # sample every symbol, centred half a period after the first crossing
    centres = phase + (np.arange(0, int((n - phase) / T)) + 0.5) * T
    centres = centres[(centres > 2) & (centres < n - 2)]
    lo = np.clip((centres - 0.30 * T).astype(int), 0, n - 1)
    hi = np.clip((centres + 0.30 * T).astype(int), 1, n)
    soft = np.array([f[a:b].mean() if b > a else 0.0 for a, b in zip(lo, hi)])
    bits = (soft > 0).astype(int)
    bstr = "".join(map(str, bits))
    ks = bstr.find(SYNC)
    return {"T": T, "rate": FS / T, "bits": bstr, "soft": soft, "sync": ks,
            "n": len(bstr), "runs": len(runs), "f_off": f_off,
            "resid": float(np.abs(resid).max())}


def main(path, fc, limit=200):
    x = load(path)
    ev = events(x)
    res = []
    for s, e in ev[:limit]:
        r = decode_burst(x, s, e)
        if r:
            r["t"] = s / FS
            r["mhz"] = fc + r["f_off"] / 1e6
            res.append(r)
    ok = [r for r in res if r["sync"] >= 0]
    print(f"{path}: {len(ev)} bursts, decoded {len(res)}, sync found in {len(ok)}")
    if not ok:
        return res
    Ts = np.array([r["T"] for r in ok])
    print(f"symbol period: {np.median(Ts)/FS*1e6:.3f} us  -> {FS/np.median(Ts):.0f} Bd "
          f"(spread {Ts.std()/FS*1e6:.3f} us)")
    print(f"preamble lengths: {sorted(set(r['sync'] for r in ok))}")
    lens = [r["n"] - r["sync"] for r in ok]
    print(f"bits after sync: {sorted(set(lens))[:10]} (median {int(np.median(lens))})")
    conf = np.concatenate([np.abs(r["soft"]) for r in ok])
    print(f"per-bit |soft|: median {np.median(conf)/1e3:.1f} kHz, p10 {np.percentile(conf,10)/1e3:.1f} kHz")
    return res


if __name__ == "__main__":
    res = main(sys.argv[1], float(sys.argv[2]))
    json.dump([{k: v for k, v in r.items() if k != "soft"} for r in res],
              open("/home/tmco/sdr/deep/grid_frames.json", "w"), indent=1)
    print("wrote deep/grid_frames.json")