'''This module contains classes that defines specific experiments. 
Each class has methods that reflect common procedures for the represented experiment.'''

from gpyb import *
from peripheral_related_classes import *
from time import sleep, time, localtime,strftime,gmtime
from scipy import gplt
import os
import MySQLdb
import sys

class EiFraExp(CommonExperiment):
    '''Common procedures with Electrochemical Interface and Frequency Response Analyser
    by Solartron Intruments'''
    def __init__(self):
        ''' Initialize GPIB Interface (interface object gpib),
        Electrochemical Interface (peripheral object ei), 
        Frequency Response Analyser (peripheral object fra)'''
        
        self.gpib=InterfaceGPIB('gpib')
        self.ei=PeripheralGPIB('EI')
        self.fra=PeripheralGPIB('FRA')
        print '\nSetting remote mode for gpib'
        self.gpib.sre(1)
        self.gpib.report_ibsta()
        self.ei.resistance_tags={'01':.1,'02':1,'03':10,'04':100,'05':1000,'06':10000}
                
    def set_eicm(self):
        self.eicm=input('''Input a list of commands to be sent to EI. 
        Type "help" for formatting problems.\n''')
        if self.eicm=="help" or self.eicm=='"help"':
            print """You should input a Python List object, containing strings.
Eg: ['command1','command2',...] 
These commands will be executed in sequence as you input them. 
If you don't want to send commands to FRA, input an empty list [] 
To read and display the machine output of the last command, give 'read'.
['command1','command-to-read','read','command3'] will display the EI bus reading after command-to-read.
Eg: ['?ER','read'] will display the last error occured on EI.
"""
            self.set_eicm()

    def set_fracm(self):
        self.fracm=input('''Input a list of commands to be sent to FRA.
       Type "help" for formatting problems.\n''')
        if self.fracm=="help" or self.fracm=='"help"':
            print """You should input a Python List object, containing strings.
Eg: ['command1','command2',...] 
If you don't want to send commands to FRA, input an empty list []
The commands will be executed in sequence as you input them. 
To read and display the machine output of the last command, give 'read'.
['command1','command-to-read','read','command3'] will display the FRA bus reading after command-to-read.
Eg: ['?ER','read'] will display the last error occured on FRA.
"""
            set_fracm()

    def setup(self,set=1,rr=None,pot=None,eicm=None,fracm=None,eOL=None,eIL=None,ePI=None,ePB=None,eVT=None,eFI=None
    ,eVX=None,eIX=None,eDG=None,eRG=None,eTR=None,ePO=None,fAM=None,fBI=None,fWF=None,fMO=None,fIS=None,fMS=None,
    fAU=None,fRA=None,fGD=None,fMA=None,fMI=None,cfg=None):
        '''Use obj.setup(0) to load eifra_cnf.py. Use obj.setup() to configure manually.
Use obj.setup(rr=ref.resistance,pot=baseDCpotential,eicm=['ei','custom','command','sequence'],fracm=['fra','custom','command','sequence']
to set the arguments inline.
In this way tou can also set a tweak (just add Tweak=Value to the arguments). 
Eg: obj.setup(fMA=32000) will set to 32000 the maximum frequence of the FRA sweep.'''
        try:
            cmd='import '+cfg; exec(cmd)
        except:
            cfg='eifra_cnf'
            cmd='import '+cfg; exec(cmd)
        
        self.ei.tweaks={'eOL':eOL,'eIL':eIL,'ePI':ePI,'ePB':ePB,'eVT':eVT,'eFI':eFI,
        'eVX':eVX,'eIX':eIX,'eDG':eDG,'eRG':eRG,'eTR':eTR,'ePO':ePO}
        
        self.fra.tweaks={'fAM':fAM,'fBI':fBI,'fWF':fWF,'fMO':fMO,'fIS':fIS,'fMS':fMS,
        'fAU':fAU,'fRA':fRA,'fGD':fGD,'fMA':fMA,'fMI':fMI}
        folder=None
        if set==0:
            
            cmd='self.rr='+cfg+'.rr'; exec(cmd)
            print 'Loaded EI Reference Resistance (RR): '+str(self.rr)
            cmd='self.pot='+cfg+'.pot'; exec(cmd)
            print 'Loaded EI DC Base Potential (pot): '+str(self.pot)
            cmd='self.eicm='+cfg+'.eicm'; exec(cmd)
            print 'Loaded Extra EI command sequence: '+repr(self.eicm)
            cmd='self.fracm='+cfg+'.fracm'; exec(cmd)
            print 'Loaded Extra FRA command sequence: '+repr(self.fracm)
            
            for tweak in self.ei.tweaks.keys():
                cmd='self.ei.'+tweak+'=str('+cfg+'.'+tweak+')'; exec(cmd)
            print 'Loaded EI tweaks'
            
            for tweak in self.fra.tweaks.keys():
                cmd='self.fra.'+tweak+'=str('+cfg+'.'+tweak+')'; exec(cmd)
            print 'Loaded FRA tweaks'
            
        elif set==1 and None in [rr,pot,eicm,fracm]:
        
            self.rr=input('Set Electrochemical Interface Reference Resistance (RR): ')
            self.pot=input('Set Base DC potential: ')
            self.set_eicm()
            self.set_fracm()
        
            #Setting EI tweaks
            for tweak in self.ei.tweaks.keys():
                if self.ei.tweaks[tweak]!=None:
                    cmd='self.ei.'+tweak+'='+str(self.ei.tweaks[tweak]); exec(cmd)
                    print 'Set EI '+tweak+' to '+str(self.ei.tweaks[tweak])
                else:
                    cmd='self.ei.'+tweak+'=str('+cfg+'.'+tweak+')'; exec(cmd)
            print 'Loaded EI tweaks'
            
            # Setting FRA tweaks
            for tweak in self.fra.tweaks.keys():
                if self.fra.tweaks[tweak]!=None:
                    cmd='self.fra.'+tweak+'='+str(self.fra.tweaks[tweak]); exec(cmd)
                    print 'Set FRA '+tweak+' to '+str(self.fra.tweaks[tweak])
                else:
                    cmd='self.fra.'+tweak+'=str('+cfg+'.'+tweak+')'; exec(cmd)
            print 'Loaded FRA tweaks'
        
        else:
            self.rr=rr
            self.pot=pot
            self.eicm=eicm
            self.fracm=fracm
        
            # Setting EI tweaks
            for tweak in self.ei.tweaks.keys():
                if self.ei.tweaks[tweak]!=None:
                    cmd='self.ei.'+tweak+'='+str(self.ei.tweaks[tweak]); exec(cmd)
                else:
                    cmd='self.ei.'+tweak+'=str('+cfg+'.'+tweak+')'; exec(cmd)
                    print 'Loaded EI '+tweak+' from einfra_cnf.py with value '+str(self.ei.tweaks[tweak])
            
            # Setting FRA tweaksx.loop()
            for tweak in self.fra.tweaks.keys():
                if self.fra.tweaks[tweak]!=None:
                    cmd='self.fra.'+tweak+'='+str(self.fra.tweaks[tweak]); exec(cmd)
                else:
                    cmd='self.fra.'+tweak+'=str('+cfg+'.'+tweak+')'; exec(cmd)
                    print 'Loaded FRA '+tweak+' from einfra_cnf.py with value '+str(self.fra.tweaks[tweak])
       
        # total Data structure:
        self.datot={'frequence':[],'real':[],'complex':[],'error':[]}
            
    def send_ei(self):
        print 'Sending commands to EI'
        for cmd in ['RR'+str(self.rr),'PX3','PY5','OL'+self.ei.eOL,'IL'+self.ei.eIL,'BY0','ON0','PI'+self.ei.ePI,'PB'+self.ei.ePB,'VT'+self.ei.eVT,'FI'+self.ei.eFI
        ,'VX'+self.ei.eVX,'IX'+self.ei.eIX,'DG'+self.ei.eDG,'RG'+self.ei.eRG,'TR'+self.ei.eTR,'PO'+self.ei.ePO,'PV'+str(self.pot),'VR'+str(self.pot)]:
            self.ei.wrt(cmd)
            self.ei.write_read('?ER')
            print 'Sent '+cmd+': GPIB status: '+str(self.ei.ibsta)+' EI error code:'
            self.ei.write_read('?ER')
            sleep(0.2)
    
    def send_custom_ei(self):
        print 'Sending custom commands to EI'
        for cmd in self.eicm:
            if cmd!='read':
                self.ei.wrt(cmd)
                print 'Sent '+cmd+': status: '+str(self.ei.ibsta)
            else: 
                red=self.ei.rd()
                print 'Read command performed on EI bus:\n'+str(red[0])
                print 'Status: '+str(self.ei.ibsta)
            sleep(0.2)

    def send_fra(self):
        print 'Sending commands to FRA'
        self.fra.wrt('AM'+self.fra.fAM)
        for cmd in ['AM'+self.fra.fAM,'BI'+self.fra.fBI,'WF'+self.fra.fWF,'MO'+self.fra.fMO,'IS'+self.fra.fIS,'MS'+self.fra.fMS,'AU'+self.fra.fAU,'RA'+self.fra.fRA,'RA2,0','IP1,1'
        ,'IP2,1','SO0201','CO0','OP2,0','OP3,1','GD'+self.fra.fGD,'MA'+self.fra.fMA,'MI'+self.fra.fMI]:
            self.fra.wrt(cmd)
            print 'Sent '+cmd+': GPIB status: '+str(self.fra.ibsta)+' FRA Error code:'
            self.fra.write_read('?ER')
            sleep(0.2)

    def send_custom_fra(self):
        print 'Sending custom commands to FRA'
        for cmd in self.fracm:
            if cmd!='read':
                self.fra.wrt(cmd)
                print 'Sent '+cmd+': GPIB status: '+str(self.fra.ibsta)
            else:
                red=self.fra.rd()
                print 'Read command performed on FRA bus:\n'+str(red[0])
                print 'Status: '+str(self.fra.ibsta)
            sleep(0.2)
    
    def resrif(self):
        '''Reads Reference Resistance from Electrochemical Interface to set the obj.ei.resistance value'''
        self.ei.wrt('GP2')
        self.ei.resistance=self.ei.resistance_tags[self.ei.write_read('?RR')[0].replace('\r\n','').replace('+','')]
    
    def loop(self,num=1):
        '''This is the measure cycle.'''
        ## Clear file:
        self.fra.wrt('FC')
         # Set reference resistance
        self.write_list({'ei':['PW1','BR1','RU1']})
        sleep(10)
        ti=time() 
        #raw_input('Press return twice to start sweeping')
        self.write_list({'fra':['SC2','RE','RG']})
        rsp=1
        print 'Waiting for service request from FRA. First loop.'
        while rsp not in ['\x06','\x07']:
            sleep(5)
            rsp=self.fra.rsp()[0]
            sys.stdout.write('Waiting for service request from FRA. Time: '+str(int(time()-ti))+' Loop: '+str(num)+'\t\t\t\r')
            
        sleep(2)
        self.write_list({'fra':['SA','SS','SG']})

    def read_buffer(self):
        # Data structure:
        self.data={'frequence':[],'real':[],'complex':[],'error':[]}
        
        self.resrif()
        self.fra.write_read('?NR',buffer_len=20)[0].replace('\r\n','')
        sleep(0.2)
        self.total_readings=self.fra.write_read('?NR',buffer_len=20)[0].replace('\r\n','')
        
        self.fra.clr() #Placebo effect!!!
        self.fra.wrt('OP2,1')

        for zero in self.total_readings:
            if zero=='0':
                self.total_readings=self.total_readings[1:]
            else:
                break
        if self.total_readings in ['0','']:
            return 'Nothing to read from FRA or access timeout'
        
        try:
            self.total_readings=int(self.total_readings)
        except:
            return 'Number of readings not valid'
        print 'Number of readings taken: '+str(self.total_readings)
        
        self.fra.write_read('FD0',buffer_len=38)
        # First reading (FD0)
        datum=self.fra.write_read('FD0',buffer_len=38)[0].replace('\r\n','')
        self.datum=datum
        if len(datum)<35:
                print '/!\\ Wrong length of the FRA data output: '+repr(datum)+'. Skipping first buffer entry.' 
                print "Try to restart GPIB output on FRA: "+repr(self.fra.wrt('OP2,1'))
        else:
            self.data['frequence']+=[datum[:11]]
            cmd="real="+datum[12:23]; exec(cmd)
            cmd="complex="+datum[24:35]; exec(cmd)
            self.data['error']+=[int(datum[36])]
            self.data['real']+=[real*self.ei.resistance]
            self.data['complex']+=[complex*self.ei.resistance]
            
        sleep(0.5)
        # Scroll down file and read (FD2)
        for measure in range(self.total_readings)[1:]:
            sleep(0.5)
            datum=self.fra.write_read('FD2',buffer_len=38)[0].replace('\r\n','')
            self.datum=datum #for debug
            if len(datum)<35:
                print '/!\\ Timeout or wrong length of the FRA data output: '+str(datum)+'. Skipping buffer entry ' +str(measure+1) 
                print "Try to restart GPIB output on FRA: "+repr(self.fra.wrt('OP2,1'))
            else:
                self.data['frequence']+=[datum[:11]]
                cmd="real="+datum[12:23]; exec(cmd)
                cmd="complex="+datum[24:35]; exec(cmd)
                self.data['error']+=[int(datum[36])]
                self.data['real']+=[real*self.ei.resistance]
                self.data['complex']+=[complex*self.ei.resistance]
                self.cmd=cmd # for debug
        self.fra.wrt('OP2,0')
        
        self.datot['error']+=self.data['error']
        self.datot['frequence']+=self.data['frequence']
        self.datot['real']+=self.data['real']
        self.datot['complex']+=self.data['complex']
        print self.datot
        return 0
            
    def ny_plot(self,plot_opt=None,output=None,target='data'):
        if target=='data': 
            x=self.data['real']; y=self.data['complex']
        else: 
            x=self.datot['real']; y=self.datot['complex']
            
        if plot_opt!=None:
            gplt.plot(x,y,plot_opt)
        else:
            gplt.plot(x,y)
        gplt.title('Nyquist plot')
        gplt.xtitle('Real')
        gplt.ytitle('Complex')
        if output!=None:
            gplt.output(output,'png')
        gplt.close()
        
        
