#!/usr/bin/env python3
import os
import time
import math
import cmath
import argparse
import errno

DEV = "/dev/usbtmc0"


def _write(dev: str, cmd: str) -> None:
    if not cmd.endswith("\n"):
        cmd += "\n"
    fd = os.open(dev, os.O_WRONLY)
    try:
        os.write(fd, cmd.encode("ascii"))
    finally:
        os.close(fd)


def _read_loop_into(dev: str, buf: bytearray, timeout: float, chunk_size: int = 4096) -> bool:
    """Keep reading (blocking) into buf until timeout; tolerate ETIMEDOUT."""
    deadline = time.monotonic() + timeout
    fd = os.open(dev, os.O_RDONLY)
    try:
        while time.monotonic() < deadline:
            try:
                chunk = os.read(fd, chunk_size)
                if chunk:
                    buf += chunk
                    return True
            except OSError as e:
                if e.errno == errno.ETIMEDOUT:
                    time.sleep(0.01)
                    continue
                raise
            time.sleep(0.01)
        return False
    finally:
        os.close(fd)


def query_line(cmd: str, timeout: float = 2.0, *, dev: str = DEV) -> str:
    _write(dev, cmd)
    time.sleep(0.03)
    buf = bytearray()
    ok = _read_loop_into(dev, buf, timeout=timeout)
    if not ok or not buf:
        raise TimeoutError(f"query_line timeout: {cmd}")
    line = bytes(buf).split(b"\n", 1)[0]
    return line.decode("ascii", errors="replace").strip()


def query_block(cmd: str, timeout: float = 5.0, *, dev: str = DEV) -> bytes:
    """
    Read IEEE488.2 block: #<N><len><payload>
    Returns payload bytes.
    """
    _write(dev, cmd)
    time.sleep(0.03)

    buf = bytearray()
    deadline = time.monotonic() + timeout

    while time.monotonic() < deadline:
        _read_loop_into(dev, buf, timeout=0.5)  # accumulate in chunks

        # parse block if possible
        if len(buf) >= 2 and buf[0:1] == b"#" and chr(buf[1]).isdigit():
            ndigits = int(chr(buf[1]))
            header_len = 2 + ndigits
            if len(buf) >= header_len:
                len_str = buf[2:header_len].decode("ascii", errors="replace")
                if len_str.isdigit():
                    payload_len = int(len_str)
                    total_len = header_len + payload_len
                    if len(buf) >= total_len:
                        return bytes(buf[header_len:total_len])

        # some responses might be plain text
        if b"\n" in buf and not buf.startswith(b"#"):
            return bytes(buf)

    raise TimeoutError(f"query_block timeout: {cmd} (got {len(buf)} bytes)")

def wrap_pi(x: float) -> float:
    return (x + math.pi) % (2.0 * math.pi) - math.pi

def time_axis(n: int, tscale: float, toffs: float) -> list[float]:
    # screen spans 12 div horizontally (approx). Works well for NORM screen data.
    t0 = toffs - 6.0 * tscale
    dt = (12.0 * tscale) / (n - 1) if n > 1 else 0.0
    return [t0 + i * dt for i in range(n)]

def phasor_at(samples: list[float], t: list[float], f0: float) -> complex:
    mean = sum(samples) / len(samples)
    w = 2.0 * math.pi * f0
    s = 0+0j
    for x, ti in zip(samples, t):
        s += (x - mean) * cmath.exp(-1j * w * ti)
    return s


def measure_phase(
    *,
    dev: str = DEV,
    timeout: float = 3.0,
    points_mode: str = "NORM",
    fetch_idn: bool = False,
) -> dict:
    """Measure CH1/CH2 phase difference using a single waveform capture.

    Returns a dict suitable for JSON serialization.
    """
    if points_mode not in {"NORM", "RAW", "MAX"}:
        raise ValueError("points_mode must be one of: NORM, RAW, MAX")

    idn = query_line("*IDN?", timeout=timeout, dev=dev) if fetch_idn else None

    # Stop & set points mode
    _write(dev, ":STOP")
    _write(dev, f":WAV:POIN:MODE {points_mode}")
    time.sleep(0.05)

    # timebase
    tscale = float(query_line(":TIM:SCAL?", timeout=timeout, dev=dev))
    toffs = float(query_line(":TIM:OFFS?", timeout=timeout, dev=dev))

    # frequency
    f0 = float(query_line(":MEAS:FREQ? CHAN1", timeout=timeout, dev=dev))

    # waveform payloads (bytes 0..255)
    b1 = query_block(":WAV:DATA? CHAN1", timeout=timeout, dev=dev)
    b2 = query_block(":WAV:DATA? CHAN2", timeout=timeout, dev=dev)
    n = min(len(b1), len(b2))
    b1, b2 = b1[:n], b2[:n]

    x1 = [float(v) for v in b1]
    x2 = [float(v) for v in b2]
    t = time_axis(n, tscale, toffs)

    amp1_pp_adc = (max(x1) - min(x1)) if x1 else float("nan")
    amp2_pp_adc = (max(x2) - min(x2)) if x2 else float("nan")

    p1 = phasor_at(x1, t, f0)
    p2 = phasor_at(x2, t, f0)
    phi1 = cmath.phase(p1)
    phi2 = cmath.phase(p2)
    dphi = wrap_pi(phi2 - phi1)

    deg = dphi * 180.0 / math.pi
    dt = dphi / (2.0 * math.pi * f0) if f0 != 0.0 else float("nan")

    _write(dev, ":RUN")

    return {
        "ts": time.time(),
        "idn": idn,
        "points_mode": points_mode,
        "n": n,
        "tscale": tscale,
        "toffs": toffs,
        "f0_hz": f0,
        "amp1_pp_adc": amp1_pp_adc,
        "amp2_pp_adc": amp2_pp_adc,
        "phi1_rad": phi1,
        "phi2_rad": phi2,
        "dphi_rad": dphi,
        "dphi_deg": deg,
        "dt_s": dt,
        "wave1": x1,
        "wave2": x2,
    }

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--dev", default=DEV)
    ap.add_argument("--mode", default="NORM", choices=["NORM", "RAW", "MAX"], help="WAV:POIN:MODE")
    ap.add_argument("--timeout", type=float, default=3.0)
    args = ap.parse_args()

    print("[+] *IDN? ->", query_line("*IDN?", timeout=args.timeout, dev=args.dev))

    r = measure_phase(dev=args.dev, timeout=args.timeout, points_mode=args.mode, fetch_idn=False)
    print(f"[+] Using f0 = {r['f0_hz']:.3f} Hz")
    print(f"[+] Got points: {r['n']} (mode={r['points_mode']})")
    print(f"[+] Phase:")
    print(f"    phi1={r['phi1_rad']:+.6f} rad, phi2={r['phi2_rad']:+.6f} rad")
    print(f"    dphi={r['dphi_rad']:+.6f} rad = {r['dphi_deg']:+.2f} deg")
    print(f"    equivalent delay dt={r['dt_s']:+.9e} s")

if __name__ == "__main__":
    main()