#This program allows to generate a waveform (analog or digital), transfer it on a file in the current folder and export it in the instrument's waveforms list.
#You can reset the "generate_waveform" flag under the section "Generate waveform settings" to disable it.
#If you try to export a waveform which is already present in the instrument's waveforms list an error will be thrown
#This program also configures the sequencer according to the parameters in section "Sequencer settings"
import pyvisa
import AWG_7000
import Waveform
import numpy as np
import matplotlib.pyplot as plt
from scipy import signal
import random
import re




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)

    #Master Encoding: Necessary if you use non-ascii characters, like accented vocals.
    master.encoding = "latin-1"

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




####################################################################################################################################
#Generate waveform settings
generate_waveform: bool = True
generated_waveform_name: str = "AsymmetricalTriangle"
generated_waveform_type = AWG_7000.WaveformType.Analog
generated_waveform_file: str = "my_waveform.txt"

#General analog waveform parameters
analog_waveform_type = Waveform.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)
phase = 0           #Sine wave phase in degrees
duty_cycle = 50     #Square wave duty cycle (percentage)
symmetry = 0.75      #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 = Waveform.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]

#Sequencer parameters
format = AWG_7000.Format.Amplitude_offset
analog_waveform_names = [["SINE", "SINE"],
                  ["RAMP", "RAMP"],
                  ["SQUARE", "SQUARE"]]
amplitudes = [[1, 1.5],
               [0.75, 0.75],
               [1, 1.2]]
offsets = [[0, 0],
            [0, 0.3],
            [0, 0]]
decreasing_strategies = [[AWG_7000.DecreasingStrategy.Decimation, AWG_7000.DecreasingStrategy.Decimation],
                         [AWG_7000.DecreasingStrategy.Decimation, AWG_7000.DecreasingStrategy.Decimation],
                         AWG_7000.DecreasingStrategy.Decimation, AWG_7000.DecreasingStrategy.Decimation]
increasing_strategies = [[AWG_7000.IncreasingStrategy.Interpolation, AWG_7000.IncreasingStrategy.Interpolation],
                         [AWG_7000.IncreasingStrategy.Interpolation, AWG_7000.IncreasingStrategy.Interpolation],
                         [AWG_7000.IncreasingStrategy.Interpolation, AWG_7000.IncreasingStrategy.Interpolation]]
sub_lengths = [[16384, 16384],
               [16384, 16384],
               [16384, 16384]]
delays = [[0, 0],
          [0, 0],
          [0, 0]]
output_formats = [[AWG_7000.OutputFormat.Waveform, AWG_7000.OutputFormat.Waveform],
                  [AWG_7000.OutputFormat.Waveform, AWG_7000.OutputFormat.Waveform],
                  [AWG_7000.OutputFormat.Waveform, AWG_7000.OutputFormat.Waveform]]
entry_lengths = [16384, 16384, 16384]
repetitions = [1, 1, 1]
sequencer_length = len(analog_waveform_names)

#Digital
digital_channels = AWG_7000.ChannelsNumber.Eight
digital_pod = AWG_7000.Pod.PodA
skew = 0
digital_voltage_level = 3
digital_waveforms = ["COUNTER", "COUNTER", "COUNTER"]
marker_custom_pattern = ["CLOCK", "CLOCK", "CLOCK"]
####################################################################################################################################




