infratek_106a.py

"""
Infratek 106A / 306A / LEM D4000 Power Analyzer GPIB Driver

Blocking driver operating in single-shot (STOP) mode.
Uses MAV serial poll for query readiness and *OPC? for acquisition complete.
write_termination='\r' is required by the device.

Tested on: IT,306A,0,2.3R6 via GPIB0::5::INSTR
"""

import time
from typing import Literal, Optional

import pyvisa


# ── Literal types for SCPI parameter completion ───────────────────────
Aperture = Literal["100M", "250M", "500M", "1", "2"]
CurrentInput = Literal["IN5", "IN30", "SH"]
SyncSource = Literal["VOLT", "CURR"]
Phase = Literal["L1", "L2", "L3"]
VoltageRange = Literal["Auto", "300M", "1", "3", "10", "30", "100", "300", "1000"]
CurrentRangeIN30 = Literal["Auto", "1", "3", "10", "30", "100"]
CurrentRangeIN5 = Literal["Auto", "15M", "50M", "150M", "500M", "1.5", "5", "15"]
ShuntRange = Literal["Auto", "60M", "200M", "600M", "2", "6"]
DisplayMode = Literal["L1", "L2", "L3", "ALL"]

# ── Aperture lookup ────────────────────────────────────────────────────
APERTURE_SECONDS: dict[Aperture, float] = {
    "100M": 0.100,
    "250M": 0.250,
    "500M": 0.500,
    "1": 1.000,
    "2": 2.000,
}


class CommandError(Exception):
    """Raised when the device signals a command error via ESB in the status byte."""


