#!/usr/bin/env python3
# /// script
# requires-python = ">=3.8"
# dependencies = ["pyserial"]
# ///
"""Chameleon Ultra crystal offset over USB, measured against an NTP server. The command-line version of index.html.

    uv run hfxo_ntp.py /dev/tty.usbmodemXXXX     # or: pipx run hfxo_ntp.py ... (both install pyserial)
    pip install pyserial
    python3 hfxo_ntp.py /dev/tty.usbmodemXXXX [--minutes 30] [--server time.cloudflare.com]
    python3 hfxo_ntp.py COM5 --minutes 15

The script needs firmware with command 3125, the HFXO tick counter.

Plug your CU in to USB, not Bluetooth, and close anything else using it, such as the CU app or the Python CLI. Keep
the CU away from readers, and run no LF commands during the run. LF code can release the crystal. The firmware
counts each release as a lapse, and the script rejects a run with any lapse. Keep the computer awake.

At the end, the script prints the crystal offset in ppm (+ means the crystal runs fast).

Method: as on the page, the script starts the counter, then reads the counter every 2 s. Each read sits
between two readings of this computer's monotonic clock. At the end, the script stops the counter and fits the count
against time. NTP gives the monotonic clock's own rate error over the same time, and the script removes that error
from the result.

Unlike the page, the script uses NTP over UDP (port 123), with one query every 16 s through the whole run. Public
NTP servers would rate-limit the page's bursts. As on the page, the fit uses only the fastest half of the queries by
round trip. A slow query carries queueing that isn't split evenly between the two directions. The script looks up
the server's name once, and every query then reaches the same machine.
"""

import argparse
import json
import math
import socket
import struct
import sys
import time

import serial

NOMINAL = 16e6  # command 3125 counts 16 MHz HFCLK ticks with the crystal held
CMD_HFXO_COUNT, CMD_GET_GIT_VERSION = 3125, 1017
STATUS_SUCCESS, STATUS_INVALID_CMD = 0x68, 0x67
START, READ, STOP = 0, 1, 2
CU_INTERVAL, NTP_INTERVAL = 2.0, 16.0
NTP_EPOCH = 2208988800  # 1900 -> 1970

# The local clock is perf_counter (mach_absolute_time on macOS, QueryPerformanceCounter on Windows), moved once onto
# a Unix-like scale to keep offsets small. Nothing steps it, and the NTP fit measures its rate error.
_BASE = time.time() - time.perf_counter()


def now():
    return _BASE + time.perf_counter()


# ---------------------------------------------------------------- CU framing (big-endian)
# SOF, LRC1, cmd u16, status u16, len u16, LRC2, data, LRC3. Each LRC covers every byte before it.


def lrc(b):
    return (0x100 - sum(b)) & 0xFF


def frame(cmd, data=b""):
    f = bytearray([0x11])
    f.append(lrc(f))
    f += struct.pack("!HHH", cmd, 0, len(data))
    f.append(lrc(f))
    f += data
    f.append(lrc(f))
    return bytes(f)


class CU:
    def __init__(self, port):
        self.s = serial.Serial(port, 115200, timeout=0.05)
        self.buf = bytearray()

    def cmd(self, cmd, data=b"", timeout=3.0):
        self.s.write(frame(cmd, data))
        end = time.monotonic() + timeout
        while time.monotonic() < end:
            self.buf += self.s.read(max(1, self.s.in_waiting))  # a fixed size would wait out the timeout
            while True:
                i = self.buf.find(0x11)
                if i < 0:
                    self.buf.clear()
                    break
                del self.buf[:i]
                b = self.buf
                if len(b) < 9:
                    break
                if b[1] != lrc(b[:1]) or b[8] != lrc(b[:8]):
                    del b[:1]
                    continue
                n = struct.unpack("!H", b[6:8])[0]
                if len(b) < 10 + n:
                    break
                if b[9 + n] != lrc(b[: 9 + n]):
                    del b[:1]
                    continue
                rcmd, status = struct.unpack("!HH", b[2:6])
                data_out = bytes(b[9 : 9 + n])
                del b[: 10 + n]
                if rcmd == cmd:
                    return status, data_out
        raise TimeoutError(f"no reply to command {cmd}")

    def version(self):
        status, d = self.cmd(CMD_GET_GIT_VERSION)
        return d.decode(errors="replace") if status == STATUS_SUCCESS else "?"

    def hfxo(self, sub):
        tb = now()
        status, d = self.cmd(CMD_HFXO_COUNT, bytes([sub]))
        ta = now()
        if status == STATUS_INVALID_CMD:
            raise SystemExit("this firmware has no command 3125 (the HFXO tick counter). Flash the test firmware.")
        if status != STATUS_SUCCESS or len(d) < 20:
            raise RuntimeError(f"3125 failed: status {status:#x}, {len(d)} bytes")
        ticks, _, _, lapses, running, hfxo = struct.unpack("!QIIHBB", d[:20])
        return {"t": (tb + ta) / 2, "ticks": ticks, "lapses": lapses, "running": bool(running), "hfxo": bool(hfxo)}


