#! /usr/bin/python



__author__ = 'Predrag Pejovic'
__license__ = 'GPLv3+'
__date__ = '23-08-2017'



from pylab import *
import serial
import time
import os



################################################################################
#
# timestamp
#
def timestamp():
    t = time.localtime()
    tx = '{:04d}'.format(t.tm_year)
    tx += '-'
    tx += '{:02d}'.format(t.tm_mon)
    tx += '-'
    tx += '{:02d}'.format(t.tm_mday)
    tx += '_'
    tx += '{:02d}'.format(t.tm_hour)
    tx += '-'
    tx += '{:02d}'.format(t.tm_min)
    tx += '-'
    tx += '{:02d}'.format(t.tm_sec)
    return tx
#
################################################################################



################################################################################
#
# fname
#
def fname(filename, ts):
    fn = filename
    if ts:
        t = timestamp()
        if fn == '':
            fn = t
        else:
            fn += '_' + t
    return fn
#
################################################################################



################################################################################
#
# scope
#
class Oscilloscope(object):
    #
    def __init__(self, address = '/dev/ttyS0', shy = False):
        try:
            self.port = serial.Serial(address, timeout = 5)
            scopeid = self.ask('*idn?')
            if scopeid == '':
                print 'cannot find oscilloscope at', address
            else:
                self.write('header 0')
                self.write('data:encdg ascii')
                self.shy = shy
                if not self.shy:
                    print scopeid[: -1]
        except:
            print 'cannot open port', address
    #
    def __del__(self):
        self.port.close()
    #
    def write(self, string):
        self.port.write(string + '\n')
    #
    def read(self):
        return self.port.readline()
    #
    def readchar(self):
        return self.port.read()
    #
    def ask(self, question):
        self.write(question)
        return self.read()
    #
    def id(self):
        return self.ask('*idn?')
    #
    #
    #
    def getwfm(self, ch, start = 1, stop = 2500):
        #
        self.write('data:start ' + str(start))
        self.write('data:stop ' + str(stop))
        #
        self.write('data:source ch' + str(ch))
        #
        astate = self.getstate()
        self.write('acquire:state 0')
        self.write('curve?')
        s = self.read()
        self.write('acquire:state ' + astate)
        if s == '':
            return float64([])
        else:
            return float64(s[: -1].split(',')) * 0.04
    #
    #
    #
    def getsamples(self, ch, start = 1, stop = 2500):
        #
        x = self.getwfm(ch, start, stop)
        self.write('ch' + str(ch) + '?')
        s = self.read()[: -1].split(';')
        probe = float(s[0])
        scale = float(s[1])
        position = float(s[2])
        return (x - position) * scale * probe
    #
    #
    #
    def savesamples(self, ch, start = 1, stop = 2500, dataformat = 'npy', 
    filename = '', ts = True):
        #
        isiton = int(self.ask('select:ch' + str(ch) + '?')[0])
        if isiton != 1:
            print 'channel', str(ch), 'is off, no action performed'
            return
        #
        if not(dataformat == 'npy' or dataformat == 'txt'):
            print 'unrecognized format', dataformat, 'no action performed'
            return
        #
        x = self.getsamples(ch, start, stop)
        #
        fn = fname(filename, ts)
        #
        if dataformat == 'npy':
            np.save(fn + '.npy', x)
        else:
            np.savetxt(fn + '.txt', x)
    #
    #
    #
    def drawfig(self, filename = '', ts = True, dots = '', fmt = 'pdf'):
        #
        astate = self.getstate()
        self.write('acquire:state 0')
        settings = ''
        #
        oneon = int(self.ask('select:ch1?')[-2 : -1])
        if oneon == 1:
            if not self.shy:
                print 'ch1 . . .'
            s1 = self.ask('ch1?')[: -1].split(';')
            settings += 'CH1:\n'
            settings += 'probe = ' + s1[0] + '\n'
            settings += 'scale = ' + s1[1] + '\n'
            settings += 'position = ' + s1[2] + '\n'
            settings += 'coupling = ' + s1[3] + '\n'
            settings += 'bandwidth = ' + s1[4] + '\n\n'
            ch1 = self.getwfm(1)
        #
        twoon = int(self.ask('select:ch2?')[-2 : -1])
        if twoon == 1:
            if not self.shy:
                print 'ch2 . . .'
            s2 = self.ask('ch2?')[: -1].split(';')
            settings += 'CH2:\n'
            settings += 'probe = ' + s2[0] + '\n'
            settings += 'scale = ' + s2[1] + '\n'
            settings += 'position = ' + s2[2] + '\n'
            settings += 'coupling = ' + s2[3] + '\n'
            settings += 'bandwidth = ' + s2[4] + '\n\n'
            ch2 = self.getwfm(2)
        #
        hrange = linspace(- 5, 5, 2500)
        hs = self.ask('horizontal:main:scale?')
        settings += 'HORIZONTAL:\nscale = ' + hs + '\n'
        #
        self.write('acquire:state ' + astate)
        #
        close('all')
        figure(1, figsize = (5, 4))
        if oneon == 1:
            plot(hrange, ch1, 'y' + dots)
        if twoon == 1:
            plot(hrange, ch2, 'c' + dots)
        xlim(- 5, 5)
        ylim(- 4, 4)
        xticks(range(-5, 6), 11 * '')
        yticks(range(-4, 5), 9 * '')
        grid()
        fn = fname(filename, ts)
        savefig(fn + '.' + fmt, bbox_inches = 'tight')
        close()
        #
        sf = open(fn + '_settings.txt', 'w')
        sf.write(settings)
        sf.close()
    #
    #
    #
    def drawxy(self, filename = '', ts = True, dots = '', fmt = 'pdf'):
        #
        oneon = int(self.ask('select:ch1?')[-2 : -1])
        twoon = int(self.ask('select:ch2?')[-2 : -1])
        if not (oneon and twoon):
            print 'both channels should be displayed, no action performed'
            return
        #
        astate = self.getstate()
        self.write('acquire:state 0')
        settings = ''
        #
        if not self.shy:
            print 'ch1 . . .'
        s1 = self.ask('ch1?')[: -1].split(';')
        settings += 'CH1:\n'
        settings += 'probe = ' + s1[0] + '\n'
        settings += 'scale = ' + s1[1] + '\n'
        settings += 'position = ' + s1[2] + '\n'
        settings += 'coupling = ' + s1[3] + '\n'
        settings += 'bandwidth = ' + s1[4] + '\n\n'
        ch1 = self.getwfm(1)
        #
        if not self.shy:
            print 'ch2 . . .'
        s2 = self.ask('ch2?')[: -1].split(';')
        settings += 'CH2:\n'
        settings += 'probe = ' + s2[0] + '\n'
        settings += 'scale = ' + s2[1] + '\n'
        settings += 'position = ' + s2[2] + '\n'
        settings += 'coupling = ' + s2[3] + '\n'
        settings += 'bandwidth = ' + s2[4] + '\n\n'
        ch2 = self.getwfm(2)
        #
        hs = self.ask('horizontal:main:scale?')
        settings += 'HORIZONTAL:\nscale = ' + hs + '\n'
        #
        self.write('acquire:state ' + astate)
        #
        close('all')
        figure(1, figsize = (5, 4))
        plot(ch1, ch2)
        xlim(- 5, 5)
        ylim(- 4, 4)
        xticks(range(-5, 6), 11 * '')
        yticks(range(-4, 5), 9 * '')
        grid()
        fn = fname(filename, ts)
        savefig(fn + '.' + fmt, bbox_inches = 'tight')
        close()
        #
        sf = open(fn + '_settings.txt', 'w')
        sf.write(settings)
        sf.close()
    #
    #
    #
    def getbmp(self, filename = '', ts = True):
        #
        astate = self.getstate()
        self.write('acquire:state 0')
        self.write('hardcopy:port rs232')
        self.write('hardcopy:format bmp')
        self.write('hardcopy start')
        #
        counter = 0
        f = ''
        while True:
            c = self.port.read()
            if c == '':
                break
            f += c
            counter += 1
            if counter % 1024 == 0:
                print '{:2d} kB'.format(counter / 1024)
        self.write('acquire:state ' + astate)
        #
        fn = fname(filename, ts)
        ff = open(fn + '.bmp', 'w')
        ff.write(f)
        ff.close()
    #
    #
    #
    def getepsstring(self):
        #
        astate = self.getstate()
        self.write('acquire:state 0')
        self.write('hardcopy:port rs232')
        self.write('hardcopy:format epsimage')
        self.write('hardcopy start')
        #
        counter = 0
        f = ''
        while True:
            c = self.readchar()
            if c == '':
                break
            f += c
            counter += 1
            if not self.shy:
                if counter % 1024 == 0:
                    print '{:2d} kB'.format(counter / 1024)
        self.write('acquire:state ' + astate)
        return f
    #
    #
    #
    def geteps(self, filename = '', ts = True):
        #
        f = self.getepsstring()
        #
        fn = fname(filename, ts)
        ff = open(fn + '.eps', 'w')
        ff.write(f)
        ff.close()
    #
    #
    #
    def getpdf(self, filename = '', ts = True):
        #
        f = self.getepsstring()
        #
        fn = fname(filename, ts)
        ff = open(fn + '.eps', 'w')
        ff.write(f)
        ff.close()
        #
        os.system('epstopdf ' + fn + '.eps')
        os.remove(fn + '.eps')
        os.system('pdfcrop ' + fn + '.pdf > /dev/null')
        os.remove(fn + '.pdf')
        os.system('mv ' + fn + '-crop.pdf ' + fn + '.pdf')
        #
        return fn
    #
    #
    #        
    def getpng(self, filename = '', ts = True, density = 600):
        #
        fn = self.getpdf(filename, ts)
        #
        if not self.shy:
            print 'Converting . . . '
        #
        os.system('convert -density ' + str(density)  + ' ' + fn +
        '.pdf ' + fn + '.png')
        os.remove(fn + '.pdf')
    #
    #
    #        
    def getpdfpng(self, filename = '', ts = True, density = 600):
        #
        fn = self.getpdf(filename, ts)
        #
        if not self.shy:
            print 'Converting . . . '
        #
        os.system('convert -density ' + str(density)  + ' ' + fn +
        '.pdf ' + fn + '.png')
    #
    #
    #   
    def getstate(self):
        return self.ask('acquire:state?')[0]
    #
    #
    #
    def run(self):
        self.write('acquire:state 1')
    #
    #
    #
    def stop(self):
        self.write('acquire:state 0')
    #
    #
    #
    def getvalue(self):
        return float(self.ask('measurement:immed:value?'))
    #
    #
    #
################################################################################