class DataStorageManager:
    def __init__(self,cfgfile,obj,target='data'):
        self.cfgfile=cfgfile
        self.obj=obj
        cfgf=cfgfile.replace('.py','')
        cmd='import '+cfgf; exec(cmd)
        for param in ['filename','filepath','image','pic',
        'myhost','mydb','myuser','mypass','mytable','mycomment']:
            try:
                cmd='self.'+param+'='+cfgf+'.'+param; exec(cmd)
            except:
                pass
        cmd='self.data=self.obj.'+target; exec(cmd)
        self.timename=str(time())[4:10]
        
    def store2f(self,wmode='w'):
        filename=self.filename
        filepath=self.filepath
        
        cwd=os.curdir
        if filename!=None and filepath==None:
            file=open(filename+'.txt',wmode)
            if self.pic==1:
                picf=open(filename+'.obj','w')
        elif filename!=None and filepath!=None and os.path.exists(filepath):
            os.chdir(filepath)
            file=open(filename+'.txt',wmode)
            if self.pic==1:
                picf=open(filename+'.obj','w')
            self.obj.ny_plot(output=filename+'.png')
            os.chdir(cwd)
        elif filename==None and filepath!=None and os.path.exists(filepath):
            filename=self.timename
            os.chdir(filepath)
            file=open(filename+'.txt',wmode)
            if self.pic==1:
                picf=open(filename+'.obj','w')
            self.obj.ny_plot(output=filename+'.png')
            os.chdir(cwd)
        else:
            file=open(self.timename+'.txt',wmode)
            if self.pic==1:
                picf=open(self.timename+'.obj','w')
            self.obj.ny_plot(output=self.timename+'.png')
        
        # Store target object
        if self.pic==1:
            pickle.dump(self.obj,picf)
        
        # Store table text
        i=0
        for freq in self.data['frequence']:
