#!/usr/bin/env python3
"""Phase-locked demodulator: the sync word itself picks the sampling phase.

Given the precise symbol period (fit from the alternating runs), every sub-symbol phase is
tried and the one that reproduces 0x2DD4 exactly is kept. That removes the last degree of
freedom, so every burst decodes to the same alignment.
"""
import numpy as np, json, sys
from scipy.signal import firwin, lfilter

FS = 2.048e6
SYNC = np.array([int(c) for c in "0010110111010100"])
CUT = 40e3


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):
    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 symbol_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)
    return T if FS / 26000 < T < FS / 20000 else None


def sample_bits(f, T, phase, nsym):
    """Integrate the discriminator over +/-0.3T around each symbol centre."""
    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)))
    soft = (cs[hi] - cs[lo]) / np.maximum(hi - lo, 1)
    return (soft > 0).astype(int), soft, centres


def best_phase(f, T):
    """Search phases over one symbol; score by exact sync matches."""
    nsym = int(len(f) / T) - 1
    best = None
    for phase in np.arange(0.0, T, T / 48):
        bits, soft, _ = sample_bits(f, T, phase, nsym)
        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))            # first exact sync occurrence
        conf = float(np.mean(np.abs(soft[max(0, idx - len(SYNC)):idx + 1])))
        # prefer an early sync (it sits at the end of the preamble) and a confident one
        score = conf - 0.05 * idx
        if best is None or score > best[0]:
            best = (score, phase, idx)
    return best


def decode_burst(x, s, e):
    f, f_off = discriminator(x, s, e)
    T = symbol_period(f)
    if T is None:
        return None
    bp = best_phase(f, T)
    if bp is None:
        return None
    conf, phase, idx = bp
    nsym = int(len(f) / T) - 1
    bits, soft, _ = sample_bits(f, T, phase, nsym)
    bstr = "".join(map(str, bits))
    return {"T": T, "rate": FS / T, "bits": bstr, "soft": soft, "sync": idx,
            "f_off": f_off, "conf": conf, "n": len(bstr)}


def payload_of(r):
    """Bits between the sync word (ends at r['sync']) and the idle alternation."""
    b = r["bits"]
    start = r["sync"] + 1                    # index in the SYNC-extended array
    start = max(0, start - 1)
    i = start
    while i < len(b) - 24:
        w = b[i:i + 24]
        if all(w[t] != w[t - 1] and w[t] == w[t - 2] for t in range(2, 24)):
            return b[start:i]
        i += 1
    return b[start:]


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
            r["payload"] = payload_of(r)
            res.append(r)
    print(f"{path}: {len(ev)} bursts, {len(res)} phase-locked")
    if not res:
        return res
    Ts = np.array([r["T"] for r in res])
    print(f"period {np.median(Ts)/FS*1e6:.3f} us ({FS/np.median(Ts):.0f} Bd), spread {Ts.std()/FS*1e6:.4f} us")
    print(f"sync index values: {sorted(set(r['sync'] for r in res))[:8]}")
    pl = [len(r["payload"]) for r in res]
    print(f"payload length: {sorted(set(pl))[:8]} (median {int(np.median(pl))})")
    import collections
    c = collections.Counter(r["payload"] for r in res)
    print(f"distinct payloads: {len(c)} of {len(res)}")
    for k, v in c.most_common(5):
        hx = "".join(f"{int(k[i:i+4],2):X}" for i in range(0, len(k)//4*4, 4))
        print(f"   x{v:2d}  {len(k):3d} bits  0x{hx}")
    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/locked_frames.json", "w"), indent=1)
    print("wrote deep/locked_frames.json")