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_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}")

    # ========== 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 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 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 read_raw(self):
        """
        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.

        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)

        # Read the data length
        length_bytes = self.socket.recv(num_digits)
        if len(length_bytes) != num_digits:
            raise RuntimeError("Failed to read data length")

        data_length = int(length_bytes)

        # Read the actual binary data
        chunks = []
        remaining = data_length
        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)

        # Read the trailing newline
        self.socket.recv(1)

        return b''.join(chunks)

    @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()