# Capture waveforms from the Tek TDS1912B oscilloscope

from usb_devices import *
from matplotlib import rc
from matplotlib import pyplot as plt

#rc('text', usetex=True)
from pylab import *
from numpy import *
import time
import string as s
import struct

class TDS1012BChannelParameters :
    def __init__(self, channel) :
        # Specify channel at 1 or 2
        self.channel = channel
        self.xunit = 0
        self.xinc = 0
        self.yunit = 0
        self.ymult = 0
        self.yoff = 0
        self.yzero = 0

class TDS1012B :
    def __init__(self, devices, sn) :
        # Devices is the dictionary of usb devices
        if devices.has_key('TDS1012B' + '-' + sn.upper()) :
            self.visa = devices['TDS1012B' + '-' + sn.upper()]
        elif devices.has_key('TDS1012B') :
            self.visa = devices['TDS1012B']     # No serial number was found.
        else :
            self.visa = ''      # Device not found
        self.identity = 'TDS1012B'
        self.sn = sn.upper()
        self.ch1 = TDS1012BChannelParameters(1)
        self.ch2 = TDS1012BChannelParameters(2)
        self.error = ''
        self.time = 0   # Used to report time taken for read and write functions

    def Write(self, stuff) :
        zero = time.time()
        try:
            self.visa.write(stuff)
        except:
            self.error = self.identity + ' write operation probably timed out.'
        self.time = time.time() - zero

    def Read(self) :
        zero = time.time()
        try:
            q = self.visa.read()
        except:
            self.error = self.identity + ' read operation probably timed out.'
            q = self.error
        self.time = time.time() - zero
        return q

    def ClearQueue(self) :
        self.Write('*esr?')
        return self.Read()

    def SetChannel(self, c) :
        error = ''
        if c > 2 | c <1 :
            error = "Incorrect channel request."
        self.Write('data:source ch'+str(c))
        return error

    def GetScaleParameters(self, ch):
        # ch = channel object
        error = ''
        if ch.channel > 2 | ch.channel <1 :
            error = "Incorrect channel request."
            return error
        self.visa.write('data:source ch'+str(ch.channel))
        try :
            self.visa.write('wfmpre:xunit?')
            ch.xunit = s.split(self.visa.read(), 'XUNIT')[1].strip().strip('"')
        except :
            error = 'Parameter read error.  Is channel ' + str(ch.channel) + ' currently displayed?'
            return error
        print ch.xunit
        self.visa.write('wfmpre:xincr?')
        ch.xinc = float(s.split(self.visa.read(), 'XINCR')[1])
        print ch.xinc
        self.visa.write('wfmpre:yunit?')
        ch.yunit = s.split(self.visa.read(), 'YUNIT')[1].strip().strip('"')
        self.visa.write('wfmpre:ymult?')
        ch.ymult = float(float(s.split(self.visa.read(), 'YMULT')[1]))
        self.visa.write('wfmpre:yoff?')
        ch.yoff = float(float(s.split(self.visa.read(), 'YOFF')[1]))
        self.visa.write('wfmpre:yzero?')
        ch.yzero = float(float(s.split(self.visa.read(), 'YZERO')[1]))
        # Make a dictionary of values
        ch.parameters = {}
        ch.parameters['xunit'] = ch.xunit
        ch.parameters['xinc'] = ch.xinc
        ch.parameters['yunit'] = ch.yunit
        ch.parameters['ymult'] = ch.ymult
        ch.parameters['yoff'] = ch.yoff
        ch.parameters['yzero'] = ch.yzero
        return error

    def GetWaveformOld(self, ch, mode) :
        # Mode is ascii or bin.
        # ch = channel object
        error = ''
        if ch.channel > 2 | ch.channel <1 :
            error = "Incorrect channel request."
            return error, []
        self.visa.write('data:source ch'+str(ch.channel))
        self.visa.write('data:start 1')
        self.visa.write('data:stop 2500')
        if s.lower(mode).find('asc') > -1 :
            self.visa.write('data:encdg ascii')
        else :
            error = 'Only ascii mode is supported.'
            return error, []
        self.visa.write('curve?')
        w = self.visa.read().split('CURVE')[1].split(',')
        x, y = self.ScaleWaveform(w, ch)
        return error, x, y

    def GetWaveform(self, ch, mode) :
        # Mode is binary or ascii.
        # ch = channel object
        error = ''
        if ch.channel > 2 | ch.channel <1 :
            error = "Incorrect channel request."
            return error, []
        self.visa.write('data:source ch'+str(ch.channel))
        self.visa.write('data:start 1')
        self.visa.write('data:stop 2500')
        if s.lower(mode).find('asc') > -1 :
            # Ascii transfer takes 4 seconds.
            mode = 'a'
            self.visa.write('data:encdg ascii')
            self.visa.write('curve?')
            w = self.visa.read().split('CURVE')[1].split(',')
        elif s.lower(mode).find('bin') > -1 :
            # Binary transfer takes a fraction of a second.
            mode = 'b'
            self.visa.write('data:encdg ribinary')  # Request signed 16 bit integer, big endian
            self.visa.write('data:width 2')
            self.visa.write('curve?')
            data_block = self.visa.read()
            print(data_block[0:13]) # This yields ':CURVE #45000' which is correct.
            # Notice that the slice function yields the characters from 0 to 13-1=12.
            raw_data = data_block[13:]  # Strip away the first 13 characters.  This seems to be right, but what are those first 13 characters?
            print(len(raw_data))
            w = struct.unpack('>2500h', raw_data) # A buffer of 2500 signed short (16 bit) integers in the big endian order, not the native little endian.
        else :
            error = 'Only ascii or binary modes are possible.'
            return error, [], []
        x, y = self.ScaleWaveform(w, ch)
        return error, x, y

    def ScaleWaveform(self, w, ch) :
        y = []
        x = []
        delx = 0
        #yoff = p['yoff']
        for z in w:
            y.append((float(z)-ch.yoff)*ch.ymult + ch.yzero)
            x.append(delx)
            delx += ch.xinc
        return x, y

