"""
Implements partial support for Siglent SDS1000X-E series oscilloscopes.
"""
import re

from lab_control.oscilloscope import Oscilloscope
from lab_control.connection.tcp_connection import TCPConnection

def _checkChannel(channel):
    assert channel in SDS1000XE.AVAILABLE_CHANNELS, "SDS1000X-E: Invalid channel {channel}"

class SDS1000XE(Oscilloscope):
    """
    Instances of this class connect to the SDS1000X-E IP
    and port and offer an API to control the device.
    """
    PORT = 5025
    TIMEOUT = 5.0
    AVAILABLE_CHANNELS = range(1, 5)

    def __init__(self, address, overrideConnection=None):
        super().__init__()
        if overrideConnection is not None:
            self._connection = overrideConnection
        else:
            self._connection = TCPConnection(address)

        self._connection.configure({"timeout" : SDS1000XE.TIMEOUT})

    def _measure(self, channel: int, code: str) -> float:
        _checkChannel(channel)
        request = f"C{channel}:PAVA? {code}\r\n"
        response = self._connection.send(request)

        if response is not None:
            pattern = r"C(?P<responseChannel>\d):PAVA .+,(?P<rawMeasurement>[\d.E+-]+)\w+"
            matches = re.search(pattern, response)
            measurement = float(matches.group("rawMeasurement"))
        else:
            measurement = None

        return measurement

    def measureAmplitude(self, channel: int) -> float:
        return self._measure(channel, "AMPL")

    def measurePeakToPeak(self, channel: int) -> float:
        return self._measure(channel, "PKPK")

    def measureRMS(self, channel: int) -> float:
        return self._measure(channel, "RMS")

    def measureFrequency(self, channel: int) -> float:
        return self._measure(channel, "FREQ")

    def setVoltsPerDivision(self, channel: int, volts: float) -> None:
        _checkChannel(channel)
        query = f"C{channel}:VDIV {volts:.2E}V\r\n"
        self._connection.send(query, responseExpected=False)

    def getDivisionsDisplayed(self) -> int:
        return 8