from scipy.fft import fft, ifft, fftfreq
import numpy as np
import sys
import matplotlib.pyplot as plt



tid_pot = np.genfromtxt(f'/nfshk2/aygun/compressible_model_v2/TEST2/elastic3/CD_0/G_145_0/tidal_potential.dat', names=True, max_rows=999)
ind_pot = np.genfromtxt(f'/nfsy3/aygun/MM_GANYMEDE/elastic/CD_0/G_20_0/love_numbers_time.dat', skip_header=15000, max_rows=999)

tidal_potential = - tid_pot['MM_RV22'] -tid_pot['Ecc_RV22'] 
print(max(tid_pot['Ecc_RV22']),max(ind_pot[:,2]) )
n  = 1000
nn = np.arange(n-1)
per = 618153.4
w   = 2.0 * np.pi / per
Delta_t = tid_pot['Times'][1] - tid_pot['Times'][0]


ecc_F = fft(-tid_pot['Ecc_RV22'])
abs_ecc = abs(ecc_F)/n
frq_ecc = 2.0 * np.pi * (nn) / (n*Delta_t)
frq_ecc= frq_ecc / w

fig = plt.figure(figsize=(10, 5))
ax = fig.add_subplot(111)
ax.plot(frq_ecc, abs_ecc, label='Amplitude Spectrum of Eccentricity')
plt.xlabel('Frequency (normalized)')
plt.ylabel('Amplitude')
plt.legend()
plt.yscale('log')
plt.grid()
plt.xlim([0, 50])  # Limit the frequency range for better visualization
plt.savefig('./figs/png/ecc_frequency_spectrum.png')
plt.savefig('./figs/pdf/ecc_frequency_spectrum.pdf')

tid_F = fft(tidal_potential)
## Remove the main frequency
tid_F = tid_F - ecc_F
amp_F = abs(tid_F)/n
frq_F = 2.0 * np.pi * (nn) / (n*Delta_t)
frq_F= frq_F / w


k22_f = fft(ind_pot[:,2])
## Remove the main frequecny
k22_f[1]  = k22_f[1]  - 1.455*0.5*n
k22_f[-1] = k22_f[-1] - 1.455*0.5*n
amp_k22 = abs(k22_f)/n
frq_k22 = 2.0 * np.pi * (nn) / (n*Delta_t)
frq_k22= frq_k22 / w


ind_pot_time_wo_1 = ifft(k22_f)
tid_pot_time_wo_1 = ifft(tid_F)


fig = plt.figure(figsize=(10, 5))
ax = fig.add_subplot(111)
ax.plot(tid_pot['TimeP'], np.real(ind_pot_time_wo_1), label='Induced Potential wo main freq')
ax.plot(tid_pot['TimeP'], np.real(tid_pot_time_wo_1), label='Tidal   Potential wo main freq')
ax.plot(tid_pot['TimeP'], -tid_pot['MM_RV22'])
plt.xlabel('Time')
plt.ylabel('Potential')
plt.legend()
plt.grid()
plt.savefig('./figs/png/ind_tid_potential_wo_1.png')
plt.savefig('./figs/pdf/ind_tid_potential_wo_1.pdf')


k2_time = np.real(ind_pot_time_wo_1) / np.real(tid_pot_time_wo_1)
fig = plt.figure(figsize=(10, 5))
ax = fig.add_subplot(111)
ax.plot(tid_pot['TimeP'], (k2_time), label='k2')
plt.xlabel('Time')
plt.ylabel('k2')
plt.legend()
plt.savefig('./figs/png/k2_time_wo_main.png')
plt.savefig('./figs/pdf/k2_time_wo_main.pdf')


k2_freq = fft(k2_time)
amp_k2 = abs(k22_f)/n
frq_k2 = 2.0 * np.pi * (nn) / (n*Delta_t)
frq_k2= frq_k2 / w


fig = plt.figure(figsize=(10, 5))
ax = fig.add_subplot(111)
ax.plot(frq_k2, amp_k2, label='Amplitude Spectrum of k2')
plt.xlabel('Frequency (normalized)')
plt.ylabel('Amplitude')
plt.legend()
plt.grid()
plt.xlim([0, 50])  # Limit the frequency range for better visualization
plt.savefig('./figs/png/k2_frequency_spectrum.png')
plt.savefig('./figs/pdf/k2_frequency_spectrum.pdf')


# sys.exit("Stopping the code execution as requested.")
## time dependent part