def QueryOscilloscopeOld(sn):
    sn = s.strip(sn)
    #sn = 'C055976'
    devices = FindUSBDevices()

    print 'Look for TDS1012B:'
    scope = TDS1012B(devices, sn)
    if scope.visa == '' :
        print "Error: TDS1012B-" + sn + " not found."
        return
    scope.Write('*IDN?')
    print "Scope identification: ", scope.Read()
    scope.ClearQueue()
    print 'Channel 1 parameters:'
    error = scope.GetScaleParameters(scope.ch1)
    if error :
        print error
        return
    print scope.ch1.parameters
    print 'Channel 2 parameters:'
    error = scope.GetScaleParameters(scope.ch2)
    if error :
        print error
        return
    print scope.ch2.parameters
    return scope

def QueryOscilloscope(sn, channel_reqs):
    error = ''
    sn = s.strip(sn)
    #sn = 'C055976'
    devices = FindUSBDevices()
    print 'Look for TDS1012B:'
    scope = TDS1012B(devices, sn)
    if scope.visa == '' :
        error = "Error: TDS1012B-" + sn + " not found."
        return error
    scope.Write('*IDN?')
    print "Scope identification: ", scope.Read()
    scope.ClearQueue()
    scope_chans = []
    channels = []
    for ch in channel_reqs :
        if ch == 1 :
            scope_chans.append(scope.ch1)
            channels.append(ch)
            print 'Channel 1 parameters:'
            error = scope.GetScaleParameters(scope.ch1)
            if error :
                return scope, error
            print scope.ch1.parameters
        elif ch == 2 :
            scope_chans.append(scope.ch2)
            channels.append(ch)
            print 'Channel 2 parameters:'
            error = scope.GetScaleParameters(scope.ch2)
            if error :
                return scope, error
            print scope.ch2.parameters
        else :
            print('Channel requested is invalid : ' + str(ch))
    if (len(scope_chans) == 0) :    # | (len(scope_chans) > 2):
        error = 'Channel declaration error: requested channels are invalid or not specified.'
        print(error)
        print(channel_reqs)
        return scope, error
    scope.channels = channels
    scope.chan_refs = scope_chans
    return scope, error

