import wave
from pyvisa.resources.resource import Resource
import pyvisa
from enum import Enum
import numpy as np


class Waveform(Enum):
    SINE = "SINE"
    SQUARE = "SQUARE"
    RAMP = "RAMP"
    PULSE = "PULSE"
    NOISE = "NOISE"
    ARB = "ARB"


class ChannelID(Enum):
    CH1 = 1
    CH2 = 2


CHANNEL_COUNT = 2


class Channel:
    def __init__(self, channel_id: ChannelID, instr: Resource):
        self.id = channel_id
        self.index = channel_id.value
        self.instr = instr

    # ---- Basic waveform ----
    def set_waveform(self, waveform: Waveform):
        self.instr.write(f"C{self.index}:BSWV WVTP,{waveform.value}")

    def set_frequency(self, freq: float):
        self.instr.write(f"C{self.index}:BSWV FRQ,{freq}")

    def set_amplitude(self, amplitude: float):
        self.instr.write(f"C{self.index}:BSWV AMP,{amplitude}")

    def set_offset(self, offset: float):
        self.instr.write(f"C{self.index}:BSWV OFST,{offset}")

    # ---- Output control (merged) ----
    def set_output(self, state: bool):
        cmd = "ON" if state else "OFF"
        self.instr.write(f"C{self.index}:OUTP {cmd}")

    # ---- Convenience ----
    def apply_sine(self, freq: float, amp: float, offset: float = 0):
        self.set_waveform(Waveform.SINE)
        self.set_frequency(freq)
        self.set_amplitude(amp)
        self.set_offset(offset)

    # ---- Query ----
    def get_frequency(self) -> float:
        response = self.instr.query(f"C{self.index}:BSWV?")
        return float(response.split(",")[1])


class GenChannels:
    def __init__(self, instr: Resource):
        self.instr = instr
        self.channels = {
            ChannelID.CH1: Channel(ChannelID.CH1, instr),
            ChannelID.CH2: Channel(ChannelID.CH2, instr),
        }

    def __getitem__(self, ch: ChannelID) -> Channel:
        return self.channels[ch]


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

        self.channels = GenChannels(self.instr)

        self.initialize()

    # ---- Initialization ----
    def initialize(self):
        self.instr.write("*RST")
        self.instr.write("*CLS")

    # ---- Global commands ----
    def identify(self) -> str:
        return self.instr.query("*IDN?")

    def close(self):
        self.instr.close()

    def get_error(self):
        return self.instr.query("SYST:ERR?")
