#!/usr/bin/env python
# *-----------------------------------------------------------------------*
# |                                                                       |
# |  Copyright (c) 2016 by Paul Scherrer Institute (http://www.psi.ch)    |
# |                                                                       |
# |              Author Thierry Zamofing (thierry.zamofing@psi.ch)        |
# *-----------------------------------------------------------------------*
'''
tuning functions for deltatau:
For now it does 10 step move and uploads the data
other available tuning progs on the powerbrick are:
  analyzerautotunemove
  autotunecalc
  autotunemove
  chirpmove
  currentautotunecalc
  currentstep
  filtercalculation
  openloopchirp +
  openloopsine +
  openlooptestmove
  othertrajectory
  parabolicmove
  randommove
  sinesweep +
  sinusoidal +
  stepmove
  user_gantry_crosscoupled.h
  usertrajectory

Modes:
  1: sine bode open loop
  2: chirp bode open loop
  3: sine bode closed loop
  4: chirp bode closed loop

TODO:
use openloopsine to create a bode diagram of the 'strecke'
'''

import os, sys, json, time
import numpy as np
import matplotlib as mpl
#mpl.use('GTKAgg')
import matplotlib.pyplot as plt
import subprocess as sprc
import telnetlib
from scipy.signal.waveforms import chirp
from utilities import *

