#!/usr/bin/env python3
# SPDX-License-Identifier: MIT
# Copyright (c) 2026 Olympian LED
"""rate_kit: waveforms a second on a Rigol DHO800 / DHO900 / MHO900, counted from outside the firmware.

One file, Python 3.8+ standard library only. It talks SCPI over the raw socket (TCP 5555). It works the same on
Rigol's stock firmware, Norbert Kiszka's mod and NovaOS, and says which one it found. Every setting it changes is
read first, written to snapshot.json before anything changes, and put back on success, on an error and on Ctrl-C
(run/stop state included).

    python3 rate_kit.py identify --host SCOPE
    python3 rate_kit.py measure --host SCOPE --method aux-self
    python3 rate_kit.py measure --host SCOPE --method aux-external --counter COUNTER
    python3 rate_kit.py restore --host SCOPE --from out/snapshot-<stamp>.json
    python3 rate_kit.py scpi-health --host SCOPE                              (read-only: latency, error queue)
    python3 rate_kit.py preset --host SCOPE --yes-factory-default ...         (*RST, count, setup block restored)
    python3 rate_kit.py trio --unit stock=HOST --unit mod=HOST --unit novaos=HOST ...

Methods (README.md has the wiring, METHOD.md the exact command sequences and the arithmetic):

  aux-self      AUX OUT set to the trigger output (:SYSTem:AOUTput TOUT), cabled into one of the unit's own inputs,
                totalled by the unit's own counter (:COUNter:MODE TOTalize). Rate = (count1 - count0) / (t1 - t0),
                t taken when each :COUNter:CURRent? answer arrives. Caveat: the counter's input is digitised even
                when hidden, so on the MHO900 the signal channel runs at 2 GSa/s instead of 4 GSa/s at 20 ns/div:
                aux-self understates the one-channel rate (clean stock MHO934: 25.4k alone, counted by a second
                scope; 13.2k with its own counter on its hidden CH4, 9.1k with CH4 shown). Use aux-external for the
                one-channel figure.
  aux-external  AUX OUT of the unit under test into another Rigol DHO/MHO (--counter HOST) that counts it. The unit
                under test runs the signal channel alone, at its full sample rate. Needs a second scope.

The kit never sends :DISPlay:DATA? (on Kiszka's mod v0.1.5 an abandoned screenshot read left SCPI unanswered until
a reboot), never saves files on the scope, and sends *RST and :SYSTem:SETup only in `preset`, behind its own flag.
No calibration and no other system commands.
"""

import argparse
import datetime
import json
import math
import os
import platform
import re
import shlex
import signal
import socket
import sys
import time

KIT_VERSION = "1.0.0"
SCPI_PORT = 5555
LEVEL_V = 1.5  # the counter's threshold (its input's trigger level) and the signal's trigger level

# Commands the kit refuses to send, whatever the caller asks (matched on the long or short form).
_REFUSED = [
    (re.compile(r"^:?DISP(LAY)?:DATA", re.I), "a screenshot (:DISPlay:DATA?) hangs SCPI on Kiszka's mod v0.1.5"),
    (re.compile(r"^\*RST", re.I), "a reset"),
    (re.compile(r"^\*(CAL|TST)", re.I), "calibration / self-test"),
    (re.compile(r"^:?CAL", re.I), "calibration"),
    (re.compile(r"^:?SAVE", re.I), "a save to the scope's storage"),
    (re.compile(r"^:?SYST(EM)?:(RES|ERAS|UPD|SET|LOCK|OPT|KEY|PON|DATE|TIME|LANG)", re.I), "a system setting"),
]


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


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)


def nrf(x):
    """A number as the bench engine writes it: 1.5 -> 1.5E0, 2e-8 -> 2E-8, -1.5 -> -1.5E0, 0 -> 0."""
    if x == 0:
        return "0"
    e = int(math.floor(math.log10(abs(x))))
    m = round(x / 10 ** e, 12)
    if abs(m) >= 10:
        m, e = m / 10, e + 1
    s = ("%.12f" % m).rstrip("0").rstrip(".")
    return "%sE%d" % (s, e)