def WaveformAcquisition(sn, channel_reqs=[1, 2], xy_pairs=[[1,2]], symbols = ['', ''], graph_title=None, show_legend=True, legends=None, legend_pos=None) :
    # If two channels are requested, they are acquired at different times.
    # Truly simultaneous waveform acquisition is under consideration.
    # This will require a GetSimultaneousWaveform function which will use the single sweep option.
    scope, error = QueryOscilloscope(sn, channel_reqs)
    if error :
        print('Fatal error encountered. ' + error)
        return
    #scope_chans = []
    #channels = []
    #for ch in channel_reqs :
        #if ch == 1 :
            #scope_chans.append(scope.ch1)
            #channels.append(ch)
        #elif ch == 2 :
            #scope_chans.append(scope.ch2)
            #channels.append(ch)
        #else :
            #print('Channel requested is invalid : ' + str(ch))
    #if (len(scope_chans) == 0) | (len(scope_chans) > 2):
        #print('Channel declaration error: requested channels are invalid or specified more than once.')
        #print(channel_reqs)
        #print('I quit.')
        #return
    data = []
    for i in range(len(scope.chan_refs)) :
        #error, wx, wy = scope.GetWaveform(scope.chan_refs[i], 'ascii')
        error, wx, wy = scope.GetWaveform(scope.chan_refs[i], 'binary')
        #print wx
        #print wy
        #return
        if error:
            print error
            return
        print('Channel ' + str(scope.channels[i]) + ' data array lengths = ' + str(len(wx)) + ', ' + str(len(wy)))
        data.append([wx, wy])
    #return
    scope.ClearQueue()
    a = raw_input('Enter plot mode as yt (the default) or xy or r or rvgs or osc :')
    aa = a.strip().lower()
    if graph_title == None :
        graph_title = raw_input('Enter the title for the graph: ')
    if (aa == '') | (aa == 'yt') :
        ys = []
        for i in range(len(data)) :
            x = data[i][0]
            y = data[i][1]
            ys.append(y)
            if legends :
                if i >= len(legends) :
                    legends.append('Channel '+str(scope.channels[i]))
                the_label = legends[i]
            else :
                the_label = 'Channel '+str(scope.channels[i])
            plot(x, y, label=the_label)
        if show_legend : legend(loc=legend_pos)
        title(graph_title)
        xlabel(scope.chan_refs[0].xunit)
        ylabel('Amplitude ('+scope.chan_refs[0].yunit+')')
        grid(True)
        show()
    if (aa == 'osc') :
        ys = []
        for i in range(len(data)) :
            x = data[i][0]
            y = data[i][1]
            ys.append(y)
            #if legends :
                #if i >= len(legends) :
                    #legends.append('Channel '+str(scope.channels[i]))
                #the_label = legends[i]
            #else :
                #the_label = 'Channel '+str(scope.channels[i])
            fty = fft.fft(y)
            ftyabs = abs(fty[0:1250])
            n = fty.size
            xstep = scope.chan_refs[0].xinc
            print n, xstep
            ftx = fft.fftfreq(n, xstep)[0:1250]/1.0e06    # in MHz
            f1 = plt.figure()
            sb1 = f1.add_subplot(111)
            sb1.plot(ftx, 20.0*log10(ftyabs))
            #if show_legend : legend(loc=legend_pos)
            plt.title(graph_title)
            plt.xlabel(scope.chan_refs[0].xunit + r'$^{-1}/10^6$')
            plt.ylabel(r'20 log10(Amplitude) in dB')
            sb1.grid(True)
            f2 = plt.figure()
            sb2 = f2.add_subplot(111)
            sb2.plot(ftx, ftyabs, label=legends[0])
            ftmaxindex = ftyabs.argmax()
            ftmax = ftyabs[ftmaxindex]
            print ftmax, ftmaxindex
            lorentz = array([])
            freqo = ftx[ftmaxindex]
            q = float(raw_input('Enter Q :'))
            if q == 0 :
                q = 200
            fwhm = freqo/q/2.0
            alpha = fwhm**2
            for ftfreq in ftx :
                lorentz = append(lorentz, ftmax*alpha/(alpha + (ftfreq - freqo)**2))
            print lorentz.size
            sb2.plot(ftx, lorentz, label=legends[1])
            if show_legend : legend(loc=legend_pos)
            plt.title(graph_title)
            plt.xlabel(scope.chan_refs[0].xunit + r'$^{-1}/10^6$')
            plt.ylabel('Amplitude (arbitrary)')
            sb2.grid(True)
        plt.show()
    elif aa == 'r' :
        if len(data) < 2 :
            print ('Two channels must be specified for r display mode.')
            return
        #if len(data) > 2 :
            #print('Using only the first two channels for the xy graph.')
        rf = float(raw_input('Enter the Rf value in Ohms: '))
        for p in xy_pairs :
            #print p
            x = array(data[p[0]-1][1])
            y = array(data[p[1]-1][1])
            r = - rf*x/y
            ys = [r]
            plot(x,r)
        title(graph_title)
        xlabel('Vd ('+scope.chan_refs[0].yunit+')')
        ylabel('Rchannel ('+'Ohms'+')')
        grid(True)
        show()
    elif aa == 'rvgs' :     # Measure jfet channel R (Vgs).
        if len(data) < 2 :
            print ('Two channels must be specified for rvgs display mode.')
            return
        #if len(data) > 2 :
            #print('Using only the first two channels for the xy graph.')
        rf = float(raw_input('Enter the Rf value in Ohms: '))
        vds = float(raw_input('Enter the Vds value in Volts: '))
        for p in xy_pairs :
            #print p
            x = array(data[p[0]-1][1])
            y = array(data[p[1]-1][1])
            r = - rf*vds/y
            #r = y/x    # test with a simple divider.  This ratio should be a constant A
            ys = [r]
            plot(x,r)
        title(graph_title)
        xlabel('Vd ('+scope.chan_refs[0].yunit+')')
        ylabel('Rchannel ('+'Ohms'+')')
        grid(True)
        show()
    elif aa == 'xy' :
        if len(data) < 2 :
            print ('Two channels must be specified for xy display mode.')
            return
        #if len(data) > 2 :
            #print('Using only the first two channels for the xy graph.')
        for p in xy_pairs :
            #print p
            x = data[p[0]-1][1]
            y = data[p[1]-1][1]
            ys = [y]
            plot(x,y)
        title(graph_title)
        xlabel('Amplitude ('+scope.chan_refs[0].yunit+')')
        ylabel('Amplitude ('+scope.chan_refs[1].yunit+')')
        grid(True)
        show()
    else :
        print('Inappropriate graph requested. Type must be either yt or xy.')
        return
    SaveResults(x, ys, scope.channels, scope.chan_refs)

