#!/usr/bin/env python3
"""2-FSK demodulator for the 868 MHz bursts: slice, clock-recover, dump bits."""
import numpy as np, sys
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from scipy.signal import firwin, lfilter

BG, DIM, ACC, ACC2 = "#0d1117", "#8b949e", "#58a6ff", "#f78166"
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=1.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(seg, fs=FS):
    ph = np.unwrap(np.angle(seg))
    return np.diff(ph) * fs / (2 * np.pi)


def demod(seg, cut=20e3):
    inst = discriminator(seg)
    b = firwin(101, cut / (FS / 2))
    return lfilter(b, 1, inst)


def symbol_rate(f, fs):
    """Fundamental of the sliced waveform -> chip/symbol clock candidates."""
    f = f - f.mean()
    S = np.abs(np.fft.rfft(f * np.hanning(len(f)))) ** 2
    fr = np.fft.rfftfreq(len(f), 1 / fs)
    m = (fr > 500) & (fr < 60e3)
    k = np.argmax(S[m])
    return fr[m][k], S, fr


def bits_from(f, fs, clk, phase=0.0):
    """Sample the sliced waveform at the recovered clock."""
    b = (f > 0).astype(int)
    period = fs / clk
    idx = np.arange(phase * period, len(b) - 1, period).astype(int)
    return b[idx]


def main(path, fc, t_on, out_prefix, cut=20e3):
    x = load(path)
    ev = events(x)
    s, e = int(t_on * FS), int(t_on * FS) + int(0.0259 * FS)
    seg = x[s:e].copy()
    n = np.arange(len(seg))
    # coarse carrier estimate, then refine
    ph = np.unwrap(np.angle(seg))
    a = np.polyfit(np.arange(len(ph)), ph, 1)[0]
    y = seg * np.exp(-2j * np.pi * a / (2 * np.pi) * FS / FS * n / FS)
    y = seg * np.exp(-1j * a * n)
    f_off = a * FS / (2 * np.pi)
    inst = demod(y, cut)
    ss = np.std(inst)

    # slicer threshold from the bimodal histogram
    hist, edges = np.histogram(inst, bins=80)
    thr = 0.0

    print(f"carrier offset {f_off/1e3:.1f} kHz, |IF| rms {ss/1e3:.1f} kHz")
    print(f"levels: p5 {np.percentile(inst,5)/1e3:.1f} kHz, p95 {np.percentile(inst,95)/1e3:.1f} kHz")

    clk, S, fr = symbol_rate(inst, FS)
    print(f"dominant modulation tone: {clk/1e3:.3f} kHz -> alternating-bit symbol rate ~{2*clk/1e3:.2f} kBd")
    b = bits_from(inst, FS, 2 * clk)
    bs = "".join(map(str, b[:400]))
    print(f"bits (symbol rate {2*clk/1e3:.2f} kBd, first 400):\n{bs}")

    # run-length histogram (Manchester-lite detection)
    runs = []
    cur, cnt = inst[0], 0
    for v in (inst > 0):
        if v == cur:
            cnt += 1
        else:
            runs.append(cnt); cur = v; cnt = 1
    runs = np.array(runs[1:-1])
    print(f"run lengths (samples @ {FS/1e6:.3f} MS/s): median {np.median(runs):.0f}, "
          f"p10 {np.percentile(runs,10):.0f}, p90 {np.percentile(runs,90):.0f}")

    fig, ax = plt.subplots(3, 1, figsize=(13, 8), facecolor=BG,
                           gridspec_kw={"height_ratios": [1.4, 1, 1]})
    for a_ in ax:
        a_.set_facecolor(BG); a_.tick_params(colors=DIM, labelsize=8)
        for sp in a_.spines.values():
            sp.set_color("#30363d")
    tt = np.arange(len(inst)) / FS * 1e3
    ax[0].plot(tt, inst / 1e3, color=ACC2, lw=0.4)
    ax[0].axhline(0, color=DIM, lw=0.6, ls="--")
    ax[0].set_xlim(0, 3); ax[0].set_ylabel("kHz"); ax[0].set_xlabel("ms")
    ax[0].set_title("discriminator, low-passed at 20 kHz — two rails = 2-FSK", color="#e6edf3", fontsize=10, loc="left")
    ax[0].set_ylim(-30, 30)
    ax[1].hist(inst / 1e3, bins=100, color=ACC, alpha=0.9)
    ax[1].set_title("amplitude histogram of the discriminator", color="#e6edf3", fontsize=9, loc="left")
    ax[1].set_xlabel("kHz")
    m = (fr > 200) & (fr < 40e3)
    ax[2].semilogy(fr[m] / 1e3, S[m] / S[m].max(), color=ACC2, lw=0.8)
    ax[2].set_xlabel("kHz"); ax[2].set_title("modulation spectrum", color="#e6edf3", fontsize=9, loc="left")
    fig.tight_layout()
    fig.savefig(f"{out_prefix}.png", dpi=130, facecolor=BG, bbox_inches="tight", pad_inches=0.15)
    print("wrote", f"{out_prefix}.png")
    return bs


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