import numpy as np
import argparse
import matplotlib.pyplot as plt
import glob

def plot_fk(k, f, sp2, f_min, f_max, k_min, k_max, record_path, channel):
    """
    Plots the F-K diagram and saves it as a PNG file.

    Parameters:
    k (numpy.ndarray): Wavenumber data.
    f (numpy.ndarray): Frequency data.
    sp2 (numpy.ndarray): Power spectral density data.
    f_OSGW (numpy.ndarray): Array for overlay plot (e.g., theoretical curve).
    vmin (float): Minimum value for color scale.
    vmax (float): Maximum value for color scale.
    record_path (str): Path to save the plot.
    ind (int): Index used in the filename.
    title (str): Title used in the filename.
    """
    plt.rcParams.update({'font.size': 35})
    fig, ax = plt.subplots(figsize=(20, 12))
    plt.pcolormesh(k, f, sp2.T, cmap='jet', vmin=-110, vmax=-50, shading='auto')
    #plt.pcolormesh(k, f, sp2.T, cmap='jet', vmin=-150, vmax=-90, shading='auto') # 8150
    #h = plt.colorbar()
    #h.set_label('Power Spectra [dB] (rel. 1 $(\epsilon/s)^2$)')
    plt.ylabel('Frequency [Hz]')
    plt.xlabel('Wavenumber [1/m]')
    #plt.ylim([0.004, 1.0])
    plt.ylim([f_min, f_max])
    #plt.xlim([-0.005, 0.005])
    plt.xlim([k_min, k_max])
    ax.set_yscale('log')
    plt.grid()
    plt.tight_layout()
    plt.savefig(f"{record_path}/PlotFK_{str(f_min)}-{str(f_max)}Hz_chn{str(channel)}.png")

def compute_energy(k, f, sp2, f_min, f_max, k_min, k_max, record_path, channel):
    k_indices = np.where((k >= k_min) & (k < k_max))
    f_indices = np.where((f >= f_min) & (f < f_max))
    f_indices = np.array(f_indices).flatten()
    k_indices = np.array(k_indices).flatten()
    f_grid, k_grid = np.meshgrid(f_indices, k_indices, indexing='ij')
    sp2_ = sp2[k_grid, f_grid]    
    k_ = k[k_indices]
    f_ = f[f_indices]
    A = np.sqrt((10 ** (sp2_ / 20)) / 2)
    A_stacked = A.sum(axis=0)/float(len(f_))
    sp2_stacked = 2*(np.abs(A_stacked)**2)
    sp2_stacked = 20*np.log10(abs(sp2_stacked))
    ## compute energy from twosides
    kp_indices = np.where(k_ > 0)
    kn_indices = np.where(k_ < 0)
    kp_energy = A_stacked[kp_indices]
    kn_energy = A_stacked[kn_indices]
    kp_percent = sum(kp_energy) * 100 / (sum(kp_energy) + sum(kn_energy))
    kn_percent = sum(kn_energy) * 100 / (sum(kp_energy) + sum(kn_energy))
    
    # plotting
    plt.figure(figsize=(10, 6))
    plt.plot(k_, sp2_stacked) #, marker='o')  # 'o' is for circle markers, you can change it
    plt.xlabel('wavenumber [1/m]')  # Replace with your actual label
    plt.ylabel('Power Spectra [dB] (rel. 1 $(\epsilon/s)^2$)')  # Replace with your actual label
    plt.fill_between(k_, -200, sp2_stacked, color='skyblue', alpha=0.4)  # Add shading below the curve
    plt.axvline(x=0, color='red', linestyle='-')
    plt.figtext(0.8, 0.2, str(f'{kp_percent:5.2f}%'), fontsize=25)
    plt.figtext(0.2, 0.2, str(f'{kn_percent:5.2f}%'), fontsize=25)
    plt.xlim([min(k_), max(k_)])
    #plt.ylim([np.amin(sp2_stacked)-10, np.amax(sp2_stacked)+10])
    plt.ylim([-140, -90])
    plt.grid(True)
    plt.tight_layout()
    plt.savefig(f"{record_path}/PlotFK_2_{str(f_min)}-{str(f_max)}Hz_chn{str(channel)}.png")

    #print(len(k), len(k_), k_indices)
    #print(len(f), len(f_), f_indices)
    #print(sp2.shape, sp2_.shape)

def main(folder_path, channel, freq_min, freq_max, wavenumber_min, wavenumber_max):
    f_path = glob.glob(folder_path+"/"+"f_values_*chn"+str(channel)+".npy")
    k_path = glob.glob(folder_path+"/"+"k_values_*chn"+str(channel)+".npy")
    sp2_path = glob.glob(folder_path+"/"+"fk_psd_*chn"+str(channel)+".npy")
    f = np.load(str(f_path)[2:-2])
    k = np.load(str(k_path)[2:-2])
    sp2 = np.load(str(sp2_path)[2:-2])

    compute_energy(k, f, sp2, float(freq_min), float(freq_max), float(wavenumber_min), float(wavenumber_max), folder_path, channel)
    plot_fk(k, f, sp2, float(freq_min), float(freq_max), float(wavenumber_min), float(wavenumber_max), folder_path, channel)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Read f, k, and sp2 data.")
    parser.add_argument("folder_path", help="Path to the .npy file containing data")
    parser.add_argument("channel", help="Channel number for reading the .npy file")
    parser.add_argument("freq_min", help="minimum frequency")
    parser.add_argument("freq_max", help="maximum frequency")
    parser.add_argument("wavenumber_min", help="minimum wavenumber")
    parser.add_argument("wavenumber_max", help="maximum wavenumber")

    args = parser.parse_args()
    main(args.folder_path, args.channel, args.freq_min, args.freq_max, args.wavenumber_min, args.wavenumber_max)