import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.signal import chirp

#------------------------------------------------
# PARAMETERS
FT = 200  # Sampling Frequency [Hz]
# Frequency of signals
f_x = 0.2  # [Hz]
f_y = 0.7  # [Hz]
f_z = 2    # [Hz]
# Ramping time and signal durations
t_tot = 120
t_start_x = 4 / f_x
t_end_x = 4 / f_x
t_cost_x = t_tot - (t_start_x + t_end_x)
t_start_y = 4 / f_y
t_end_y = 4 / f_y
t_cost_y = t_tot - (t_start_y + t_end_y)
t_start_z = 4 / f_z
t_end_z = 4 / f_z
t_cost_z = t_tot - (t_start_z + t_end_z)
# Amplitude of position in [mm]
Ax = 30  # [mm]
Ay = 0   # [mm]
Az = 20  # [mm]
#------------------------------------------------
# Time vector
t = np.arange(0, t_tot + 1 / FT, 1 / FT)
#------------------------------------------------
# X-direction
Ax = Ax * 1e-3  # Conversion from mm to m
Amp_array_x = Ax * np.ones(len(t))

value_to_extract_up_to_start = t_start_x
t_start = t[t <= value_to_extract_up_to_start]
m_start = Ax / t_start_x
Amp_array_x[:len(t_start)] = m_start * t_start

value_to_extract_up_to_end = t_end_x
t_end = t[t >= (t_tot - value_to_extract_up_to_end)]
m_end = Ax / t_end_x
Amp_array_x[-len(t_end):]=m_end*(t_end-(t_tot - t_end_x)*np.ones(len(t_end)))
Amp_array_x[-len(t_end):] = Amp_array_x[-len(t_end):][::-1]

plt.figure()
plt.plot(t, Amp_array_x * 1000)
plt.xlabel('Time [s]')
plt.ylabel('Amplitude X Position [mm]')
plt.title('Trend of Amplitude X [mm] in function of time [s]')
#------------------------------------------------
# Y-direction
Ay = Ay * 1e-3  # Conversion from mm to m
Amp_array_y = Ay * np.ones(len(t))

value_to_extract_up_to_start = t_start_y
t_start = t[t <= value_to_extract_up_to_start]
m_start = Ay / t_start_y
Amp_array_y[:len(t_start)] = m_start * t_start

value_to_extract_up_to_end = t_end_y
t_end = t[t >= (t_tot - value_to_extract_up_to_end)]
m_end = Ay / t_end_y
Amp_array_y[-len(t_end):]=m_end*(t_end-(t_tot-t_end_y)*np.ones(len(t_end)))
Amp_array_y[-len(t_end):] = Amp_array_y[-len(t_end):][::-1]

plt.figure()
plt.plot(t, Amp_array_y * 1000)
plt.xlabel('Time [s]')
plt.ylabel('Amplitude Y Position [mm]')
plt.title('Trend of Amplitude Y [mm] in function of time [s]')
#-------------------------------------------------
# Z-direction
Az = Az * 1e-3  # Conversion from mm to m
Amp_array_z = Az * np.ones(len(t))

value_to_extract_up_to_start = t_start_z
t_start = t[t <= value_to_extract_up_to_start]
m_start = Az / t_start_z
Amp_array_z[:len(t_start)] = m_start * t_start

value_to_extract_up_to_end = t_end_z
t_end = t[t >= (t_tot - value_to_extract_up_to_end)]
m_end = Az / t_end_z
Amp_array_z[-len(t_end):]=m_end*(t_end-(t_tot-t_end_z)*np.ones(len(t_end)))
Amp_array_z[-len(t_end):] = Amp_array_z[-len(t_end):][::-1]

plt.figure()
plt.plot(t, Amp_array_z * 1000)
plt.xlabel('Time [s]')
plt.ylabel('Amplitude Z Position [mm]')
plt.title('Trend of Amplitude Z [mm] in function of time [s]')
#------------------------------------------------
# Sine sweep in X direction
sine_signal_x = chirp(t, f0=f_x, f1=f_x, t1=t_tot, method='linear')
sine_signal_x = sine_signal_x * Amp_array_x

plt.figure()
plt.plot(t, sine_signal_x * 1000, t, np.zeros(len(t)))
plt.title('Position along X')
plt.xlabel('Time [s]')
plt.ylabel('Position [mm]')
#------------------------------------------------
# Sine sweep in Y direction
sine_signal_y = chirp(t, f0=f_y, f1=f_y, t1=t_tot, method='linear')
sine_signal_y = sine_signal_y * Amp_array_y

plt.figure()
plt.plot(t, sine_signal_y * 1000, t, np.zeros(len(t)))
plt.title('Position along Y')
plt.xlabel('Time [s]')
plt.ylabel('Position [mm]')
#------------------------------------------------
# Sine sweep in Z direction
sine_signal_z = chirp(t, f0=f_z, f1=f_z, t1=t_tot, method='linear')
sine_signal_z = sine_signal_z * Amp_array_z

plt.figure()
plt.plot(t, sine_signal_z * 1000, t, np.zeros(len(t)))
plt.title('Position along Z')
plt.xlabel('Time [s]')
plt.ylabel('Position [mm]')
#------------------------------------------------
# Assemble signal
Signal = np.vstack((sine_signal_x * 1000, sine_signal_y * 1000, sine_signal_z * 1000)).T

