#!/usr/bin/env python3
"""Capture one single-shot CH3 frame from Hantek DPO7204C, save raw + PNG.

Flow: query settings -> :SINGle -> poll :TRIGger:STATus? until STOP ->
WAVeform:DATA:ALL? CHANnel3 (drained completely, multi-packet aware).
Scope is left in STOP so the on-screen frame matches the saved data.

Usage: .venv/bin/python tools/capture.py <name> [device]
Saves captures/<name>.bin, captures/<name>.npz, captures/<name>.png
"""
import os
import sys
import time

import matplotlib

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

DEV = sys.argv[2] if len(sys.argv) > 2 else "/dev/usbtmc2"
NAME = sys.argv[1] if len(sys.argv) > 1 else "capture"
OUTDIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "captures")
os.makedirs(OUTDIR, exist_ok=True)

FIRST_HDR = 11 + 117  # '#9'+9-digit len, then 18 bytes counters + 99 bytes info
NEXT_HDR = 11 + 18    # follow-up packets: counters only


def read_exact(fd, n):
    buf = b""
    while len(buf) < n:
        chunk = os.read(fd, min(1 << 20, n - len(buf)))
        if not chunk:
            raise IOError("short read from scope")
        buf += chunk
    return buf


def read_packet(fd):
    head = read_exact(fd, 11)
    assert head[:2] == b"#9", f"bad packet start: {head!r}"
    pkt_len = int(head[2:11])
    body = read_exact(fd, pkt_len)
    return head + body


def query(fd, cmd):
    os.write(fd, cmd.encode() + b"\n")
    return os.read(fd, 256).decode(errors="replace").strip()


fd = os.open(DEV, os.O_RDWR)
print("IDN:", query(fd, "*IDN?"))

tdiv = float(query(fd, ":TIMebase:SCALe?"))
vdiv = float(query(fd, ":CHANnel3:SCALe?"))
voff = float(query(fd, ":CHANnel3:OFFSet?"))
print(f"tdiv={tdiv} s/div  vdiv={vdiv} V/div  offset={voff} V")

os.write(fd, b":SINGle\n")
for _ in range(100):
    time.sleep(0.1)
    st = query(fd, ":TRIGger:STATus?")
    if st == "STOP":
        break
else:
    sys.exit(f"scope did not reach STOP (last status: {st})")
print("status: STOP (frame captured)")

srate = float(query(fd, ":ACQuire:SRATe?"))
print(f"srate={srate:.3e} Sa/s")

os.write(fd, b"WAVeform:DATA:ALL? CHANnel3\n")
pkt = read_packet(fd)
total_len = int(pkt[11:20])
info_hdr = pkt[29:FIRST_HDR]
data = pkt[FIRST_HDR:]
raw_all = pkt
while len(data) < total_len:
    os.write(fd, b"WAVeform:DATA:ALL? CHANnel3\n")
    p = read_packet(fd)
    raw_all += p
    data += p[NEXT_HDR:]
print(f"received {len(data)} samples (declared {total_len})")
print("info header hex:", info_hdr.hex(" "))
os.close(fd)

raw = np.frombuffer(data[:total_len], dtype=np.uint8).astype(np.float32)
# unsigned 8-bit, 25.6 levels/div, mid-screen = code 128, offset shifts zero
volts = (raw - 128.0) / 25.6 * vdiv - voff
t = np.arange(len(volts)) / srate * 1e3  # ms

base = os.path.join(OUTDIR, NAME)
with open(base + ".bin", "wb") as f:
    f.write(raw_all)
np.savez(base + ".npz", volts=volts, srate=srate, vdiv=vdiv, voff=voff, tdiv=tdiv)

fig, axes = plt.subplots(2, 1, figsize=(16, 8))
axes[0].plot(t, volts, lw=0.5)
axes[0].set_title(f"{NAME} — full frame ({srate:.0e} Sa/s, {vdiv} V/div, off {voff} V)")
# zoom on activity: region where signal deviates from its median
dev_idx = np.where(np.abs(volts - np.median(volts)) > 0.5)[0]
if len(dev_idx):
    lo = max(0, dev_idx[0] - int(0.05 * (dev_idx[-1] - dev_idx[0] + 1)) - 100)
    hi = min(len(volts), dev_idx[-1] + int(0.05 * (dev_idx[-1] - dev_idx[0] + 1)) + 100)
    axes[1].plot(t[lo:hi], volts[lo:hi], lw=0.7)
    axes[1].set_title("zoom on activity")
for ax in axes:
    ax.set_xlabel("t, ms")
    ax.set_ylabel("U, V")
    ax.grid(True, alpha=0.3)
fig.tight_layout()
fig.savefig(base + ".png", dpi=110)
print("saved:", base + ".png")
print(f"range: min={volts.min():.3f} V  max={volts.max():.3f} V  median={np.median(volts):.3f} V")