import matplotlib.gridspec as gridspec

from rapid_fgm import *
from datetime import datetime, timedelta
import numpy as np
import matplotlib.pyplot as plt
import cartopy.crs as ccrs

#USER INPUT

SC_ID = 'C3'

TIME_PERIOD = (datetime(2002,6,27,00), datetime(2002,6,29,0))

T_PLOT = datetime(2002,6,28,19,15)

PARTICLE_TYPE = "electron"  #'electron' oder "ion"

instrument = "IES" if PARTICLE_TYPE == "electron" else "IIMS"


#-----------------------------------------------------
r = RAPID(sc_id=SC_ID, time_period=TIME_PERIOD, particle_type=PARTICLE_TYPE)

data_rapid = r.load_rapid()
r.get_rapid_var(r.data_rapid)
print('rapid time:', r.time_rapid[0], r.time_rapid[-1])




flow = FLOW(
    sc_id=SC_ID,
    time_period=TIME_PERIOD,
    particle_type=PARTICLE_TYPE
)

flow.get_flow_var(flow.data_flow)
flow.load_flow()

#font in Times New Roman
plt.rcParams['font.family'] = 'Times New Roman'
plt.rcParams['mathtext.fontset'] = 'stix'


def vec_to_lonlat(n):
    x, y, z = n[..., 0], n[..., 1], n[..., 2]
    lon = np.degrees(np.arctan2(y, x))
    lat = np.degrees(np.arcsin(z))
    return lon, lat


# Load Data 
r = RAPID(sc_id=SC_ID, time_period=TIME_PERIOD)
r.get_rapid_var(r.data_rapid)

e = FLOW(sc_id=SC_ID, time_period=TIME_PERIOD)
e.load_flow()
e.get_flow_var(e.data_flow)

f = FGM(sc_id=SC_ID, time_period=TIME_PERIOD)
f.load_fgm()
f.get_fgm_var(f.data_fgm)

w = WORKSPACE()
sync = w.sync_all_by_resolution(r, f, e)

w.spacecrafts(
    sc_id=SC_ID,
    t_sync=sync['time'],
    rapid_sync=sync['rapid'],
    fgm_sync=sync['fgm'],
    sc2gse=sync['sc2gse'],
    time_period=None
)

w.rapid_n_sc2gse(SC_ID)
w.calc_pitch_angle(SC_ID)



idx = w.time_to_index(w.SC_data[SC_ID]['time'], T_PLOT)


B = w.SC_data[SC_ID]['B_vec'][idx]
B_hat = B / np.linalg.norm(B)
phi = np.linspace(0, 2*np.pi, 200)


fig = plt.figure(figsize=(16, 9))
fig.subplots_adjust(hspace=1.5)
gs = gridspec.GridSpec(2, 4, figure=fig)

