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

w=40.  # ms step between samples
ts=.2 # sampling time
x = np.arange(0, 400, w)
y=np.cos(x)

xx = np.arange(0, 400, ts)

ax=plt.gca()
ax.xaxis.set_ticks(x)
markerline, stemlines, baseline = ax.stem(x, y, '-')

yf=np.fft.fft(y)


#best trajectory with lowest frequency
y_iftf=np.hstack((yf,np.zeros(len(xx)-len(x))))
y_ift=np.fft.ifft(y_iftf)*w/ts
ax.plot(xx,y_ift,'-b',label='ift')


#plt.figure()
#ax=plt.gca()
#ax.xaxis.set_ticks(x)
#markerline, stemlines, baseline = ax.stem(x, y, '-')

#PVT move
t=np.hstack((y[-1:],y,y[:1]))

n=int(w/ts)
v=(t[2:]-t[:-2])/(w*2)

y_pvt=np.ndarray(len(xx))*0
xx1=xx[:n]
for i in range(len(x)-1):
  d=y[i]
  c=v[i]
  a=( -2*(y[i+1]-y[i]-v[i]*w)+   w*(v[i+1]-v[i]))/w**3
  b=(3*w*(y[i+1]-y[i]-v[i]*w)-w**2*(v[i+1]-v[i]))/w**3
  y_pvt[i*n:(i+1)*n]=a*xx1**3+b*xx1**2+c*xx1+d

ax.plot(xx,y_pvt,'-g',label='pvt')

#PVT move with stop 
v*=0
y_p0t=np.ndarray(len(xx))*0
for i in range(len(x)-1):
  d=y[i]
  c=v[i]
  a=( -2*(y[i+1]-y[i]-v[i]*w)+   w*(v[i+1]-v[i]))/w**3
  b=(3*w*(y[i+1]-y[i]-v[i]*w)-w**2*(v[i+1]-v[i]))/w**3
  y_p0t[i*n:(i+1)*n]=a*xx1**3+b*xx1**2+c*xx1+d

ax.plot(xx,y_p0t,'-r',label='p0t')

ax.legend(loc='best')
plt.show(block=False)


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

y_iftf=np.fft.fft(y_ift)
y_pvtf=np.fft.fft(y_pvt)
y_p0tf=np.fft.fft(y_p0t)


#f=np.arange(0,1E3/(2*ts),1E3/(2*ts*(len(xx)-1)))
f=np.linspace(0,1E3/(2*ts),len(xx))

db_mag=20*np.log10(abs(y_iftf))
ax.semilogx(f,db_mag,'-b',label='ift')  # Bode magnitude plot
db_mag=20*np.log10(abs(y_pvtf))
ax.semilogx(f,db_mag,'-g',label='pvt')  # Bode magnitude plot
db_mag=20*np.log10(abs(y_p0tf))
ax.semilogx(f,db_mag,'-r',label='p0t')  # Bode magnitude plot
ax.yaxis.set_label_text('dB ampl')
ax.xaxis.set_label_text('frequency [Hz]')
plt.grid(True)

ax.legend(loc='best')
plt.show(block=False)