scanengine-3/hardware/tektronix_base.py

"""
Base class for Tektronix oscilloscope control with raw SCPI command support.
"""

import socket


class TektronixOscilloscopeBase:
    """
    Base class for Tektronix oscilloscope instruments.

    Provides low-level SCPI command interface for fast data transfer.
    """

    # Valid acquisition modes (short and long forms)
    # SCPI uses capital letters for short form, scope responds with full uppercase
    ACQUIRE_MODES = {
        'SAMPLE': ['SAM', 'SAMPLE', 'SAMple'],
        'PEAKDETECT': ['PEAK', 'PEAKDETECT', 'PEAKdetect'],
        'HIRES': ['HIR', 'HIRES', 'HIRes'],
        'AVERAGE': ['AVE', 'AVERAGE', 'AVErage'],
        'ENVELOPE': ['ENV', 'ENVELOPE', 'ENVelope']
    }

    # Valid trigger coupling modes
    TRIGGER_COUPLING = {
        'DC': ['DC'],
        'HFREJ': ['HFRej', 'HFREJ'],
        'LFREJ': ['LFRej', 'LFREJ'],
        'NOISEREJ': ['NOISErej', 'NOISEREJ']
    }

    # Valid trigger slope options
    TRIGGER_SLOPE = {
        'RISE': ['RISe', 'RISE'],
        'FALL': ['FALL'],
        'EITHER': ['EITher', 'EITHER']
    }

    # Valid trigger modes
    TRIGGER_MODE = {
        'AUTO': ['AUTO'],
        'NORMAL': ['NORMal', 'NORMAL']
    }

    # Valid channel coupling modes
    CHANNEL_COUPLING = {
        'AC': ['AC'],
        'DC': ['DC']
    }

    # Valid channel termination values (ohms)
    CHANNEL_TERMINATION = [50, 1000000]

    # Valid data encoding formats
    DATA_ENCODING = {
        'ASCII': ['ASCIi', 'ASCII'],
        'RIBINARY': ['RIBinary', 'RIBINARY'],  # Signed integer, MSB first (recommended)
        'RPBINARY': ['RPBinary', 'RPBINARY'],  # Positive integer, MSB first
        'FPBINARY': ['FPBinary', 'FPBINARY'],  # Floating point binary
        'SRIBINARY': ['SRIbinary', 'SRIBINARY'],  # Signed integer, byte-swapped
        'SRPBINARY': ['SRPbinary', 'SRPBINARY'],  # Positive integer, byte-swapped
        'SFPBINARY': ['SFPbinary', 'SFPBINARY']   # Floating point binary, byte-swapped
    }

    # Valid waveform output preamble encoding
    WFMOUTPRE_ENCODING = {
        'BINARY': ['BINary', 'BINARY'],
        'ASCII': ['ASCii', 'ASCII']
    }

    # Valid byte order options
    BYTE_ORDER = {
        'LSB': ['LSB'],  # Least significant byte first
        'MSB': ['MSB']   # Most significant byte first
    }

    def __init__(self, resource_name=None, port=4000, timeout=5.0, terminator='\n'):
        """
        Initialize the oscilloscope base class.

        Args:
            resource_name: IP address or hostname of the instrument
            port: TCP port for SCPI communication (default: 4000)
            timeout: Socket timeout in seconds (default: 5.0)
            terminator: Line terminator for SCPI commands (default: '\n')
        """
        self.resource_name = resource_name
        self.port = port
        self.timeout = timeout
        self.terminator = terminator
        self.socket = None
        self._connected = False

    def connect(self):
        """Establish TCP socket connection to the oscilloscope."""
        if not self.resource_name:
            raise ValueError("resource_name (IP address/hostname) must be provided")

        if self._connected:
            return

        try:
            self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.socket.settimeout(self.timeout)
            self.socket.connect((self.resource_name, self.port))
            self._connected = True
        except socket.error as e:
            self.socket = None
            self._connected = False
            raise ConnectionError(
                f"Failed to connect to {self.resource_name}:{self.port} - {e}") from e

    def disconnect(self):
        """
        Close TCP socket connection to the oscilloscope.
        """
        if self.socket:
            try:
                self.socket.close()
            except Exception:
                pass
            finally:
                self.socket = None
                self._connected = False

    def _normalize_channel(self, channel):
        """Normalize channel input to integer (1-4)."""
        if isinstance(channel, str):
            channel_upper = channel.upper()
            if not (channel_upper.startswith('CH') and len(channel_upper) == 3 and channel_upper[2].isdigit()):
                raise ValueError(f"Invalid channel: {channel}. Must be CH1, CH2, CH3, or CH4")
            channel = int(channel_upper[2])

        if not isinstance(channel, int) or channel < 1 or channel > 4:
            raise ValueError(f"Invalid channel: {channel}. Must be 1-4")

        return channel

    def set_acquire_mode(self, mode):
        """Set the acquisition mode."""
        # Normalize mode to uppercase for comparison
        mode_upper = mode.upper()

        # Check if mode is valid (any short or long form)
        valid = False
        for variants in self.ACQUIRE_MODES.values():
            if mode_upper in [v.upper() for v in variants]:
                valid = True
                break

        if not valid:
            valid_options = []
            for variants in self.ACQUIRE_MODES.values():
                valid_options.extend(variants)
            raise ValueError(f"Invalid acquisition mode: {mode}. "
                           f"Valid options: {', '.join(valid_options)}")

        self.write(f"ACQuire:MODe {mode}")

    def get_fastframe_state(self):
        """Query the current FastFrame state."""
        response = self.query("HORizontal:FASTframe:STATE?")
        return int(response)

    def set_fastframe_state(self, state):
        """Enable or disable FastFrame mode."""
        # Convert boolean to int if needed
        if isinstance(state, bool):
            state = 1 if state else 0

        # Validate state
        if state not in (0, 1):
            raise ValueError(f"Invalid FastFrame state: {state}. Must be 0 (off) or 1 (active)")

        self.write(f"HORizontal:FASTframe:STATE {state}")

    def set_fastframe_count(self, count):
        """Set the number of FastFrame frames."""
        # Validate count
        if not isinstance(count, int) or count <= 0:
            raise ValueError(f"Invalid frame count: {count}. Must be a positive integer")

        self.write(f"HORizontal:FASTframe:COUNt {count}")

    def get_fastframe_max_frames(self):
        """Query how many FastFrame frames the current horizontal settings allow."""
        return int(self.query("HORizontal:FASTframe:MAXFRames?"))

    def get_record_length(self):
        """Query the current horizontal record length."""
        response = self.query("HORizontal:MODe:RECOrdlength?")
        return int(response)

    def set_sample_rate(self, rate):
        """Set the horizontal sample rate."""
        # Validate rate
        if not isinstance(rate, (int, float)) or rate <= 0:
            raise ValueError(f"Invalid sample rate: {rate}. Must be a positive number")

        self.write(f"HORizontal:MODe:SAMPLERate {rate}")

    def set_trigger_slope(self, slope):
        """Set the trigger slope."""
        slope_upper = slope.upper()

        # Check if slope is valid
        valid = False
        for variants in self.TRIGGER_SLOPE.values():
            if slope_upper in [v.upper() for v in variants]:
                valid = True
                break

        if not valid:
            valid_options = []
            for variants in self.TRIGGER_SLOPE.values():
                valid_options.extend(variants)
            raise ValueError(f"Invalid trigger slope: {slope}. "
                           f"Valid options: {', '.join(valid_options)}")

        self.write(f"TRIGger:A:EDGE:SLOpe {slope}")

    def set_trigger_source(self, source):
        """Set the trigger source."""
        # Convert integer to channel string if needed
        if isinstance(source, int):
            if source < 1 or source > 4:
                raise ValueError(f"Invalid trigger source: {source}. Must be 1-4")
            source = f"CH{source}"

        # Validate channel format
        source_upper = source.upper()
        if not (source_upper.startswith('CH') and len(source_upper) == 3 and source_upper[2].isdigit()):
            raise ValueError(f"Invalid trigger source: {source}. Must be CH1, CH2, CH3, or CH4")

        channel_num = int(source_upper[2])
        if channel_num < 1 or channel_num > 4:
            raise ValueError(f"Invalid trigger source: {source}. Channel must be 1-4")

        self.write(f"TRIGger:A:EDGE:SOUrce {source}")

    def set_trigger_level(self, channel, level):
        """Set the trigger level for a specific channel."""
        # Convert to channel number if string
        if isinstance(channel, str):
            channel_upper = channel.upper()
            if not (channel_upper.startswith('CH') and len(channel_upper) == 3 and channel_upper[2].isdigit()):
                raise ValueError(f"Invalid channel: {channel}. Must be CH1, CH2, CH3, or CH4")
            channel = int(channel_upper[2])

        # Validate channel number
        if not isinstance(channel, int) or channel < 1 or channel > 4:
            raise ValueError(f"Invalid channel: {channel}. Must be 1-4")

        # Validate level
        if not isinstance(level, (int, float)):
            raise ValueError(f"Invalid trigger level: {level}. Must be a number")

        self.write(f"TRIGger:A:LEVel:CH{channel} {level}")

    def set_trigger_mode(self, mode):
        """Set the trigger mode."""
        mode_upper = mode.upper()

        # Check if mode is valid
        valid = False
        for variants in self.TRIGGER_MODE.values():
            if mode_upper in [v.upper() for v in variants]:
                valid = True
                break

        if not valid:
            valid_options = []
            for variants in self.TRIGGER_MODE.values():
                valid_options.extend(variants)
            raise ValueError(f"Invalid trigger mode: {mode}. "
                           f"Valid options: {', '.join(valid_options)}")

        self.write(f"TRIGger:A:MODe {mode}")

    # ========== Channel Control Methods ==========

    def set_channel_bandwidth(self, channel, bandwidth):
        """Set the bandwidth for a specific channel."""
        channel = self._normalize_channel(channel)
        self.write(f"CH{channel}:BANdwidth {bandwidth}")

    def set_channel_coupling(self, channel, coupling):
        """Set the coupling mode for a specific channel."""
        channel = self._normalize_channel(channel)
        coupling_upper = coupling.upper()

        # Validate coupling
        valid = False
        for variants in self.CHANNEL_COUPLING.items():
            if coupling_upper in [v.upper() for v in variants]:
                valid = True
                break

        if not valid:
            valid_options = []
            for variants in self.CHANNEL_COUPLING.values():
                valid_options.extend(variants)
            raise ValueError(f"Invalid coupling mode: {coupling}. "
                           f"Valid options: {', '.join(valid_options)}")

        self.write(f"CH{channel}:COUPling {coupling}")

    def set_channel_label_name(self, channel, name):
        """Set the label name for a specific channel."""
        channel = self._normalize_channel(channel)

        if not isinstance(name, str):
            raise ValueError(f"Invalid label name: {name}. Must be a string")

        # Need to quote the string for SCPI
        self.write(f'CH{channel}:LABel:NAMe "{name}"')

    def set_channel_position(self, channel, position):
        """Set the vertical position for a specific channel."""
        channel = self._normalize_channel(channel)

        if not isinstance(position, (int, float)):
            raise ValueError(f"Invalid position: {position}. Must be a number")

        self.write(f"CH{channel}:POSition {position}")

    def set_channel_scale(self, channel, scale):
        """Set the vertical scale for a specific channel."""
        channel = self._normalize_channel(channel)

        if not isinstance(scale, (int, float)) or scale <= 0:
            raise ValueError(f"Invalid scale: {scale}. Must be a positive number")

        self.write(f"CH{channel}:SCAle {scale}")

    def set_channel_termination(self, channel, termination):
        """Set the termination for a specific channel."""
        channel = self._normalize_channel(channel)

        if termination not in self.CHANNEL_TERMINATION:
            raise ValueError(f"Invalid termination: {termination}. "
                           f"Valid options: {', '.join(map(str, self.CHANNEL_TERMINATION))}")

        self.write(f"CH{channel}:TERmination {termination}")

    # ========== Measurement Methods ==========

    # Immediate measurements the alignment/inspection code asks for.  The
    # instrument accepts many more; this list is what has been exercised here,
    # and an unlisted type is far more likely to be a typo than a deliberate
    # choice.
    MEASUREMENT_TYPES = {
        'MEAN': ['MEAN'],
        'AMPLITUDE': ['AMPlitude', 'AMPLITUDE'],
        'MAXIMUM': ['MAXimum', 'MAXIMUM'],
        'MINIMUM': ['MINImum', 'MINIMUM'],
        'PK2PK': ['PK2pk', 'PK2PK'],
        'RMS': ['RMS'],
    }

    # Tektronix returns this sentinel when a measurement cannot be made (no
    # acquisition yet, source off, signal outside the graticule).  It is a
    # valid float, so it has to be caught explicitly or it reads as a
    # 1e38 V measurement.
    MEASUREMENT_INVALID = 9.9e37

    def measure_immediate(self, channel, measurement_type='MEAN'):
        """Take an immediate measurement on one channel and return it in volts.

        "Immediate" measurements are computed on demand and are not added to
        the scope's on-screen measurement badges, so this leaves whatever the
        operator has set up on the front panel untouched.

        Raises ValueError if the instrument reports the measurement as
        unavailable, which on a triggered-acquisition scope usually means it
        has not acquired anything yet.
        """
        channel = self._normalize_channel(channel)

        if measurement_type.upper() not in self.MEASUREMENT_TYPES:
            raise ValueError(
                f"Invalid measurement type: {measurement_type}. "
                f"Valid options: {', '.join(self.MEASUREMENT_TYPES)}")

        self.write(f"MEASUrement:IMMed:SOUrce1 CH{channel}")
        self.write(f"MEASUrement:IMMed:TYPe {measurement_type}")
        response = self.query("MEASUrement:IMMed:VALue?")
        try:
            value = float(response)
        except ValueError as exc:
            raise ValueError(
                f"Unparseable {measurement_type} measurement on CH{channel}: "
                f"{response!r}") from exc

        if abs(value) >= self.MEASUREMENT_INVALID:
            raise ValueError(
                f"CH{channel} {measurement_type} is unavailable (the scope "
                f"returned its no-measurement sentinel). Check that the "
                f"channel is on and that the acquisition is triggering.")
        return value

    def get_acquisition_count(self):
        """Number of acquisitions since the acquisition was last started.

        A caller polling a free-running scope uses this to tell a fresh
        reading from a stale one: if the count has not moved, the record has
        not changed and every measurement taken off it is the previous
        answer.
        """
        return int(float(self.query("ACQuire:NUMACq?")))

    # ========== Waveform Transfer Methods ==========

    def set_data_source(self, source):
        """Set the data source for waveform transfer."""
        # Convert integer to channel string if needed
        if isinstance(source, int):
            if source < 1 or source > 4:
                raise ValueError(f"Invalid data source: {source}. Must be 1-4")
            source = f"CH{source}"

        # Validate channel format
        source_upper = source.upper()
        if not (source_upper.startswith('CH') and len(source_upper) == 3 and source_upper[2].isdigit()):
            raise ValueError(f"Invalid data source: {source}. Must be CH1, CH2, CH3, or CH4")

        channel_num = int(source_upper[2])
        if channel_num < 1 or channel_num > 4:
            raise ValueError(f"Invalid data source: {source}. Channel must be 1-4")

        self.write(f"DATa:SOUrce {source}")

    def set_data_encoding(self, encoding):
        """Set the curve transfer encoding (e.g. RIBinary = signed int, MSB first)."""
        valid = ('ASCII', 'RIBinary', 'RPBinary', 'FPBinary',
                 'SRIbinary', 'SRPbinary', 'SFPbinary')
        if encoding.upper() not in [v.upper() for v in valid]:
            raise ValueError(f"Invalid data encoding: {encoding}. "
                             f"Valid options: {', '.join(valid)}")
        self.write(f"DATa:ENCdg {encoding}")

    def set_data_width(self, width):
        """Set bytes per sample for curve transfers."""
        if width not in (1, 2):
            raise ValueError(f"Invalid data width: {width}. Must be 1 or 2")
        self.write(f"DATa:WIDth {width}")

    def query_wfmoutpre(self):
        """Query all waveform output preamble parameters."""
        return self.query("WFMOutpre?")

    def transfer_curve(self):
        """
        Transfer waveform curve data from the oscilloscope.

        This method reads binary curve data according to the current DATa:SOUrce,
        DATa:ENCdg, and WFMOutpre settings.

        Returns:
            bytes: Raw binary curve data

        Raises:
            RuntimeError: If not connected or transfer fails
        """
        # Send the CURVe? query
        self.write("CURVe?")


        # Read the binary data using IEEE 488.2 format
        return self.read_raw()

    def transfer_fastframe(self, parse=True, byte_count=1, signed=True, byte_order='MSB'):
        """
        Transfer all FastFrame waveform data from the oscilloscope.

        When FastFrame is enabled, a single CURVe? query returns all frames
        as sequential IEEE 488.2 binary blocks. This method reads all frames
        efficiently in a single operation.

        Args:
            parse: If True, parse the raw bytes into integer arrays (default: True)
            byte_count: Number of bytes per sample for parsing (1 or 2)
            signed: True for signed integer parsing, False for unsigned
            byte_order: 'MSB' or 'LSB' for byte order when parsing

        Returns:
            list: List of waveforms. If parse=True, each waveform is a list of
                  integer ADC values. If parse=False, each waveform is raw bytes.

        Raises:
            RuntimeError: If not connected, FastFrame not enabled, or transfer fails
            ValueError: If parsing parameters are invalid

        Example:
            # Enable FastFrame with 1000 frames
            scope.set_fastframe_state(True)
            scope.set_fastframe_count(1000)

            # Acquire data
            scope.write("ACQuire:STATE RUN")
            time.sleep(2)
            scope.write("ACQuire:STATE STOP")

            # Transfer all frames at once
            waveforms = scope.transfer_fastframe()
            print(f"Got {len(waveforms)} frames, {len(waveforms[0])} points each")
        """
        # Verify FastFrame is enabled
        if not self.get_fastframe_state():
            raise RuntimeError("FastFrame is not enabled. Enable it with set_fastframe_state(True)")

        # Use the number of frames actually acquired, not the configured maximum.
        # If the stage stops early, fewer triggers arrive and the scope captures
        # fewer frames than configured — reading the configured count would block.
        acquired = int(self.query("ACQuire:NUMFRAMESACQuired?"))
        if acquired <= 0:
            raise RuntimeError("Scope acquired 0 FastFrame frames — no data to read")
        frame_count = acquired

        # Send a single CURVe? query - scope will return all frames
        self.write("CURVe?")

        # Read all frames as sequential IEEE 488.2 binary blocks
        waveforms = []
        for _ in range(frame_count):
            curve_bytes = self.read_raw()

            if parse:
                waveform = self.parse_curve_data(
                    curve_bytes,
                    byte_count=byte_count,
                    signed=signed,
                    byte_order=byte_order
                )
            else:
                waveform = curve_bytes

            waveforms.append(waveform)

        return waveforms

    def transfer_fastframe_bulk(self, frame_count, samples_per_frame,
                                bytes_per_sample=1):
        """Transfer a whole FastFrame burst as one contiguous buffer.

        Unlike transfer_fastframe this does not care how the scope frames the
        response — it accumulates blocks until it has the expected byte count,
        so one large IEEE block and one block per frame both work.  Returns a
        bytearray of frame_count * samples_per_frame * bytes_per_sample bytes.
        """
        if not self.get_fastframe_state():
            raise RuntimeError("FastFrame is not enabled. Enable it with set_fastframe_state(True)")

        expected = frame_count * samples_per_frame * bytes_per_sample
        if expected <= 0:
            raise RuntimeError(
                f"Nothing to transfer: {frame_count} frames × "
                f"{samples_per_frame} samples × {bytes_per_sample} bytes"
            )

        self.write("CURVe?")

        buf = bytearray()
        while len(buf) < expected:
            block = self.read_raw(expected_bytes=expected - len(buf))
            if not block:
                raise RuntimeError(
                    f"Scope returned an empty block {len(buf)}/{expected} bytes "
                    "into the burst transfer"
                )
            buf += block

        if len(buf) != expected:
            raise RuntimeError(
                f"Burst transfer overran: got {len(buf)} bytes, expected {expected}"
            )
        return buf

    def parse_curve_data(self, curve_bytes, byte_count=1, signed=True, byte_order='MSB'):
        """Parse raw curve data into integer array."""
        if byte_count not in (1, 2):
            raise ValueError(f"Invalid byte count: {byte_count}. Must be 1 or 2")

        if byte_order not in ('MSB', 'LSB'):
            raise ValueError(f"Invalid byte order: {byte_order}. Must be 'MSB' or 'LSB'")

        values = []

        if byte_count == 1:
            # 8-bit data
            for byte in curve_bytes:
                if signed:
                    # Convert unsigned byte to signed
                    value = byte if byte < 128 else byte - 256
                else:
                    value = byte
                values.append(value)
        else:
            # 16-bit data
            num_samples = len(curve_bytes) // 2
            for i in range(num_samples):
                byte1 = curve_bytes[i * 2]
                byte2 = curve_bytes[i * 2 + 1]

                if byte_order == 'MSB':
                    # Most significant byte first
                    value = (byte1 << 8) | byte2
                else:
                    # Least significant byte first
                    value = (byte2 << 8) | byte1

                if signed:
                    # Convert unsigned to signed
                    if value >= 32768:
                        value = value - 65536

                values.append(value)

        return values

    def write(self, command):
        """Send a raw SCPI command to the instrument."""
        if not self._connected or not self.socket:
            raise RuntimeError("Not connected to instrument")

        if not command.endswith(self.terminator):
            command += self.terminator

        self.socket.sendall(command.encode('ascii'))

    def query(self, command):
        """Send a SCPI query and return the response."""
        if not self._connected or not self.socket:
            raise RuntimeError("Not connected to instrument")

        self.write(command)

        response = b''
        terminator_bytes = self.terminator.encode('ascii')
        while True:
            chunk = self.socket.recv(4096)
            if not chunk:
                break
            response += chunk
            if terminator_bytes in chunk:
                break

        return response.decode('ascii').strip()

    def _recv_exact(self, count):
        """Read exactly `count` bytes; recv() is free to return fewer."""
        chunks = []
        remaining = count
        while remaining > 0:
            chunk = self.socket.recv(min(remaining, 65536))
            if not chunk:
                raise RuntimeError("Connection closed while reading data")
            chunks.append(chunk)
            remaining -= len(chunk)
        return b''.join(chunks)

    def read_raw(self, expected_bytes=None):
        """
        Read raw binary data from the instrument.

        Useful for fast waveform data transfer. Handles IEEE 488.2 binary block format.
        Format: #<N><digits><data><newline>
        where N is a digit indicating how many digits follow,
        and those digits specify the length of the data block.

        `#0` announces an indeterminate-length block, normally delimited by EOI
        — which a raw socket never sees.  Pass expected_bytes to say how many
        bytes to take in that case.

        Returns:
            bytes: Raw binary data (without IEEE 488.2 header)

        Raises:
            RuntimeError: If not connected or invalid format
            socket.error: If communication fails
        """
        if not self._connected or not self.socket:
            raise RuntimeError("Not connected to instrument")

        # Read the '#' character
        header = self.socket.recv(1)
        if header != b'#':
            # Read whatever else is buffered so the error shows the actual response
            self.socket.setblocking(False)
            try:
                more = self.socket.recv(100)
            except OSError:
                more = b''
            self.socket.setblocking(True)
            raise RuntimeError(f"Invalid IEEE 488.2 header: expected '#', got {header}, followed by: {more[:50]}")

        # Read the digit indicating length of length field
        length_of_length = self.socket.recv(1)
        if not length_of_length.isdigit():
            raise RuntimeError(f"Invalid length specifier: {length_of_length}")

        num_digits = int(length_of_length)

        if num_digits == 0:
            # Indeterminate length: no byte count follows, and the EOI that
            # would delimit it does not exist on a raw socket.
            if expected_bytes is None:
                raise RuntimeError(
                    "Scope returned an indeterminate-length block (#0); "
                    "read_raw needs expected_bytes to size it over a socket"
                )
            data_length = expected_bytes
        else:
            data_length = int(self._recv_exact(num_digits))

        data = self._recv_exact(data_length)

        # Read the trailing newline / block separator
        self._recv_exact(1)

        return data

    @property
    def is_connected(self):
        """Check if instrument is connected."""
        return self._connected

    def __enter__(self):
        """Context manager entry."""
        self.connect()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit."""
        self.disconnect()