for ch in range(8):

    ax = fig.add_subplot(gs[ch], projection=ccrs.Mollweide(central_longitude=0))
    ax.set_global()


    flux = w.SC_data[SC_ID]['flux'][idx][ch,...]
    flux = np.ma.masked_where(flux < 0, flux)
    #print('flux: ', flux)
    n_gse = w.SC_data[SC_ID]['n_gse'][idx]

    
    lon, lat = vec_to_lonlat(-1*n_gse)


    lon = lon[:flux.shape[0], :flux.shape[1]]
    lat = lat[:flux.shape[0], :flux.shape[1]]
    

    if np.any(~flux.mask):
        mean_flux = flux[~flux.mask].mean()
    else:
        mean_flux = np.nan

    mask = ~flux.mask
        
    
    ax.gridlines(
        crs=ccrs.PlateCarree(),
        draw_labels=False,
        linewidth=0.5,
        linestyle='--',
        color='gray'
    )
    

    
    # x-axis
    for lon_tick in np.arange(-180, 181, 60):
        ax.text(lon_tick, 0, f"{lon_tick}°",
                transform=ccrs.PlateCarree(),
                ha='center', va='center', fontsize=12, fontname='Times New Roman')

    # y-axis
    for lat_tick in np.arange(-90, 91, 30):
        ax.text(-180, lat_tick, f"{lat_tick}°",
                transform=ccrs.PlateCarree(),
                ha='center', va='center', fontsize=12, fontname='Times New Roman')
    
    sc = ax.scatter(
        lon[mask],
        lat[mask],
        c=flux[mask],
        s=30,
        cmap='jet',
        transform=ccrs.PlateCarree(),
        zorder=10
    )
    #assumption: same GF for all channels
    
    GF = 2.2E-3
    dE = np.array([1.130E+01, 1.760E+01, 2.640E+01, 3.300E+01,
               4.840E+01, 6.820E+01, 9.240E+01, 7.000E+01])
    dt = 0.25

    factor = GF * dE[ch] * dt
    cbar_flux = fig.colorbar(sc, ax=ax, orientation='horizontal',fraction=0.05, pad=0.19)


    cbar_flux.set_label(r'Flux / $(\mathrm{cm^{2}\,s\,sr\,keV})^{-1}$', fontsize = 12)

    cbar_counts = cbar_flux.ax.secondary_xaxis("top")
    cbar_counts.tick_params(labelsize=12)

    flux_ticks = cbar_flux.get_ticks()

    factor = GF * dE[7] * dt  

    cbar_counts.set_xlim(cbar_flux.ax.get_xlim())
    cbar_counts.set_xticks(flux_ticks)
    cbar_counts.set_xticklabels([f"{t * factor:.0f}" for t in flux_ticks])

    cbar_counts.set_xlabel("Counts", fontsize = 12)
    cbar_counts.xaxis.set_ticks_position("top")
    cbar_counts.xaxis.set_label_position("top")

    cbar_flux.update_ticks()
    flux_ticks = cbar_flux.get_ticks()

    cbar_flux.ax.tick_params(labelsize=12)

    #  Pitch Angle Isolines
    B = B_hat 

    lon_grid = np.linspace(-180, 180, 200)
    lat_grid = np.linspace(-90, 90, 100)

    LON, LAT = np.meshgrid(lon_grid, lat_grid)


    LONG_rad = np.deg2rad(LON)
    LAT_rad = np.deg2rad(LAT)


    x = -1*np.cos(LAT_rad) * np.cos(LONG_rad)
    y = -1*np.cos(LAT_rad) * np.sin(LONG_rad)
    z = -1*np.sin(LAT_rad)

    v_grid = np.stack([x, y, z], axis=-1)

# calc pitch angle
    pa = np.degrees(np.arccos(np.clip(np.sum(v_grid * B, axis=-1), -1, 1)))

# Contour lines
    levels = np.arange(0, 181, 15)

    cs = ax.contour(
    LON, LAT, pa,
    levels=levels,
    colors='black',
    linewidths=0.6,
    transform=ccrs.PlateCarree(),
    zorder=2
)

   
    labels = ax.clabel(cs, inline=True, fontsize=10, fmt='%d°')

    for label in labels:
        label.set_fontweight('bold')

    


    from matplotlib.lines import Line2D
    contour_handle = Line2D([0], [0], color='black', lw=0.8, label='Pitch angle')


   
    IES_energy = r.energy_table[ch]
    ax.set_title(f'{SC_ID} - Ch {ch+1} ({IES_energy:.1f} keV)', fontsize=14, fontname='Times New Roman')


# global title

Position = w.SC_data[SC_ID]['fgm_pos'][idx]
RE_km = 6371
POSITION = Position / RE_km
fig.suptitle(
    f'{SC_ID} RAPID {instrument} View Directions\n {w.SC_data[SC_ID]["time"][idx]} - Position GSE ($R_E$) {POSITION}',
    fontsize=16, fontname='Times New Roman'
)


plt.tight_layout(rect=[0, 0.08, 1, 0.95])

timestamp = T_PLOT.strftime("%Y%m%d_%H%M%S")
filename = 'C:/Users/lview/OneDrive/Dokumente/Desktop/C3_RAPID_20020628/191500/Mollweide/'+ f'{SC_ID}_{instrument}_'+f'{timestamp}'+'_MW_all.pdf'
print(filename)
plt.savefig(filename)

plt.show()