#!/usr/bin/env python3
# SPDX-License-Identifier: MIT
# Copyright (c) 2026 Olympian LED
"""step_kit: the absolute bandwidth of one oscilloscope channel, from a fast-edge pulser, over SCPI only.

One file, Python 3.8+ with numpy. It talks SCPI over the raw socket (TCP 5555 by default). It is written for, and
tested against, the SCPI dialect of the Rigol DHO800, DHO900 and MHO900.

    python3 step_kit.py measure  --host SCOPE --ch 1 [--imp 50] [--band OFF] [--records 8] [--out run.json]
    python3 step_kit.py analyse  run.npz [--edge auto|rise|fall]       (the same arithmetic on saved records)
    python3 step_kit.py selftest                                         (synthetic records of known responses)
    python3 step_kit.py report   run.json [--csv PREFIX]                 (the table)

The method (METHOD.md has it step by step): the pulser free-runs on its own clock, so its edges fall at every phase
of the scope's sample clock. Each record's edges are found, their times fitted to one period, and every sample is
placed by its time from its edge on a grid of 1/os of a sample interval (equivalent time). The folded step's
derivative is the channel's impulse response; its spectrum over its own DC level is |H(f)|, absolute. From it: the
-3 dB frequency, the 10-90 % rise time, the overshoot, a noise floor (the same fold a quarter period later), and an
upper bound of the edge jitter that smears the fold.

`measure` reads every setting it will change first and puts each one back at the end (on success, on an error and
on Ctrl-C), then reads them all again and lists any that still differ. It never sends *RST, a calibration, a
screenshot or a file command. Every read has a timeout; after a timeout it opens a new connection and checks the
scope still answers *IDN?.

Exit codes: 0 measured and restored; 1 refused or failed, restored; 2 usage error or an untested model;
3 the restore is incomplete; 4 the scope stopped answering; 130 Ctrl-C, restored.
"""
import argparse
import csv
import datetime
import hashlib
import json
import math
import re
import socket
import sys
import time
from pathlib import Path

import numpy as np

VERSION = "1.0.0"
SCPI_PORT = 5555
PAGE = 250_000  # points per :WAVeform:DATA? page


class KitError(Exception):
    """A failure the kit reports and exits on (after restoring)."""


class Refused(KitError):
    """The measurement refused a record or a setting (exit 1)."""


class ScpiTimeout(KitError):
    pass


class ScpiHung(KitError):
    """The scope stopped answering *IDN? on a fresh connection: nothing more can be sent, nor restored."""


def utc_iso(t):
    return datetime.datetime.fromtimestamp(t, datetime.timezone.utc).isoformat(timespec="milliseconds").replace("+00:00", "Z")


class Clock:
    """Time for everything but socket deadlines; the tests swap in a fake one that never really sleeps."""

    def now(self):
        return time.monotonic()

    def wall(self):
        return time.time()

    def sleep(self, s):
        if s > 0:
            time.sleep(s)


# ================================================================================================ the analysis
# Transcribed from our bench tool's `step` method, arithmetic unchanged.
#
# The pulser's square, sampled at dt, is folded on its own period: each record's rising 50 % crossings are fitted
# as t0 + n P + q (n - n_mean)^2, and every sample within the window about an edge is placed by its time from that
# edge on a grid of 1/os sample. The pulser's clock is not locked to the scope's, so its edges fall at every phase
# and the grid fills from many records, so nothing above the ADC's Nyquist frequency aliases (each sample is the
# analog step at a known instant). The step's derivative is the channel's impulse response, its spectrum the
# channel's |H(f)| over its own DC level. A Gaussian edge of rise time tr is down 0.035 dB at 1 GHz at 30 ps
# (exp(-(2 pi f tr / 2.563)^2 / 2)): --source-rise-ps divides it out when it matters.

def square_levels(x):
    """The two levels of a square record (each half's median about the 1 / 99 % percentiles' middle)."""
    p1, p99 = np.percentile(x, [1, 99])
    mid = 0.5 * (p1 + p99)
    lo, hi = x[x < mid], x[x >= mid]
    if not len(lo) or not len(hi):
        return float(p1), float(p99)
    return float(np.median(lo)), float(np.median(hi))


def rising_edges(x, lo, hi):
    """Sub-sample times (in samples) of the rising 50 % crossings, with 30 / 70 % hysteresis so noise near the level
    is one crossing; the last crossing before the record leaves the 30 % band."""
    mid = 0.5 * (lo + hi)
    st = np.zeros(len(x), np.int8)
    st[x < lo + 0.3 * (hi - lo)] = -1
    st[x > lo + 0.7 * (hi - lo)] = 1
    nz = np.nonzero(st)[0]
    s = st[nz]
    turn = (s[1:] == 1) & (s[:-1] == -1)
    out = []
    for i0, i1 in zip(nz[:-1][turn], nz[1:][turn]):
        seg = x[i0:i1 + 1]
        k = np.nonzero((seg[:-1] < mid) & (seg[1:] >= mid))[0]
        if not len(k):
            continue
        k = int(i0 + k[-1])
        out.append(k + (mid - x[k]) / (x[k + 1] - x[k]))
    return np.array(out, float)


def fit_edges(t):
    """The pulser's edge times t (samples) as t0 + n P + q (n - n_mean)^2 over each edge's cycle number n (a slow
    frequency drift between the pulser's and the scope's clocks in q); outliers past 5 sigma dropped once. Returns
    (n, fitted times, residual rms in samples, period in samples)."""
    if len(t) < 8:
        raise ValueError(f"{len(t)} edges in the record: too few to fold (a deeper record or a faster pulser)")
    p0 = float(np.median(np.diff(t)))
    n = np.round((t - t[0]) / p0).astype(np.int64)
    keep = np.r_[True, np.diff(n) > 0]  # one edge per cycle (a false crossing shares its cycle: dropped)
    t, n = t[keep], n[keep]
    for _ in range(2):
        A = np.column_stack([np.ones(len(n)), n.astype(float), (n - n.mean()) ** 2])
        c, *_ = np.linalg.lstsq(A, t, rcond=None)
        r = t - A @ c
        sd = float(np.std(r))
        ok = np.abs(r) <= max(5 * sd, 1e-6)
        if ok.all():
            break
        t, n = t[ok], n[ok]
    return n, A @ c, sd, float(c[1])


def fold_step(x, T, pre, post, os, sel=None):
    """The samples of x about each fitted edge time T (samples), on a grid of 1 / os samples from -pre to +post
    samples: per bin the sum of the values, the sum of their times from the edge and the count. sel(J) keeps a subset
    of sample indices (one ADC core's)."""
    L = int(np.ceil(pre + post)) + 2
    J = np.floor(T - pre).astype(np.int64)[:, None] + np.arange(L)[None, :]
    tau = J - T[:, None]
    ok = (tau >= -pre) & (tau < post) & (J >= 0) & (J < len(x))
    if sel is not None:
        ok &= sel(J)
    nb = int(round((pre + post) * os))
    tv = tau[ok]
    idx = np.floor((tv + pre) * os).astype(np.int64)
    good = (idx >= 0) & (idx < nb)
    s = np.bincount(idx[good], weights=x[J[ok]][good], minlength=nb)
    st = np.bincount(idx[good], weights=tv[good], minlength=nb)
    c = np.bincount(idx[good], minlength=nb).astype(float)
    return np.stack([s, st, c])


def fold_wave(F, pre, os):
    """The folded step on the regular grid (-pre + (i + 0.5) / os samples): each bin's mean value at its mean time,
    interpolated -- a bin's samples need not sit at its centre (a pulser near a rational multiple of the sample
    interval puts them at a few phases only), and the box-car's centre would shift the step by up to half a bin."""
    s, st, c = F
    ok = c > 0
    tm, ym = st[ok] / c[ok], s[ok] / c[ok]
    grid = -pre + (np.arange(len(c)) + 0.5) / os
    return np.interp(grid, tm, ym)


def gaussian_edge_db(f, rise_s):
    """A Gaussian edge's own spectrum (dB) at f for its 10-90 % rise time."""
    sig = rise_s / 2.563
    return -10 * np.log10(np.e) * (2 * np.pi * np.asarray(f) * sig) ** 2


# summarise(): a point more than this far (dB) below both its neighbours is a one-bin dip, never the -3 dB corner
SPUR_DB = 1.0