class PBTuning:
    tuneDir='/opt/ppmac/tune/'

    def __init__(self,args):
      self.args=args

    def do_command(self,cmd,*param):
      args=self.args
      host=args.host
      cmd=('ssh','root@'+host, self.tuneDir+cmd)+tuple(map(str,param))
      print(' '.join(cmd))
      p=sprc.Popen(cmd, shell=False, stdin=sprc.PIPE, stdout=sprc.PIPE, stderr=sprc.PIPE)
      res=p.wait()
      print(res)
      print(p.stderr.read())
      print(p.stdout.read())
      PBGatherPlot='/home/zamofing_t/scripts/PBGatherPlot'
      try:
        fnLoc=self.fnLoc
      except AttributeError:
        fnLoc = '/tmp/gather.txt'
      cmd=(PBGatherPlot,'-m24','-v0','--host',host,'--dat',fnLoc)
      p = sprc.Popen(cmd, shell=False, stdin=sprc.PIPE, stdout=sprc.PIPE, stderr=sprc.PIPE)
      retval = p.wait()
      print(p.stderr.read())
      #print(p.stdout.read())

      #self.meta = {'timebase': ServoPeriod*gather['Period']}
      self.data = np.genfromtxt(fnLoc, delimiter=' ')
      return self.data

    def phase_amp(self,frq,rep):
      try:
        ax1=self.ax1
        ax2=self.ax2
        ax1.clear()
        ax2.clear()
      except AttributeError:
        fig = plt.figure(); self.ax1=ax1 = fig.add_subplot(1, 1, 1)
        fig = plt.figure(); self.ax2=ax2 = fig.add_subplot(1, 1, 1)
      n = self.data.shape[0]
      w=np.hamming(n)
      res=np.ndarray((2,2))
      col=('b-','g-')
      for i in (0,1):
        data=self.data[:,i]
        avg=data.mean(0)
        #data=data-data[0]
        data=data-avg
        data*=w
        ax1.plot(data, col[i])
        f = np.fft.fft(data)
        ax2.plot(np.absolute(f) , col[i])
        idx = int(rep)
        bias = np.absolute(f[0] / n)
        phase = np.angle(f[idx])
        #ampl = np.absolute(f[idx]) * 2 / n
        ampl = np.absolute(f[idx]) * 2 / w.sum()
        res[i,:]=(ampl,phase)
      plt.pause(.05)
      return (res[1,0]/res[0,0],res[1,1]-res[0,1])

    def bode_sine(self,openloop=True,motor=1,minFrq=1,maxFrq=300,numFrq=150,amp=10,file='/tmp/bode.npz'):
      if os.path.isfile(file):
        f=np.load(file)
        bode=f['bode']
        meta=f['meta'].item()
        meta['file']=file
      else:
        #motor 1 maximum: 13750
        #amp= percentage of maximum amplitude
        fig = plt.figure()
        ax = fig.add_subplot(1, 1, 1)
        frqLst=np.logspace(np.log10(minFrq),np.log10(maxFrq),numFrq)
        n=len(frqLst)
        #frqLst=(10,15,20,25,30)
        bode=np.ndarray((n,3))
        bode[:, 0]=frqLst
        #for i in range(n):
        for i in range(n-1,-1,-1):
          frq=frqLst[i]
          t=1
          rep=max(1,frq*t)
          self.fnLoc='/tmp/gather%d.txt'%i
          if openloop:
            data=self.do_command('openloopsine',motor,amp,frq,rep,0)
          else:
            data=self.do_command('sinusoidal',motor,amp,frq,rep,0)
            data=data[:,(1,2)]
          ax.clear()
          avg=data.mean(0)
          print(avg)
          ax.plot(data[:, 0]-avg[0] , 'b-', label='input')
          ax.plot(data[:, 1]-avg[1], 'g-', label='output')
          plt.pause(.05)
          bode[i,1:]=self.phase_amp(frq, rep)
          print('frq %g ampl %g phase %g'%tuple(bode[i,:]))
          plt.draw_all()

        meta={'motor':motor,'date':time.asctime()}
        np.savez_compressed(file, bode=bode, meta=meta)
        meta['file']=file
      self.bode_sine_plot(bode, meta)

    def bode_sine_plot(self,bode,meta):
      frq=bode[:,0]
      db_mag=20*np.log10(bode[:,1])
      phase=np.degrees(np.unwrap(bode[:,2]))# numpy.unwrap(p, discont=3.141592653589793, axis=-1)
      if phase[0]>0: phase-=360
      fig = plt.figure()
      fig.canvas.set_window_title(os.path.basename(meta['file']))
      ax = fig.add_subplot(2, 1, 1)
      ax.semilogx(frq, db_mag,'.-')  # Bode magnitude plot
      ax.xaxis.set_label_text('dB Mag.')
      plt.grid(True)
      #ax.loglog(frqLst, bode[:,0],'.-')  # Bode magnitude plot
      ax = fig.add_subplot(2, 1, 2)
      ax.semilogx(frq, phase,'.-')  # Bode phase plot
      ax.xaxis.set_label_text('phase')
      plt.grid(True)

    def bode_chirp(self,openloop=True,motor=1,minFrq=10,maxFrq=150,tSec=30.,amp=10,mode=1,file='/tmp/gather.npz'):
      #amp= percentage of maximum amplitud
      if os.path.isfile(file):
        f=np.load(file)
        data=f['data']
        #try:
        meta=f['meta'].item()
        #except KeyError:
        #  meta={'motor': motor, 'date': time.asctime(), 'minFrq': 10, 'maxFrq': 150, 'tSec': 30, 'amp': amp}
        #  np.savez_compressed(file, data=data, meta=meta)
        meta['file']=file
      else:
        mode=1
        #Sys.ServoPeriod=0.2ms = 5kHz -> therefore factor 2000 for the time second->servoCnt
        #don't know why the frequency must be scaled...
        if openloop:
          data=self.do_command('openloopchirp',motor,amp,minFrq/2.,maxFrq/2.,tSec*2000,mode,0)
        else:
          data=self.do_command('sinesweep', motor, amp, minFrq/2., maxFrq/2., tSec*2000, mode, 0)
          data = data[:, (1, 2)]

        meta={'motor':motor,'date':time.asctime(),'minFrq':minFrq,'maxFrq':maxFrq,'tSec':tSec,'amp':amp}
        np.savez_compressed(file, data=data, meta=meta)
        meta['file'] = file
      self.bode_chirp_plot(data, meta,openloop)

    def bode_chirp_plot(self, data, meta,openloop):
      tSec=meta['tSec']
      minFrq=meta['minFrq']
      maxFrq=meta['maxFrq']

      #n=25000
      #t=np.linspace(0, tSec, 2*n+1, endpoint=True)
      #c=chirp(t, f0=minFrq, f1=maxFrq, t1=tSec, phi=-90, method='logarithmic')

      #plt.figure()
      #plt.clf()
      #plt.subplot(2, 1, 1)
      #plt.plot(t, c)
      #tstr="Logarithmic Chirp, f(0)=%g, f(%g)=%g"%(minFrq, tSec, maxFrq)
      #plt.title(tstr)

      #plt.subplot(2, 1, 2)
      #plt.plot(t, minFrq*(maxFrq/minFrq)**(t/tSec), 'r')
      ## yscale('log')
      #plt.grid(True)
      #plt.ylabel('Frequency (Hz)')
      #plt.xlabel('time (sec)')

      if openloop:
        n=1000
        d=np.concatenate((np.ones(n-1)*data[0, 1],data[:, 1]))
        d=np.convolve(d,np.ones(n),'valid')/n
        data[:, 1]-=d
        # 5..150 Hz
        # 15 sec
        c = data[:, 0]
        o = data[:, 1]
        o -= o[0]
      else:
        data=data-data[0,:]
        c = data[:,0]
        o=data[:,1]
      t = np.linspace(0, tSec, data.shape[0], endpoint=True)
      n=data.shape[0]/2
      fig=plt.figure()
      fig.canvas.set_window_title(os.path.basename(meta['file']+' raw'))
      plt.plot(t, c)
      plt.plot(t, o)
      #bode_plot(o)
      #plt.show()

      #n samples
      #t1 seconds
      #ts=t1/(2*n)#samplingperiode =t[1]-t[0]
      #w_k=2*pi*k/T
      #w=2*np.pi*np.arange(n+1)/t1*(2*n)
      f=np.arange(n+1)/tSec #Hz
      #w*=2.*np.pi #rad/sec
      fig=plt.figure()
      fig.canvas.set_window_title(os.path.basename(meta['file']+' bode'))
      ftX=np.fft.rfft(c)
      ftY=np.fft.rfft(o)

      i=int(minFrq*tSec); j=int(maxFrq*tSec); #print(w[i],w[j])
      f=f[i:j+1]
      ftX=ftX[i:j+1]
      ftY=ftY[i:j+1]

      ft=ftY/ftX
      phase=np.angle(ft)
      phase=np.degrees(np.unwrap(phase))

      #magDb=10*np.log10(np.abs(ft)) #in decibel
      mag=np.abs(ftY)/np.abs(ftX)
      magDb=20*np.log10(mag) #in decibel (20=10*2: factor 2 because rfft only half)
      #magDb=np.abs(ftY)/np.abs(ftX)
      ax=plt.subplot(2,1,1)
      ax.semilogx(f,magDb,'b')    # Bode magnitude plot
      #ax.plot(w,magDb)    # Bode magnitude plot
      ax.axvline(minFrq,c='k');ax.axvline(maxFrq,c='k')
      ax.grid(True)
      ax=plt.subplot(2,1,2)

      ax.semilogx(f,phase,'b')  # Bode phase plot
      ax.yaxis.set_label_text('Amplitude [dB]')

      #ax.set_ylim(phase[i],phase[j])
      ax.grid(True)
      ax.xaxis.set_label_text('Frequency [Hz]')
      ax.yaxis.set_label_text('Phase [degree]')
      #plt.axvline(x=0.22058956)
      #plt.show()

      #meta={'motor':motor,'date':time.asctime()}
      #np.savez_compressed(file, bode=bode, meta=meta)

    def bode_full_plot(self, mot,base):
      fn='sine_ol_mot%d_ext.npz'%mot
      file=os.path.join(base,fn)
      f=np.load(file)
      bode=f['bode']
      meta=f['meta'].item()
      frq=bode[:,0]
      mag=bode[:,1]
      phase=bode[:,2]

      for fn in ('chirp_ol_mot%d_exta.npz','chirp_ol_mot%d_extb.npz','chirp_ol_mot%d_extc.npz'):
        fn=fn%mot
        file=os.path.join(base,fn)
        f=np.load(file)
        data=f['data']
        meta=f['meta'].item()
        tSec=meta['tSec']
        minFrq=meta['minFrq']
        maxFrq=meta['maxFrq']
        amp=meta['amp']
        n=1000
        d=np.concatenate((np.ones(n-1)*data[0, 1],data[:, 1]))
        d=np.convolve(d,np.ones(n),'valid')/n
        data[:, 1]-=d
        # 5..150 Hz
        # 15 sec
        c = data[:, 0]
        o = data[:, 1]
        o -= o[0]
        n=data.shape[0]/2
        ftX=np.fft.rfft(c)
        ftY=np.fft.rfft(o)
        i=int(minFrq*tSec); j=int(maxFrq*tSec); #print(w[i],w[j])
        f=np.arange(n+1)/tSec #Hz
        f=f[i:j+1]
        ftX=ftX[i:j+1]
        ftY=ftY[i:j+1]

        ft=ftY/ftX
        frq=np.concatenate((frq,f))
        phase=np.concatenate((phase,np.angle(ft)))
        #if amp==50:amp*=np.sqrt(2)
        mag=np.concatenate((mag,(np.abs(ftY)/np.abs(ftX))))

      db_mag=20*np.log10(mag)
      phase=np.degrees(np.unwrap(phase))# numpy.unwrap(p, discont=3.141592653589793, axis=-1)
      fig = plt.figure()
      fig.canvas.set_window_title('full bode of motor %d'%mot)
      ax = fig.add_subplot(2, 1, 1)
      ax.semilogx(frq, db_mag,'-')  # Bode magnitude plot
      ax.xaxis.set_label_text('dB Mag.')
      plt.grid(True)
      #ax.loglog(frqLst, bode[:,0],'.-')  # Bode magnitude plot
      ax = fig.add_subplot(2, 1, 2)
      ax.semilogx(frq, phase,'-')  # Bode phase plot
      ax.xaxis.set_label_text('phase')
      ax.set_ylim(-360,0)
      plt.grid(True)


