#!/usr/bin/env python
# *-----------------------------------------------------------------------*
# |                                                                       |
# |  Copyright (c) 2019 by Paul Scherrer Institute (http://www.psi.ch)    |
# |                                                                       |
# |              Author Thierry Zamofing (thierry.zamofing@psi.ch)        |
# *-----------------------------------------------------------------------*
'''
Trajectory comparison:
pvt: position velocity time
p0t: position velocity=0 time
ift: inverse fourier transformation

-> look at trajectory and frequency components
'''
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import os

np.set_printoptions(precision=3, suppress=True)

def derivate_fft_test():
  n=32.
  frq=1.
  t=np.arange(n)/n*2*np.pi
  p=np.sin(t*frq)  # position array of trajectory

  pf=np.fft.fft(p)
  print (pf)
  f=np.fft.fftfreq(pf.shape[0], d=1/n)

  pfd=pf*f*1j #differentiate in fourier
  print (pfd)
  pd=np.fft.ifft(pfd)
  print (pd)

  ax=plt.figure().add_subplot(1, 1, 1)
  ax.plot(t, p, '.-b', label='p')
  ax.plot(t, pd, '.-r', label='pd')
  ax.grid(True)
  plt.show(block=False)

  pass

def gen_pvt(p,v,t,ts):
  '''generates a pvt motion
     p: position array
     v: velocity array
     t: time array
     ts: servo cycle time
     !!! it is assumed, that the time intervals are constant !!!
     '''
  pvt=np.ndarray(int(t[-1]/ts))*0


  tt1=np.arange(0,t[1]-t[0],ts)
  for i in range(len(t)-1):
    d=p[i]
    c=v[i]
    a=(-2*(p[i+1]-p[i]-v[i]*w)+w*(v[i+1]-v[i]))/w**3
    b=(3*w*(p[i+1]-p[i]-v[i]*w)-w**2*(v[i+1]-v[i]))/w**3
    pvt[i*n:(i+1)*n]=a*tt1**3+b*tt1**2+c*tt1+d

  return pvt


(a,b)=os.path.split(__file__)
fnBase=os.path.abspath(os.path.join(a,'../MXfastStageDoc'))

#derivate_fft_test()
w=40.  # ms step between samples
ts=.2 # sampling time
n=int(w/ts)# servo cycle between samples
k=320  #number of unique samples

t = np.arange(0, w*(k+1), w)  #time array of trajectory

#p=3.*np.cos(t)+4.             #position array of trajectory
#p=3.*np.sin(1.3+2.*t/(w*k)*2.*np.pi)+10.
#p+=np.cos(1.5*t/(w*k)*2.*np.pi)
#p=np.cos(8*t*np.pi*2./(k*w))  #eine schwingung
np.random.seed(10)
p=np.random.random(k+1)*4.  #position array of trajectory


p[-1]=p[0]                # put the first position at the end
#p=np.array([1,-1]*16+[1,]);
#p+=3

tt = np.arange(t[0],t[-1], ts) #time array of servo cycles
fig=plt.figure()
ax=fig.add_subplot(1, 1, 1)
ax.xaxis.set_ticks(t)
markerline, stemlines, baseline = ax.stem(t, p, '-')


#### best trajectory with lowest frequency ###
p_iftf=np.fft.fft(p[:-1])
ft=np.hstack((p_iftf[:k/2],np.zeros((n-1)*k),p_iftf[k/2:]))
pp_ift=np.fft.ifft(ft)*n
ax.plot(tt,pp_ift,'-b',label='ift')

### PVT move ###
p2=np.hstack((p[-2],p,p[1]))
v=(p2[2:]-p2[:-2])/(w*2)

pp_pvt=gen_pvt(p,v,t,ts)
ax.plot(tt,pp_pvt,'-g',label='pvt')

### PVT move with stop ###
v*=0
pp_p0t=gen_pvt(p,v,t,ts)
ax.plot(tt,pp_p0t,'-r',label='p0t')
ax.set_xlim(0,400)