def summarise(freqs, db, snr=None, min_snr=10.0, d0=None):
    """-3 dB (relative to the lowest frequency's level, or to d0: the step's 0 dB is its own DC level), peaking,
    monotone, flatness to half the -3 dB frequency."""
    f = np.array(freqs, float)
    d = np.array(db, float)
    ok = np.ones(len(d), bool) if snr is None else np.array(snr, float) >= min_snr
    f, d = f[ok], d[ok]
    if len(d) < 3:
        return None
    d0 = d[0] if d0 is None else d0
    rel = d - d0

    # the first point at or below -3 dB that is not a one-bin dip: a bin more than SPUR_DB below both its neighbours
    # is a spur or interference, not the corner
    def dip(k):
        return 0 < k < len(rel) - 1 and rel[k - 1] - rel[k] > SPUR_DB and rel[k + 1] - rel[k] > SPUR_DB

    below = [k for k in range(len(rel)) if rel[k] <= -3.0 and not dip(k)]
    if len(below) == 0:
        f3 = None
    else:
        i = below[0]
        f3 = float(f[0]) if i == 0 else float(np.exp(np.interp(-3.0, [rel[i], rel[i - 1]], [np.log(f[i]), np.log(f[i - 1])])))
    upto = f <= (f3 if f3 else f[-1])
    peak = float(np.max(rel[upto])) if upto.any() else 0.0
    half = f <= (0.5 * f3 if f3 else f[-1])
    flat = float(np.ptp(rel[half])) if half.sum() > 1 else 0.0
    # monotone past the peak to the -20 dB point (0.3 dB tolerance against measurement noise)
    stop = np.where(rel <= -20.0)[0]
    j = stop[0] if len(stop) else len(rel)
    seg = rel[int(np.argmax(rel[:j])) if j else 0:j]
    mono = bool(np.all(np.diff(seg) <= 0.3)) if len(seg) > 1 else True
    return dict(dc_db=round(float(d0), 2), f3db=f3, peaking_db=round(peak, 2), flat_db=round(flat, 2), monotone=mono)