# ---------------------------------------------------------------- NTP


def ntp_query(addr, timeout=2.0):
    """One SNTP query. Returns (local mid time, server - local offset, round trip), all in seconds."""
    with socket.socket(socket.AF_INET6 if ":" in addr[0] else socket.AF_INET, socket.SOCK_DGRAM) as s:
        s.settimeout(timeout)
        t1 = now()
        s.sendto(b"\x23" + 47 * b"\0", addr)  # LI 0, version 4, mode 3 (client)
        d, _ = s.recvfrom(512)
        t4 = now()
    if len(d) < 48 or d[1] == 0:  # stratum 0: kiss-o'-death, e.g. rate limited
        raise RuntimeError(f"server refused ({d[12:16]!r})")

    def ts(i):
        sec, frac = struct.unpack("!II", d[i : i + 8])
        return sec - NTP_EPOCH + frac / 2**32

    t2, t3 = ts(32), ts(40)
    return (t1 + t4) / 2, ((t2 - t1) + (t3 - t4)) / 2, (t4 - t1) - (t3 - t2)


# ---------------------------------------------------------------- fits


def linfit(x, y):
    n = len(x)
    mx, my = sum(x) / n, sum(y) / n
    sxx = sum((v - mx) ** 2 for v in x)
    k = sum((a - mx) * (b - my) for a, b in zip(x, y)) / sxx
    b0 = my - k * mx
    res = [b - (k * a + b0) for a, b in zip(x, y)]
    sig = math.sqrt(sum(r * r for r in res) / (n - 2) / sxx) if n > 2 else float("nan")
    return k, b0, res, sig


def median(a):
    s = sorted(a)
    m = len(s) // 2
    return s[m] if len(s) % 2 else (s[m - 1] + s[m]) / 2


def robust_fit(x, y):
    """Least squares, refit once without points more than 5 MADs off the line. Returns (k, sigma k, kept, n)."""
    k, _, res, sig = linfit(x, y)
    med = median(res)
    mad = median([abs(r - med) for r in res]) or 1e-12
    keep = [abs(r - med) <= 5 * 1.4826 * mad for r in res]
    if len(x) >= 4 and 3 <= sum(keep) < len(x):
        k, _, res, sig = linfit([v for v, ok in zip(x, keep) if ok], [v for v, ok in zip(y, keep) if ok])
    return k, sig, sum(keep), len(x)


# A path's delay wanders by a fraction of a millisecond, and the scatter of the queries can't show that.
# with_path_floor adds PATH_FLOOR (in seconds) to each end of an NTP fit, in quadrature. Without it, the halves check
# fails on wander that biases the result by less than the +- allows.
PATH_FLOOR = 0.3e-3


def with_path_floor(sig, x):
    """The NTP fit's slope sigma, with PATH_FLOOR added for the span of x."""
    return math.hypot(sig, math.sqrt(2) * PATH_FLOOR / (max(x) - min(x)))