def trim_num(x):
    """3.0 -> 3, 10e6 -> 10000000, 0.5 -> 0.5"""
    if float(x).is_integer():
        return str(int(x))
    return repr(float(x))


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(reply):
    """:ACQuire:MDEPth? answer -> a settable value (1.0000E+03 -> 1k; AUTO stays)."""
    v = parse_num(reply)
    if v is None:
        return reply
    n = int(round(v))
    if n >= 1000000 and n % 1000000 == 0:
        return "%dM" % (n // 1000000)
    if n >= 1000 and n % 1000 == 0:
        return "%dk" % (n // 1000)
    return str(n)


# ---------------------------------------------------------------------------------------------------- SCPI


class Scpi:
    """SCPI on the raw socket. 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."""

    def __init__(self, host, port=SCPI_PORT, timeout=5.0, clock=None, trace=None, label="uut", echo=False,
                 redact=None, pace=0.075):
        self.host, self.port, self.timeout = host, port, timeout
        self.clock = clock or Clock()
        self.trace = trace if trace is not None else []
        self.label, self.echo = label, echo
        self.redact = redact or (lambda s: s)
        self.pace = pace
        self.sock = None
        self.buf = b""
        self.t0 = self.clock.now()
        self.sent = []  # every line sent, for the tests

    # -- connection
    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, allow=False):
        """`allow`: only the commands that exist to send one of the refused forms pass it, each behind its own
        explicit flag (preset: *RST and :SYSTem:SETup)."""
        for rx, why in _REFUSED:
            if rx.search(cmd.strip()) and not allow:
                raise KitError("refused to send %r: %s" % (cmd, why))
        if self.sock is None:
            self.connect()
        self.sent.append(cmd)
        try:
            self.sock.sendall((cmd + "\n").encode("ascii"))
        except OSError as e:
            raise KitError("send %r failed: %s" % (cmd, e))

    def _send_raw_header(self, header):
        if self.sock is None:
            self.connect()
        self.sent.append(header)
        try:
            self.sock.sendall((header + " ").encode("ascii"))
        except OSError as e:
            raise KitError("send %r failed: %s" % (header, e))

    def _readline(self, timeout):
        deadline = time.monotonic() + timeout
        while b"\n" not in self.buf:
            left = deadline - time.monotonic()
            if left <= 0:
                raise ScpiTimeout("no answer")
            self.sock.settimeout(left)
            try:
                chunk = self.sock.recv(65536)
            except socket.timeout:
                raise ScpiTimeout("no answer")
            except OSError as e:
                raise KitError("read failed: %s" % e)
            if not chunk:
                raise KitError("the scope closed the connection")
            self.buf += chunk
        line, _, self.buf = self.buf.partition(b"\n")
        return line.decode("ascii", "replace").strip()

    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), "host": self.label, "send": cmd}
        if reply is not None:
            e["reply"] = self.redact(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("[%s %8.3f] %s%s\n" % (self.label, e["s"], cmd, "" if reply is None and not note else
                                                    "  -> " + (e.get("reply") if reply is not None else note)))

    # -- the three operations
    def write(self, cmd, pace=None, allow=False):
        t = self.clock.now()
        self._send(cmd, allow)
        self._log(cmd, None, t, None)
        self.clock.sleep(self.pace if pace is None else pace)

    def block(self, cmd, timeout=20.0, allow=False):
        """An IEEE 488.2 definite block answer (#<n><len><bytes>), as bytes."""
        t = self.clock.now()
        self._send(cmd, allow)
        deadline = time.monotonic() + timeout

        def need(n):
            while len(self.buf) < n:
                left = deadline - time.monotonic()
                if left <= 0:
                    raise ScpiTimeout("%s: block incomplete after %.0f s" % (cmd, timeout))
                self.sock.settimeout(left)
                try:
                    chunk = self.sock.recv(65536)
                except socket.timeout:
                    continue
                if not chunk:
                    raise KitError("the scope closed the connection")
                self.buf += chunk
        try:
            need(2)
            if self.buf[:1] != b"#":
                raise KitError("%s: not a block: %r" % (cmd, self.buf[:20]))
            nd = int(self.buf[1:2])
            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 = 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 write_block(self, header, data, pace=1.0):
        """`header #<n><len><bytes>`: a setting that takes a block (:SYSTem:SETup)."""
        t = self.clock.now()
        self._send_raw_header(header)
        self.sent[-1] = "%s <%d-byte block>" % (header, len(data))
        ln = str(len(data))
        try:
            self.sock.sendall(b"#" + str(len(ln)).encode() + ln.encode() + data + b"\n")
        except OSError as e:
            raise KitError("send %s block failed: %s" % (header, e))
        self._log("%s <%d-byte block>" % (header, len(data)), None, t, None)
        self.clock.sleep(pace)

    def query(self, cmd, timeout=None, allow=False):
        """The answer, or ScpiTimeout (the connection is checked and replaced first) / ScpiHung."""
        t = self.clock.now()
        self._send(cmd, allow)
        try:
            r = self._readline(self.timeout if timeout is None else timeout)
        except ScpiTimeout:
            self._log(cmd, None, t, self.clock.now(), "TIMEOUT")
            self._recover()
            raise ScpiTimeout("%s: no answer within %.1f s" % (cmd, self.timeout if timeout is None else timeout))
        t1 = self.clock.now()
        self._log(cmd, r, t, t1)
        self.last_reply_at = t1
        self.last_reply_wall = self.clock.wall()
        return r

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

    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("%s:%d stopped answering *IDN? on a fresh connection" % (self.host, self.port))
        self._log("*IDN?", r, t, self.clock.now(), "connection replaced after a timeout")

    def drain_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


# ---------------------------------------------------------------------------------------------------- identity


class Identity:
    def __init__(self, idn, novaos):
        parts = [p.strip() for p in idn.split(",")]
        while len(parts) < 4:
            parts.append("")
        self.idn = idn
        self.maker, self.model, self.serial, self.fw = parts[:4]
        m = re.match(r"^(DHO|MHO)(\d)(\d)(\d)", self.model.upper())
        if not m:
            raise KitError("%s %s is not a DHO800, DHO900 or MHO900" % (self.maker, self.model))
        series = m.group(1) + m.group(2) + "00"
        if series not in ("DHO800", "DHO900", "MHO900"):
            raise KitError("%s %s is not a DHO800, DHO900 or MHO900" % (self.maker, self.model))
        self.series = series
        self.family = "mho" if series == "MHO900" else "dho"
        self.channels = int(m.group(4))
        if "NK TECHNOLOGIES" in self.maker.upper():
            self.firmware = "nk-mod"
        elif novaos:
            self.firmware = "novaos"
        elif "RIGOL" in self.maker.upper():
            self.firmware = "rigol-stock"
        else:
            self.firmware = "unknown"
        self.novaos_version = novaos

    def redact(self, show_serial):
        if show_serial or not self.serial:
            return lambda s: s
        ser = self.serial
        return lambda s: s.replace(ser, "<serial hidden>") if isinstance(s, str) else s

    def as_dict(self, show_serial):
        r = self.redact(show_serial)
        return {"idn": r(self.idn), "maker": self.maker, "model": self.model, "serial": r(self.serial),
                "idn_firmware": self.fw, "series": self.series, "family": self.family, "channels": self.channels,
                "firmware": self.firmware, "novaos_version": self.novaos_version}

    # bench defaults (ours; say yours with --signal-input / --aux-input)
    def default_signal_input(self):
        return 1 if self.family == "mho" else 4

    def default_aux_input(self):
        return 4 if self.family == "mho" else 3

    def default_generator(self):
        return "SOURce2" if self.family == "mho" else "SOURce"


def identify(s):
    idn = s.query("*IDN?")
    if "NK TECHNOLOGIES" in idn.upper():
        nova = None  # the mod's maker string says it all; no unknown query sent
    else:
        nova = s.try_query(":NOVAos:VERSion?", min(2.0, s.timeout))
        if nova is None:
            s.drain_errors(3)  # the stock queues -100 for the unknown query
    return Identity(idn, nova)


# ---------------------------------------------------------------------------------------------------- snapshot


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


class Snapshot:
    """What the kit may touch on one scope, read before anything changes."""

    def __init__(self, s, ident, gen, level_sources):
        self.s, self.ident, self.gen = s, ident, gen
        self.values = {}
        self.levels = {}  # edge trigger level per source channel
        self.level_sources = sorted(set(level_sources))
        self.state = None
        self.edge_source = None

    def keys(self):
        n = self.ident.channels
        k = [":TRIGger:SWEep"]
        k += [":CHANnel%d:DISPlay" % i for i in range(1, n + 1)]
        for i in range(1, n + 1):
            k += [":CHANnel%d:%s" % (i, x) for x in ("SCALe", "OFFSet", "PROBe", "COUPling", "BWLimit")]
        k += [":TIMebase:MAIN:SCALe", ":TIMebase:MAIN:OFFSet", ":ACQuire:MDEPth", ":ACQuire:TYPE",
              ":TRIGger:MODE", ":TRIGger:EDGE:SOURce", ":TRIGger:EDGE:SLOPe", ":TRIGger:EDGE:LEVel",
              ":TRIGger:SWEep"]
        if self.gen:
            k += [":%s:%s" % (self.gen, x) for x in ("FUNCtion", "FREQuency", "VOLTage:AMPLitude", "VOLTage:OFFSet",
                                                     "OUTPut:STATe")]
        k += [":SYSTem:AOUTput", ":DISPlay:TYPE", ":DISPlay:GRADing:TIME", ":DISPlay:COLor",
              ":COUNter:ENABle", ":COUNter:SOURce", ":COUNter:MODE", ":COUNter:TOTalize:ENABle"]
        return k

    def take(self):
        s = self.s
        s.drain_errors(5)
        self.state = s.query(":TRIGger:STATus?")
        for k in self.keys():
            v = s.try_query(k + "?")
            if v is not None and k not in self.values:
                self.values[k] = v
        missing = [k for k in self.keys() if k not in self.values]
        s.drain_errors(5)
        # the edge level is kept per source: read it on each source the kit will set it on, then put the source back
        for c in self.level_sources:
            s.write(":TRIGger:EDGE:SOURce " + ch(c))
            v = s.try_query(":TRIGger:EDGE:LEVel?")
            if v is not None:
                self.levels[c] = v
            s.drain_errors(3)
        self.edge_source = self.values.get(":TRIGger:EDGE:SOURce")
        if self.edge_source:
            s.write(":TRIGger:EDGE:SOURce " + self.edge_source)
        s.drain_errors(3)
        return missing

    def has(self, key):
        return key in self.values

    def as_dict(self):
        return {"state": self.state, "values": self.values, "edge_levels": {str(k): v for k, v in self.levels.items()},
                "generator": self.gen}

    @classmethod
    def from_dict(cls, s, ident, d):
        sn = cls(s, ident, d.get("generator"), [int(k) for k in d.get("edge_levels", {})])
        sn.values = dict(d["values"])
        sn.levels = {int(k): v for k, v in d.get("edge_levels", {}).items()}
        sn.state = d.get("state")
        sn.edge_source = sn.values.get(":TRIGger:EDGE:SOURce")
        return sn


def _bool_key(k):
    return k.endswith((":DISPlay", ":STATe", ":ENABle", ":COLor"))


def settable(key, reply):
    """A query's answer as a value to write back: 1/0 -> ON/OFF on the switches, depths as 1k/10k; the rest as
    answered (the answers are the commands' short forms, e.g. CHAN1, NORM, TOT, SQU, which the scope accepts)."""
    r = reply.strip()
    if _bool_key(key) and r in ("0", "1"):
        return "ON" if r == "1" else "OFF"
    if key.endswith(":MDEPth"):
        return depth_word(r)
    return r


def same(key, a, b):
    if a is None or b is None:
        return a == b
    if _bool_key(key):
        norm = {"1": "ON", "0": "OFF"}
        return norm.get(a.strip().upper(), a.strip().upper()) == norm.get(b.strip().upper(), b.strip().upper())
    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()


def restore(sn):
    """Put every snapshot value back (only those that differ now), in the order the stock accepts them: channels on
    before their scale and offset (a hidden channel ignores them on the stock), :RUN before the depth, trigger levels
    per source, the counter last, then the channels' display and the run state. Returns a report."""
    s, v = sn.s, sn.values
    rep = {"written": [], "failed": [], "still_different": [], "not_readable": []}

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

    def put(k, want=None):
        if k not in v:
            return
        now = cur(k)
        if now is None or same(k, now, v[k]):
            return
        val = want if want is not None else settable(k, v[k])
        try:
            s.write("%s %s" % (k, val), pace=_pace(k))
            rep["written"].append("%s %s" % (k, val))
        except KitError as e:
            rep["failed"].append("%s %s (%s)" % (k, val, e))

    n = sn.ident.channels
    chans = range(1, n + 1)
    # 1. a channel whose settings must come back is shown first
    for i in chans:
        k = ":CHANnel%d:DISPlay" % i
        if k not in v:
            continue
        now = cur(k)
        if now is None or same(k, now, "ON"):
            continue
        # the mod ignores a hidden channel's scale, offset and trigger level: show it while they go back
        need = settable(k, v[k]) == "ON" or i in sn.levels
        if not need:
            for x in (":CHANnel%d:SCALe" % i, ":CHANnel%d:OFFSet" % i):
                if x in v:
                    c = cur(x)
                    need = need or (c is not None and not same(x, c, v[x]))
        if need:
            s.write(k + " ON")
            rep["written"].append(k + " ON")
    for part in ("SCALe", "OFFSet", "COUPling", "BWLimit", "PROBe"):
        for i in chans:
            put(":CHANnel%d:%s" % (i, part))
    put(":TIMebase:MAIN:SCALe")
    put(":TIMebase:MAIN:OFFSet")
    put(":ACQuire:TYPE")
    if ":ACQuire:MDEPth" in v:
        now = cur(":ACQuire:MDEPth")
        if now is not None and not same(":ACQuire:MDEPth", now, v[":ACQuire:MDEPth"]):
            s.write(":RUN")  # the depth is taken while running
            s.try_query(":TRIGger:STATus?")
            put(":ACQuire:MDEPth")
    for k in (":DISPlay:TYPE", ":DISPlay:GRADing:TIME", ":DISPlay:COLor"):
        put(k)
    if sn.levels:
        for c in sorted(sn.levels):
            s.write(":TRIGger:EDGE:SOURce " + ch(c))
            now = s.try_query(":TRIGger:EDGE:LEVel?")
            if now is None:
                rep["not_readable"].append(":TRIGger:EDGE:LEVel (%s)" % ch(c))
            elif not same(":TRIGger:EDGE:LEVel", now, sn.levels[c]):
                s.write(":TRIGger:EDGE:LEVel " + sn.levels[c])
                rep["written"].append(":TRIGger:EDGE:LEVel %s (%s)" % (sn.levels[c], ch(c)))
        put(":TRIGger:EDGE:SOURce")
        if sn.edge_source and ":TRIGger:EDGE:SOURce" not in [w.split(" ")[0] for w in rep["written"]]:
            # the level loop moved the source; put it back even though the first read matched
            now = s.try_query(":TRIGger:EDGE:SOURce?")
            if now is not None and not same(":TRIGger:EDGE:SOURce", now, sn.edge_source):
                s.write(":TRIGger:EDGE:SOURce " + sn.edge_source)
                rep["written"].append(":TRIGger:EDGE:SOURce " + sn.edge_source)
    else:
        put(":TRIGger:EDGE:SOURce")
    put(":TRIGger:EDGE:SLOPe")
    put(":TRIGger:SWEep")
    put(":TRIGger:MODE")
    if sn.gen:
        for x in ("FUNCtion", "FREQuency", "VOLTage:AMPLitude", "VOLTage:OFFSet", "OUTPut:STATe"):
            put(":%s:%s" % (sn.gen, x))
    # the counter is switched off before its mode or source changes back: on our stock MHO934, switching it to FREQ
    # while its source was still a shown, unconnected input was followed within a second by the app aborting
    # (std::out_of_range in stoi, 2026-10-16 21:26:47; METHOD.md run J)
    ck = (":COUNter:TOTalize:ENABle", ":COUNter:MODE", ":COUNter:SOURce")
    if ":COUNter:ENABle" in v:
        on_now = cur(":COUNter:ENABle")
        if on_now is not None and same(":COUNter:ENABle", on_now, "ON"):
            moves = [k for k in ck if k in v and not same(k, cur(k) or "", v[k])]
            if moves or settable(":COUNter:ENABle", v[":COUNter:ENABle"]) == "OFF":
                s.write(":COUNter:ENABle OFF", pace=0.5)
                rep["written"].append(":COUNter:ENABle OFF")
    for k in ck + (":COUNter:ENABle", ":SYSTem:AOUTput"):
        put(k)
    for i in chans:
        put(":CHANnel%d:DISPlay" % i)
    # 2. the run state
    st = s.try_query(":TRIGger:STATus?")
    if st is not None and sn.state:
        was_stopped = sn.state.upper().startswith("STOP")
        if was_stopped and not st.upper().startswith("STOP"):
            s.write(":STOP")
            rep["written"].append(":STOP")
        elif not was_stopped and st.upper().startswith("STOP"):
            s.write(":RUN")
            rep["written"].append(":RUN")
    # 3. read everything back
    for k in sn.keys():
        if k in v:
            now = s.try_query(k + "?", 4.0)
            if now is None:
                if k not in rep["not_readable"]:
                    rep["not_readable"].append(k)
            elif not same(k, now, v[k]):
                rep["still_different"].append({"key": k, "was": v[k], "now": now})
    rep["errors"] = s.drain_errors(10)
    rep["complete"] = not rep["failed"] and not rep["still_different"]
    return rep


def _pace(cmd):
    """Waits after a write, as the bench engine paced them: the counter's setup 0.5 s, the generator 0.18 s."""
    c = cmd.upper()
    if c.startswith((":COUNTER:ENABLE ON", ":COUNTER:SOURCE", ":COUNTER:SOUR ", ":COUNTER:MODE", ":COUN:MODE")):
        return 0.5
    if re.match(r"^:SOURCE\d?:(FREQ|VOLT|OUTP)", c):
        return 0.175
    return 0.075


# ---------------------------------------------------------------------------------------------------- the run


def scenes(args):
    out = []
    for chans in args.channels:
        for tdiv in args.tdiv:
            for depth in args.depth:
                out.append({"tdiv": tdiv, "depth": depth, "channels": chans})
    return out


def scene_label(sc):
    t = sc["tdiv"]
    tl = ("%g ns/div" % (t * 1e9)) if t < 1e-6 else ("%g us/div" % (t * 1e6)) if t < 1e-3 else ("%g ms/div" % (t * 1e3))
    return "%s, %s, %s" % (tl, sc["depth"], "one input" if sc["channels"] == "one" else "every input")


class Runner:
    def __init__(self, args, clock=None):
        self.a = args
        self.clock = clock or Clock()
        self.trace = []
        self.report = {"kit": "rate_kit", "kit_version": KIT_VERSION, "argv": list(args.argv),
                       "started": utc_iso(self.clock.wall()), "host_pc": {"python": platform.python_version(),
                                                                          "os": platform.system(), "machine": platform.machine()},
                       "methods": args.methods, "scopes": {}, "scenes": [], "notes": [], "restore": {}, "trace": self.trace}
        self.snaps = {}
        self.scopes = {}

    # -- helpers
    def say(self, msg):
        if not self.a.quiet:
            sys.stderr.write(msg + "\n")
            sys.stderr.flush()

    def open_scope(self, label, hostport):
        host, port = split_host(hostport)
        s = Scpi(host, port, timeout=self.a.timeout, clock=self.clock, trace=self.trace, label=label, echo=self.a.trace)
        s.connect()
        ident = identify(s)
        s.redact = ident.redact(self.a.show_serial)
        for e in self.trace:
            if "reply" in e:
                e["reply"] = s.redact(e["reply"])
        d = ident.as_dict(self.a.show_serial)
        d["address"] = hostport if self.a.show_address else "<address hidden>"
        self.report["scopes"][label] = d
        self.scopes[label] = (s, ident)
        self.say("%s: %s %s, %s firmware %s%s" % (label, ident.maker, ident.model, ident.firmware, ident.fw,
                                                   (" (NovaOS %s)" % ident.novaos_version) if ident.novaos_version else ""))
        return s, ident

    def snapshot(self, label, gen, level_sources):
        s, ident = self.scopes[label]
        sn = Snapshot(s, ident, gen, level_sources)
        missing = sn.take()
        self.snaps[label] = sn
        if missing:
            self.report["notes"].append("%s: not answered, so neither changed nor restored: %s" % (label, ", ".join(missing)))
        self.save_snapshots()
        return sn

    def save_snapshots(self):
        d = {lab: {"identity": self.report["scopes"][lab], "snapshot": sn.as_dict()} for lab, sn in self.snaps.items()}
        with open(self.path("snapshot", "json"), "w", encoding="utf-8", newline="\n") as f:
            json.dump(d, f, indent=1)

    def path(self, kind, ext):
        return os.path.join(self.a.out, "%s-%s.%s" % (kind, self.a.stamp, ext))

    # -- setup in the bench engine's order (METHOD.md "aux-self, command by command")
    def preamble_counter(self, s, sn, aux, display_bench):
        a = ch(aux)
        disp = sn.values.get(":%s:DISPlay" % a)
        if disp is not None and settable(":X:DISPlay", disp) == "OFF":
            # Kiszka's mod ignores a hidden channel's scale, offset and trigger level (no error): set them shown
            s.write(":%s:DISPlay ON" % a)
        s.write(":%s:SCALe 1" % a)
        s.write(":%s:OFFSet -2" % a)
        s.write(":TRIGger:EDGE:SOURce " + a)
        s.write(":TRIGger:EDGE:LEVel " + nrf(LEVEL_V))
        for k, val in (display_bench or ()):
            if sn.has(k):
                s.write("%s %s" % (k, val))
        s.write(":COUNter:ENABle ON", pace=0.5)
        s.write(":COUNter:SOURce " + a, pace=0.5)
        s.write(":COUNter:MODE TOTalize", pace=0.5)
        s.write(":COUNter:TOTalize:ENABle ON")

    def setup_scene(self, s, ident, sc, sig, aux_shown, gen):
        """The scene, as the bench engine's `ctx.setup(View)` wrote it."""
        a = self.a
        s.write(":RUN")
        s.try_query(":TRIGger:STATus?")
        if sc["depth"].lower() != "as-is":
            s.write(":ACQuire:MDEPth %s" % sc["depth"])
        s.try_query(":ACQuire:MDEPth?")
        if a.scene == "minimal":
            # only what the scene needs: the signal, the time base, the trigger; channels, probes, the acquisition
            # type and the display stay as the scope has them (e.g. right after its own preset)
            if gen:
                s.write(":%s:FUNCtion %s" % (gen, a.function))
                s.write(":%s:FREQuency %s" % (gen, trim_num(a.frequency)), pace=0.175)
                s.write(":%s:VOLTage:AMPLitude %s" % (gen, trim_num(a.amplitude)), pace=0.175)
                s.write(":%s:VOLTage:OFFSet %s" % (gen, nrf(a.amplitude / 2.0)), pace=0.175)
                s.write(":%s:OUTPut:STATe ON" % gen, pace=0.175)
            # the signal's input at 1 V/div, centred: a 1.5 V level is clamped to the screen (50 mV/div after *RST)
            s.write(":CHANnel%d:DISPlay ON" % sig)
            s.write(":CHANnel%d:SCALe 1" % sig)
            s.write(":CHANnel%d:OFFSet %s" % (sig, nrf(-a.amplitude / 2.0)))
            s.write(":TIMebase:MAIN:SCALe %s" % nrf(sc["tdiv"]))
            s.write(":TRIGger:EDGE:SOURce " + ch(sig))
            s.write(":TRIGger:EDGE:LEVel " + nrf(LEVEL_V))
            s.write(":TRIGger:SWEep NORMal")
            return
        s.write(":ACQuire:TYPE NORMal")
        if gen:
            s.write(":%s:FUNCtion %s" % (gen, a.function))
            s.write(":%s:FREQuency %s" % (gen, trim_num(a.frequency)), pace=0.175)
            s.write(":%s:VOLTage:AMPLitude %s" % (gen, trim_num(a.amplitude)), pace=0.175)
            s.write(":%s:VOLTage:OFFSet %s" % (gen, nrf(a.amplitude / 2.0)), pace=0.175)
            s.write(":%s:OUTPut:STATe ON" % gen, pace=0.175)
        n = ident.channels
        if sc["channels"] == "all":
            on = [sig] + [i for i in range(1, n + 1) if i != sig]
        else:
            on = [sig] + ([aux_shown] if aux_shown else [])
        for i in range(1, n + 1):
            if i not in on:
                s.write(":CHANnel%d:DISPlay OFF" % i)
        for i in on:
            s.write(":CHANnel%d:DISPlay ON" % i)
        # the channels the scene draws: the signal, and with every input on, the others (AUX input at offset -2)
        view = [(sig, "1", nrf(-a.amplitude / 2.0))]
        if sc["channels"] == "all":
            view += [(i, "1", "-2" if i == aux_shown else "0") for i in range(1, n + 1) if i != sig]
        for i, scale, off in view:
            s.write(":CHANnel%d:PROBe 1" % i)
            s.write(":CHANnel%d:COUPling DC" % i)
            s.write(":CHANnel%d:BWLimit OFF" % i)
            s.write(":CHANnel%d:SCALe %s" % (i, scale))
            s.write(":CHANnel%d:OFFSet %s" % (i, off))
        s.write(":TIMebase:MAIN:SCALe %s" % nrf(sc["tdiv"]))
        s.write(":TIMebase:MAIN:OFFSet 0")
        s.write(":TRIGger:MODE EDGE")
        s.write(":TRIGger:EDGE:SOURce " + ch(sig))
        s.write(":TRIGger:EDGE:SLOPe POSitive")
        s.write(":TRIGger:EDGE:LEVel " + nrf(LEVEL_V))
        s.write(":TRIGger:SWEep NORMal")

    def readback(self, s, ident, gen, extra=()):
        keys = [":ACQuire:MDEPth", ":ACQuire:SRATe", ":ACQuire:TYPE", ":TIMebase:MAIN:SCALe", ":TIMebase:MAIN:OFFSet",
                ":TRIGger:MODE", ":TRIGger:EDGE:SOURce", ":TRIGger:EDGE:LEVel", ":TRIGger:EDGE:SLOPe", ":TRIGger:SWEep"]
        for i in range(1, ident.channels + 1):
            keys += [":CHANnel%d:DISPlay" % i]
        keys += list(extra)
        if gen:
            keys += [":%s:%s" % (gen, x) for x in ("FUNCtion", "FREQuency", "VOLTage:AMPLitude", "VOLTage:OFFSet", "OUTPut:STATe")]
        keys += [":SYSTem:AOUTput", ":DISPlay:TYPE", ":DISPlay:GRADing:TIME", ":DISPlay:COLor",
                 ":COUNter:ENABle", ":COUNter:SOURce", ":COUNter:MODE", ":COUNter:TOTalize:ENABle"]
        return {k: s.try_query(k + "?") for k in keys}

    # -- one window
    def window(self, s, counter, uut):
        """Clear, wait, read the count, poll for --seconds, read again. Returns the raw readings."""
        a = self.a
        res = {}
        if counter is not None:
            counter.write(":COUNter:TOTalize:CLEar")
        self.clock.sleep(a.clear_wait)
        polls = []
        if counter is not None:
            c0 = parse_num(counter.query(":COUNter:CURRent?"))
            t0, w0 = counter.last_reply_at, counter.last_reply_wall
        if uut is not None and self.scopes["uut"][1].firmware == "novaos":
            res["novaos_density_rate"] = uut.try_query(":NOVAos:DENSity:RATE?", 2.0)
        k = max(1, int(round(a.seconds)))
        if a.poll == "quiet":  # nothing on the wire between the two counter reads
            self.clock.sleep(a.seconds)
            k = 0
        for _ in range(k):
            self.clock.sleep(a.seconds / k)
            p = {"t": utc_iso(self.clock.wall()), "status": (uut.try_query(":TRIGger:STATus?") or "TIMEOUT").upper()}
            if a.poll == "counter" and counter is not None:
                p["count"] = parse_num(counter.query(":COUNter:CURRent?"))
                p["s"] = round(counter.last_reply_at - t0, 4)
            polls.append(p)
        if counter is not None:
            c1 = parse_num(counter.query(":COUNter:CURRent?"))
            t1, w1 = counter.last_reply_at, counter.last_reply_wall
            dt = t1 - t0
            res.update({"count0": c0, "t0": utc_iso(w0), "count1": c1, "t1": utc_iso(w1), "seconds": round(dt, 4),
                        "rate": (c1 - c0) / dt if (c0 is not None and c1 is not None and dt > 0) else None})
            pts = [(0.0, c0)] + [(p["s"], p["count"]) for p in polls if p.get("count") is not None] + [(dt, c1)]
            if len(pts) > 2 and all(c is not None for _, c in pts):
                res["rate_fit"] = slope(pts)
                res["distinct_counts"] = len(set(c for _, c in pts))
        res["polls"] = polls
        return res


    # -- the whole run
    def run(self):
        a = self.a
        methods = a.methods
        uut, ident = self.open_scope("uut", a.host)
        sig = a.signal_input or ident.default_signal_input()
        aux = a.aux_input or ident.default_aux_input()
        gen = None if a.signal == "external" else (a.generator or ident.default_generator())
        counter = None
        cident = None
        if "aux-external" in methods:
            if not a.counter:
                raise KitError("aux-external needs --counter HOST (the scope that counts the unit's AUX OUT)")
            counter, cident = self.open_scope("counter", a.counter)
        if "aux-self" in methods:
            if sig == aux:
                raise KitError("the signal (CH%d) and AUX OUT (CH%d) need two different inputs" % (sig, aux))
            counter = uut
            cident = ident
        self.report["wiring"] = {"signal_input": sig, "aux_input": aux if "aux-self" in methods else None,
                                 "counter_input": a.counter_input if "aux-external" in methods else None,
                                 "generator": gen or "external", "counter_channel": a.counter_channel}
        # the snapshots before anything changes
        load = a.load_counter_input
        self.snapshot("uut", gen, [sig] + ([aux] if "aux-self" in methods else []) + ([load] if load else []))
        if "aux-external" in methods:
            self.snapshot("counter", None, [a.counter_input])
        failed = None
        try:
            if a.clear_measurements:
                # a measurement on screen costs the stock about 21 % (METHOD.md run G); its items are not restorable
                uut.write(":MEASure:DELete" if ident.family == "mho" else ":MEASure:CLEar", pace=0.5)
                self.report["notes"].append("measurement items cleared (--clear-measurements); not restored")
            # preambles
            uut.write(":SYSTem:AOUTput TOUT")
            if "aux-self" in methods:
                self.preamble_counter(uut, self.snaps["uut"], aux, display_values(a))
            elif display_values(a):
                for k, val in display_values(a):
                    if self.snaps["uut"].has(k):
                        uut.write("%s %s" % (k, val))
            if "aux-external" in methods:
                self.preamble_counter(counter, self.snaps["counter"], a.counter_input, None)
                counter.write(":TRIGger:SWEep AUTO")  # the counting scope runs free; its own trigger does not matter
                counter.write(":RUN")
            if load:
                # the unit's own counter on one of its inputs, set up while shown, not read: what counting on the
                # unit's own input costs it, measured by the external counter
                self.preamble_counter(uut, self.snaps["uut"], load, None)
                if a.load_counter_channel == "hidden":
                    uut.write(":CHANnel%d:DISPlay OFF" % load)
                self.report["notes"].append("the unit's own counter runs on CH%d (%s), not read" % (load, a.load_counter_channel))
            counting = "aux-self" in methods or "aux-external" in methods
            aux_shown = aux if ("aux-self" in methods and a.counter_channel == "shown") else None
            if load and a.load_counter_channel == "shown":
                aux_shown = load
            for sc in scenes(a):
                label = scene_label(sc)
                self.say("scene: " + label)
                self.setup_scene(uut, ident, sc, sig, aux_shown, gen)
                pre = self.readback(uut, ident, gen)
                # the bench engine's signal check; on the stock and the mod it may put a live measurement on screen
                vpp = uut.try_query(":MEASure:ITEM? VPP,%s" % ch(sig)) if a.signal_check else None
                uut.write(":RUN")
                self.clock.sleep(a.settle)
                post = self.readback(uut, ident, gen)
                errs = uut.drain_errors(10)
                n_on = sum(1 for i in range(1, ident.channels + 1) if settable(":X:DISPlay", post.get(":CHANnel%d:DISPlay" % i) or "0") == "ON")
                rec = {"scene": sc, "label": label, "settings_after_setup": pre, "settings_while_counting": post,
                       "signal_vpp": vpp, "scope_errors": errs, "channels_on": n_on, "repeats": []}
                if counter is not None and counter is not uut:
                    rec["counter_settings"] = self.readback(counter, cident, None)
                for r in range(a.repeats):
                    raw = self.window(uut, counter if counting else None, uut)
                    raw["repeat"] = r + 1
                    rec["repeats"].append(raw)
                    msg = "  repeat %d:" % (r + 1)
                    if raw.get("rate") is not None:
                        msg += " %.0f wfm/s (AUX OUT counted)" % raw["rate"]
                    self.say(msg)
                rec["srate"] = uut.try_query(":ACQuire:SRATe?")
                rec["mdepth"] = uut.try_query(":ACQuire:MDEPth?")
                rec["summary"] = summarize(rec)
                if counting and rec["summary"].get("aux_rate", {}).get("max") in (0, 0.0):
                    trig = set(p["status"] for r in rec["repeats"] for p in r["polls"])
                    if trig & {"TD", "RUN", "AUTO"}:
                        rec["warning"] = ("the counter read 0 while the scope triggered (%s): no AUX OUT pulses reached "
                                          "the counter's input, or its threshold (that input's trigger level) is not "
                                          "inside the pulses" % "/".join(sorted(trig)))
                self.report["scenes"].append(rec)
            if counter is not None:
                counter.write(":COUNter:TOTalize:ENABle OFF")
                self.report["errors_at_end"] = counter.drain_errors(10)
        except KeyboardInterrupt:
            failed = "interrupted (Ctrl-C)"
        except ScpiHung as e:
            failed = "hung: %s" % e
        except Exception as e:  # restore whatever happened
            failed = "%s: %s" % (type(e).__name__, e)
        finally:
            self.restore_all(failed)
        return failed

    def restore_all(self, failed):
        old = None
        try:
            old = signal.signal(signal.SIGINT, signal.SIG_IGN)  # a second Ctrl-C must not stop the restore
        except ValueError:
            pass
        try:
            for label, sn in self.snaps.items():
                if failed and failed.startswith("hung") and label == "uut":
                    self.report["restore"][label] = {"complete": False, "why": "SCPI hung: nothing could be sent. When "
                                                     "it answers again: rate_kit.py restore --host ... --from " +
                                                     self.path("snapshot", "json")}
                    continue
                try:
                    self.report["restore"][label] = restore(sn)
                except Exception as e:
                    self.report["restore"][label] = {"complete": False, "why": "%s: %s" % (type(e).__name__, e)}
        finally:
            if old is not None:
                signal.signal(signal.SIGINT, old)
        self.report["finished"] = utc_iso(self.clock.wall())
        self.report["failed"] = failed


def display_values(a):
    """--display: the writes it implies, or None to leave the display alone."""
    return {"bench": ((":DISPlay:TYPE", "VECTors"), (":DISPlay:GRADing:TIME", "INFinite"), (":DISPlay:COLor", "ON")),
            "default": ((":DISPlay:TYPE", "VECTors"), (":DISPlay:GRADing:TIME", "MIN"), (":DISPlay:COLor", "OFF")),
            }.get(a.display)


def slope(pts):
    n = len(pts)
    mx = sum(p[0] for p in pts) / n
    my = sum(p[1] for p in pts) / n
    sxx = sum((p[0] - mx) ** 2 for p in pts)
    return sum((p[0] - mx) * (p[1] - my) for p in pts) / sxx if sxx > 0 else None


def stats(xs):
    xs = [x for x in xs if x is not None]
    if not xs:
        return {}
    return {"n": len(xs), "mean": sum(xs) / len(xs), "min": min(xs), "max": max(xs)}


def summarize(rec):
    s = {}
    s["aux_rate"] = stats([r.get("rate") for r in rec["repeats"]])
    fit = stats([r.get("rate_fit") for r in rec["repeats"]])
    if fit:
        s["aux_rate_fit"] = fit
    nv = [parse_num((r.get("novaos_density_rate") or "").split(",")[0]) for r in rec["repeats"]]
    if any(v is not None for v in nv):
        s["novaos_density_rate"] = stats(nv)
    return {k: v for k, v in s.items() if v}


# ---------------------------------------------------------------------------------------------------- output


def fmt(v, nd=0):
    if v is None:
        return "-"
    if isinstance(v, float):
        return "{:,.{}f}".format(v, nd)
    return str(v)


def markdown(rep):
    L = ["# rate_kit run %s" % rep["started"], ""]
    L.append("- kit %s, Python %s on %s; command: `%s`" % (rep["kit_version"], rep["host_pc"]["python"],
                                                          rep["host_pc"]["os"], " ".join(shlex.quote(x) for x in rep["argv"])))
    for lab, d in rep["scopes"].items():
        L.append("- %s: `%s` -> %s, %s firmware %s%s" % (lab, d["idn"], d["series"], d["firmware"], d["idn_firmware"],
                                                         (" (NovaOS %s)" % d["novaos_version"]) if d.get("novaos_version") else ""))
    L.append("- methods: %s; wiring: %s" % (", ".join(rep["methods"]), json.dumps(rep.get("wiring", {}))))
    if rep.get("failed"):
        L.append("- **stopped early: %s**" % rep["failed"])
    L += ["", "## Results", "", "| scene | sample rate | depth read | channels on | AUX OUT wfm/s mean | min | max | "
          "slope fit | NovaOS engine |", "|---|---:|---:|---:|---:|---:|---:|---:|---:|"]
    for sc in rep["scenes"]:
        s = sc["summary"]
        ar, fit, nv = s.get("aux_rate", {}), s.get("aux_rate_fit", {}), s.get("novaos_density_rate", {})
        L.append("| %s | %s | %s | %s | %s | %s | %s | %s | %s |" % (
            sc["label"], sc.get("srate") or "-", sc.get("mdepth") or "-", sc["channels_on"], fmt(ar.get("mean")),
            fmt(ar.get("min")), fmt(ar.get("max")), fmt(fit.get("mean")), fmt(nv.get("mean")) if nv else "-"))
    for sc in rep["scenes"]:
        if sc.get("warning"):
            L.append("")
            L.append("**%s**: %s" % (sc["label"], sc["warning"]))
    L += ["", "## Raw counter readings", "", "| scene | repeat | t0 (UTC) | count0 | t1 (UTC) | count1 | seconds | rate | "
          "trigger status polled |", "|---|---:|---|---:|---|---:|---:|---:|---|"]
    for sc in rep["scenes"]:
        for r in sc["repeats"]:
            sts = ",".join(sorted(set(p["status"] for p in r["polls"])))
            L.append("| %s | %d | %s | %s | %s | %s | %s | %s | %s |" % (
                sc["label"], r["repeat"], r.get("t0", "-"), fmt(r.get("count0")), r.get("t1", "-"), fmt(r.get("count1")),
                fmt(r.get("seconds"), 3) if r.get("seconds") is not None else "-", fmt(r.get("rate"), 1), sts))
    L += ["", "## Settings read back while counting", ""]
    for sc in rep["scenes"]:
        L.append("- %s: %s" % (sc["label"], "; ".join("%s=%s" % (k, v) for k, v in sc["settings_while_counting"].items())))
        if sc.get("counter_settings"):
            L.append("  - counter scope: %s" % "; ".join("%s=%s" % (k, v) for k, v in sc["counter_settings"].items()))
    L += ["", "## Restore", ""]
    for lab, r in rep["restore"].items():
        if r.get("complete"):
            L.append("- %s: complete (%d settings written back)" % (lab, len(r.get("written", []))))
        else:
            L.append("- %s: **incomplete**: %s" % (lab, json.dumps({k: r.get(k) for k in ("why", "failed", "still_different") if r.get(k)})))
    for n in rep["notes"]:
        L.append("- note: %s" % n)
    L += ["", "%d SCPI lines in the JSON's `trace` (time, command, answer, ms)." % len(rep["trace"]), ""]
    return "\n".join(L)


def split_host(hp):
    if hp.count(":") == 1:
        h, p = hp.split(":")
        return h, int(p)
    return hp, SCPI_PORT


def csv_floats(s):
    return [float(x) for x in s.split(",") if x.strip()]


def csv_words(s):
    return [x.strip() for x in s.split(",") if x.strip()]


HELP_EPILOG = """
caveats, per method:
  aux-self      The counter's channel is digitised, shown or hidden: on the MHO900 at 20 ns/div the signal channel runs
                at 2 GSa/s instead of 4 GSa/s while counting, so this understates the one-channel rate by about half.
                Showing the counter's channel costs more again (clean stock MHO934, mean of 1k / 10k / AUTO: 25.4k
                alone, counted by a second scope; 13.2k with its own counter on its hidden CH4; 9.1k with CH4 shown).
                Kiszka's mod ignores a hidden channel's scale, offset and trigger level, so the kit sets them while
                the channel is shown, then hides it if asked. AUX OUT pulses once per accepted trigger; the stock's
                UltraAcquire pulses once per sequence.
  aux-external  Only the other scope's counter is used; it must be a Rigol DHO/MHO with :COUNter. The unit under test
                runs its signal channel alone at its full sample rate. Cable length and the counter's threshold
                (1.5 V, its trigger level on that input) matter; AUX OUT is 3.3 V logic: never into 50 Ohm.
"""


def build_parser():
    p = argparse.ArgumentParser(prog="rate_kit.py", description=__doc__.split("\n\n")[0],
                                formatter_class=argparse.RawDescriptionHelpFormatter, epilog=HELP_EPILOG)
    sub = p.add_subparsers(dest="cmd")

    def common(sp):
        sp.add_argument("--host", required=True, help="unit under test, HOST or HOST:PORT (SCPI, default port 5555)")
        sp.add_argument("--timeout", type=float, default=5.0, help="seconds to wait for any SCPI answer (default 5)")
        sp.add_argument("--trace", action="store_true", help="print every SCPI line as it goes (always kept in the JSON)")
        sp.add_argument("--show-serial", action="store_true", help="keep the serial number in the output (hidden by default)")
        sp.add_argument("--show-address", action="store_true", help="keep IP addresses in the output (hidden by default)")
        sp.add_argument("--quiet", action="store_true")

    sp = sub.add_parser("identify", help="read-only: *IDN?, family, firmware (stock / Kiszka mod / NovaOS)")
    common(sp)

    sp = sub.add_parser("measure", help="count waveforms a second", formatter_class=argparse.RawDescriptionHelpFormatter,
                        epilog=HELP_EPILOG)
    common(sp)
    measure_args(sp)

    sp = sub.add_parser("restore", help="write a snapshot.json back (after a hang or a kill)")
    common(sp)
    sp.add_argument("--from", dest="src", required=True, help="snapshot-<stamp>.json from an earlier run")
    sp.add_argument("--label", default="uut", help="which scope of the snapshot (uut or counter)")

    sp = sub.add_parser("scpi-health", help="read-only: SCPI latency and the error queue",
                        formatter_class=argparse.RawDescriptionHelpFormatter, epilog=HEALTH_EPILOG)
    common(sp)
    sp.add_argument("--n", type=int, default=20, help="queries timed (default 20)")
    sp.add_argument("--out", default=None, help="output directory (default ./rate-kit-<UTC stamp>)")

    sp = sub.add_parser("preset", help="count on a factory-default scope: save the setup, *RST, count, restore the setup",
                        formatter_class=argparse.RawDescriptionHelpFormatter, epilog=PRESET_EPILOG)
    common(sp)
    measure_args(sp)
    sp.add_argument("--yes-factory-default", action="store_true",
                    help="required: *RST resets the scope to factory defaults before counting (the setup is restored "
                         "from :SYSTem:SETup? afterwards; measurement items cleared by the reset do not come back)")


    sp = sub.add_parser("trio", help="the same scene on several units (stock, mod, NovaOS), one comparison table",
                        formatter_class=argparse.RawDescriptionHelpFormatter, epilog=TRIO_EPILOG)
    sp.add_argument("--unit", action="append", required=True, metavar="LABEL=HOST[,method=M[+M]][,counter=HOST@INPUT]",
                    help="repeat per unit, e.g. stock=STOCK,counter=COUNTER@1 mod=MOD,method=aux-self")
    sp.add_argument("--timeout", type=float, default=5.0)
    sp.add_argument("--trace", action="store_true")
    sp.add_argument("--show-serial", action="store_true")
    sp.add_argument("--show-address", action="store_true")
    sp.add_argument("--quiet", action="store_true")
    measure_args(sp, host=False)
    return p


def measure_args(sp, host=True):
    sp.add_argument("--method", default="aux-self",
                    help="aux-self (the unit's own counter) | aux-external (a second scope's counter)")
    sp.add_argument("--counter", help="aux-external: the counting scope, HOST[:PORT]")
    sp.add_argument("--counter-input", type=int, default=1, help="aux-external: the counting scope's input with AUX OUT (default 1)")
    sp.add_argument("--signal-input", type=int, help="input carrying the signal (default: CH1 on an MHO900, CH4 on a DHO; our bench)")
    sp.add_argument("--aux-input", type=int, help="aux-self: input carrying AUX OUT (default: CH4 on an MHO900, CH3 on a DHO)")
    sp.add_argument("--counter-channel", choices=("shown", "hidden"), default="shown",
                    help="aux-self: show the AUX input while counting (default) or hide it (as our bench engine did on "
                         "2026-10-16; its scale, offset and level are set while shown either way)")
    sp.add_argument("--signal", choices=("builtin", "external"), default="builtin",
                    help="builtin: the scope's generator drives the signal input; external: your generator, set by you")
    sp.add_argument("--generator", help="generator subsystem (default SOURce2 on an MHO900, SOURce on a DHO900 'S')")
    sp.add_argument("--frequency", type=float, default=10e6, help="square-wave frequency, Hz (default 10e6)")
    sp.add_argument("--amplitude", type=float, default=3.0, help="Vpp, from 0 V up (default 3: a 0..3 V square)")
    sp.add_argument("--tdiv", type=csv_floats, default=[20e-9, 2e-6], help="time bases, s/div, comma list (default 20e-9,2e-6)")
    sp.add_argument("--depth", type=csv_words, default=["1k", "10k", "AUTO"], help="memory depths (default 1k,10k,AUTO); as-is: leave the depth alone")
    sp.add_argument("--channels", type=csv_words, default=["one"], help="one (the signal alone) and/or all (every input on)")
    sp.add_argument("--scene", choices=("full", "minimal"), default="full",
                    help="full: every channel, probe, coupling and trigger setting written as our bench engine did; "
                         "minimal: only the generator, the time base, the edge source / level and NORMal sweep (and "
                         "the depth unless --depth as-is): the rest as the scope has it")
    sp.add_argument("--display", choices=("as-is", "bench", "default"), default="as-is",
                    help="as-is: leave the display alone; bench: vectors, infinite persistence, colour grade on, as our "
                         "bench engine's 2026-10-16 run; default: vectors, minimum persistence, colour grade off")
    sp.add_argument("--seconds", type=float, default=10.0, help="counting window (default 10)")
    sp.add_argument("--repeats", type=int, default=3, help="windows per scene (default 3)")
    sp.add_argument("--settle", type=float, default=2.0, help="seconds after :RUN before counting (default 2)")
    sp.add_argument("--clear-wait", type=float, default=3.0, help="seconds between the counter's CLEar and its first read (default 3)")
    sp.add_argument("--poll", choices=("status", "counter", "quiet"), default="status",
                    help="in the window, poll :TRIGger:STATus? once a second (default, the bench's traffic), or also "
                         ":COUNter:CURRent? (adds a least-squares slope and shows how often the count moves), or send "
                         "nothing between the two counter reads (quiet)")
    sp.add_argument("--signal-check", dest="signal_check", action="store_true",
                    help="send :MEASure:ITEM? VPP on the signal input before counting, as our bench engine did. Off by "
                         "default: the query leaves a live measurement on screen, and that alone cost the stock MHO934 "
                         "21 %% (24,954 -> 19,683 wfm/s at 1k, 2026-10-16); the kit does not remove it")
    sp.add_argument("--no-signal-check", dest="signal_check", action="store_false", help=argparse.SUPPRESS)
    sp.add_argument("--out", default=None, help="output directory (default ./rate-kit-<UTC stamp>)")
    sp.add_argument("--function", default="SQUare", help="generator function (default SQUare; e.g. SINusoid)")
    sp.add_argument("--load-counter-input", type=int,
                    help="aux-external: also run the unit's OWN counter (totalize) on this input, unread, to measure what "
                         "counting on the unit's own input costs it")
    sp.add_argument("--load-counter-channel", choices=("hidden", "shown"), default="hidden",
                    help="that input hidden (default) or shown while counting; its scale, offset and level are set shown")
    sp.add_argument("--clear-measurements", action="store_true",
                    help="delete every measurement item before counting (:MEASure:DELete on an MHO900, :MEASure:CLEar on "
                         "a DHO): one on screen costs the stock about 21 %%. Not restorable")


HEALTH_EPILOG = """
scpi-health, read-only by default:
  - *IDN? and :TRIGger:STATus? timed --n times, the error queue read.
"""

PRESET_EPILOG = """
preset: what we ran on our stock MHO934 on 2026-10-16 (METHOD.md runs G and J):
  1. :SYSTem:SETup? (a block) saved to setup-<stamp>.bin, with the run state, the generator and 47 settings read;
  2. *RST (the programming guide: "Restores the instrument to its factory default settings"; the firmware has no
     :SYSTem:PRESet), then *OPC? until it answers 1;
  3. the measurement, with --scene minimal unless you pass --scene full;
  4. :SYSTem:SETup <block>, the generator and the run state written back; those settings read again and compared.
  Exit 3, with the block's path, if anything differs. Measurement items the reset cleared are not restored.
"""


TRIO_EPILOG = """
trio: one scene, several units, one table. Every measure option applies to every unit; per unit, method= and
counter=HOST@INPUT override --method / --counter / --counter-input. Each unit's full run lands in OUT/<label>/;
OUT/trio-<stamp>.md / .json hold the comparison. The units run one after another, not at once.
"""


def stamp_of(clock):
    return datetime.datetime.fromtimestamp(clock.wall(), datetime.timezone.utc).strftime("%Y%m%dT%H%M%SZ")


def write_outputs(a, kind, rep, md, idents=()):
    text_json, text_md = json.dumps(rep, indent=1), md
    if not getattr(a, "show_serial", False):
        for ident in idents:
            if ident and ident.serial:
                text_json = text_json.replace(ident.serial, "<serial hidden>")
                text_md = text_md.replace(ident.serial, "<serial hidden>")
    if not getattr(a, "show_address", False):
        for hp in [getattr(a, "host", None), getattr(a, "counter", None)]:
            if hp:
                h = split_host(hp)[0]
                text_json, text_md = text_json.replace(h, "<address hidden>"), text_md.replace(h, "<address hidden>")
    pj = os.path.join(a.out, "%s-%s.json" % (kind, a.stamp))
    pm = os.path.join(a.out, "%s-%s.md" % (kind, a.stamp))
    with open(pj, "w", encoding="utf-8", newline="\n") as f:
        f.write(text_json)
    with open(pm, "w", encoding="utf-8", newline="\n") as f:
        f.write(text_md)
    return pj, pm


def say(a, msg):
    if not getattr(a, "quiet", False):
        sys.stderr.write(msg + "\n")
        sys.stderr.flush()


# ------------------------------------------------------------------------------------------------ scpi-health


def scpi_health(a, clock):
    host, port = split_host(a.host)
    trace = []
    s = Scpi(host, port, timeout=a.timeout, clock=clock, trace=trace, label="uut", echo=a.trace)
    rep = {"kit": "rate_kit scpi-health", "kit_version": KIT_VERSION, "argv": a.argv, "started": utc_iso(clock.wall()),
           "trace": trace}
    ident = None
    try:
        s.connect()
        ident = identify(s)
        s.redact = ident.redact(a.show_serial)
        rep["scope"] = ident.as_dict(a.show_serial)
        lat = {"*IDN?": [], ":TRIGger:STATus?": []}
        for i in range(a.n):
            for q in lat:
                t0 = time.monotonic()
                r = s.try_query(q, a.timeout)
                lat[q].append(None if r is None else round((time.monotonic() - t0) * 1000.0, 2))
        rep["latency_ms"] = {q: stats(v) for q, v in lat.items()}
        rep["unanswered"] = {q: sum(1 for x in v if x is None) for q, v in lat.items()}
        rep["errors"] = s.drain_errors(10)
    except ScpiHung as e:
        rep["hung"] = str(e)
    except KitError as e:
        rep["error"] = str(e)
    s.close()
    rep["finished"] = utc_iso(clock.wall())
    lines = ["# scpi-health %s" % rep["started"], ""]
    if rep.get("scope"):
        lines.append("- `%s` (%s, %s)" % (rep["scope"]["idn"], rep["scope"]["series"], rep["scope"]["firmware"]))
    for k in ("hung", "error"):
        if rep.get(k):
            lines.append("- **%s**: %s" % (k, rep[k]))
    for q, st in rep.get("latency_ms", {}).items():
        lines.append("- `%s` x%d: mean %.1f ms, min %.1f, max %.1f, unanswered %d" % (
            q, a.n, st.get("mean", 0), st.get("min", 0), st.get("max", 0), rep["unanswered"][q]))
    if rep.get("errors"):
        lines.append("- error queue: %s" % "; ".join(rep["errors"]))
    a.out = a.out or "rate-kit-%s" % a.stamp
    os.makedirs(a.out, exist_ok=True)
    pj, pm = write_outputs(a, "scpi-health", rep, "\n".join(lines) + "\n", [ident])
    say(a, "\n".join(lines))
    say(a, "wrote %s and %s" % (pj, pm))
    if rep.get("hung"):
        return 4
    if rep.get("error") or any(rep.get("unanswered", {}).values()):
        return 1
    return 0


# ------------------------------------------------------------------------------------------------ preset

PRESET_KEYS = [":ACQuire:MDEPth", ":ACQuire:TYPE", ":TIMebase:MAIN:SCALe", ":TIMebase:MAIN:OFFSet", ":TIMebase:MODE",
               ":TRIGger:MODE", ":TRIGger:SWEep", ":TRIGger:EDGE:SOURce", ":TRIGger:EDGE:LEVel", ":TRIGger:EDGE:SLOPe",
               ":TRIGger:COUPling", ":TRIGger:HOLDoff", ":TRIGger:NREJect", ":DISPlay:TYPE", ":DISPlay:GRADing:TIME",
               ":DISPlay:COLor", ":DISPlay:WBRightness", ":DISPlay:GBRightness", ":DISPlay:GRID", ":DISPlay:RULers",
               ":SYSTem:AOUTput", ":COUNter:ENABle", ":COUNter:SOURce", ":COUNter:MODE", ":MEASure:STATistic:DISPlay",
               ":MEASure:AMSource", ":CHANnel1:DISPlay", ":CHANnel2:DISPlay", ":CHANnel3:DISPlay", ":CHANnel4:DISPlay",
               ":CHANnel1:SCALe", ":CHANnel1:OFFSet", ":CHANnel1:PROBe", ":CHANnel1:COUPling", ":CHANnel1:BWLimit",
               ":CHANnel1:IMPedance", ":CHANnel1:INVert", ":CHANnel1:VERNier", ":MATH1:DISPlay", ":BUS1:DISPlay",
               ":MASK:ENABle", ":HISTogram:ENABle"]


def preset(a, clock):
    if not a.yes_factory_default:
        print("error: preset resets the scope to factory defaults (*RST); pass --yes-factory-default")
        return 2
    host, port = split_host(a.host)
    trace = []
    s = Scpi(host, port, timeout=a.timeout, clock=clock, trace=trace, label="uut", echo=a.trace)
    s.connect()
    ident = identify(s)
    s.redact = ident.redact(a.show_serial)
    gen = None if a.signal == "external" else (a.generator or ident.default_generator())
    gen_keys = [":%s:%s" % (gen, x) for x in ("FUNCtion", "FREQuency", "VOLTage:AMPLitude", "VOLTage:OFFSet",
                                               "OUTPut:STATe")] if gen else []
    a.out = a.out or "rate-kit-%s" % a.stamp
    os.makedirs(a.out, exist_ok=True)
    rep = {"kit": "rate_kit preset", "kit_version": KIT_VERSION, "argv": a.argv, "started": utc_iso(clock.wall()),
           "scope": ident.as_dict(a.show_serial), "trace": trace}
    state = s.query(":TRIGger:STATus?")
    blob = s.block(":SYSTem:SETup?", 30.0, allow=True)
    bpath = os.path.join(a.out, "setup-%s.bin" % a.stamp)
    with open(bpath, "wb") as f:
        f.write(blob)
    # :CHANnel<n>:IMPedance is in the MHO900 programming guide (3.6.7) only: the DHO800/DHO900 inputs are 1 MOhm
    keys = [k for k in PRESET_KEYS if ident.family == "mho" or not k.endswith(":IMPedance")] + gen_keys
    qt = min(3.0, a.timeout)
    before = {k: s.try_query(k + "?", qt) for k in keys}
    s.drain_errors(10)
    rep.update({"state_before": state, "setup_block_bytes": len(blob), "setup_block_file": os.path.basename(bpath),
                "before": before})
    code = 1
    try:
        s.write("*RST", pace=1.0, allow=True)
        for _ in range(40):
            if s.try_query("*OPC?", 3.0) == "1":
                break
            clock.sleep(1.0)
        clock.sleep(2.0)
        after = {k: s.try_query(k + "?", qt) for k in keys}
        s.drain_errors(10)
        rep["after_reset"] = after
        rep["reset_changed"] = {k: (before[k], after[k]) for k in keys if not same(k, before[k], after[k])}
        s.close()
        if a.scene == "full" and "--scene" not in a.argv:
            a.scene = "minimal"
        m = argparse.Namespace(**vars(a))
        m.out = os.path.join(a.out, "measure")
        code, mrep, mpath = run_measure(m, clock)
        rep["measure_report"] = os.path.relpath(mpath, a.out) if mpath else None
        rep["measure_summary"] = [{"label": sc["label"], "summary": sc["summary"]} for sc in mrep.get("scenes", [])]
    finally:
        # the app may be restarting (the stock's watchdog relaunches it after a crash): wait for SCPI, up to 3 min
        for _ in range(36):
            try:
                s.connect()
                if s.try_query("*IDN?", 5.0):
                    break
            except KitError:
                pass
            clock.sleep(5.0)
        s.write_block(":SYSTem:SETup", blob, pace=1.0)
        for _ in range(40):
            if s.try_query("*OPC?", 3.0) == "1":
                break
            clock.sleep(1.0)
        clock.sleep(3.0)
        for k in gen_keys:
            v = before.get(k)
            if v is not None:
                s.write("%s %s" % (k, settable(k, v)), pace=0.2)
        now_state = s.try_query(":TRIGger:STATus?")
        if state.upper().startswith("STOP") and now_state and not now_state.upper().startswith("STOP"):
            s.write(":STOP")
        elif not state.upper().startswith("STOP") and now_state and now_state.upper().startswith("STOP"):
            s.write(":RUN")
        now = {k: s.try_query(k + "?", qt) for k in keys}
        rep["restore_diff"] = {k: (before[k], now[k]) for k in keys if not same(k, before[k], now[k])}
        rep["state_after"] = s.try_query(":TRIGger:STATus?")
        s.close()
    rep["finished"] = utc_iso(clock.wall())
    lines = ["# rate_kit preset %s" % rep["started"], "",
             "- `%s`; setup block %d bytes saved as `%s`" % (rep["scope"]["idn"], len(blob), rep["setup_block_file"]),
             "- *RST changed: %s" % ("; ".join("%s %s -> %s" % (k, x, y) for k, (x, y) in rep.get("reset_changed", {}).items()) or "nothing read"),
             "", "| scene | AUX wfm/s mean | min | max |", "|---|---:|---:|---:|"]
    for sc in rep.get("measure_summary", []):
        ar = sc["summary"].get("aux_rate", {})
        lines.append("| %s | %s | %s | %s |" % (sc["label"], fmt(ar.get("mean")), fmt(ar.get("min")), fmt(ar.get("max"))))
    lines += ["", "- restore: %s; state %s (was %s)" % ("complete" if not rep["restore_diff"] else
                                                        "**differs: %s**" % json.dumps(rep["restore_diff"]),
                                                        rep["state_after"], state)]
    pj, pm = write_outputs(a, "preset", rep, "\n".join(lines) + "\n", [ident])
    say(a, "\n".join(lines))
    if rep["restore_diff"]:
        say(a, "RESTORE DIFFERS: write %s back with :SYSTem:SETup" % bpath)
        return 3
    return code


# ------------------------------------------------------------------------------------------------ trio

def parse_unit(spec):
    label, _, rest = spec.partition("=")
    parts = rest.split(",")
    u = {"label": label.strip(), "host": parts[0].strip()}
    for p in parts[1:]:
        k, _, v = p.partition("=")
        u[k.strip()] = v.strip()
    if not u["label"] or not u["host"]:
        raise KitError("--unit %r: want LABEL=HOST[,method=M][,counter=HOST@INPUT]" % spec)
    return u


def trio(a, clock):
    units = [parse_unit(x) for x in a.unit]
    a.out = a.out or "rate-kit-%s" % a.stamp
    os.makedirs(a.out, exist_ok=True)
    results = []
    worst = 0
    for u in units:
        m = argparse.Namespace(**vars(a))
        m.host = u["host"]
        m.method = u.get("method", a.method).replace("+", ",")  # method=M1+M2 inside a --unit
        if "counter" in u:
            ch_, _, inp = u["counter"].partition("@")
            m.counter, m.counter_input = ch_, int(inp or a.counter_input)
        m.out = os.path.join(a.out, u["label"])
        say(a, "== %s" % u["label"])
        code, mrep, mpath = run_measure(m, clock)
        worst = max(worst, code)
        results.append({"label": u["label"], "code": code, "report": os.path.relpath(mpath, a.out) if mpath else None,
                        "scope": mrep.get("scopes", {}).get("uut"), "methods": mrep.get("methods"),
                        "scenes": [{"label": sc["label"], "srate": sc.get("srate"), "summary": sc["summary"]}
                                   for sc in mrep.get("scenes", [])], "failed": mrep.get("failed")})
    labels = []
    for r in results:
        for sc in r["scenes"]:
            if sc["label"] not in labels:
                labels.append(sc["label"])
    lines = ["# rate_kit trio %s" % utc_iso(clock.wall()), ""]
    for r in results:
        sc = r["scope"] or {}
        lines.append("- **%s**: `%s`, %s firmware %s%s, methods %s%s" % (
            r["label"], sc.get("idn"), sc.get("firmware"), sc.get("idn_firmware"),
            (" (NovaOS %s)" % sc["novaos_version"]) if sc.get("novaos_version") else "", ",".join(r["methods"] or []),
            ("; **stopped: %s**" % r["failed"]) if r["failed"] else ""))
    lines += ["", "| scene | " + " | ".join(r["label"] for r in results) + " |",
              "|---|" + "---:|" * len(results)]

    def cell(r, label):
        for sc in r["scenes"]:
            if sc["label"] == label:
                sm = sc["summary"]
                parts = []
                if sm.get("aux_rate"):
                    ar = sm["aux_rate"]
                    parts.append("AUX %s (%s-%s)" % (fmt(ar["mean"]), fmt(ar["min"]), fmt(ar["max"])))
                if sm.get("novaos_density_rate"):
                    parts.append("engine %s" % fmt(sm["novaos_density_rate"]["mean"]))
                parts.append(str(sc.get("srate") or "-") + " Sa/s")
                return "<br>".join(parts)
        return "-"
    for label in labels:
        lines.append("| %s | %s |" % (label, " | ".join(cell(r, label) for r in results)))
    rep = {"kit": "rate_kit trio", "kit_version": KIT_VERSION, "argv": a.argv, "units": results}
    pj = os.path.join(a.out, "trio-%s.json" % a.stamp)
    pm = os.path.join(a.out, "trio-%s.md" % a.stamp)
    text_json, text_md = json.dumps(rep, indent=1), "\n".join(lines) + "\n"
    if not a.show_address:
        for u in units:
            for h in [u["host"]] + ([u["counter"].partition("@")[0]] if "counter" in u else []):
                hh = split_host(h)[0]
                text_json, text_md = text_json.replace(hh, "<address hidden>"), text_md.replace(hh, "<address hidden>")
    with open(pj, "w", encoding="utf-8", newline="\n") as f:
        f.write(text_json)
    with open(pm, "w", encoding="utf-8", newline="\n") as f:
        f.write(text_md)
    say(a, text_md)
    say(a, "wrote %s and %s" % (pj, pm))
    return worst


# ------------------------------------------------------------------------------------------------ measure

def run_measure(a, clock):
    """One measure run; returns (exit code, report, JSON path)."""
    a.methods = csv_words(a.method)
    bad = [m for m in a.methods if m not in ("aux-self", "aux-external")]
    if bad or not a.methods:
        print("error: unknown method %s" % ", ".join(bad))
        return 2, {}, None
    if "aux-self" in a.methods and "aux-external" in a.methods:
        print("error: aux-self and aux-external both count AUX OUT; pick one")
        return 2, {}, None
    for c in a.channels:
        if c not in ("one", "all"):
            print("error: --channels takes one and/or all")
            return 2, {}, None
    a.stamp = getattr(a, "stamp", None) or stamp_of(clock)
    a.out = a.out or "rate-kit-%s" % a.stamp
    os.makedirs(a.out, exist_ok=True)
    r = Runner(a, clock=clock)
    try:
        failed = r.run()
    except KitError as e:  # before anything changed (connect, identity, snapshot)
        failed = str(e)
        r.report["failed"] = failed
        if r.snaps:
            r.restore_all(failed)
    for sc, _ in r.scopes.values():
        sc.close()
    rep = r.report
    pj, pm = write_outputs(a, "rate-kit", rep, markdown(rep), [i for _, i in r.scopes.values()])
    r.say("wrote %s and %s" % (pj, pm))
    incomplete = [lab for lab, x in rep["restore"].items() if not x.get("complete")]
    if incomplete:
        r.say("RESTORE INCOMPLETE on %s: see the report's restore section" % ", ".join(incomplete))
    if failed:
        r.say("stopped: %s" % failed)
    if failed and failed.startswith("hung"):
        return 4, rep, pj
    if incomplete:
        return 3, rep, pj
    if failed == "interrupted (Ctrl-C)":
        return 130, rep, pj
    if failed:
        return 1, rep, pj
    return 0, rep, pj


def main(argv=None, clock=None):
    argv = list(sys.argv[1:] if argv is None else argv)
    p = build_parser()
    a = p.parse_args(argv)
    if not a.cmd:
        p.print_help()
        return 2
    a.argv = ["rate_kit.py"] + argv
    clock = clock or Clock()
    if a.cmd == "identify":
        host, port = split_host(a.host)
        s = Scpi(host, port, timeout=a.timeout, clock=clock, echo=a.trace)
        try:
            ident = identify(s)
        except KitError as e:
            print("error: %s" % e)
            return 2
        finally:
            s.close()
        d = ident.as_dict(a.show_serial)
        print(json.dumps(d, indent=1))
        return 0
    if a.cmd == "restore":
        with open(a.src, encoding="utf-8") as f:
            d = json.load(f)[a.label]
        host, port = split_host(a.host)
        s = Scpi(host, port, timeout=a.timeout, clock=clock, echo=a.trace)
        ident = identify(s)
        sn = Snapshot.from_dict(s, ident, d["snapshot"])
        try:
            rep = restore(sn)
        finally:
            s.close()
        print(json.dumps(rep, indent=1))
        return 0 if rep["complete"] else 3
    a.stamp = stamp_of(clock)
    if a.cmd == "scpi-health":
        return scpi_health(a, clock)
    if a.cmd == "preset":
        return preset(a, clock)
    if a.cmd == "trio":
        return trio(a, clock)
    return run_measure(a, clock)[0]


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