import time
from pyvisa.resources.resource import Resource
import pyvisa
from enum import Enum

import numpy as np

CHANNEL_COUNT = 4
# Rigol DHO4804


class TrigState(Enum):
    TD = 0
    WAIT = (1,)
    RUN = (2,)
    AUTO = (3,)
    STOP = (4,)


class Channel:
    def __init__(self, index, instr):
        self.index = index
        self.instr = instr
        self.enabled = False
        self.probe = 10
        self.scale_V = 5

        self.setProbe(self.probe)
        self.setVScale(self.scale_V)

        self.last_scale = 1  # V/div

    def setProbe(self, probe: float):
        self.instr.write(f":CHANnel{self.index}:PROBe {probe}")

    def setVScale(self, scale: float):
        self.instr.write(f":CHANnel{self.index}:SCALe {scale}")

    def getWaveformRaw(self):
        self.instr.write(f":WAVeform:SOURce CHANnel{self.index}")
        self.instr.write(":WAV:DATA?")
        response = self.instr.read_raw()
        data_start_offset = int(response[1:2].decode("ASCII")) + 2
        data_len = (len(response) - data_start_offset) // 2 * 2

        dt = np.dtype("<u2")
        datapoints = np.frombuffer(
            response[(data_start_offset) : (data_start_offset + data_len)], dtype=dt
        )

        return datapoints

    def getOffset(self):

        ret = self.instr.query(f":CHANnel{self.index}:OFFSet?")
        return float(ret)

    def setOffset(self, offset: float):

        self.instr.write(f":CHANnel{self.index}:OFFSet {offset}")

    def getWaveform(self):

        waveform = self.getWaveformRaw()
        offset = self.getYReference()

        waveform = (waveform.astype("int") - offset) * self.getYincrement()

        return waveform

    #
    def getYincrement(self):
        self.instr.write(f":WAVeform:SOURce CHANnel{self.index}")

        ret = self.instr.query(":WAVeform:YINCrement?")
        return float(ret)

    def getXincrement(self):

        self.instr.write(f":WAVeform:SOURce CHANnel{self.index}")

        ret = self.instr.query(":WAVeform:XINCrement?")
        return float(ret)

    def getYOrigin(self):

        self.instr.write(f":WAVeform:SOURce CHANnel{self.index}")

        ret = self.instr.query(":WAVeform:YORigin?")
        return int(ret)

    def getYReference(self):

        self.instr.write(f":WAVeform:SOURce CHANnel{self.index}")

        ret = self.instr.query(":WAVeform:YREFerence?")
        return int(ret)

    def genTimeArray(self, waveform):
        x_delta = self.getXincrement()
        tim = np.arange(0, len(waveform) * x_delta, x_delta)

        return tim

    def getVpp(self):
        self.instr.write(f":MEASure:SOURce CHANnel{self.index}")
        ret = self.instr.query(":MEASure:ITEM? VPP")
        saturated = False

        vpp = float(ret)

        if vpp > 1e10:
            saturated = True

        return vpp, saturated

    def getVrms(self):
        self.instr.write(f":MEASure:SOURce CHANnel{self.index}")
        ret = self.instr.query(":MEASure:ITEM? VRMS")
        saturated = False

        rms = float(ret)

        if rms > 1e10:
            saturated = True

        return rms, saturated

    def clampVscale(self, val):
        if val > 100:
            val = 100
        if val < 0.01:
            val = 0.01
        return val

    def getSpan(self):  # display the difference between min and max voltage in dump
        waveform = self.getWaveformRaw()
        return np.max(waveform) - np.min(waveform)


class RigolOsc:
    def __init__(self, visa_address: str):
        rm = pyvisa.ResourceManager()
        self.instr: Resource = rm.open_resource(visa_address)

        self.channels = [Channel(i + 1, self.instr) for i in range(CHANNEL_COUNT)]
        self.initialize()

    def getChannel(self, channel_id: int) -> Channel:
        if 0 < channel_id <= CHANNEL_COUNT:
            return self.channels[channel_id - 1]
        else:
            raise Exception("Invalid channel id")

    def single(self):  # arm single mode
        self.instr.write(":SINGle")
        self.instr.query("*OPC?")
        time.sleep(0.01)  # Either this, or RUN before triggering
        # print(self.getTrigStatus())

        while self.getTrigStatus() != TrigState.STOP:
            continue
            # print(self.getTrigStatus())

            # continue

    def getTrigStatus(self):
        ret = self.instr.query(":TRIGger:STATus?")
        if "TD" in ret:
            return TrigState.TD
        elif "WAIT" in ret:
            return TrigState.WAIT
        elif "RUN" in ret:
            return TrigState.RUN
        elif "AUTO" in ret:
            return TrigState.AUTO
        elif "STOP" in ret:
            return TrigState.STOP
        else:
            print("error")
            # print(ret)
            raise (Exception("Unknown state"))

    def setPoints(self, points):
        self.instr.write(f":WAVeform:POINts {points}")

    def initialize(self):
        self.instr.write(":WAVeform:FORMat WORD")
        self.instr.write(":WAVeform:MODE RAW")
        self.enable10MIn()

    def getMemoryDepth(self):

        ret = self.instr.query(":ACQuire:MDEPth?")
        depth = int(float(ret))
        # print(depth)
        return depth

    def run(self):
        self.instr.write(":RUN")

    def setTimescale(self, timebase: float):
        self.instr.write(f":TIMebase:MAIN:SCALe {timebase}")

    def enable10MIn(self):
        self.instr.write(":SYSTem:RCLock CINput")
