# -*- coding: utf-8 -*-
"""
This Python program processes the previously imported N2 and N3d data to construct and visualise 
Saturn Kilometric Radiation (SKR) spectrograms. The datasets are merged, filtered according to 
frequency and signal-to-noise criteria, and transformed into two-dimensional arrays representing 
polarization properties as a function of time and frequency. The resulting spectrograms are used as 
diagnostic plots in the subsequent analysis. 

@author: Dorian Jost
"""

import numpy as np
import pandas as pd
from matplotlib import pyplot as plt
from matplotlib.dates import AutoDateLocator, DateFormatter
from datetime import datetime, timedelta
from pathlib import Path
from tqdm import tqdm
import readafile

def get_dataframe(year, doy1, hour1, doy2, hour2, data_dir_n2, data_dir_n3d):
    start = datetime(year, 1, 1) + timedelta(days=doy1 - 1, hours=hour1)
    end = datetime(year, 1, 1) + timedelta(days=doy2 - 1, hours=hour2)
    current = start
    print("Fetching freqs and times")
    
    while current <= end:
        doy = current.timetuple().tm_yday
        hour = current.hour
        doy_str = f"{doy:03d}"
        hour_str = f"{hour:02d}"
        
        n2_file = Path(data_dir_n2) / f"P{year}{doy_str}.{hour_str}"
        n3d_file = Path(data_dir_n3d) / f"N3d_dsq{year}{doy_str}.{hour_str}"
        
        try:
            n2 = readafile.read_n2(str(n2_file))
            n3d = readafile.read_n3d(str(n3d_file))
        except FileNotFoundError:
            print("Missing file", f"{n3d_file}")
            current += timedelta(hours=1)
            start = current
            continue    
        

        n2columns = n2.dtype.names
        n2_df = pd.DataFrame.from_records({ n: list(n2[n]) for n in n2columns},columns=n2columns)        

        n3dcolumns = n3d.dtype.names
        n3d_df = pd.DataFrame.from_records({ n: list(n3d[n]) for n in n3dcolumns},columns=n3dcolumns)
        
        DF_hour = pd.merge(n2_df,n3d_df, on="num", how ="inner")
        if current == start:
            DF = DF_hour.copy() 
            current += timedelta(hours=1)
            continue
             
        DF = pd.concat([DF, DF_hour], ignore_index = True)        
        current += timedelta(hours=1)
        
    ### splitting sn into two columns ###
    DF["sn0"] = DF["sn"].apply(lambda x: x[0])
    DF["sn1"] = DF["sn"].apply(lambda x: x[1])
    DF.pop("sn")
  
    return DF

def plot_beta(ax, data):   
    timestamps = data["timestamps"]
    beta_xw = data["beta_xw"]
    beta_uw = data["beta_uw"]
    beta_vw = data["beta_vw"]
    
       
#     plt.figure(figsize=(10, 6))
    ax.plot(timestamps, beta_xw, label=r'$\beta_{xw}$')
    ax.plot(timestamps, beta_uw, label=r'$\beta_{uw}$')
    ax.plot(timestamps, beta_vw, label=r'$\beta_{vw}$')
    
    ax.set_ylim(-90, 90)
    ax.margins(x=0)
    
    ax.set_xlabel('Time', fontsize = 16)
    ax.set_ylabel(r'$\beta$ (degrees)', fontsize = 16)
    ax.set_title(r'$\beta$-Angle Between Vector and Reference Planes', fontsize = 20)
    
    # Set custom format for x-axis ticks
    ax.xaxis.set_major_formatter(DateFormatter('%H'))

    # Automatically manage the tick positions and labels
    ax.xaxis.set_major_locator(AutoDateLocator())

    plt.legend()
    plt.grid(True)