try:
    #Generate new waveform only if not already present in the instrument's waveforms list
    if generate_waveform and not re.findall(generated_waveform_name, AWG.GetWaveformsList()):
        #Analog waveforms
        if generated_waveform_type == AWG_7000.WaveformType.Analog:

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

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

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

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

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

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

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

            elif analog_waveform_type == Waveform.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 in the file Waveform")

            #Generate instance of class waveform to graph generated function
            wf = Waveform.AnalogWaveform(samples, dt, generated_waveform_name)
            t, y = wf.get_timeseries()
            plt.plot(t, y)
            plt.xlabel("Time (s)")
            plt.ylabel("Amplitude")
            plt.title(wf.name)
            plt.grid(True)
            plt.show()

            #Transfer generated analog waveform into file
            with open(generated_waveform_file, 'w') as file:
                for point in y: file.write(f"{str(point)}\n")


        #Digital waveforms
        elif generated_waveform_type == AWG_7000.WaveformType.Digital:
            if digital_waveform_type == Waveform.DigitalType.ArbitraryPattern: data = arbitrary_pattern

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

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

            elif digital_waveform_type == Waveform.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 in the file Waveform")

            #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)]

            #Generate instance of class waveform to graph generated function
            wf = Waveform.DigitalWaveform(data_dig, dt)
            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("Digital Waveform")
            plt.legend()
            plt.grid(True)
            plt.show()

            #Create and transfer data into file
            with open(generated_waveform_file, 'w') as file:
                for pattern in data:
                    file.write(str(pattern) + "\n")

        else: raise Exception("Parameter waveform_type must be an instance of the class WaveformType defined in the file AWG_7000")


        #Transfer waveform file to instrument and import in waveform list
        AWG.TransferFile(generated_waveform_file, generated_waveform_file)
        AWG.ImportFileInWaveformList(generated_waveform_name, generated_waveform_file, generated_waveform_type)
        AWG.DeleteFile(generated_waveform_file)
    



    #Instrument configuration (sequencer, channel and general settings)
    AWG.Initialize(True)

    for channel in AWG.channels:
        AWG.ConfigureChannelSettings(channel, 100, 0, 0, 0, 0, AWG_7000.KeepLastMode.CustomOnStop, AWG_7000.Polarity.Normal, 0, 0,
                                     AWG_7000.ChannelOptimization.High_Linearity, 0, AWG_7000.State.OFF, AWG_7000.KeepLastMode.CustomOnStop, 0)
        
    AWG.ConfigureGeneralSettings(AWG_7000.OperatingMode.Full_rate, AWG_7000.IncreasingStrategy.Interpolation, AWG_7000.DecreasingStrategy.Decimation, AWG_7000.RunMode.Continuous,
                                 AWG_7000.WaitTriggerState.First_sample, 1, AWG_7000.BurstSelection.NumberCycles, AWG_7000.JumpMode.AfterRepetitions)
    
    for name in digital_waveforms:
        if name != "":
            AWG.SetDigitalChannelsNumber(AWG_7000.ChannelsNumber.Eight)
            AWG.ConfigureDigitalChannels(digital_channels, digital_pod, digital_voltage_level, AWG_7000.SetValue.PRESENT, skew, AWG_7000.SetValue.PRESENT)
            AWG.SetDigitalsState(AWG_7000.State.ON)

    AWG.SetSequencerLength(sequencer_length, AWG_7000.SetValue.PRESENT)

    #Program all wanted sequencer entries
    for i in range(sequencer_length):
        entry = AWG_7000.SequencerEntry(entry_number=i+1, analog_waveform_names=analog_waveform_names[i], amplitudes=amplitudes[i], offsets=offsets[i],
                                        decreasing_strategies=decreasing_strategies[i], increasing_strategies=increasing_strategies[i], sub_lengths=sub_lengths[i],
                                        delays=delays[i], digital_waveform_name=digital_waveforms[i], custom_pattern=marker_custom_pattern[i], length=entry_lengths[i],
                                        repetitions=repetitions[i], output_formats=output_formats[i])
        AWG.ConfigureSequencerElement(entry)

    for channel in AWG.channels: AWG.SetOutputState(channel, AWG_7000.State.ON)

    while(True):
        if input("Press enter to start") == "":
            AWG.AWGStart()
            break
    
    while(True):
        if input("Press enter to stop") == "":
            AWG.AWGStop()
            break

except Exception: print(f"The following exception has occured: {e}")