def SaveResults(x, ys, channels, scope_chans) :
    # y is an array of y values
    # channels is an array of channels
    outfile = raw_input('Enter a file name or nothing to skip: ')
    if s.strip(outfile) :
        o = open(outfile, 'w')
        if channels :
            o.write('# Channel information:\n')
            for i in range(len(scope_chans)) :
                o.write('#\tChannel '+str(channels[i])+' : '+scope_chans[i].xunit+', '+scope_chans[i].yunit+'\n')
        o.write('\n# x value followed by y for each channel:\n')
        for i in range(len(x)) :
            dat = str(x[i])
            for y in ys:
                dat += ',' + str(y[i])
            dat += '\n'
            o.write(dat)
            #o.write(str(x[i]) + ',' + str(amplitude[i]) + ',' + str(phase[i]) + '\n')
        o.close()
    return

#  Execute -------------------------------------------------------------------------------------------------------------
if __name__ == '__main__':
    # channel_reqs is a list of channels requested: [1,2], [1], [2,1], [1,1,2,1,2]
    # One or more are allowed, and they can be the same to compare consecutive waveforms.
    # If a channel is requested more than once then waveforms from that channel are acquired at different times.
    # Truly simultaneous waveform acquisition is under consideration.
    # This will require a GetSimultaneousWaveform function which will use the single sweep option.
    # Graphs can be of the yt (y versus time) or xy (y versus x) type.  A query will appear in the io window.
    # The title can be specified in response to a query in the io window.
    # Data can be saved as a csv file.
    # To see a legend for multiple waveforms on a single yt graph, show_legend=True.
    # A TDS1012B oscilloscope is assumed, and the serial number must be set in WaveformAcquisition.
    symbols = ['', '', '', '', '', '']
    # The next call to WaveformAcquistion is intended for yt mode.
    #legends = [r'$I_b \times 10^5\Omega$',  r'$I_c \times 10^3\Omega$'] # Use the r'$abc$' form for complete functionality.
    #the_title = r'$I_c=V_c/ 10^3\Omega$ and $I_b=V_b/ 10^5\Omega$ for 2N4123 NPN $\beta=110$ Sine Wave Input'     # If this title is desired, set graph_title=the_title in WaveformAcquistion
    #WaveformAcquisition('C055976', channel_reqs=[1,2], xy_pairs=[[1,2]], symbols=symbols, graph_title=the_title, show_legend=True, legends=legends, legend_pos='lower right')
    # The next call to WaveformAcquistion is intended for xy mode.
    #the_title = r'$I_c=V_c/ 10^3\Omega$ vs $I_b=V_b/ 10^5\Omega$ for 2N4123 NPN $\beta=110$ Sine Wave Input'     # If this title is desired, set graph_title=the_title in WaveformAcquistion
    #WaveformAcquisition('C055976', channel_reqs=[1,2], xy_pairs=[[1,2]], symbols=symbols, graph_title=the_title, show_legend=True, legends=legends)
    # The next call to WaveformAcquistion is intended for r mode.
    #the_title = 'Channel Resistance'     # If this title is desired, set graph_title=the_title in WaveformAcquistion
    #WaveformAcquisition('C055976', channel_reqs=[1,2], xy_pairs=[[1,2]], symbols=symbols, graph_title=the_title, show_legend=False)
    # The next call to WaveformAcquistion is intended for osc mode.
    #the_title = r'JFET Oscillator Spectrum'     # If this title is desired, set graph_title=the_title in WaveformAcquistion
    #WaveformAcquisition('C055976', channel_reqs=[2], symbols=symbols, graph_title=the_title, show_legend=True, legends=legends, legend_pos='lower right')
    # The next call to WaveformAcquistion is intended for xy mode.
    """
    legends = [r'Input',  r'Output'] # Use the r'$abc$' form for complete functionality.
    the_title = r'74LS04 Output vs. Input at 3 MHz'     # If this title is desired, set graph_title=the_title in WaveformAcquistion
    WaveformAcquisition('C055976', channel_reqs=[1,2], xy_pairs=[[1,2]], symbols=symbols, graph_title=the_title, show_legend=True, legends=legends)
    """