# -*- coding: utf-8 -*-
"""逻辑分析仪波形 .bin → 时频曲线拟合。

文件格式（自动探测）:
  16 通道逻辑分析仪导出：每采样 2 字节（小端 uint16），bit0..15 = CH0..CH15。
  信号通道 = 值 0x0100 的位（bit8），即字节流的奇数索引字节承载 CH8。
  采样率默认 100MS/s（可用 --fs 覆盖；也可自动验证 2000/5000Hz 命中）。

用法:
    python plot_waveform.py <xxx.bin> [--fs 100e6] [--ch 8] [--out <png>]
"""
import argparse
import sys

import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt


def extract_channel(m, ch):
    """从 2 字节/采样（小端 uint16）格式中提取通道 ch 的数字序列。"""
    # 目标位在 16bit 值中的字节位置：bit<8 → 低字节(偶索引)；>=8 → 高字节(奇索引)
    if ch < 8:
        lo = m[0::2]
        return (lo >> ch) & 1
    hi = m[1::2]
    return (hi >> (ch - 8)) & 1


def rising_edges(sig):
    """上升沿索引（数字序列）。"""
    return np.nonzero((sig[1:] == 1) & (sig[:-1] == 0))[0] + 1


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("bin_path")
    ap.add_argument("--fs", type=float, default=None, help="采样率 Hz")
    ap.add_argument("--ch", type=int, default=0, help="信号通道号")
    ap.add_argument("--out", default=None, help="输出png路径")
    args = ap.parse_args()

    m = np.memmap(args.bin_path, dtype=np.uint8, mode="r")
    if len(m) % 2 != 0:
        sys.exit("odd file size, not 2-byte/sample format?")
    nsamp = len(m) // 2
    print("samples:", nsamp)

    ch = args.ch
    sig = extract_channel(m, ch)
    ones = int(sig.sum())
    print("CH%d: %d ones / %d samples (%.2f%%)"
          % (ch, ones, nsamp, 100.0 * ones / nsamp))

    nz = np.nonzero(sig)[0]
    if len(nz) == 0:
        sys.exit("channel has no signal")
    lo, hi = int(nz[0] * 0.97), int(nz[-1] * 1.03)
    print("signal range: sample %d .. %d" % (lo, hi))

    # ---- 采样率：自动验证（检查匀速段 2000/5000Hz 命中） ----
    if args.fs is None:
        candidates = [100e6, 50e6, 200e6, 25e6]
        best = []
        for fs in candidates:
            down = max(1, int(fs // 5000000))
            s = sig[lo:hi:down]
            r = rising_edges(s)
            if len(r) < 20:
                continue
            per = np.diff(r) * down / fs
            f = 1.0 / np.maximum(per, 1e-9)
            f = f[(f > 100) & (f < 20000)]
            if len(f) < 10:
                continue
            h, e = np.histogram(f, bins=400, range=(0, 20000))
            e2 = h[(e[:-1] >= 1940) & (e[:-1] <= 2060)].sum()
            e5 = h[(e[:-1] >= 4850) & (e[:-1] <= 5150)].sum()
            best.append(((e2 + e5) / max(h.sum(), 1), fs))
        best.sort(reverse=True)
        fs = best[0][1]
        print("sample-rate ranking:", ["%g:%.3f" % (f, s) for s, f in best])
        print("using fs = %g Hz" % fs)
    else:
        fs = float(args.fs)

    # ---- 上升沿 → 时频 ----
    r = rising_edges(sig[lo:hi])
    r = r + lo
    t_rise = r / fs
    per = np.diff(t_rise)
    freq = 1.0 / np.maximum(per, 1e-9)
    t_mid = (t_rise[:-1] + t_rise[1:]) / 2.0

    print("rising edges: %d" % len(r))
    print("first pulse @ %.3fs, last @ %.3fs" % (t_rise[0], t_rise[-1]))
    print("total pulses: %d" % len(r))
    print("average freq: %.1f Hz"
          % (len(r) / max(t_rise[-1] - t_rise[0], 1e-9)))

    # ---- 三段结构（按 10ms 桶取中位频率） ----
    buckets = np.arange(t_rise[0] - 0.05, t_rise[-1] + 0.05, 0.01)
    med = []
    bt = []
    for i in range(len(buckets) - 1):
        sel = (t_mid >= buckets[i]) & (t_mid < buckets[i + 1])
        if sel.sum() > 0:
            med.append(np.median(freq[sel]))
            bt.append((buckets[i] + buckets[i + 1]) / 2)
    med = np.array(med)
    bt = np.array(bt)
    active = med > 100
    edges = np.diff(active.astype(int))
    starts = bt[1:][edges == 1]
    stops = bt[1:][edges == -1]
    if len(bt) > 0 and active[0]:
        starts = np.concatenate(([bt[0]], starts))
    if len(bt) > 0 and active[-1]:
        stops = np.concatenate((stops, [bt[-1]]))
    print("segments:")
    for s, e in zip(starts, stops):
        sel = (t_rise >= s) & (t_rise <= e)
        n = sel.sum()
        print("  %.3f ~ %.3fs  dur=%.3fs  pulses=%d  avg=%.0fHz"
              % (s, e, e - s, n, n / max(e - s, 1e-9)))
    for i in range(len(stops) - 1):
        print("  gap seg%d->%d: %.3fms" % (i + 1, i + 2,
                                           (starts[i + 1] - stops[i]) * 1e3))

    # ---- 画图 ----
    out = args.out or (args.bin_path.rsplit(".", 1)[0] + "_timefreq.png")
    fig, axes = plt.subplots(2, 1, figsize=(14, 9), sharex=True,
                             gridspec_kw={"height_ratios": [1, 3]})

    ax = axes[0]
    step = max(1, int(fs // 200000))
    xs = np.arange(lo, hi, step) / fs
    ax.plot(xs, sig[lo:hi:step], lw=0.4, color="steelblue")
    ax.set_ylabel("CH%d raw" % ch)
    ax.set_yticks([0, 1])
    ax.grid(alpha=0.3)

    ax = axes[1]
    ax.plot(t_mid, freq, lw=0.8, color="crimson")
    ax.axhline(2000, color="gray", ls="--", lw=0.8)
    ax.axhline(5000, color="gray", ls="--", lw=0.8)
    ax.set_ylabel("instantaneous freq (Hz)")
    ax.set_xlabel("time (s)")
    ax.set_ylim(0, 6000)
    ax.grid(alpha=0.3)
    ax.legend(["freq", "2000Hz", "5000Hz"], loc="upper right", fontsize=8)

    fig.suptitle("CH%d time-frequency: %d pulses, fs=%g Hz, ch=%d"
                 % (ch, len(r), fs, ch))
    fig.tight_layout()
    fig.savefig(out, dpi=130)
    print("saved:", out)


if __name__ == "__main__":
    main()