from scipy.fft import fft, ifft, fftfreq
"""
This script performs frequency analysis on tidal and induced potentials using Fourier Transform.
It reads data from specified files, computes the Fourier Transform, and visualizes the results.

Modules:
    - scipy.fft: For computing the Fast Fourier Transform (FFT).
    - numpy: For numerical operations.
    - matplotlib.pyplot: For plotting the results.

Functions:
    - None

Variables:
    - tid_pot: Tidal potential data read from a file.
    - ind_pot: Induced potential data read from a file.
    - fig, ax: Matplotlib figure and axis objects for plotting.
    - n: Number of data points for FFT.
    - nn: Array of indices for FFT frequencies.
    - per: Period of the signal.
    - w: Angular frequency of the signal.
    - Delta_t: Time step between consecutive data points.
    - tid_F: FFT of the tidal potential.
    - amp_F: Amplitude spectrum of the tidal potential.
    - frq_F: Frequency spectrum of the tidal potential (normalized by angular frequency).
    - k22_f: FFT of the induced potential.
    - amp_k22: Amplitude spectrum of the induced potential.
    - frq_k22: Frequency spectrum of the induced potential (normalized by angular frequency).
    - eps: Threshold for filtering small amplitudes.
    - k22: Complex ratio of induced potential to tidal potential in the frequency domain.

Steps:
    1. Load tidal potential data (`tid_pot`) and induced potential data (`ind_pot`) from files.
    2. Plot the time-domain tidal and induced potentials.
    3. Compute the FFT of the tidal and induced potentials.
    4. Plot the frequency-domain amplitude spectra of the tidal and induced potentials.
    5. Compute the frequency-dependent ratio (`k22`) of induced to tidal potential.
    6. Plot the frequency-dependent ratio (`k22`).

Output:
    - Time-domain plot of tidal and induced potentials saved as:
        - './figs/tidal_potential.png'
        - './figs/tidal_potential.pdf'
    - Frequency-domain amplitude spectra saved as:
        - './figs/fft_tidal_induced.png'
        - './figs/fft_tidal_induced.pdf'
    - Frequency-dependent ratio (`k22`) plot saved as:
        - './figs/k22.png'
        - './figs/k22.pdf'

Notes:
    - Ensure the input file paths are correct and accessible.
    - The script assumes specific column names in the input files.
    - The frequency range for plots is limited to [0, 50] for amplitude spectra and [0, 500] for `k22`.
"""
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'] 
fig = plt.figure(figsize=(10, 5))
ax = fig.add_subplot(111)
ax.plot(tid_pot['TimeP'], tidal_potential , label='Tidal Potential')
ax.plot(tid_pot['TimeP'], ind_pot[:,2] , label='Induced Potential')
plt.xlabel('Time')
plt.ylabel('Potential')
plt.legend()
plt.savefig('./figs/tidal_potential.png')
plt.savefig('./figs/tidal_potential.pdf')

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]


tid_F = fft(tidal_potential)
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 first frequecny
k22_f[1] = complex(0.0)
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)

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



fig = plt.figure(figsize=(10, 5))
ax = fig.add_subplot(111)
ax.plot(frq_F, amp_F, label='Tidal Potential')
ax.plot(frq_k22, amp_k22, label='Induced Potential')
plt.xlabel('Frequency')
plt.ylabel('Amplitude')
plt.legend()
plt.yscale('log')
plt.xlim([0, 50])
plt.savefig('./figs/fft_tidal_induced.png')
plt.savefig('./figs/fft_tidal_induced.pdf')


eps=1e-6
k22 = np.zeros((n-1), dtype=np.complex_)
for it in range(n-1):
    if amp_F[it]>eps:
        k22[it] = k22_f[it]/tid_F[it]

fig = plt.figure(figsize=(10, 5))
ax = fig.add_subplot(111)
ax.plot(frq_F, abs(k22), label='k22')
plt.xlabel('Frequency')
plt.ylabel('Amplitude')
plt.legend()
# plt.yscale('log')
plt.xlim([-1, 100])
plt.savefig('./figs/k22.png')
plt.savefig('./figs/k22.pdf')

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

k22_t = ind_pot[:,2]/(tidal_potential)

fig = plt.figure(figsize=(10, 5))
ax = fig.add_subplot(111)
ax.plot(tid_pot['TimeP'], abs(k22_t), label='k22')
plt.xlabel('Time')
plt.ylabel('Amplitude')
plt.legend()
plt.savefig('./figs/k22_time.png')
plt.savefig('./figs/k22_time.pdf')

k22_Four = fft(k22_t)
amp_k22_Four = abs(k22_Four)/n
frq_k22_Four = 2.0 * np.pi * (nn) / (n*Delta_t)
frq_k22_Four= frq_k22_Four / w


fig = plt.figure(figsize=(10, 5))
ax = fig.add_subplot(111)
ax.plot(frq_k22_Four, amp_k22_Four, label='k22')
plt.xlabel('Frequency')
plt.ylabel('Amplitude')
plt.legend()
plt.yscale('log')
plt.xlim([0, 50])
plt.savefig('./figs/k22_Four.png')
plt.savefig('./figs/k22_Four.pdf')

cutoff = 10
k22_Four[cutoff:1000-cutoff] = 0.0
k22_t2 = ifft(k22_Four)
fig = plt.figure(figsize=(10, 5))
ax = fig.add_subplot(111)
ax.plot(tid_pot['TimeP'], np.abs(k22_t2)*(-(tid_pot['Ecc_RV22'][:] + tid_pot['MM_RV22'][:])), label='computed', marker='o', markersize=2)
ax.plot(tid_pot['TimeP'], ind_pot[:,2], label='induced')
plt.xlabel('Time')
plt.ylabel('Amplitude')
# plt.xlim([0.4,0.6])
# plt.ylim([-1.4,-1.3])
plt.legend()
plt.savefig('./figs/k22_time2.png')
plt.savefig('./figs/k22_time2.pdf')

fig = plt.figure(figsize=(10, 5))
ax = fig.add_subplot(111)
ax.plot(tid_pot['TimeP'], np.abs(k22_t2), label=r'$k_{22}$ from IFFT', marker='o', markersize=2)
ax.plot(tid_pot['TimeP'], k22_t, label=r'$k_{22}$ by division')
plt.ylim([0.51,0.56])
plt.xlabel('Time')
plt.ylabel('Amplitude')
plt.legend()
plt.savefig('./figs/k22_time_itself.png')
plt.savefig('./figs/k22_time_itself.pdf')


# print(ind_pot[0,2]/(-(tid_pot['Ecc_RV22'][0] + tid_pot['MM_RV22'][0])))
# k22_t2 = np.zeros(n-1)
# for it in range(0, n-1):
#     x0 = 0.4
#     for i in range(1, 10):
#         x = x0 - (ind_pot[it,2] - x0*(-(tid_pot['Ecc_RV22'][it] + tid_pot['MM_RV22'][it])))/((tid_pot['Ecc_RV22'][it] + tid_pot['MM_RV22'][it]))
#         x0 = x
#     k22_t2[it] = x

# fig = plt.figure(figsize=(10, 5))
# ax = fig.add_subplot(111)
# ax.plot(tid_pot['TimeP'], abs(k22_t2), label='k22', marker='x')
# ax.plot(tid_pot['TimeP'], abs(k22_t), label='k22', marker='o')
# plt.xlabel('Time')
# plt.ylabel('Amplitude')
# plt.legend()
# plt.savefig('./figs/k22_time2.png')
# plt.savefig('./figs/k22_time2.pdf')