def step_analyse(records, dt, os=16, pre_ns=15.0, post_ns=45.0, cores=0, points=80, fmin=0.0, fmax=0.0, source_rise_ps=0.0):
    """The absolute response of one channel from records of the pulser's square (volts, sample interval dt): the
    equivalent-time step, its derivative's spectrum over its DC level (sinc of the os grid divided out), a noise
    floor from the same fold a quarter period away (the flat top), rise time, overshoot, the edge-time residual (an
    upper bound of the jitter that smears the fold) and, with cores=N, each of N interleaved ADC cores' own step: its
    50 % crossing against the whole fold's (the core's timing skew) and its height. Raises ValueError when the record
    cannot be folded (few edges, or the pulser's period locked to the sample clock so the grid is not covered).

    The defaults are the ones `measure` runs (a 15 ns + 45 ns window, 80 log-spaced frequencies)."""
    pre, post = pre_ns * 1e-9 / dt, post_ns * 1e-9 / dt
    S = C = N = None
    per, resid, nedges, levels = [], [], 0, []
    fits, turns = [], []
    for x in records:
        lo, hi = square_levels(x)
        if hi - lo <= 0:
            raise ValueError("a flat record: no square on the channel")
        t = rising_edges(x, lo, hi)
        n, T, sd, P = fit_edges(t)
        if post + pre > 0.45 * P:
            raise ValueError(f"the window ({pre_ns:g} + {post_ns:g} ns) does not fit in half the pulser's period "
                             f"({P * dt * 1e9:.1f} ns): a slower pulser or a shorter --post-ns")
        F = fold_step(x, T, pre, post, os)
        # the noise floor: the same fold a quarter period later (inside the flat top, away from both edges)
        Fn = fold_step(x, T + 0.25 * P, pre, post, os) if (pre + post) < 0.2 * P else None
        S = F if S is None else S + F
        if Fn is not None:
            N = Fn if N is None else N + Fn
        per.append(P)
        resid.append(sd)
        nedges += len(T)
        levels.append((lo, hi))
        fits.append((x, T))
        turns.append(phase_turns(P, len(T)))
    C = S[2]
    empty = int(np.sum(C == 0))
    if empty:
        raise ValueError(f"{empty} of {len(C)} grid bins never sampled: the pulser's period ({np.mean(per):.4f} samples) is "
                         f"locked to the sample clock or the records hold too few edges -- change its frequency by a few %")
    y = fold_wave(S, pre, os)
    dtau = dt / os
    n0 = int(round(pre * os))
    base = float(np.mean(y[:max(1, n0 // 2)]))
    top = float(np.mean(y[-max(1, len(y) // 8):]))
    amp = top - base
    lo_m, hi_m = (float(np.mean([v[i] for v in levels])) for i in (0, 1))
    ys = (y - base) / amp

    def cross(z, lvl):
        k = int(np.argmax(z >= lvl))
        return (k - 1 + (lvl - z[k - 1]) / (z[k] - z[k - 1])) if k > 0 else float(k)

    rise_ps = (cross(ys, 0.9) - cross(ys, 0.1)) * dtau * 1e12
    # each edge's own time against the folded step (one Gauss-Newton step of a shift-only fit over its samples from
    # 1.5 ns before to 3 ns after): the spread is the jitter that smears the fold (pulser, sample clock) plus the
    # noise's share, an upper bound; the 50 % crossings the fold was fitted on carry each edge's sample-phase error
    grid = -pre + (np.arange(len(y)) + 0.5) / os
    dy = np.gradient(y, 1.0 / os)
    shifts = []
    a0, a1 = 1.5e-9 / dt, 3e-9 / dt
    for x, T in fits:
        L = int(np.ceil(a0 + a1)) + 2
        J = np.floor(T - a0).astype(np.int64)[:, None] + np.arange(L)[None, :]
        tau = J - T[:, None]
        ok = (tau >= -a0) & (tau < a1) & (J >= 0) & (J < len(x))
        m = np.interp(tau, grid, y)
        g = np.interp(tau, grid, dy)
        r = np.where(ok, x[np.clip(J, 0, len(x) - 1)] - m, 0.0)
        g = np.where(ok, g, 0.0)
        shifts.append(-np.sum(r * g, axis=1) / np.maximum(np.sum(g * g, axis=1), 1e-30))
    sh = np.concatenate(shifts)
    jitter_ps = float(np.std(sh)) * dt * 1e12
    overshoot = float(np.max(ys[n0:]) - 1.0) * 100.0
    h = np.diff(y)
    w = np.ones(len(h))
    q = len(h) // 4
    w[-q:] = 0.5 * (1 + np.cos(np.linspace(0, np.pi, q)))
    nfft = 1 << int(np.ceil(np.log2(len(h) * 4)))
    Hs = np.fft.rfft(h * w, nfft)
    f = np.fft.rfftfreq(nfft, dtau)
    corr = np.sinc(f * dtau)  # the grid's box-car average
    mag = np.abs(Hs) / np.maximum(np.abs(corr), 1e-6)
    H0 = abs(float(np.sum(h * w)))
    db = 20 * np.log10(np.maximum(mag / H0, 1e-12))
    if source_rise_ps:
        db = db - gaussian_edge_db(f, source_rise_ps * 1e-12)
    if N is not None and np.all(N[2] > 0):
        yn = fold_wave(N, pre, os)
        Hn = np.abs(np.fft.rfft(np.diff(yn) * w, nfft)) / np.maximum(np.abs(corr), 1e-6)
        kern = np.ones(9) / 9
        floor = np.sqrt(np.convolve(Hn ** 2, kern, mode="same"))
        snr = 20 * np.log10(np.maximum(mag, 1e-30) / np.maximum(floor, 1e-30))
    else:
        snr = None
    f_lo = max(fmin, 2.0 / ((pre_ns + post_ns) * 1e-9))
    f_hi = min(fmax or 1e12, 0.4 / dtau)
    pick = np.unique(np.searchsorted(f, np.geomspace(f_lo, f_hi, points)))
    pick = pick[pick < len(f)]
    out = dict(freqs=[round(float(f[k]), 1) for k in pick], db=[round(float(db[k]), 3) for k in pick],
               snr=[round(float(snr[k]), 1) for k in pick] if snr is not None else None,
               rise_ps=round(float(rise_ps), 1), overshoot_pct=round(overshoot, 2), step_v=round(amp, 6),
               square_v=round(hi_m - lo_m, 6), settled_db=round(float(20 * np.log10(abs(amp) / max(hi_m - lo_m, 1e-12))), 3),
               edges=nedges, records=len(records), period_ns=round(float(np.mean(per)) * dt * 1e9, 4),
               edge_resid_ps=round(float(np.sqrt(np.mean(np.square(resid)))) * dt * 1e12, 2), jitter_ps=round(jitter_ps, 2),
               grid_ps=round(dtau * 1e12, 2),
               min_bin_count=int(C.min()),
               step_dt=dtau, step=[round(float(v), 5) for v in ys[::max(1, os // 8)]], step_stride=max(1, os // 8),
               step_t0_ps=round(float(grid[0]) * dt * 1e12, 4), phase_turns=round(float(min(turns)), 2))
    s3 = summarise(out["freqs"], out["db"], out["snr"], 10.0, d0=0.0)
    out["f3db"] = s3["f3db"] if s3 else None
    out["peaking_db"] = s3["peaking_db"] if s3 else None
    if out["f3db"]:
        sig = out["jitter_ps"] * 1e-12
        out["jitter_bound_db_at_f3"] = round(float(-10 * np.log10(np.e) * (2 * np.pi * out["f3db"] * sig) ** 2), 3)
        if source_rise_ps == 0.0:
            out["edge_30ps_db_at_f3"] = round(float(gaussian_edge_db(out["f3db"], 30e-12)), 4)
    if cores:
        half = 0.5
        mids, heights = [], []
        for m in range(cores):
            Sm = None
            for x, T in fits:
                Fm = fold_step(x, T, pre, post, os // 2 or 1, sel=lambda J, m=m: (J % cores) == m)
                Sm = Fm if Sm is None else Sm + Fm
            if np.any(Sm[2] == 0):
                mids.append(None)
                heights.append(None)
                continue
            ym = fold_wave(Sm, pre, os // 2 or 1)
            n0m = int(round(pre * (os // 2 or 1)))
            bm, tm = float(np.mean(ym[:max(1, n0m // 2)])), float(np.mean(ym[-max(1, len(ym) // 8):]))
            zm = (ym - bm) / (tm - bm)
            mids.append(cross(zm, half) / (os // 2 or 1))
            heights.append(tm - bm)
        ok = [v for v in mids if v is not None]
        ref_mid = float(np.mean(ok)) if ok else 0.0
        # a core that samples LATE sees the step EARLIER: its sample-time skew is minus the shift of its own fold
        out["cores"] = dict(n=cores, sample_late_ps=[None if v is None else round(float(ref_mid - v) * dt * 1e12, 2) for v in mids],
                            height_db=[None if v is None else round(float(20 * np.log10(v / np.mean([h for h in heights if h]))), 3) for v in heights])
    return out


# below this many phase turns a record (phase_turns), the report warns: the fit cannot average out the crossings'
# sample-phase error, and the fold reads the -3 dB low (METHOD.md, "The pulser's frequency")
MIN_TURNS = 5.0


def phase_turns(P, m, qmax=8):
    """How many times the pulser's edges sweep the sample interval in one record: m edges of period P samples. Over
    each q <= qmax, a period near N + p/q samples puts the edges in q clusters, each advancing by |qP - round(qP)| / q
    a cycle; they cover the interval once every 1 / |qP - round(qP)| edges. The smallest m |qP - round(qP)| over q
    is the record's count. A diagnostic beside the analysis; it changes none of its numbers."""
    return min(m * abs(q * P - round(q * P)) for q in range(1, qmax + 1))


def fast_edge(x, lo, hi):
    """Which edge of a square record is the fast one: 'rise' or 'fall', by the median slope (V a sample) at the 50 %
    crossings of each. A negative-going pulser (the LBE-1322) has its fast edge falling (0 V down to its level) and
    its return rising slower."""
    mid = 0.5 * (lo + hi)
    up = np.nonzero((x[:-1] < mid) & (x[1:] >= mid))[0]
    dn = np.nonzero((x[:-1] >= mid) & (x[1:] < mid))[0]
    su = float(np.median(x[up + 1] - x[up])) if len(up) else 0.0
    sd = float(np.median(x[dn] - x[dn + 1])) if len(dn) else 0.0
    return "fall" if sd > su else "rise"


def pulser_verdict(x):
    """Is this record a free-running pulser? A square of more than 20 mV with a regular period and at least 10
    rising edges. Returns (verdict, lo, hi, text)."""
    lo, hi = square_levels(x)
    t = rising_edges(x, lo, hi) if hi - lo > 0.02 else np.array([])
    reg = float(np.std(np.diff(t)) / np.median(np.diff(t))) if len(t) > 3 else 1.0
    ok = len(t) >= 10 and reg < 0.01 and hi - lo > 0.02
    return ok, lo, hi, f"{hi - lo:.3f} V square, {len(t)} rising edges, period spread {reg:.2g}"


def turns_warning(x, lo, hi):
    """A warning (text) when the pulser's edges sweep the sample interval fewer than MIN_TURNS times in record x."""
    try:
        _, T, _, P = fit_edges(rising_edges(x, lo, hi))
    except ValueError:
        return None
    n = phase_turns(P, len(T))
    if n >= MIN_TURNS:
        return None
    return ("the pulser's period is %.4f samples: its edges sweep the sample interval only %.1f times a record "
            "(%d edges). The fold then reads the -3 dB low. Change the pulser's frequency by a few hundred Hz and run "
            "again (METHOD.md, \"The pulser's frequency\")." % (P, n, len(T)))


def ladder_up(v):
    """The 1-2-5 step at or above v."""
    e = math.floor(math.log10(v))
    for m in (1, 2, 5, 10):
        if m * 10 ** e >= v * 0.999:
            return m * 10 ** e
    return 10 ** (e + 1)


def ladder_down(tb):
    """The next 1-2-5 time base below tb."""
    e = 10.0 ** math.floor(math.log10(tb * 1.0001))
    mant = round(tb / e, 3)
    return (0.5 * e) if mant <= 1.0 else (1.0 * e if mant <= 2.0 else 2.0 * e)


# ================================================================================================ the self-test

def erf_np(z):
    """erf, Abramowitz & Stegun 7.1.26 (|error| < 1.5e-7): numpy has none, and the kit needs only numpy."""
    s = np.sign(z)
    z = np.abs(z)
    t = 1.0 / (1.0 + 0.3275911 * z)
    y = 1.0 - (((((1.061405429 * t - 1.453152027) * t) + 1.421413741) * t - 0.284496736) * t + 0.254829592) * t * np.exp(-z * z)
    return s * y


def synth_pulser(n, dt, period, resp, f3, lo=-0.25, hi=0.25, noise=1e-3, jitter=2e-12, drift_ppm=2.0, skew=None, lsb=None, seed=1,
                 fall=None):
    """A pulser's square (50 % duty) through a known response, sampled at dt: `pole` a single pole at f3, `gauss` a
    Gaussian of -3 dB at f3; edge jitter (s rms), a slow frequency drift between the two clocks, white noise (V rms),
    quantisation (lsb V) and per-core sample-time skews (s; core j % len(skew)). `fall` = (resp, f3) gives the falling
    edge its own response (a pulser whose two edges differ). For the self-test and the tests."""
    rng = np.random.default_rng(seed)
    j = np.arange(n)
    t = j * dt + (np.asarray(skew)[j % len(skew)] if skew is not None else 0.0)
    t0 = 0.37 * period
    ncyc = int(n * dt / period) + 4
    k_all = np.arange(-2, ncyc)
    drift = drift_ppm * 1e-6 * period * (k_all - ncyc / 2) ** 2 / ncyc
    Tr = t0 + k_all * period + drift + rng.normal(0, jitter, len(k_all))
    Tf = Tr + 0.5 * period + rng.normal(0, jitter, len(k_all))

    def edge(resp, f3):
        if resp == "pole":
            tau = 1.0 / (2 * np.pi * f3)
            return lambda u: np.where(u > 0, -np.expm1(-np.clip(u, 0, None) / tau), 0.0)
        sig = np.sqrt(np.log(2)) / (2 * np.pi * f3)
        return lambda u: 0.5 * (1 + erf_np(u / (sig * np.sqrt(2))))
    s = edge(resp, f3)
    sf = edge(*fall) if fall else s
    k = np.floor((t - t0) / period).astype(np.int64) + 2  # index into k_all of the cycle each sample is in
    y = np.zeros(n)
    for d in (-1, 0, 1):
        kk = np.clip(k + d, 0, len(k_all) - 1)
        y += s(t - Tr[kk]) - sf(t - Tf[kk])
    x = lo + (hi - lo) * y + rng.normal(0, noise, n)
    if lsb:
        x = np.round(x / lsb) * lsb
    return x


def selftest(echo=print, tol_pct=1.0):
    """The analysis on synthetic pulser records of known responses. Returns (failures, results)."""
    fails, results = [], []

    def check(name, cond, msg, **kw):
        echo(("ok    " if cond else "FAIL  ") + name + ": " + msg)
        results.append(dict(name=name, ok=bool(cond), msg=msg, **kw))
        if not cond:
            fails.append(name)

    def f3_case(name, recs, dt, f3, **kw):
        out = step_analyse(recs, dt, **kw)
        got = out["f3db"]
        err = (got / f3 - 1) * 100 if got else float("nan")
        check(name, got is not None and abs(err) < tol_pct,
              f"-3 dB {got / 1e6 if got else 0:.1f} MHz ({err:+.2f} %), rise {out['rise_ps']} ps, {out['edges']} edges, "
              f"jitter {out['jitter_ps']} ps (2 ps put in), min bin {out['min_bin_count']}",
              f3_true=f3, f3_read=got, err_pct=None if got is None else round(err, 3), rise_ps=out["rise_ps"],
              jitter_ps=out["jitter_ps"], edges=out["edges"])
        check(name + " jitter", out["jitter_ps"] < 15.0, f"{out['jitter_ps']} ps, bound {out.get('jitter_bound_db_at_f3')} dB at -3 dB")
        return out

    # 1.25 GSa/s, 12 bits over 8 div at 100 mV/div, 1.5 mV noise, 2 ps edge jitter, 2 ppm clock drift
    dt = 0.8e-9
    lsb = 8 * 0.1 / 4096
    for resp, f3, skew in (("pole", 400e6, None), ("gauss", 900e6, None), ("pole", 300e6, [0, 0, 10e-12, 0])):
        recs = [synth_pulser(400_000, dt, 731.3e-9 * (1 + 3e-4 * r), resp, f3, noise=1.5e-3, lsb=lsb, skew=skew, seed=7 + r)
                for r in range(4)]
        label = f"{resp} {f3 / 1e6:g} MHz at 1.25 GSa/s" + (" (Nyquist 625 MHz)" if f3 > 625e6 else "")
        out = f3_case(label, recs, dt, f3, cores=4 if skew else 0)
        if skew:
            sk = out["cores"]["sample_late_ps"]
            others = [sk[0], sk[1], sk[3]]
            late = sk[2] - sum(others) / 3
            check("core skew", abs(late - 10.0) < 2.5 and max(others) - min(others) < 2.5,
                  f"cores sample late by {sk} ps (core 2 is 10 ps late: {late:.1f} ps found)")
    # 4 GSa/s, 250,000 points a record (what a 1 Mpts single gives over SCPI on the stock firmware), 8 records, each
    # from its own start, a 1,092,233 Hz pulser through a Gaussian of 1.3 GHz, 12 bits over 8 div at 200 mV/div
    dt4 = 0.25e-9
    lsb4 = 8 * 0.2 / 4096

    def recs4(fp, f3):
        # one pulser frequency for every record (a real pulser's), each record from its own start phase
        out = []
        rng = np.random.default_rng(3)
        for r in range(8):
            x = synth_pulser(250_000 + 4000, dt4, 1 / fp, "gauss", f3, lo=-0.5, hi=0.0, noise=1.5e-3, lsb=lsb4, seed=21 + r)
            off = int(rng.integers(0, 4000))
            out.append(x[off:off + 250_000])
        return out
    out = f3_case("gauss 1300 MHz at 4 GSa/s, 250,000 points x 8, pulser 1,092,233 Hz", recs4(1_092_233, 1.3e9), dt4, 1.3e9)
    check("1,092,233 Hz sweeps the sample interval", out["phase_turns"] >= MIN_TURNS,
          f"{out['phase_turns']} turns a record (warned below {MIN_TURNS:g})")
    # a pulser whose period is close to a whole number of samples (1,003,017 Hz: 3,988.03 samples at 4 GSa/s) must be
    # warned about: its edges sweep the sample interval about twice a 250,000-point record
    out = step_analyse(recs4(1_003_017, 1.3e9), dt4)
    check("slow phase walk warned", out["phase_turns"] < MIN_TURNS,
          f"1,003,017 Hz: {out['phase_turns']} turns a record; -3 dB {out['f3db'] / 1e6:.1f} MHz of 1300 "
          f"({(out['f3db'] / 1.3e9 - 1) * 100:+.2f} %)", f3_true=1.3e9, f3_read=out["f3db"], turns=out["phase_turns"])
    # a negative-going pulser whose fast edge falls and whose return rises slowly: --edge auto must pick the falling
    # edge, and the inverted record must fold to the falling edge's response
    recs = [synth_pulser(400_000, dt, 731.3e-9 * (1 + 3e-4 * r), "pole", 150e6, lo=-0.5, hi=0.0, noise=1.5e-3, lsb=lsb,
                         seed=11 + r, fall=("gauss", 900e6)) for r in range(4)]
    lo, hi = square_levels(recs[0])
    picked = fast_edge(recs[0], lo, hi)
    check("fast falling edge picked", picked == "fall", f"--edge auto picks {picked!r} (fall 900 MHz Gaussian, rise 150 MHz pole)")
    f3_case("falling edge folded (gauss 900 MHz)", [-x for x in recs], dt, 900e6)
    # a pulser whose period is a whole number of samples: every edge at the same phase, the grid not covered
    rec = synth_pulser(400_000, dt, 1000 * dt, "pole", 400e6, jitter=0.0, drift_ppm=0.0, noise=1e-3)
    try:
        step_analyse([rec], dt)
        check("locked period refused", False, "folded a period of exactly 1000 samples")
    except ValueError as e:
        check("locked period refused", "never sampled" in str(e), str(e)[:100])
    return fails, results


# ================================================================================================ SCPI

def parse_num(s):
    try:
        v = float(str(s).strip())
    except (TypeError, ValueError):
        return None
    return None if (math.isnan(v) or abs(v) >= 9.9e37) else v


def depth_word(n):
    """Points as the depth word the scope takes back: 1000000 -> 1M, 10000 -> 10k. AUTO and other words stay."""
    v = parse_num(n)
    if v is None:
        return str(n).strip()
    n = int(round(v))
    if n >= 1_000_000 and n % 1_000_000 == 0:
        return "%dM" % (n // 1_000_000)
    if n >= 1000 and n % 1000 == 0:
        return "%dk" % (n // 1000)
    return str(n)


def parse_depth(s):
    """1M, 100k, 1e6, 1000000 -> points."""
    m = re.fullmatch(r"\s*([0-9.]+)\s*([kKmM]?)\s*", str(s))
    if not m:
        raise argparse.ArgumentTypeError(f"not a depth: {s!r} (say 1M, 100k or 1000000)")
    v = float(m.group(1)) * {"": 1, "k": 1e3, "K": 1e3, "m": 1e6, "M": 1e6}[m.group(2)]
    return int(round(v))


# Lines the kit never sends, whatever a caller asks: a reset, the setup block, calibration, a screenshot, files.
_REFUSED = [re.compile(p, re.I) for p in (r"^\*RST\b", r"^:?SYST(em)?:SET", r"^[:*]?CAL", r"^:?DISP(lay)?:DATA",
                                           r"^:?SAVE", r"^:?LOAD", r"^:?SYST(em)?:(RES|REB|SHUT)")]


class Scpi:
    """SCPI on the raw socket, LF-terminated. Every read has a deadline. After a timeout the connection is replaced
    (a late answer would otherwise be read as the next query's) and *IDN? checks the scope still answers; if not,
    ScpiHung. Every line sent is logged with its time, its answer and its latency."""

    def __init__(self, host, port=SCPI_PORT, timeout=5.0, clock=None, echo=False, pace=0.05):
        self.host, self.port, self.timeout = host, port, timeout
        self.clock = clock or Clock()
        self.echo, self.pace = echo, pace
        self.trace = []
        self.sock = None
        self.buf = b""
        self.t0 = self.clock.now()
        self.hide = []  # strings never written to the trace (the serial number)

    def connect(self):
        self.close()
        try:
            self.sock = socket.create_connection((self.host, self.port), timeout=min(self.timeout, 6.0))
        except OSError as e:
            raise KitError("cannot connect to %s:%d (%s)" % (self.host, self.port, e))
        self.buf = b""

    def close(self):
        if self.sock is not None:
            try:
                self.sock.close()
            except OSError:
                pass
        self.sock = None

    def _send(self, cmd):
        for rx in _REFUSED:
            if rx.search(cmd.strip()):
                raise KitError("refused to send %r: the kit never resets, calibrates, saves or takes a screenshot" % cmd)
        if self.sock is None:
            self.connect()
        try:
            self.sock.sendall((cmd + "\n").encode("ascii"))
        except OSError as e:
            raise KitError("send %r failed: %s" % (cmd, e))

    def _fill(self, deadline, what):
        left = deadline - time.monotonic()
        if left <= 0:
            raise ScpiTimeout(what)
        self.sock.settimeout(left)
        try:
            chunk = self.sock.recv(1 << 20)
        except socket.timeout:
            raise ScpiTimeout(what)
        except OSError as e:
            raise KitError("read failed: %s" % e)
        if not chunk:
            raise KitError("the scope closed the connection")
        self.buf += chunk

    def _readline(self, timeout):
        deadline = time.monotonic() + timeout
        while b"\n" not in self.buf:
            self._fill(deadline, "no answer")
        line, _, self.buf = self.buf.partition(b"\n")
        return line.decode("ascii", "replace").strip()

    def _clean(self, s):
        for h in self.hide:
            if h:
                s = s.replace(h, "<serial not recorded>")
        return s

    def _log(self, cmd, reply, t_send, t_reply, note=None):
        e = {"t": utc_iso(self.clock.wall()), "s": round(t_send - self.t0, 4), "send": cmd}
        if reply is not None:
            e["reply"] = self._clean(reply)
        if t_reply is not None:
            e["ms"] = round((t_reply - t_send) * 1000.0, 2)
        if note:
            e["note"] = note
        self.trace.append(e)
        if self.echo:
            sys.stderr.write("[%8.3f] %s%s\n" % (e["s"], cmd, "" if reply is None and not note else
                                                 "  -> " + (e.get("reply") if reply is not None else note)))

    def write(self, cmd, pace=None):
        t = self.clock.now()
        self._send(cmd)
        self._log(cmd, None, t, None)
        self.clock.sleep(self.pace if pace is None else pace)

    def query(self, cmd, timeout=None):
        tmo = self.timeout if timeout is None else timeout
        t = self.clock.now()
        self._send(cmd)
        try:
            r = self._readline(tmo)
        except ScpiTimeout:
            self._log(cmd, None, t, self.clock.now(), "TIMEOUT")
            self._recover()
            raise ScpiTimeout("%s: no answer within %.1f s" % (cmd, tmo))
        self._log(cmd, r, t, self.clock.now())
        return r

    def try_query(self, cmd, timeout=None):
        try:
            return self.query(cmd, timeout)
        except ScpiTimeout:
            return None

    def block(self, cmd, timeout=30.0):
        """An IEEE 488.2 definite-length block answer (#<n><len><bytes>, then an optional LF), as bytes."""
        t = self.clock.now()
        self._send(cmd)
        deadline = time.monotonic() + timeout

        def need(n):
            while len(self.buf) < n:
                self._fill(deadline, "%s: block incomplete after %.0f s" % (cmd, timeout))
        try:
            need(2)
            if self.buf[:1] != b"#":
                raise KitError("%s: not a block: %r" % (cmd, self.buf[:20]))
            nd = self.buf[1:2]
            if not nd.isdigit() or nd == b"0":
                raise KitError("%s: not a definite-length block: %r" % (cmd, self.buf[:20]))
            nd = int(nd)
            need(2 + nd)
            ln = int(self.buf[2:2 + nd])
            need(2 + nd + ln)
        except ScpiTimeout:
            self._log(cmd, None, t, self.clock.now(), "TIMEOUT (block)")
            self._recover()
            raise
        data = bytes(self.buf[2 + nd:2 + nd + ln])
        self.buf = self.buf[2 + nd + ln:]
        if self.buf[:1] == b"\n":
            self.buf = self.buf[1:]
        self._log(cmd, "<%d-byte block>" % ln, t, self.clock.now())
        return data

    def _recover(self):
        self.connect()
        t = self.clock.now()
        self._send("*IDN?")
        try:
            r = self._readline(6.0)
        except (ScpiTimeout, KitError):
            self._log("*IDN?", None, t, self.clock.now(), "TIMEOUT (after a timeout): the scope's SCPI is hung")
            raise ScpiHung("the scope stopped answering *IDN? on a fresh connection")
        self._log("*IDN?", r, t, self.clock.now(), "connection replaced after a timeout")

    def errors(self, limit=20):
        errs = []
        for _ in range(limit):
            r = self.try_query(":SYSTem:ERRor?", 3.0)
            if r is None or r.startswith("0,") or r.startswith("+0,"):
                break
            errs.append(r)
        return errs


class Identity:
    """*IDN?: maker, model, serial, firmware. The kit records the model and the firmware field only."""

    def __init__(self, idn):
        parts = [p.strip() for p in idn.split(",")]
        parts += [""] * (4 - len(parts))
        self.maker, self.model, self.serial = parts[:3]
        self.firmware = ",".join(parts[3:]).strip(",")
        m = re.match(r"^(DHO|MHO)(\d)(\d)(\d)", self.model.upper())
        self.series = None
        if m and (m.group(1) + m.group(2)) in ("DHO8", "DHO9", "MHO9"):
            self.series = m.group(1) + m.group(2) + "00"
        self.channels = int(m.group(4)) if m else 4
        self.has_50 = self.series == "MHO900"  # the DHO800 / DHO900 inputs are 1 MOhm only
        self.max_rate = 4e9 if self.series == "MHO900" else 1.25e9

    def as_dict(self):
        return {"model": self.model, "firmware": self.firmware, "series": self.series}


# ================================================================================================ the scope

def ch(n):
    return "CHANnel%d" % n


def same(a, b):
    """Two answers of one query are the same setting: numbers by value, words by their text."""
    if a is None or b is None:
        return a == b
    x, y = parse_num(a), parse_num(b)
    if x is not None and y is not None:
        return abs(x - y) <= 1e-9 * max(1.0, abs(x), abs(y))
    return a.strip().upper() == b.strip().upper()


class Scope:
    """One channel of one scope, set up for the pulser's step and put back after."""

    def __init__(self, s, ident, chan, imp, clock):
        self.s, self.id, self.ch, self.imp, self.clock = s, ident, chan, imp, clock
        self.tb = 0.0
        self.depth = 0
        self.snap = {}
        self.levels = {}
        self.state = None
        self.edge_source = None
        self.unanswered = []
        self.last_hash = None
        self.preamble = None
        self.points = None

    # -- the settings the kit may change, in the order they are put back
    def chan_keys(self):
        k = [":CHANnel%d:IMPedance" % self.ch] if self.id.has_50 else []
        return k + [":CHANnel%d:%s" % (self.ch, x) for x in ("PROBe", "COUPling", "BWLimit", "SCALe", "OFFSet")]

    def keys(self):
        k = [":CHANnel%d:DISPlay" % i for i in range(1, self.id.channels + 1)]
        k += self.chan_keys()
        k += [":ACQuire:TYPE", ":ACQuire:MDEPth", ":TIMebase:MAIN:SCALe", ":TIMebase:MAIN:OFFSet",
              ":TRIGger:MODE", ":TRIGger:EDGE:SOURce", ":TRIGger:SWEep",
              ":WAVeform:SOURce", ":WAVeform:MODE", ":WAVeform:FORMat", ":WAVeform:STARt", ":WAVeform:STOP"]
        return k

    def take_snapshot(self):
        s = self.s
        s.errors(5)
        self.state = s.query(":TRIGger:STATus?")
        for k in self.keys():
            v = s.try_query(k + "?")
            if v is None:
                self.unanswered.append(k)
                s.errors(3)  # an unanswered query may queue an error: read it off, so later checks are clean
            else:
                self.snap[k] = v
        # the edge level belongs to the selected source: read it on the found source and on the test channel
        self.edge_source = self.snap.get(":TRIGger:EDGE:SOURce")
        v = s.try_query(":TRIGger:EDGE:LEVel?")
        if v is not None and self.edge_source:
            self.levels[self.edge_source] = v
        mine = "CHAN%d" % self.ch
        if self.edge_source and not same(self.edge_source, mine):
            s.write(":TRIGger:EDGE:SOURce " + ch(self.ch))
            v = s.try_query(":TRIGger:EDGE:LEVel?")
            if v is not None:
                self.levels[mine] = v
            s.write(":TRIGger:EDGE:SOURce " + self.edge_source)
        s.errors(3)
        return dict(state=self.state, values=dict(self.snap), edge_levels=dict(self.levels), unanswered=list(self.unanswered))

    # -- set up
    def setup(self, scale, depth, tb, band, offset=0.0):
        s = self.s
        s.write(":RUN")  # the depth is taken while the scope runs
        for n in range(1, self.id.channels + 1):
            s.write(":CHANnel%d:DISPlay %d" % (n, 1 if n == self.ch else 0))
        c = self.ch
        s.write(":CHANnel%d:PROBe 1" % c)
        if self.id.has_50:
            s.write(":CHANnel%d:IMPedance %s" % (c, "FIFTy" if self.imp == "50" else "OMEG"))
        s.write(":CHANnel%d:COUPling DC" % c)
        s.write(":CHANnel%d:BWLimit %s" % (c, band))
        s.write(":CHANnel%d:SCALe %s" % (c, scale))
        s.write(":CHANnel%d:OFFSet %s" % (c, offset))
        s.write(":ACQuire:TYPE NORMal")
        s.write(":ACQuire:MDEPth %s" % depth_word(depth))
        s.write(":TIMebase:MAIN:SCALe %s" % tb)
        s.write(":TIMebase:MAIN:OFFSet 0")
        s.write(":TRIGger:MODE EDGE")
        s.write(":TRIGger:EDGE:SOURce %s" % ch(c))
        s.write(":TRIGger:EDGE:LEVel 0")
        s.write(":TRIGger:SWEep AUTO")
        s.write(":WAVeform:MODE RAW")
        s.write(":WAVeform:FORMat WORD")
        if self.id.has_50 and self.imp == "50":
            got = s.query(":CHANnel%d:IMPedance?" % c)
            if not got.upper().startswith("FIFT"):
                raise Refused("CH%d: 50 Ohm not taken (%s)" % (c, got))
        errs = s.errors()
        if errs:
            raise Refused("settings refused: %s" % errs)
        # the depth and the time base once more, after the channel changes settled, and read back
        self.depth = depth
        self.clock.sleep(0.5)
        s.write(":ACQuire:MDEPth %s" % depth_word(depth))
        s.write(":TIMebase:MAIN:SCALe %s" % tb)
        for _ in range(50):
            md, ts = parse_num(s.query(":ACQuire:MDEPth?")), parse_num(s.query(":TIMebase:MAIN:SCALe?"))
            if md == depth and ts and abs(ts / tb - 1) < 0.01:
                break
            self.clock.sleep(0.1)
        else:
            raise Refused("depth %d / time base %g not taken" % (depth, tb))
        self.tb = tb
        self.clock.sleep(0.3)

    def set_tb(self, tb):
        self.s.write(":TIMebase:MAIN:SCALe %.4e" % tb)
        for _ in range(50):
            got = parse_num(self.s.query(":TIMebase:MAIN:SCALe?")) or 0.0
            if got > 0 and abs(got / tb - 1) < 0.3:  # the scope's ladder rounds a scale to its 1-2-5 steps
                self.tb = got
                break
            self.clock.sleep(0.1)
        self.clock.sleep(0.3)

    def full_rate(self, depth, rate):
        """The time base at which `depth` points are taken at `rate`, read back: a record can span more than the
        screen's ten divisions, so depth / rate / 10 is not enough. Down the 1-2-5 ladder until :ACQuire:SRATe?
        reads the rate with the depth kept; refuses when no time base gives both. Running: a stopped scope may not
        recompute its sample rate."""
        tb = depth / rate / 10.0
        tried = []
        self.s.write(":RUN")
        for _ in range(10):
            self.s.write(":ACQuire:MDEPth %s" % depth_word(depth))
            self.set_tb(tb)
            self.clock.sleep(0.8)
            sr = parse_num(self.s.query(":ACQuire:SRATe?")) or 0.0
            md = parse_num(self.s.query(":ACQuire:MDEPth?")) or 0.0
            tried.append("%.3g s/div: %g MSa/s, %g pts" % (self.tb, sr / 1e6, md))
            if sr >= rate * 0.99 and md == depth:
                return self.tb, tried
            if md != depth and sr >= rate * 0.99:
                break
            tb = ladder_down(self.tb)
        raise Refused("no time base takes %d points at %g MSa/s: %s" % (depth, rate / 1e6, "; ".join(tried)))

    def capture(self):
        """One single acquisition of the test channel: (codes, volts, sample interval, preamble)."""
        s = self.s
        for attempt in range(3):
            s.write(":SINGle", pace=0.2)
            t0 = self.clock.now()
            forced = False
            while True:
                st = s.query(":TRIGger:STATus?")
                if st.upper().startswith("STOP"):
                    break
                # forced only once the record has had time to fill (twelve divisions)
                if self.clock.now() - t0 > max(0.6, 12 * self.tb + 0.3) and not forced:
                    s.write(":TFORce")
                    forced = True
                if self.clock.now() - t0 > 20 + 12 * self.tb:
                    raise Refused("the single did not stop (%s)" % st)
                self.clock.sleep(0.01)
            codes, x, dt, pre = self.read_record()
            h = hashlib.sha256(codes.tobytes()).hexdigest()
            if h != self.last_hash:
                self.last_hash = h
                return codes, x, dt, pre
            # byte for byte the previous record: the single had not started when the status was read
        raise Refused("three singles in a row gave the previous record again")

    def read_record(self):
        s, c = self.s, self.ch
        s.write(":WAVeform:SOURce %s" % ch(c))
        s.write(":WAVeform:MODE RAW")
        s.write(":WAVeform:FORMat WORD")
        raw_pre = s.query(":WAVeform:PREamble?")
        pre = raw_pre.split(",")
        if len(pre) < 10:
            raise Refused("a preamble of %d fields: %r" % (len(pre), raw_pre))
        pts = int(float(pre[2]))
        if pts <= 0 or pts > self.depth:
            raise Refused("CH%d: the record holds %d points, not the %d set" % (c, pts, self.depth))
        # Rigol's stock MHO900 firmware answers a 1 Mpts single's RAW preamble with 250,000 points: the samples it
        # gives, at the record's own interval (the caller checks the rate)
        xinc, yinc, yorig, yref = float(pre[4]), float(pre[7]), float(pre[8]), float(pre[9])
        raw = bytearray()
        start = 1
        while start <= pts:
            stop = min(pts, start + PAGE - 1)
            s.write(":WAVeform:STARt %d" % start)
            s.write(":WAVeform:STOP %d" % stop)
            b = s.block(":WAVeform:DATA?")
            if len(b) != 2 * (stop - start + 1):
                raise Refused("points %d..%d: %d bytes, not %d" % (start, stop, len(b), 2 * (stop - start + 1)))
            raw += b
            start = stop + 1
        codes = np.frombuffer(bytes(raw), dtype="<u2").copy()
        x = to_volts(codes, yinc, yorig, yref)
        self.points = pts
        self.preamble = dict(raw=raw_pre, points=pts, xincrement=xinc, xorigin=float(pre[5]), xreference=float(pre[6]),
                             yincrement=yinc, yorigin=yorig, yreference=yref)
        return codes, x, xinc, self.preamble

    def read_back(self):
        """Every setting the kit set, as the scope answers it now, with the sample rate."""
        out = {}
        for k in self.keys() + [":ACQuire:SRATe", ":TRIGger:EDGE:LEVel"]:
            v = self.s.try_query(k + "?")
            out[k] = v
        return out

    # -- put back
    def restore(self):
        s, v = self.s, self.snap
        rep = {"written": [], "still_different": [], "not_readable": [], "errors": []}

        def cur(k):
            r = s.try_query(k + "?", 4.0)
            if r is None and k not in rep["not_readable"]:
                rep["not_readable"].append(k)
            return r

        def put(k, val=None):
            if k not in v:
                return
            now = cur(k)
            if now is None or same(now, v[k]):
                return
            want = val if val is not None else (depth_word(v[k]) if k.endswith(":MDEPth") else v[k])
            s.write("%s %s" % (k, want))
            rep["written"].append("%s %s" % (k, want))

        def one_pass():
            c = self.ch
            dk = ":CHANnel%d:DISPlay" % c
            # the stock ignores a hidden channel's scale and offset: the test channel is shown while they go back
            if dk in v:
                now = cur(dk)
                if now is not None and not same(now, "1"):
                    s.write(dk + " 1")
                    rep["written"].append(dk + " 1")
            for k in self.chan_keys():
                put(k)
            md = ":ACQuire:MDEPth"
            if md in v:
                now = cur(md)
                if now is not None and not same(now, v[md]):
                    st = s.try_query(":TRIGger:STATus?")
                    if st is not None and st.upper().startswith("STOP"):
                        s.write(":RUN")  # the depth is taken while running; the run state is put back last
                        rep["written"].append(":RUN")
            for k in (":ACQuire:TYPE", md, ":TIMebase:MAIN:SCALe", ":TIMebase:MAIN:OFFSet"):
                put(k)
            # each source's own level, then the source found
            src_now = cur(":TRIGger:EDGE:SOURce")
            for src in sorted(self.levels):
                s.write(":TRIGger:EDGE:SOURce " + src)
                now = s.try_query(":TRIGger:EDGE:LEVel?")
                if now is None:
                    rep["not_readable"].append(":TRIGger:EDGE:LEVel (%s)" % src)
                elif not same(now, self.levels[src]):
                    s.write(":TRIGger:EDGE:LEVel " + self.levels[src])
                    rep["written"].append(":TRIGger:EDGE:LEVel %s (%s)" % (self.levels[src], src))
            if self.edge_source:
                s.write(":TRIGger:EDGE:SOURce " + self.edge_source)
                if src_now is not None and not same(src_now, self.edge_source):
                    rep["written"].append(":TRIGger:EDGE:SOURce " + self.edge_source)
            for k in (":TRIGger:SWEep", ":TRIGger:MODE"):
                put(k)
            # a waveform source whose channel is off is refused (the DHO): shown for the moment
            wk = ":WAVeform:SOURce"
            if wk in v:
                now = cur(wk)
                if now is not None and not same(now, v[wk]):
                    m = re.match(r"CHAN(?:nel)?(\d)$", v[wk].strip(), re.I)
                    shown = None
                    if m:
                        shown = cur(":CHANnel%s:DISPlay" % m.group(1))
                        if shown is not None and not same(shown, "1"):
                            s.write(":CHANnel%s:DISPlay 1" % m.group(1))
                    s.write("%s %s" % (wk, v[wk]))
                    rep["written"].append("%s %s" % (wk, v[wk]))
                    if m and shown is not None and not same(shown, "1"):
                        s.write(":CHANnel%s:DISPlay %s" % (m.group(1), shown))
            for k in (":WAVeform:MODE", ":WAVeform:FORMat", ":WAVeform:STARt", ":WAVeform:STOP"):
                put(k)
            for n in range(1, self.id.channels + 1):
                put(":CHANnel%d:DISPlay" % n)

        def compare():
            diff = []
            for k in self.keys():
                if k in v:
                    now = s.try_query(k + "?", 4.0)
                    if now is not None and not same(now, v[k]):
                        diff.append({"setting": k, "was": v[k], "now": now})
            for src in sorted(self.levels):
                s.write(":TRIGger:EDGE:SOURce " + src)
                now = s.try_query(":TRIGger:EDGE:LEVel?")
                if now is not None and not same(now, self.levels[src]):
                    diff.append({"setting": ":TRIGger:EDGE:LEVel (%s)" % src, "was": self.levels[src], "now": now})
            if self.edge_source:
                s.write(":TRIGger:EDGE:SOURce " + self.edge_source)
            return diff

        for _ in range(2):
            # twice: a value the first pass clamped against a setting not yet back takes on the second
            one_pass()
            self.clock.sleep(0.5)
            diff = compare()
            if not diff:
                break
        # the run state last
        st = s.try_query(":TRIGger:STATus?")
        if self.state and st is not None:
            if not self.state.upper().startswith("STOP"):
                if st.upper().startswith("STOP"):
                    s.write(":RUN")
                    rep["written"].append(":RUN")
            else:
                # found stopped: one forced single at the restored settings, so :ACQuire:SRATe? (which the stock
                # recomputes only by acquiring) reads the restored time base's rate again. It replaces the held trace.
                s.write(":SINGle", pace=0.5)
                s.write(":TFORce")
                rep["written"].append(":SINGle, :TFORce (found stopped)")
                for _ in range(50):
                    r = s.try_query(":TRIGger:STATus?")
                    if r is None or r.upper().startswith("STOP"):
                        break
                    self.clock.sleep(0.1)
        rep["still_different"] = compare()
        rep["unanswered_in_snapshot"] = list(self.unanswered)
        rep["errors"] = s.errors(10)
        rep["complete"] = not rep["still_different"]
        return rep


def to_volts(codes, yinc, yorig, yref):
    """WORD codes to volts, as the preamble defines them: (code - (yorigin + yreference)) * yincrement."""
    return (codes.astype(float) - (yorig + yref)) * yinc


# ================================================================================================ commands

def parse_host(h):
    if h.count(":") == 1:
        host, port = h.rsplit(":", 1)
        return host, int(port)
    return h, SCPI_PORT


def analysis_args(a):
    return dict(os=a.os, pre_ns=a.pre_ns, post_ns=a.post_ns, cores=a.cores, points=a.points, fmax=a.fmax,
                source_rise_ps=a.source_rise_ps)


def write_json(path, doc):
    Path(path).parent.mkdir(parents=True, exist_ok=True)
    Path(path).write_text(json.dumps(doc, indent=1), encoding="utf-8", newline="\n")


def write_csv(prefix, res):
    """PREFIX-step.csv: the folded step (time from the 50 % fit, ps; level, 0 = base, 1 = top).
    PREFIX-response.csv: |H(f)| (Hz, dB, SNR dB)."""
    p = str(prefix)
    if p.lower().endswith(".csv"):
        p = p[:-4]
    step_path, resp_path = Path(p + "-step.csv"), Path(p + "-response.csv")
    step_path.parent.mkdir(parents=True, exist_ok=True)
    dt_ps = res["step_dt"] * 1e12 * res["step_stride"]
    with open(step_path, "w", newline="", encoding="utf-8") as fh:
        w = csv.writer(fh, lineterminator="\n")
        w.writerow(["time_ps", "level"])
        for i, v in enumerate(res["step"]):
            w.writerow(["%.3f" % (res["step_t0_ps"] + i * dt_ps), v])
    with open(resp_path, "w", newline="", encoding="utf-8") as fh:
        w = csv.writer(fh, lineterminator="\n")
        w.writerow(["freq_hz", "db", "snr_db"])
        snr = res.get("snr") or [None] * len(res["freqs"])
        for f, d, n in zip(res["freqs"], res["db"], snr):
            w.writerow([f, d, "" if n is None else n])
    return [str(step_path), str(resp_path)]


def report_lines(doc):
    r = doc["results"]
    unit = doc.get("unit") or {}
    head = "%s %s" % (unit.get("model", "records"), unit.get("firmware", ""))
    f3 = "%.1f MHz" % (r["f3db"] / 1e6) if r.get("f3db") else "beyond %.0f MHz" % (r["freqs"][-1] / 1e6)
    edge = {"fall": "falling", "rise": "rising"}.get(doc.get("edge"), doc.get("edge"))
    rows = [("-3 dB", f3),
            ("rise 10-90 %", "%.1f ps" % r["rise_ps"]),
            ("overshoot", "%.2f %%" % r["overshoot_pct"]),
            ("peaking", "%+.2f dB" % r["peaking_db"] if r.get("peaking_db") is not None else "-"),
            ("edge folded", edge),
            ("edges", "%d in %d records" % (r["edges"], r["records"])),
            ("jitter (upper bound)", "%.2f ps" % r["jitter_ps"]),
            ("jitter bound at -3 dB", "%.3f dB" % r["jitter_bound_db_at_f3"] if r.get("jitter_bound_db_at_f3") is not None else "-"),
            ("30 ps source edge at -3 dB", "%.4f dB (not divided out)" % r["edge_30ps_db_at_f3"] if r.get("edge_30ps_db_at_f3") is not None else
             ("divided out (%g ps)" % doc.get("analysis", {}).get("source_rise_ps", 0))),
            ("sample rate", "%.4g GSa/s" % (doc["rate"] / 1e9) if doc.get("rate") else "-"),
            ("grid", "%.2f ps (os %s)" % (r["grid_ps"], doc.get("analysis", {}).get("os", "?"))),
            ("step / square", "%.4f V / %.4f V (%+.3f dB)" % (r["step_v"], r["square_v"], r["settled_db"])),
            ("pulser period", "%.4f ns" % r["period_ns"]),
            ("fewest samples in a bin", "%d" % r["min_bin_count"]),
            ("phase turns a record", "%.2f%s" % (r["phase_turns"], "  (fewer than %g: read low)" % MIN_TURNS
                                                 if r["phase_turns"] < MIN_TURNS else "") if "phase_turns" in r else "-")]
    title = "## %s CH%s %s" % (head.strip(), doc.get("channel", "?"), doc.get("settings", {}).get("asked", {}).get("band", ""))
    out = [title.rstrip(),
           "", "| quantity | value |", "|---|---|"]
    out += ["| %s | %s |" % row for row in rows]
    if r.get("cores"):
        out += ["", "cores (sample late, ps): %s; height (dB): %s" % (r["cores"]["sample_late_ps"], r["cores"]["height_db"])]
    return out


def cmd_selftest(a):
    fails, _ = selftest()
    if fails:
        print("selftest FAILED:", ", ".join(fails))
        return 1
    print("selftest: every case passed")
    return 0


def load_records(path):
    z = np.load(path, allow_pickle=False)
    if "codes" in z.files:
        codes = z["codes"]
        recs = [to_volts(codes[i], float(z["yincrement"][i]), float(z["yorigin"][i]), float(z["yreference"][i]))
                for i in range(codes.shape[0])]
    else:
        recs = [np.asarray(r, float) for r in z["volts"]]
    dt = float(z["dt"])
    meta = json.loads(str(z["meta"])) if "meta" in z.files else {}
    return recs, dt, meta


def pick_edge(edge, x):
    if edge != "auto":
        return edge
    lo, hi = square_levels(x)
    return fast_edge(x, lo, hi)


def cmd_analyse(a):
    recs, dt, meta = load_records(a.records_file)
    edge = pick_edge(a.edge, recs[0])
    try:
        res = step_analyse([-x if edge == "fall" else x for x in recs], dt, **analysis_args(a))
    except ValueError as e:
        print("step_kit: %s" % e, file=sys.stderr)
        return 1
    doc = dict(tool="step_kit", version=VERSION, command="analyse", date=utc_iso(time.time()), source=str(a.records_file),
               unit=meta.get("unit"), channel=meta.get("channel"), rate=1.0 / dt, edge=edge,
               settings=meta.get("settings", {}), analysis=analysis_args(a), results=res)
    for line in report_lines(doc):
        print(line)
    if a.out:
        write_json(a.out, doc)
        print("wrote", a.out)
    if a.csv:
        print("wrote", ", ".join(write_csv(a.csv, res)))
    return 0


def cmd_report(a):
    for p in a.files:
        doc = json.loads(Path(p).read_text(encoding="utf-8"))
        if "results" not in doc:
            print("%s: no results" % p)
            continue
        print("\n".join(report_lines(doc)))
        print()
        if a.csv:
            pre = a.csv if len(a.files) == 1 else "%s-%s" % (a.csv, Path(p).stem)
            print("wrote", ", ".join(write_csv(pre, doc["results"])))
    return 0


def cmd_measure(a, clock=None):
    clock = clock or Clock()
    host, port = parse_host(a.host)
    s = Scpi(host, port, timeout=a.timeout, clock=clock, echo=a.verbose)
    stamp = datetime.datetime.now(datetime.timezone.utc).strftime("%Y%m%d-%H%M%S")
    out_path = a.out or "step-kit-%s.json" % stamp
    doc = dict(tool="step_kit", version=VERSION, command="measure", date=utc_iso(clock.wall()))
    s.connect()
    idn = s.query("*IDN?")
    ident = Identity(idn)
    if ident.serial:
        s.hide.append(ident.serial)
        for e in s.trace:
            if "reply" in e:
                e["reply"] = s._clean(e["reply"])
    doc["unit"] = ident.as_dict()
    if ident.series is None and not a.untested_model:
        print("step_kit: %s is not a DHO800, DHO900 or MHO900, whose SCPI the kit is tested on (--untested-model to "
              "try anyway)" % ident.model, file=sys.stderr)
        s.close()
        return 2
    if not 1 <= a.ch <= ident.channels:
        print("step_kit: --ch %d: the %s has %d channels" % (a.ch, ident.model, ident.channels), file=sys.stderr)
        s.close()
        return 2
    imp = a.imp or ("50" if ident.has_50 else "1M")
    if ident.series is None:
        ident.has_50 = imp == "50"  # an untested model: :CHANnel<n>:IMPedance is sent only when 50 Ohm is asked
    if imp == "50" and not ident.has_50 and ident.series is not None:
        print("step_kit: the %s's inputs are 1 MOhm only: feed the pulser through a 50 Ohm feed-through and pass "
              "--imp 1M" % ident.model, file=sys.stderr)
        s.close()
        return 2
    rate = a.rate or ident.max_rate
    depth = a.depth
    opc = s.query("*OPC?")
    if opc.strip() != "1":
        print("step_kit: *OPC? answers %r: the connection is one answer behind; nothing sent" % opc, file=sys.stderr)
        s.close()
        return 1
    sc = Scope(s, ident, a.ch, imp, clock)
    doc["snapshot"] = sc.take_snapshot()
    asked = dict(channel=a.ch, impedance=imp, band=a.band, rate=rate, depth=depth, records=a.records, edge=a.edge)
    doc["channel"] = a.ch
    doc["settings"] = {"asked": asked}
    doc["analysis"] = analysis_args(a)
    rc = 0
    recs_codes, recs_pre, recs = [], [], []
    try:
        # 1. the pulser on the channel, alone, at the full rate, 0.5 V/div
        sc.setup(0.5, depth, depth / rate / 10.0, a.band)
        _, tried = sc.full_rate(depth, rate)
        _, x, dt, _ = sc.capture()
        ok, lo, hi, what = pulser_verdict(x)
        doc["pulser_check"] = dict(found=ok, what=what, full_rate_ladder=tried)
        print("CH%d: %s%s" % (a.ch, what, " <- pulser" if ok else ""), flush=True)
        if not ok:
            raise Refused("no free-running pulser on CH%d (a square of more than 20 mV, at least 10 regular edges a record)" % a.ch)
        if abs(1.0 / dt / rate - 1) > 0.01:
            raise Refused("CH%d records at %.1f MSa/s, not the %g asked" % (a.ch, 1 / dt / 1e6, rate / 1e6))
        warn = turns_warning(x, lo, hi)
        if warn:
            doc.setdefault("warnings", []).append(warn)
            print("WARNING: " + warn, flush=True)
        # 2. the edge folded: --edge, or the faster of the two
        edge = a.edge if a.edge != "auto" else fast_edge(x, lo, hi)
        print("folding the %s edge" % ("falling" if edge == "fall" else "rising"), flush=True)
        # 3. scaled so the step spans about 5 div, centred, NORMal edge trigger at its middle
        scale = ladder_up(max(hi - lo, 1e-3) / 5.0)
        offset = -(lo + hi) / 2
        mid = (lo + hi) / 2
        sc.setup(scale, depth, depth / rate / 10.0, a.band, offset=offset)
        tb, tried = sc.full_rate(depth, rate)
        s.write(":TRIGger:EDGE:SOURce %s" % ch(a.ch))
        s.write(":TRIGger:EDGE:LEVel %.4f" % mid)
        s.write(":TRIGger:SWEep NORMal")
        clock.sleep(0.5)
        got_band = s.try_query(":CHANnel%d:BWLimit?" % a.ch)
        errs = s.errors()
        if errs:
            raise Refused("CH%d: settings refused: %s" % (a.ch, errs))
        asked.update(scale=scale, offset=offset, trigger_level=round(mid, 4), timebase=tb, edge_folded=edge)
        doc["settings"]["read_back"] = sc.read_back()
        doc["settings"]["band_read_back"] = got_band
        doc["edge"] = edge
        # 4. the records
        for i in range(a.records):
            codes, x, dt_i, pre = sc.capture()
            if abs(1.0 / dt_i / rate - 1) > 0.01:
                raise Refused("record %d at %.1f MSa/s, not the %g asked" % (i + 1, 1 / dt_i / 1e6, rate / 1e6))
            if np.max(x) > mid + 3.9 * scale or np.min(x) < mid - 3.9 * scale:
                raise Refused("CH%d reaches the screen's edge (%.3f .. %.3f V at %g V/div): a clipped step is not a step"
                              % (a.ch, np.min(x), np.max(x), scale))
            recs_codes.append(codes)
            recs_pre.append(pre)
            recs.append(x)
            dt = dt_i
            print("  record %d: %d points at %.4g GSa/s" % (i + 1, len(x), 1 / dt_i / 1e9), flush=True)
        doc["preamble"] = recs_pre[-1]
        doc["rate"] = 1.0 / dt
        doc["points_per_record"] = int(len(recs[-1]))
        try:
            res = step_analyse([-x if edge == "fall" else x for x in recs], dt, **analysis_args(a))
        except ValueError as e:
            raise Refused("CH%d: %s" % (a.ch, e))
        doc["results"] = res
    except Refused as e:
        doc["refused"] = str(e)
        print("step_kit: refused: %s" % e, file=sys.stderr)
        rc = 1
    except ScpiHung as e:
        doc["error"] = str(e)
        print("step_kit: %s; nothing can be restored (the snapshot is in %s)" % (e, out_path), file=sys.stderr)
        rc = 4
    except KeyboardInterrupt:
        doc["error"] = "interrupted"
        rc = 130
    except KitError as e:
        doc["error"] = str(e)
        print("step_kit: %s" % e, file=sys.stderr)
        rc = 1
    finally:
        try:
            if rc == 4:
                raise ScpiHung("the scope stopped answering before the restore")
            rep = sc.restore()
        except ScpiHung as e:
            rep = {"complete": False, "error": str(e)}
            rc = 4
        except KitError as e:
            rep = {"complete": False, "error": str(e)}
        doc["restore"] = rep
        doc["trace"] = s.trace
        s.close()
    if not rep.get("complete") and rc in (0, 1):
        rc = 3
    write_json(out_path, doc)
    print("restore: %s" % ("every setting back" if rep.get("complete") else "INCOMPLETE: %s" % (rep.get("still_different") or rep.get("error"))))
    if a.save_records and recs_codes:
        meta = dict(unit=doc["unit"], channel=a.ch, settings=doc["settings"], edge=doc.get("edge"))
        np.savez_compressed(a.save_records, codes=np.stack(recs_codes),
                            yincrement=np.array([p["yincrement"] for p in recs_pre]),
                            yorigin=np.array([p["yorigin"] for p in recs_pre]),
                            yreference=np.array([p["yreference"] for p in recs_pre]),
                            dt=np.array(recs_pre[-1]["xincrement"]), meta=np.array(json.dumps(meta)))
        print("wrote", a.save_records)
    print("wrote", out_path)
    if "results" in doc:
        print("\n".join(report_lines(doc)))
        if a.csv:
            print("wrote", ", ".join(write_csv(a.csv, doc["results"])))
    return rc


def add_analysis(p):
    p.add_argument("--edge", default="auto", choices=["auto", "rise", "fall"],
                   help="the pulser edge folded (auto: the faster; a negative-going pulser's fast edge falls)")
    p.add_argument("--os", type=int, default=16, help="the equivalent-time grid, this many points a sample interval")
    p.add_argument("--pre-ns", type=float, default=15.0, help="the fold's window before the edge, ns")
    p.add_argument("--post-ns", type=float, default=45.0, help="the fold's window after the edge, ns")
    p.add_argument("--points", type=int, default=80, help="log-spaced frequencies in the response")
    p.add_argument("--fmax", type=float, default=0.0, help="the highest frequency reported, Hz (default 0.4 / grid)")
    p.add_argument("--cores", type=int, default=0, help="also fold each of N interleaved ADC cores alone (skew, gain)")
    p.add_argument("--source-rise-ps", type=float, default=0.0,
                   help="divide out a Gaussian source edge of this 10-90 %% rise, ps (default 0: nothing divided out)")
    p.add_argument("--csv", default="", help="write PREFIX-step.csv and PREFIX-response.csv")


def main(argv=None, clock=None):
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    sub = ap.add_subparsers(dest="cmd", required=True)
    m = sub.add_parser("measure", help="measure one channel over SCPI")
    m.add_argument("--host", required=True, help="HOST or HOST:PORT (port %d by default)" % SCPI_PORT)
    m.add_argument("--ch", type=int, required=True, help="the channel the pulser is on")
    m.add_argument("--imp", choices=["50", "1M"], help="input impedance (default: 50 Ohm where the model has it, else 1M)")
    m.add_argument("--band", default="OFF", help="the band limit word, as the scope takes it (OFF, 20M, 250M, ...)")
    m.add_argument("--rate", type=float, default=0.0, help="sample rate, Sa/s (default: the model's one-channel rate)")
    m.add_argument("--depth", type=parse_depth, default=1_000_000, help="memory depth (default 1M)")
    m.add_argument("--records", type=int, default=8, help="single acquisitions folded together (default 8)")
    m.add_argument("--out", default="", help="the result JSON (default step-kit-<UTC stamp>.json)")
    m.add_argument("--save-records", default="", help="also save the raw records (.npz) for `analyse`")
    m.add_argument("--timeout", type=float, default=5.0, help="seconds to wait for each answer")
    m.add_argument("--untested-model", action="store_true", help="run on a model the kit is not tested on")
    m.add_argument("--verbose", action="store_true", help="print every SCPI line")
    add_analysis(m)
    an = sub.add_parser("analyse", help="the analysis on saved records")
    an.add_argument("records_file")
    an.add_argument("--out", default="", help="write the result JSON")
    add_analysis(an)
    sub.add_parser("selftest", help="the analysis on synthetic records of known responses")
    r = sub.add_parser("report", help="the table of result JSONs")
    r.add_argument("files", nargs="+")
    r.add_argument("--csv", default="", help="write PREFIX-step.csv and PREFIX-response.csv")
    a = ap.parse_args(argv)
    if a.cmd == "selftest":
        return cmd_selftest(a)
    if a.cmd == "analyse":
        return cmd_analyse(a)
    if a.cmd == "report":
        return cmd_report(a)
    try:
        return cmd_measure(a, clock)
    except ScpiHung as e:
        print("step_kit: %s" % e, file=sys.stderr)
        return 4
    except KitError as e:
        print("step_kit: %s" % e, file=sys.stderr)
        return 1


if __name__ == "__main__":
    sys.exit(main())