"""
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 (single trigger) just before the Agilent LP capture.
The LP→HS current step droops the 1.8 V rail, triggering the Rigol.
The waveform is then read over SCPI and written directly 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 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)

    rigol.write(":CHANnel1:DISPlay 1")
    rigol.write(":CHANnel2:DISPlay 0")
    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(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: 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_waveform_csv(path: Path) -> int:
    """
    Read Ch1 waveform from Rigol over SCPI and write to CSV.
    Sends :STOP first to ensure acquisition is complete before reading —
    this is reliable regardless of trigger/status state.

    Returns the number of samples written, or 0 on error.
    """
    try:
        rigol.write(":STOP")
        time.sleep(0.3)

        rigol.write(":WAVeform:SOURce CHANnel1")
        rigol.write(":WAVeform:FORMat ASC")   # Rigol DS1000Z uses ASC not ASCII
        time.sleep(0.1)

    except Exception as e:
        print(f"[RIGOL] 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] 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] Data read error: {e}")
        return 0

    if not vals:
        print("[RIGOL] No samples parsed — check scope 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] CSV write error: {e}")
        return 0