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

#------------------------------------------------
# Parameters
FT = 200  # Sampling Frequency [Hz]
f_x_exc = 1.0413  # [Hz]
f_z = 2.2293  # [Hz]
Az = 30  # [mm]
Ax_exc = 15  # [mm]
t_tot = 80  # [s]
t_excitation = 20  # [s]
t_start_x = 4 / f_x_exc  # Duration of starting ramp [s]
t_end_x = 4 / f_x_exc  # Duration of ending ramp [s]
t_cost_x=t_excitation-(t_start_x + t_end_x)  # Duration of the constant signal
t_start_z = 4 / f_z  # Duration of starting ramp [s]
t_end_z = 4 / f_z  # Duration of ending ramp [s]
t_cost_z=t_tot-(t_start_z + t_end_z)  # Duration of the constant signal [s]
#------------------------------------------------
# Time vector
t = np.arange(0, t_tot + 1 / FT, 1 / FT)  # Time vector [s]
t_array_exct = np.arange(0, t_excitation + 1 / FT, 1 / FT)  # [s]
#------------------------------------------------
#AMPLITUDE FOR HORIZONTAL INITIAL EXCITATION
# Amplitude of the signal (X-direction excitation)
Ax_exc = Ax_exc * 1e-3  # [m] conversion from mm to m
Amp_array_x = Ax_exc * np.ones(len(t_array_exct))

# Calculate amplitudes with ramps
value_to_extract_up_to_start = t_start_x
t_start = t_array_exct[t_array_exct <= value_to_extract_up_to_start]
m_start = Ax_exc / 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_array_exct[t_array_exct>=(t_excitation-value_to_extract_up_to_end)]
m_end = Ax_exc / t_end_x
Amp_array_x[-len(t_end):]=m_end*(t_end-(t_excitation - t_end_x)*np.ones(len(t_end)))
Amp_array_x[-len(t_end):] = Amp_array_x[-len(t_end):][::-1]
#------------------------------------------------
#AMPLITUDE FOR VERTICAL EXCITATION
# Amplitude of the signal (Z-direction)
Az = Az * 1e-3  # [m] conversion from mm to m
Amp_array_z = Az * np.ones(len(t))

# Calculate amplitudes with ramps
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]
#-------------------------------------------------
# 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
sine_signal_z = sine_signal_z * 1000  # Conversion to [mm]

# Plotting
plt.figure()
plt.plot(t, sine_signal_z)
plt.plot(t, np.zeros(len(sine_signal_z)), 'r--')
plt.title('Position along Z')
plt.xlabel('time [s]')
plt.ylabel('Position [mm]')
plt.ylim(-40, 40)
plt.grid(True)
plt.minorticks_on()
plt.grid(which='minor', linestyle=':', linewidth=0.5, color='black')
#------------------------------------------------
#ASSEMBLING WITH HORIZONTAL INITIAL EXCITATION
# Section for vertical excitation with horizontal start
sine_signal_exc = chirp(t_array_exct, f0=f_x_exc, f1=f_x_exc, t1=t_excitation, method='linear')
sine_signal_exc = sine_signal_exc * Amp_array_x

# Assembling the signal
sine_signal_exc = sine_signal_exc[:-1]
sine_signal_z = sine_signal_z[:-1]

sine_signal_z = np.concatenate((np.zeros(len(sine_signal_exc) // 2), sine_signal_z))
sine_signal_exc = np.concatenate((sine_signal_exc, np.zeros(len(sine_signal_z) - len(sine_signal_exc))))
sine_signal_exc = sine_signal_exc * 1000

# Effective time, considering the initial horizontal excitation
t_tot_effective = t_excitation / 2 + t_tot
effective_time = np.arange(0, t_tot_effective + 1 / FT, 1 / FT)[:-1]

# Plotting the final signals
plt.figure()
plt.subplot(2, 1, 1)
plt.plot(effective_time, sine_signal_exc)
plt.title('Position along X')
plt.xlabel('time [s]')
plt.ylabel('Position [mm]')
plt.xlim(0, 20)
plt.ylim(-40, 40)
plt.grid(True)
plt.minorticks_on()
plt.grid(which='minor', linestyle=':', linewidth=0.5, color='black')
plt.subplot(2, 1, 2)
plt.plot(effective_time, sine_signal_z)
plt.title('Position along Z')
plt.xlabel('time [s]')
plt.ylabel('Position [mm]')
plt.xlim(0, 20)
plt.ylim(-40, 40)
plt.grid(True)
plt.minorticks_on()
plt.grid(which='minor', linestyle=':', linewidth=0.5, color='black')
plt.tight_layout()  # Add this line for proper layout

plt.figure()
plt.plot(effective_time, sine_signal_exc, effective_time, sine_signal_z)
plt.xlabel('time [s]')
plt.ylabel('Position [mm]')
plt.legend(['Sine signal in X', 'Sine Signal in Z'])
plt.ylim(-40, 40)
plt.grid(True)
plt.minorticks_on()
plt.grid(which='minor', linestyle=':', linewidth=0.5, color='black')

# Final signal
Signal = np.column_stack([sine_signal_exc, np.zeros(len(sine_signal_exc)), sine_signal_z])
