#!/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

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 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))
      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]
        data=(data-data[0])*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_gather(self,motor=1,minFrq=1,maxFrq= 100,numFrq= 100,amp=10,file='/tmp/bode.npz'):
      #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):
        frq=frqLst[i]
        t=1
        rep=max(1,frq*t)
        self.fnLoc='/tmp/gather%d.txt'%i
        data=self.do_command('openloopsine',motor,amp,frq,rep,0)
        ax.clear()
        ax.plot(data[:, 0]-data[0, 0] , 'b-', label='input')
        ax.plot(data[:, 1]-data[0, 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()

      fig = plt.figure()
      ax = fig.add_subplot(2, 1, 1)
      ax.semilogx(frqLst, bode[:,0])  # Bode magnitude plot
      #ax.loglog(frqLst, bode[:,0],'.-')  # Bode magnitude plot
      ax = fig.add_subplot(2, 1, 2)
      ax.semilogx(frqLst, bode[:,1])  # Bode phase plot
      plt.show()
      meta={'motor':motor,'date':time.asctime()}
      np.savez_compressed(file, bode=bode, meta=meta)

    def bode_plot(self,file):
      f=np.load(file)
      bode=f['bode']
      meta=f('meta')
      frq=bode[:,0]
      db_mag=20*np.log10(bode[:,1])
      phase=np.unwrap(bode[:,2])# numpy.unwrap(p, discont=3.141592653589793, axis=-1)
      fig = plt.figure()
      ax = fig.add_subplot(2, 1, 1)
      ax.semilogx(frq, db_mag,'.-')  # Bode magnitude plot
      ax.xaxis.set_label_text('dB Mag.')
      #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.show()


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('--host', help='hostname', metavar='HOST', required=True)
    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')
    args=parser.parse_args()

    plt.ion()
    tune=PBTuning(args)
    base='/home/zamofing_t/Documents/prj/SwissFEL/epics_ioc_modules/ESB_MX/python/data/'
    file=sys.path.join(base,'bode1.npz')
    tune.bode(file=file,motor=1)
    tune.bode_plot(file)
    #tune.do_command('stepmove',1,100,500,0,0)
  #------------------ Main Code ----------------------------------
  #ssh_test()
  ret=parse_args()
  exit(ret)