##            try:
            real=self.data['real'][i]
            complex=self.data['complex'][i]
            error=self.data['error'][i]
            file.write('%s\t%e\t%e\t%s\n' % (freq, real, complex, error))
            i+=1
##            except:
##                file.write('inconsistent line\n')

        file.close()
        picf.close()
        os.chdir(cwd)

    def store2my(self):
        mydb=self.mydb
        myhost=self.myhost
        myuser=self.myuser
        mypass=self.mypass
        connection=MySQLdb.connect(myhost,myuser,mypass,mydb)
        cursor=connection.cursor()
        filename=self.filename
        mytable=self.mytable
        if filename==None:
            filename=self.timename
        
        txt=''
        i=0
        for freq in self.data['frequence']:
            try:
                real=self.data['real'][i]
                complex=self.data['complex'][i]
                error=self.data['error'][i]
                txt=txt+'%s \t %e \t %e \t  %s \n' % (freq, real, complex, error)
            except:
                txt=txt+'inconsistent line\n'
            i+=1
        
        self.obj.ny_plot(output='ny_plot.temp.png')
        op=open('ny_plot.temp.png','r')
        op.seek(0)
        # Img string
        img=op.read()
        op.close()
        os.remove('ny_plot.temp.png')
        # Obj string
        self.objs=pickle.dumps(self.obj)
        oggi=strftime("%d %b %Y %H:%M:%S", gmtime())
        # Data struct:
        kys={'Nome':filename,'Data':oggi,'.txt':txt,
        '.img':img,'.obj':self.objs}
        # Insert data:
        cmd="""cursor.execute("INSERT INTO """+mytable+""" (Nome, Data, txt, png, obj, Commento) VALUES (%s,%s,%s,%s,%s,%s)", (filename, oggi, txt, img, self.objs,self.mycomment))"""; exec(cmd)
        cursor.close()
        connection.close()