import pyvisa
import numpy as np
from datetime import datetime
import matplotlib.pyplot as plt
from enum import Enum
from scipy import signal
import random
from AWG_7000 import AWG7000
import re
from Waveform import AnalogType
from Waveform import DigitalType
from Waveform import AnalogWaveform
from Waveform import DigitalWaveform
from Waveform import WaveformType




try:
    # Initialize VISA
    rm = pyvisa.ResourceManager()

    # List connected resources
    resources = rm.list_resources()
    print(f"\n\n\nConnected devices: {resources}")

    #Instrument TCP/IP address
    address_listed = resources[1]
    address_not_listed = "TCPIP0::192.168.0.55::inst0::INSTR"

    #If you use NI MAX or similar, you can also use one of the adresses in the list 
    master = rm.open_resource(address_not_listed)

    #Initialize AWG_5062
    AWG = AWG7000(master)
except Exception as e: print(f"The following exception has occured while connecting with the instrument: {e}")




#################################################################################################################################################
waveform_type = WaveformType.Analog
waveform_name = "MyTriangle"
waveform_file = "MyWaveform.txt"

#Analog waveform parameters
analog_waveform_type = AnalogType.Triangle
sampling_frequency = 1000000000  #Hz
analog_n_samples = 1000
frequency = 1000000  #Hz
amplitude = 0.5  #V
offset = 0  #V
dt = 1 / (sampling_frequency)
t = np.linspace(0, 1/frequency, int(sampling_frequency/frequency), endpoint=False)
#Specific waveform parameters
phase = 0           #Sine wave phase in degrees
duty_cycle = 50     #Square wave duty cycle (percentage)
symmetry = 0.2      #Triangle wave symmetry: set at 0.5 for ideal triangle; set at 0 or 1 for sawtooth
sigma = 1           #Gaussian wave standard deviation
gamma = 0.5         #Lorentz wave parameter
tau = 1             #Exponential wave time constant in seconds

#Digital waveform parametrs
digital_waveform_type = DigitalType.Counter
dt = 0.001                  #Sampling period (just for graphic representation)
data = []                   #Leave this list empty
counter_start_value = 0
counter_stop_value = 255
digital_n_samples = 10              #Number of random bits to generate for random waveform
fixed_decimal_value = 100   #Integer between 0 and 255
arbitrary_pattern = [1, 4, 7, 109, 44, 123]
#################################################################################################################################################




if waveform_type == WaveformType.Analog:
    if analog_waveform_type == AnalogType.Sine: samples = np.sin(np.linspace(0, 2 * np.pi, analog_n_samples) + phase) + offset

    elif analog_waveform_type == AnalogType.Square: samples = amplitude * signal.square(2 * np.pi * frequency * t, duty_cycle/100) + offset
        
    elif analog_waveform_type == AnalogType.Triangle: samples = amplitude * signal.sawtooth(2 * np.pi * frequency * t, symmetry) + offset

    elif analog_waveform_type == AnalogType.Gaussian: samples = amplitude * np.exp(-(2 * np.pi * frequency * t - 1/(2*frequency))**2 / (2 * sigma**2)) + offset

    elif analog_waveform_type == AnalogType.Sinc: samples = amplitude * np.sinc(2 * np.pi * frequency * (t-1/(2*frequency))) + offset

    elif analog_waveform_type == AnalogType.Lorentz: samples = amplitude / (1 + ((2 * np.pi * frequency * (t - 1 / (2*frequency))) / gamma)**2) + offset

    elif analog_waveform_type == AnalogType.Exp_Decay: samples = amplitude * np.exp(-2 * np.pi * frequency * t/tau) + offset

    elif analog_waveform_type == AnalogType.Exp_Rise: samples = amplitude * np.exp(2 * np.pi * frequency * t/tau) + offset

    elif analog_waveform_type == AnalogType.Haversine: samples = 0.5 * (1 - np.cos(2 * np.pi * frequency * t))

    else: raise Exception("Parameter analog_waveform_type must be an instance of the class AnalogType defined above in this file")

    try:
        wf = AnalogWaveform(samples, dt, waveform_name)
        t, y = wf.get_timeseries()

        #Waveform graph
        plt.plot(t, y)
        plt.xlabel("Time (s)")
        plt.ylabel("Amplitude")
        plt.title(wf.name)
        plt.grid(True)
        plt.show()

        #Transfer file
        with open(waveform_file, 'w') as file:
            for point in y: file.write(f"{str(point)}\n")
    except Exception as e: print(f"The following exception has occured: {e}")




elif waveform_type == WaveformType.Digital:
    if digital_waveform_type == DigitalType.ArbitraryPattern: data = arbitrary_pattern

    elif digital_waveform_type == DigitalType.Counter:
        for i in range(counter_start_value, counter_stop_value + 1): data.append(i)

    elif digital_waveform_type == DigitalType.Random:
        for i in range(digital_n_samples): data.append(random.randint(0, 255))

    elif digital_waveform_type == DigitalType.FixedDecimalValue:
        for i in range(digital_n_samples): data.append(fixed_decimal_value)
    else: raise Exception("Parameter digital_waveform_type must be an instance of the class DigitalType defined above in this file")

    try:
        #Generate matrix data_dig to realize the graph
        data_dig = []
        for number in data:
            #Converts pattern into a bynary list of 8 bits
            binary_str = bin(number)[2:]
            list_bit = [int(b) for b in binary_str]
            for i in range(8 - len(list_bit)): list_bit = [0] + list_bit
            data_dig.append(list_bit)
        #Transpose the matrix to graph the waveform correctly (each line contains the values of each digital channel in a sampling instant)
        data_dig = [list(line) for line in zip(*data_dig)]

        #Create waveform associated with data_dig
        wf = DigitalWaveform(data_dig, dt)

        #Waveform graph
        time_axis = [wf.dt * i for i in range(wf.data.shape[1])]
        offset = 2
        for i, line in enumerate(wf.data):
            plt.step(time_axis, line + i * offset, where='post', label=wf.line_names[i])
        plt.xlabel("Time (s)")
        plt.ylabel("Digital Lines")
        plt.title(waveform_name)
        plt.legend()
        plt.grid(True)
        plt.show()

        #Create and transfer data into file
        with open(waveform_file, 'w') as file:
            for pattern in data:
                file.write(str(pattern) + "\n")
    except Exception as e: raise Exception(f"The followiung exception has occured: {e}")


#Transfer waveform file to instrument and import in waveform list if not already present in the instrument's waveforms list
#if not re.findall(waveform_name, AWG.GetWaveformsList()):
AWG.TransferFile(waveform_file, waveform_file)
AWG.ImportFileInWaveformList(waveform_name, waveform_file, waveform_type)
AWG.DeleteFile(waveform_file)