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

def plot_fk(k, f, sp2, vmin, vmax, 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': 20})
    fig, ax = plt.subplots(figsize=(20, 12))
    plt.pcolormesh(k, f, sp2.T, cmap='jet', vmin=vmin, vmax=vmax, shading='auto')
    h = plt.colorbar()
    h.set_label('Power Spectra [dB] (rel. 1 $(\epsilon/s)^2$)')
    plt.ylabel('frequency [1/s]')
    plt.xlabel('wavenumber [1/m]')
    plt.ylim([0.004, 1.0])
    plt.xlim([-0.005, 0.005])
    ax.set_yscale('log')
    plt.grid()
    plt.tight_layout()
    plt.savefig(f"{record_path}/PlotFK_chn{str(channel)}.png")

def main(folder_path, channel):
    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")
    print(f_path)
    f = np.load(str(f_path)[2:-2])
    k = np.load(str(k_path)[2:-2])
    sp2 = np.load(str(sp2_path)[2:-2])

    plot_fk(k, f, sp2, -110, -50, 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")

    args = parser.parse_args()
    main(args.folder_path, args.channel)