class Infratek106A:
    """Blocking driver for Infratek 106A/306A in single-shot (STOP) mode.

    Usage:
        with Infratek106A() as dev:
            dev.configure(current_input="IN30", sync_source="CURR", aperture="250M", phase="L1")
            dev.trigger()
            v = dev.voltage_rms()        # reads L1 (current phase)
            v2 = dev.voltage_rms("L2")   # switches to L2, then reads

    For GUI use: run trigger() and read methods in a worker thread.
    """

    _MAV = 0x10  # bit 4: Message Available
    _ESB = 0x20  # bit 5: Event Status Bit (summary of ESR via ESE mask)

    def __init__(self, resource: str = "GPIB0::5::INSTR", timeout_ms: int = 10000,
                 warmup_count: int = 1):
        self._resource_str = resource
        self._timeout_ms = timeout_ms
        self._aperture_s: float = 0.250  # default 250ms
        self._phase: Phase = "L1"        # tracks current FORM:PHAS
        self._needs_warmup: bool = True   # first trigger discards garbage
        self._warmup_count: int = warmup_count
        self._rm: Optional[pyvisa.ResourceManager] = None
        self._inst: Optional[pyvisa.resources.Resource] = None

    # ── Connection ─────────────────────────────────────────────────────

    def open(self) -> "Infratek106A":
        self._rm = pyvisa.ResourceManager()
        self._inst = self._rm.open_resource(
            self._resource_str, write_termination="\r"
        )
        self._inst.timeout = self._timeout_ms
        self._init_device()
        return self

    def close(self):
        if self._inst:
            self._inst.close()
            self._inst = None
        if self._rm:
            self._rm.close()
            self._rm = None

    def __enter__(self):
        return self.open()

    def __exit__(self, *args):
        self.close()

    def _init_device(self):
        """Reset, configure ESE for CME detection, enter STOP mode."""
        self._write("*RST")
        time.sleep(1.0)
        self._write("*CLS")
        time.sleep(0.1)
        self._write("*ESE 32")    # ESB = CME only
        time.sleep(0.05)
        self._write("DIS:M ALL")
        time.sleep(0.05)
        self._write("ACQ:H Stop")
        time.sleep(0.3)

    # ── Low-level I/O ──────────────────────────────────────────────────

    def _write(self, cmd: str):
        self._inst.write(cmd)

    def _write_acq_batch(self, cmds: list[str]):
        """Send multiple ACQ config commands, then re-enter STOP mode once.

        Commands like ACQ:IN, ACQ:SYNC, ACQ:APER implicitly put the
        device back into RUN. This batches them with a single ACQ:H Stop.
        """
        if not cmds:
            return
        for cmd in cmds:
            self._inst.write(cmd)
        time.sleep(1)
        self._inst.write("ACQ:H Stop")
        time.sleep(0.3)
        self._needs_warmup = True

    def _wait_mav(self, timeout_s: float = 5.0):
        """Poll STB for MAV."""
        deadline = time.perf_counter() + timeout_s
        while time.perf_counter() < deadline:
            if self._inst.read_stb() & self._MAV:
                return
            time.sleep(0.005)
        raise TimeoutError(f"MAV not set within {timeout_s:.1f}s")

    def _query(self, cmd: str, timeout_s: float = 5.0) -> str:
        self._inst.write(cmd)
        self._wait_mav(timeout_s)
        return self._inst.read().strip()

    def _query_float(self, cmd: str, timeout_s: float = 5.0) -> float:
        return float(self._query(cmd, timeout_s))

    # ── Trigger / Acquire ──────────────────────────────────────────────

    def trigger(self):
        """*CLS → *TRG → sleep(aperture) → *OPC?. Blocking.

        If acquisition settings changed since the last trigger,
        warmup_count extra triggers are performed first (results discarded).
        """
        if self._needs_warmup:
            for _ in range(self._warmup_count):
                self._write("*CLS")
                self._write("*TRG")
                time.sleep(self._aperture_s)
                self._query("*OPC?", timeout_s=self._aperture_s + 5.0)
            self._needs_warmup = False
        self._write("*CLS")
        self._write("*TRG")
        time.sleep(self._aperture_s)
        self._query("*OPC?", timeout_s=self._aperture_s + 5.0)

    # ── Configuration ──────────────────────────────────────────────────

    def identify(self) -> str:
        return self._query("*IDN?")

    def configure(
        self,
        current_input: Optional[CurrentInput] = None,
        sync_source: Optional[SyncSource] = None,
        aperture: Optional[Aperture] = None,
        voltage_range: Optional[VoltageRange] = None,
        current_range: Optional[CurrentRangeIN30 | CurrentRangeIN5 | ShuntRange] = None,
        phase: Optional[Phase] = None,
        warmup_count: Optional[int] = None,
    ):
        """Configure device settings in one batch.

        Only sends commands for parameters that are provided.
        All ACQ commands are batched with a single sleep + ACQ:H Stop.
        """
        if warmup_count is not None:
            self._warmup_count = warmup_count
        acq_cmds: list[str] = []
        if current_input is not None:
            acq_cmds.append(f"ACQ:IN {current_input}")
        if sync_source is not None:
            acq_cmds.append(f"ACQ:SYNC {sync_source}")
        if aperture is not None:
            self._aperture_s = APERTURE_SECONDS[aperture]
            acq_cmds.append(f"ACQ:APER {aperture}")
        if voltage_range is not None:
            acq_cmds.append(f"ACQ:RAN:VOLT {voltage_range}")
        if current_range is not None:
            acq_cmds.append(f"ACQ:RAN:CURR {current_range}")
        self._write_acq_batch(acq_cmds)
        if phase is not None:
            self._write(f"FORM:PH {phase}")
            time.sleep(0.05)
            self._phase = phase

    def error(self) -> int:
        return int(self._query("ERR?"))

    # ── Phase switching ────────────────────────────────────────────────

    def _ensure_phase(self, phase: Optional[Phase]):
        """Switch FORM:PHAS only if different from current."""
        if phase is not None and phase != self._phase:
            self._write(f"FORM:PH {phase}")
            time.sleep(0.05)
            self._phase = phase

    # ── Read methods (call after trigger) ──────────────────────────────

    # Voltage
    def voltage_rms(self, phase: Optional[Phase] = None) -> float:
        self._ensure_phase(phase)
        return self._query_float("VOLT:RMS?")

    def voltage_rms_ac(self, phase: Optional[Phase] = None) -> float:
        self._ensure_phase(phase)
        return self._query_float("VOLT:RMS:AC?")

    def voltage_rect(self, phase: Optional[Phase] = None) -> float:
        self._ensure_phase(phase)
        return self._query_float("VOLT:RECT?")

    def voltage_mean(self, phase: Optional[Phase] = None) -> float:
        self._ensure_phase(phase)
        return self._query_float("VOLT:MEAN?")

    def voltage_min(self, phase: Optional[Phase] = None) -> float:
        self._ensure_phase(phase)
        return self._query_float("VOLT:MIN?")

    def voltage_max(self, phase: Optional[Phase] = None) -> float:
        self._ensure_phase(phase)
        return self._query_float("VOLT:MAX?")

    def voltage_peak(self, phase: Optional[Phase] = None) -> float:
        self._ensure_phase(phase)
        return self._query_float("VOLT:PEAK?")

    def voltage_crest(self, phase: Optional[Phase] = None) -> float:
        self._ensure_phase(phase)
        return self._query_float("VOLT:CREST?")

    def voltage_form(self, phase: Optional[Phase] = None) -> float:
        self._ensure_phase(phase)
        return self._query_float("VOLT:FORM?")

    def voltage_thd(self, phase: Optional[Phase] = None) -> float:
        self._ensure_phase(phase)
        return self._query_float("VOLT:THD?")

    # Current
    def current_rms(self, phase: Optional[Phase] = None) -> float:
        self._ensure_phase(phase)
        return self._query_float("CURR:RMS?")

    def current_rms_ac(self, phase: Optional[Phase] = None) -> float:
        self._ensure_phase(phase)
        return self._query_float("CURR:RMS:AC?")

    def current_rect(self, phase: Optional[Phase] = None) -> float:
        self._ensure_phase(phase)
        return self._query_float("CURR:RECT?")

    def current_mean(self, phase: Optional[Phase] = None) -> float:
        self._ensure_phase(phase)
        return self._query_float("CURR:MEAN?")

    def current_min(self, phase: Optional[Phase] = None) -> float:
        self._ensure_phase(phase)
        return self._query_float("CURR:MIN?")

    def current_max(self, phase: Optional[Phase] = None) -> float:
        self._ensure_phase(phase)
        return self._query_float("CURR:MAX?")

    def current_peak(self, phase: Optional[Phase] = None) -> float:
        self._ensure_phase(phase)
        return self._query_float("CURR:PEAK?")

    def current_crest(self, phase: Optional[Phase] = None) -> float:
        self._ensure_phase(phase)
        return self._query_float("CURR:CREST?")

    def current_form(self, phase: Optional[Phase] = None) -> float:
        self._ensure_phase(phase)
        return self._query_float("CURR:FORM?")

    def current_thd(self, phase: Optional[Phase] = None) -> float:
        self._ensure_phase(phase)
        return self._query_float("CURR:THD?")

    def current_int(self, phase: Optional[Phase] = None) -> float:
        """Average charge (short-time integration)."""
        self._ensure_phase(phase)
        return self._query_float("CURR:RECT:INT?")

    # Power
    def power_active(self, phase: Optional[Phase] = None) -> float:
        self._ensure_phase(phase)
        return self._query_float("POW:ACT?")

    def power_apparent(self, phase: Optional[Phase] = None) -> float:
        self._ensure_phase(phase)
        return self._query_float("POW:APP?")

    def power_reactive(self, phase: Optional[Phase] = None) -> float:
        self._ensure_phase(phase)
        return self._query_float("POW:REA?")

    def power_factor(self, phase: Optional[Phase] = None) -> float:
        self._ensure_phase(phase)
        return self._query_float("POW:FAC?")

    def power_active_ac(self, phase: Optional[Phase] = None) -> float:
        self._ensure_phase(phase)
        return self._query_float("POW:ACT:AC?")

    def power_apparent_ac(self, phase: Optional[Phase] = None) -> float:
        self._ensure_phase(phase)
        return self._query_float("POW:APP:AC?")

    def power_reactive_ac(self, phase: Optional[Phase] = None) -> float:
        self._ensure_phase(phase)
        return self._query_float("POW:REA:AC?")

    def power_factor_ac(self, phase: Optional[Phase] = None) -> float:
        self._ensure_phase(phase)
        return self._query_float("POW:FAC:AC?")

    # Energy
    def energy_active(self, phase: Optional[Phase] = None) -> float:
        self._ensure_phase(phase)
        return self._query_float("EN:ACT?")

    def energy_apparent(self, phase: Optional[Phase] = None) -> float:
        self._ensure_phase(phase)
        return self._query_float("EN:APP?")

    def energy_reset(self):
        self._write("EN:RES")

    def current_reset(self):
        """Reset charge integration values."""
        self._write("CURR:RES")

    # Frequency
    def frequency(self, phase: Optional[Phase] = None) -> float:
        self._ensure_phase(phase)
        return self._query_float("FREQ?")

    # ── FFT harmonics ─────────────────────────────────────────────────

    def _query_fft(self, cmd: str, start: int, end: int,
                   phase: Optional[Phase] = None) -> list[float]:
        """Set FORM:START/END, query FFT, then read remaining values with '?'."""
        self._ensure_phase(phase)
        self._write(f"FORM:START {start}")
        self._write(f"FORM:END {end}")
        count = end - start + 1
        values: list[float] = []
        # First value comes from the actual command
        values.append(self._query_float(cmd))
        # Remaining values come from consecutive '?' queries
        for _ in range(count - 1):
            values.append(self._query_float("?"))
        return values

    def voltage_fft(self, start: int = 1, end: int = 50,
                    phase: Optional[Phase] = None) -> list[float]:
        """Voltage harmonics from start to end (1–99)."""
        return self._query_fft("VOLT:FFT?", start, end, phase)

    def current_fft(self, start: int = 1, end: int = 50,
                    phase: Optional[Phase] = None) -> list[float]:
        """Current harmonics from start to end (1–99)."""
        return self._query_fft("CURR:FFT?", start, end, phase)

    def power_fft(self, start: int = 1, end: int = 50,
                  phase: Optional[Phase] = None) -> list[float]:
        """Power harmonics from start to end (1–99)."""
        return self._query_fft("POW:FFT?", start, end, phase)

    def analog_inputs(self, start: int = 0, end: int = 7) -> list[float]:
        """Read analog input channels (0–7)."""
        return self._query_fft("AINP?", start, end)

    # Acquisition quality
    def acquisition_quality(self) -> int:
        """Returns overload/underload bitmask."""
        return int(self._query("ACQ:QU?"))