"""
rigol_scope.py

Controls the Rigol DS1202Z-E at 192.168.45.5 for 1.8 V supply rail monitoring.
Called from dual_capture() in mipi_test.py during the LP pass.

The scope is armed just before the Agilent LP capture.
The LP→HS current step droops the 1.8 V rail, triggering the Rigol.
The CH1 waveform is read over SCPI and written to the local data/ folder.
"""

import csv
import time
import vxi11
from pathlib import Path

RIGOL_HOST     = "192.168.45.5"
V18_SCALE      = 0.1       # V/div  — 100 mV/div; 10 divs = ±500 mV around 1.8 V
V18_OFFSET     = -1.8      # V      — shifts zero reference so 1.8 V sits at screen centre
V18_TIMEBASE   = 1e-6      # s/div  — 1 µs/div = 10 µs total window
V18_TRIG_LEVEL = 1.76      # V      — falling-edge trigger on supply droop > 40 mV
TRIG_TIMEOUT_S = 15.0      # s      — wait this long for Rigol to capture after arming

rigol: vxi11.Instrument | None = None


# ---------------------------------------------------------------------------
# Connection
# ---------------------------------------------------------------------------

def connect() -> bool:
    global rigol
    try:
        rigol = vxi11.Instrument(RIGOL_HOST)
        rigol.timeout = 10
        idn = rigol.ask("*IDN?").strip()
        print(f"[RIGOL] Connected: {idn}")
        return True
    except Exception as e:
        print(f"[RIGOL] Connection failed — 1.8 V monitoring disabled: {e}")
        rigol = None
        return False


def disconnect():
    global rigol
    if rigol:
        try:
            rigol.close()
        except Exception:
            pass
        rigol = None


def is_connected() -> bool:
    return rigol is not None


# ---------------------------------------------------------------------------
# Setup
# ---------------------------------------------------------------------------

def configure():
    """
    Configure Rigol CH1 for 1.8 V supply monitoring.
    AUTO trigger sweep: if no droop occurs, scope still captures on timeout
    so we always get a supply snapshot even when the rail is healthy.
    """
    rigol.write(":STOP")
    time.sleep(0.2)

    # CH1 — 1.8 V supply rail
    rigol.write(":CHANnel1:DISPlay 1")
    rigol.write(":CHANnel1:COUPling DC")
    rigol.write(":CHANnel1:PROBe 10")
    rigol.write(f":CHANnel1:SCALe {V18_SCALE:.3f}")
    rigol.write(f":CHANnel1:OFFSet {V18_OFFSET:.3f}")

    rigol.write(":CHANnel2:DISPlay 0")

    rigol.write(f":TIMebase:MAIN:SCALe {V18_TIMEBASE:.2E}")
    rigol.write(":TRIGger:MODE EDGE")
    rigol.write(":TRIGger:EDGe:SOURce CHANnel1")
    rigol.write(":TRIGger:EDGe:SLOPe NEGative")
    rigol.write(f":TRIGger:EDGe:LEVel {V18_TRIG_LEVEL:.3f}")
    rigol.write(":TRIGger:SWEep AUTO")   # auto: captures even without a droop trigger
    time.sleep(0.3)
    rigol.write(":RUN")                  # start acquiring immediately after configure
    time.sleep(0.2)

    print(f"[RIGOL] Configured: CH1=1.8 V rail, {int(V18_TIMEBASE*1e6)} µs/div, "
          f"trigger <{V18_TRIG_LEVEL} V falling (AUTO sweep, running)")


# ---------------------------------------------------------------------------
# Acquisition
# ---------------------------------------------------------------------------

def arm():
    """Ensure scope is running so it is actively acquiring when the LP event occurs.
    The waveform is frozen with :STOP inside read_waveform_csv() at collection time."""
    rigol.write(":RUN")


def wait_captured(timeout_s: float = TRIG_TIMEOUT_S) -> bool:
    """
    Poll until the scope has completed its single acquisition.
    DS1000Z reports STOP when done (triggered or auto-timed-out).
    Returns True when ready, False if timeout exceeded.
    """
    deadline = time.time() + timeout_s
    while time.time() < deadline:
        try:
            status = rigol.ask(":TRIGger:STATus?").strip().upper()
            if status in ("STOP", "TD"):
                return True
        except Exception:
            pass
        time.sleep(0.1)
    return False


def _read_channel_csv(channel: str, path: Path, stop_first: bool = True) -> int:
    """
    Read one Rigol channel waveform over SCPI and write to CSV.
    stop_first=False skips :STOP when the scope was already stopped by a prior read.
    Returns the number of samples written, or 0 on error.
    """
    try:
        if stop_first:
            rigol.write(":STOP")
            time.sleep(0.3)
        rigol.write(f":WAVeform:SOURce {channel}")
        rigol.write(":WAVeform:FORMat ASC")   # Rigol DS1000Z uses ASC not ASCII
        time.sleep(0.1)
    except Exception as e:
        print(f"[RIGOL] {channel} waveform setup error: {e}")
        return 0

    try:
        preamble = rigol.ask(":WAVeform:PREamble?").strip().split(",")
        # [0]=fmt [1]=type [2]=points [3]=count [4]=x_incr [5]=x_orig [6]=x_ref
        # [7]=y_incr [8]=y_orig [9]=y_ref
        x_incr = float(preamble[4])
        x_orig = float(preamble[5])
        x_ref  = float(preamble[6])
    except Exception as e:
        print(f"[RIGOL] {channel} preamble error: {e}")
        return 0

    try:
        raw = rigol.ask(":WAVeform:DATA?").strip()

        # Strip TMC binary header (#<n_digits><byte_count>...) if present
        if raw.startswith("#"):
            n_digits = int(raw[1])
            raw = raw[2 + n_digits:]

        vals = [float(v) for v in raw.split(",") if v.strip()]
    except Exception as e:
        print(f"[RIGOL] {channel} data read error: {e}")
        return 0

    if not vals:
        print(f"[RIGOL] {channel}: no samples parsed — check channel and format settings")
        return 0

    try:
        path.parent.mkdir(exist_ok=True)
        with open(path, "w", newline="") as f:
            writer = csv.writer(f)
            writer.writerow(["Time (s)", "Voltage (V)"])
            for i, v in enumerate(vals):
                t = x_orig + (i - x_ref) * x_incr
                writer.writerow([f"{t:.9f}", f"{v:.6f}"])
        return len(vals)
    except Exception as e:
        print(f"[RIGOL] {channel} CSV write error: {e}")
        return 0


def read_waveform_csv(path: Path) -> int:
    """Read CH1 (1.8 V supply) waveform from Rigol and write to CSV."""
    return _read_channel_csv("CHANnel1", path, stop_first=True)