if __name__=='__main__':
  from argparse import ArgumentParser,RawDescriptionHelpFormatter
  def parse_args():
    'main command line interpreter function'
    #usage: gpasciiCommunicator.py --host=PPMACZT84 myPowerBRICK.cfg
    (h, t)=os.path.split(sys.argv[0]);cmd='\n  '+(t if len(h)>3 else sys.argv[0])+' '
    exampleCmd=('-n',
                '-v15'
               )
    epilog=__doc__+'''
Examples:'''+''.join(map(lambda s:cmd+s, exampleCmd))+'\n '

    parser=ArgumentParser(epilog=epilog,formatter_class=RawDescriptionHelpFormatter)

    parser.add_argument('--plot', nargs='*')
    parser.add_argument('--host', help='hostname', metavar='HOST')
    parser.add_argument('--verbose', '-v', type=int, help='verbosity bits (see below)', default=0)
    parser.add_argument('--dryrun', '-n', action='store_true', help='dryrun to stdout')
    parser.add_argument('--mode', '-m', type=int, help='modes (see below)', default=1)
    parser.add_argument('--mot', type=int, help='motor', default=1)
    parser.add_argument('--tag', help='tag', default='')

    args=parser.parse_args()

    #plt.ion()
    tune=PBTuning(args)
    #data=self.do_command('openloopsine',motor,amp,frq,rep,0)
    # #data=self.do_command('openloopchirp',motor,amp,minFrq,maxFrq,tSec*1000,mode,0)

    #THIS IS A TEST TO DO 5 Hz during 10 seconds
    #frq=5
    #frq2=frq*1.1
    #ts=10 #seconds

    #5 kHz Servo
    #ds=tune.do_command('openloopsine', 2, 20, frq, frq*ts, 0)
    #dc=tune.do_command('openloopchirp', 2, 20, frq/2., frq2/2., ts*2000, 1, 0) #THIS IS NOT 10 Seconds!!!

    #dc=tune.do_command('openloopchirp', 2, 10, 5, 6, 10*1000, 1, 0) #THIS IS NOT 10 Seconds!!!
    #ds=ds-ds[0,:]
    #dc=dc-dc[0,:]
    #plt.figure()
    #plt.plot(ds[:,0],'b')
    #plt.plot(ds[:,1],'g')
    #plt.figure()
    #plt.plot(dc[:,0],'b')
    #plt.plot(dc[:,1],'g')
    #plt.show()
    #return
    base='/home/zamofing_t/Documents/prj/SwissFEL/epics_ioc_modules/ESB_MX/python/PBTuning/'
    #file=os.path.join(base, 'sine_ol_mot_tst.npz')
    #tune.bode_sine(openloop=True,file=file,motor=2, minFrq=20,maxFrq=200,numFrq=30)
    #file=os.path.join(base, 'chirp_ol_mot_tst.npz')
    #tune.bode_chirp(openloop=True, file=file, motor=2, minFrq=20,maxFrq=200,tSec=30.)
    #plt.show()
    #return

    if args.plot:
      #display bode plots
      if args.plot[0]=='ext':
        tune.bode_full_plot(mot=1,base=base)
        tune.bode_full_plot(mot=2,base=base)
        plt.show()
        return
      for fn in args.plot:
        if os.path.basename(fn).startswith('sine_ol_mot'):
          tune.bode_sine(openloop=True,file=fn)
        if os.path.basename(fn).startswith('chirp_ol_mot'):
          tune.bode_chirp(openloop=True,file=fn)
        if os.path.basename(fn).startswith('sine_cl_mot'):
          tune.bode_sine(openloop=False, file=fn)
        if os.path.basename(fn).startswith('chirp_cl_mot'):
          tune.bode_chirp(openloop=False, file=fn)
      plt.show()
      return

    mode=args.mode
    tag=args.tag
    mot=args.mot
    if mode==1:
      ol=True
      file=os.path.join(base,'sine_ol_mot%d_%s.npz'%(mot,tag))
      #def bode_sine(self, openloop=True, motor=1, minFrq=1, maxFrq=300, numFrq=150, amp=10, file='/tmp/bode.npz'):
      tune.bode_sine(openloop=True,file=file,motor=mot)
    elif mode==2:
      file=os.path.join(base,'chirp_ol_mot%d_%s.npz'%(mot,tag))
      #def bode_chirp(self,openloop=True,motor=1,minFrq=10,maxFrq=150,tSec=30.,amp=10,mode=1,file='/tmp/gather.npz'):
      #tune.bode_chirp(openloop=True,file=file,motor=mot, minFrq=10, maxFrq=300, tSec=30)
      tune.bode_chirp(openloop=True,file=file[:-4]+'a.npz',motor=mot,       minFrq=10, maxFrq=300, tSec=30)
      tune.bode_chirp(openloop=True,file=file[:-4]+'b.npz',motor=mot,amp=50,minFrq=100,maxFrq=500, tSec=30)
      tune.bode_chirp(openloop=True,file=file[:-4]+'c.npz',motor=mot,amp=50,minFrq=300,maxFrq=1500,tSec=10)
    elif mode==3:
      ol=True
      file=os.path.join(base,'sine_cl_mot%d_%s.npz'%(mot,tag))
      tune.bode_sine(openloop=False,file=file,motor=mot)
    elif mode==4:
      file=os.path.join(base,'chirp_cl_mot%d_%s.npz'%(mot,tag))
      tune.bode_chirp(openloop=False,file=file,motor=mot)
    #tune.do_command('stepmove',1,100,500,0,0)
    plt.show()
  #------------------ Main Code ----------------------------------
  #ssh_test()
  ret=parse_args()
  exit(ret)

#enable plc1
#./PBTuning.py --host SAR-CPPM-EXPMX1 --mode 1 --mot 1 --tag ext
#enable plc1
#./PBTuning.py --host SAR-CPPM-EXPMX1 --mode 2 --mot 1 --tag ext
#enable plc1
#./PBTuning.py --host SAR-CPPM-EXPMX1 --mode 1 --mot 2 --tag ext
#enable plc1
#./PBTuning.py --host SAR-CPPM-EXPMX1 --mode 2 --mot 2 --tag ext
#enable plc1
#./PBTuning.py --host SAR-CPPM-EXPMX1 --plot ext