#!/usr/bin/env python3
"""Clock recovery + payload extraction for the Uster 868 MHz sensor.

Sweeps the symbol rate and keeps the one that makes every burst yield the same
payload length, then extracts the bytes and tests weather-value encodings.
"""
import numpy as np, json, collections
from scipy.signal import firwin, lfilter

FS = 2.048e6
PATH = "/home/tmco/sdr/deep/a_8676_g40.iq"
SYNC = "0010110111010100"          # 0x2D 0xD4


def load(path):
    raw = np.fromfile(path, 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].copy()
    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)
    return lfilter(firwin(101, 20e3 / (FS / 2)), 1, inst)


def bits_at(f, rate):
    bb = f > 0
    period = FS / rate
    edges = np.where(np.diff(bb.astype(np.int8)) != 0)[0]
    if len(edges) < 10:
        return None
    # phase-lock: start half a period after the first transition
    idx = np.arange(edges[0] + period / 2, len(bb) - 1, period).astype(int)
    return "".join(map(str, bb[idx].astype(int)))


def payload_after_sync(b):
    k = b.find(SYNC)
    if k < 0:
        return None
    rest = b[k + len(SYNC):]
    # payload ends where the idle alternating pattern resumes
    for i in range(len(rest) - 24):
        w = rest[i:i + 24]
        if all(w[t] != w[t - 1] and w[t] == w[t - 2] for t in range(2, 24)):
            return rest[:i]
    return rest


def main():
    x = load(PATH)
    ev = events(x)
    print(f"{len(ev)} bursts")
    feats = [discriminator(x, s, e) for s, e in ev[:24]]

    best = None
    print("\nrate scan — want one payload length dominating")
    for rate in np.arange(22400, 23800, 40):
        lens = []
        for f in feats:
            b = bits_at(f, rate)
            if not b:
                continue
            p = payload_after_sync(b)
            if p is not None and len(p) >= 8:
                lens.append(len(p))
        if len(lens) < 8:
            continue
        c = collections.Counter(lens)
        top, n = c.most_common(1)[0]
        score = n / len(lens)
        if best is None or score > best[0]:
            best = (score, rate, top, n, len(lens))
        if score > 0.55:
            print(f"  {rate:6.0f} Bd -> length {top:3d} bits in {n}/{len(lens)} ({score*100:.0f}%)")
    score, rate, top, n, tot = best
    print(f"\nbest: {rate:.0f} Bd, payload {top} bits in {n}/{tot} bursts ({score*100:.0f}%)")

    # extract with the winning rate
    pays = []
    for f in feats:
        b = bits_at(f, rate)
        if not b:
            continue
        p = payload_after_sync(b)
        if p is not None and abs(len(p) - top) <= 1:
            pays.append(p[:top])
    print(f"extracted {len(pays)} payloads of {top} bits")
    c = collections.Counter(pays)
    print(f"distinct: {len(c)}")
    for k, v in c.most_common(4):
        hx = "".join(f"{int(k[i:i+4], 2):X}" for i in range(0, len(k) // 4 * 4, 4))
        print(f"   x{v}: {k[:64]}  = 0x{hx}")

    out = {"rate_bd": float(rate), "payload_bits": int(top),
           "payloads": [{"bits": k, "n": v,
                         "hex": "".join(f"{int(k[i:i+4], 2):X}" for i in range(0, len(k) // 4 * 4, 4))}
                        for k, v in c.most_common(10)]}
    json.dump(out, open("/home/tmco/sdr/deep/clock.json", "w"), indent=1)
    print("wrote deep/clock.json")


if __name__ == "__main__":
    main()