def plot_latitude(ax, data):
    timestamps = data["timestamps"]
    lat = data["lat"]
    
    ax.plot(timestamps,lat)
    ax.set_ylim(-90, 90)
    ax.margins(x=0)
    
    ax.set_xlabel('Time', fontsize = 16)
    ax.set_ylabel('Latitude (degrees)', fontsize = 16)
    ax.set_title('Spacecraft Latitude', fontsize = 20)
    
    
    # Set custom format for x-axis ticks
    ax.xaxis.set_major_formatter(DateFormatter('%#H'))
    
    # Automatically manage the tick positions and labels
    ax.xaxis.set_major_locator(AutoDateLocator())
    ticks = ax.get_xticks()
    ax.set_xticks(ticks[:-1])
    plt.grid(True)



if __name__ == "__main__":
    # Specify time to visualize
    year, doy1, doy2, hour1, hour2 = 2006, 362, 362, 9, 10
    
    # Set up directories
    if doy1 <= 90: quarter = "001_090"
    elif doy1 > 90 and doy1 <= 180: quarter = "091_180"
    elif doy1 > 180 and doy1 <= 270: quarter = "181_270"
    elif doy1 > 270 and doy1 <= 366: quarter = "271_366"
    else: print("DOY not a valid number")

    data_dir_n2 = f"E:/HFRdata/{year}_{quarter}/n2"
    data_dir_n3d = f"E:/HFRdata/{year}_{quarter}/n3d"

    beta_dir = f"E:/CassiniData/ephemeris/{year}/beta_{year}.dat"
    lat_dir = r"C:\Users\doria\Desktop\Cassini\cassini-main\casephem_04_17"
    
    # Get the dataframe and calculate 2D arrays
    DF = get_dataframe(year, doy1, hour1, doy2, hour2, data_dir_n2, data_dir_n3d)
    
    snr_threshold = 10 
    DF = DF[DF["f"]<=1000] # restrict to frequencies below 1KHz
    DF = DF[(DF["sn0"] >= snr_threshold) & (DF["sn1"] >= snr_threshold)] #ignore data below certain signal to noise ratio
    values = list(DF[["f","t97","v","q","u","s"]].drop_duplicates().itertuples(index=False, name=None))
    
    y = sorted(set(v[0] for v in values))
    x = sorted(set(v[1] for v in values))
    V = np.zeros((len(y), len(x)))
    Q,U,S = V.copy(), V.copy(), V.copy()
    
    for yi, xi, vi, qi, ui, si in tqdm(values):
        V[y.index(yi), x.index(xi)] = vi
        Q[y.index(yi), x.index(xi)] = qi
        U[y.index(yi), x.index(xi)] = ui
        S[y.index(yi), x.index(xi)] = si
        
    tvec = 24*(x- np.floor(x))
    freqs = y
    ### plotting ###
    fig = plt.figure(figsize = (10,6),layout="constrained")

    # --- Plot spectograms ---
    ax1 = fig.add_subplot(211)
    mesh1 = ax1.pcolormesh(tvec, freqs, V, shading ="auto", cmap ="jet", vmin = -1, vmax = 1)
    ax1.set_yscale("log")
    
    ax2 = fig.add_subplot(212)
    mesh2 = ax2.pcolormesh(tvec, freqs, np.sqrt(Q**2 + U**2), shading ="auto", cmap ="jet", vmin = 0, vmax = 1)
    ax2.set_yscale("log")
    
    time_delta = 24
    labels = [str(h) if h % 3 == 0 else "" for h in range(time_delta)]
    for ax in [ax1, ax2]:
        ax.set_yscale('log')
        ax.set_xlabel('Time [hours]', fontsize = 16)
        ax.set_ylabel('Frequency [KHz]', fontsize = 16)
        ax.grid(False)
    
    ax1.set_title("Circular polarization v", fontsize = 20)
    ax2.set_title("Linear polarization $\sqrt{q^2+u^2}$", fontsize = 20)
    fig.colorbar(mesh1, ax=ax1)
    fig.colorbar(mesh2, ax=ax2)
    
    # --- Plot latitude ---
    lat_data = readafile.read_ephemeris_data(lat_dir,year,doy1,hour1,year,doy2,hour2+1,coord_sys = "geographic")
    plt.show()