def main():
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("port", help="the CU's serial port (/dev/tty.usbmodem..., /dev/ttyACM0, COM5)")
    ap.add_argument(
        "--minutes", type=float, default=30, help="how long to measure. 30 by default (15 gives about +-0.5 ppm)"
    )
    ap.add_argument("--server", default="time.cloudflare.com", help="the NTP server, looked up once")
    ap.add_argument("--log", help="write every sample to this JSON-lines file")
    a = ap.parse_args()

    info = socket.getaddrinfo(a.server, 123, type=socket.SOCK_DGRAM)[0]
    addr = info[4][:2]
    cu = CU(a.port)
    print(f"firmware {cu.version()}, NTP {a.server} ({addr[0]}), {a.minutes:g} min")
    log = open(a.log, "w") if a.log else None  # noqa: SIM115 (closed in the finally below)
    cu_pts, ref_pts = [], []
    p = cu.hfxo(START)
    t0 = p["t"]
    end = t0 + a.minutes * 60
    next_ntp = t0
    try:
        while True:
            t = now()
            if t >= next_ntp:
                try:
                    mid, off, rtt = ntp_query(addr)
                    ref_pts.append((mid, off, rtt))
                    if log:
                        log.write(json.dumps({"ntp": [mid, off, rtt]}) + "\n")
                except (OSError, RuntimeError) as e:
                    print(f"  NTP query failed: {e}")
                next_ntp += NTP_INTERVAL
            p = cu.hfxo(READ)
            if p["lapses"] or not p["running"]:
                raise SystemExit("the CU released its crystal during the run, probably for LF activity. Run again.")
            cu_pts.append((p["t"], p["ticks"]))
            if log:
                log.write(json.dumps({"cu": [p["t"], p["ticks"]]}) + "\n")
            if p["t"] >= end:
                break
            done = (p["t"] - t0) / 60
            print(f"\r  {done:5.1f} / {a.minutes:g} min, {len(cu_pts)} reads, {len(ref_pts)} NTP", end="", flush=True)
            time.sleep(max(0.0, CU_INTERVAL - (now() - t)))
    except KeyboardInterrupt:
        print("\n  stopped early")
    finally:
        try:
            cu.hfxo(STOP)
        except (OSError, RuntimeError) as e:
            print(f"  STOP failed: {e}")
        if log:
            log.close()
    print()
    if len(cu_pts) < 10 or len(ref_pts) < 12:
        raise SystemExit("too short to fit")

    # The CU: ticks per local second. Residual scatter in us shows USB stalls (and laptop sleep).
    x = [t - cu_pts[0][0] for t, _ in cu_pts]
    y = [c - cu_pts[0][1] for _, c in cu_pts]
    f, fsig, _, _ = robust_fit(x, y)
    rms_us = math.sqrt(sum(r * r for r in linfit(x, y)[2]) / len(x)) / f * 1e6
    if rms_us > 20000:
        raise SystemExit(f"CU reads scatter {rms_us / 1e3:.0f} ms rms. The computer probably slept. Run again.")

    # The local clock against NTP: offset(t) = a + k t, where offset is server minus local, through the fastest half
    # of the queries. k > 0 means the local clock runs slow.
    cut = sorted(r for _, _, r in ref_pts)[len(ref_pts) // 2]
    fast = [p for p in ref_pts if p[2] <= cut]
    rx = [m - ref_pts[0][0] for m, _, _ in fast]
    ry = [o for _, o, _ in fast]
    k, ksig, kept, _ = robust_fit(rx, ry)
    ksig = with_path_floor(ksig, rx)

    # The halves check fits each half of the run on its own and compares the crystal's offset in each. The local
    # clock's rate can change between the halves, and the device's rate against the local clock then moves the
    # opposite way.
    half = len(fast) // 2
    split = fast[half][0]
    halves = []
    for queries, keep in ((fast[: half + 1], lambda t: t <= split), (fast[half:], lambda t: t >= split)):
        hx = [m - ref_pts[0][0] for m, _, _ in queries]
        hk, hk_sig, _, _ = robust_fit(hx, [o for _, o, _ in queries])
        hk_sig = with_path_floor(hk_sig, hx)
        pts = [(t, c) for t, c in cu_pts if keep(t)]
        hf, hf_sig, _, _ = robust_fit([t - pts[0][0] for t, _ in pts], [c - pts[0][1] for _, c in pts])
        halves.append(((hf / (1 + hk) / NOMINAL - 1) * 1e6, math.hypot(hf_sig / NOMINAL, hk_sig) * 1e6, -hk * 1e6))

    ppm = (f / (1 + k) / NOMINAL - 1) * 1e6
    sig = math.hypot(fsig / NOMINAL, ksig) * 1e6
    print(f"crystal offset   {ppm:+.2f} +- {sig:.2f} ppm   (+ = fast)")
    print(f"  uncorrected    {(f / NOMINAL - 1) * 1e6:+.2f} ppm vs this computer's clock")
    print(
        f"  computer clock {-k * 1e6:+.2f} ppm vs NTP (fastest {kept} of {len(ref_pts)} queries, min round trip "
        f"{min(r for _, _, r in ref_pts) * 1e3:.1f} ms)"
    )
    (p1, s1, c1), (p2, s2, c2) = halves
    ok = abs(p2 - p1) <= 3 * math.hypot(s1, s2)
    print(
        f"  halves         {p1:+.2f} then {p2:+.2f} ppm (computer clock {c1:+.2f} then {c2:+.2f}): "
        f"{'ok' if ok else 'DISAGREE (path to the server changed?): re-run'}"
    )
    print(f"  CU reads       {len(cu_pts)}, {rms_us:.0f} us rms")


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