### PVT with ift velocity move -> PFT ###
f=np.fft.fftfreq(k, d=1./k)
p_pftf=np.fft.fft(p[:-1])
p_pftfd=p_pftf*f*1j  # differentiate in fourier
#print (p_pftfd)
p_pftd=np.fft.ifft(p_pftfd)
#print (p_pftd)

p_pftd=np.hstack((p_pftd,p_pftd[0]))
#ax2=plt.figure().add_subplot(1,1,1)
#ax2.plot(t,p_pftd,'-b',label='dift')
#ax2.grid(True)
v=p_pftd.real/(k*2*np.pi)
pp_pft=gen_pvt(p,v,t,ts)
ax.plot(tt,pp_pft,'-c',label='pft')

ax.xaxis.set_label_text('time(ms)')
ax.yaxis.set_label_text('position')

ax.legend(loc='best')
plt.show(block=False)
fig.savefig(os.path.join(fnBase,'traj1.eps'))


### frequency plots ###

fig=plt.figure()
ax=fig.add_subplot(1,1,1)#ax=plt.gca()

#normalize with l -> value of k means amplitude of k at a given frequency
pp_iftf=np.fft.rfft(pp_ift)/(2*n)
pp_pvtf=np.fft.rfft(pp_pvt)/(2*n)
pp_p0tf=np.fft.rfft(pp_p0t)/(2*n)
pp_pftf=np.fft.rfft(pp_pft)/(2*n)

f=np.fft.rfftfreq(pp_ift.shape[0], d=ts*1E-3)
f=f[1:] #remove dc value frequency

def FreqSpecDb(spec,w=20.):
  #spec = abs(spec)
  spec = np.convolve(spec, np.ones(w) / w, 'same');
  spec = 20 * np.log10(spec)
  return spec

def FreqSpec(spec,w=20.):
  #spec = abs(spec)
  spec = np.convolve(spec, np.ones(w) / w, 'same');
  return spec

mag=FreqSpecDb(abs(pp_iftf[1:]))
ax.semilogx(f,mag,'-b',label='ift')  # Bode magnitude plot
mag=FreqSpecDb(abs(pp_pvtf[1:]))
ax.semilogx(f,mag,'-g',label='pvt')  # Bode magnitude plot
mag=FreqSpecDb(abs(pp_p0tf[1:]))
ax.semilogx(f,mag,'-r',label='p0t')  # Bode magnitude plot
mag=FreqSpecDb(abs(pp_pftf[1:]))
ax.semilogx(f,mag,'-c',label='pft')  # Bode magnitude plot
ax.yaxis.set_label_text('dB ampl')
#ax.yaxis.set_label_text('ampl')
ax.xaxis.set_label_text('frequency [Hz]')
plt.grid(True)

ax.legend(loc='best')
ax.set_xlim(1,100)
plt.show(block=False)
fig.savefig(os.path.join(fnBase,'traj2.eps'))

fig=plt.figure()
ax=fig.add_subplot(1,1,1)#ax=plt.gca()
magift=abs(pp_iftf[1:])
mag=FreqSpec(abs(pp_pvtf[1:])-magift)
ax.semilogx(f,mag,'-g',label='pvt-ift')  # Bode magnitude plot
mag=FreqSpec(abs(pp_p0tf[1:])-magift)
ax.semilogx(f,mag,'-r',label='p0t-ift')  # Bode magnitude plot
mag=FreqSpec(abs(pp_pftf[1:])-magift)
ax.semilogx(f,mag,'-c',label='pft-ift')  # Bode magnitude plot

#ax.yaxis.set_label_text('dB ampl')
ax.yaxis.set_label_text('ampl')
ax.xaxis.set_label_text('frequency [Hz]')
ax.set_xlim(1,100)
plt.grid(True)
ax.legend(loc='best')

plt.show(block=False)
fig.savefig(os.path.join(fnBase,'traj3.eps'))
plt.show()