import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from rapid_fgm import *
from datetime import datetime, timedelta
from matplotlib.lines import Line2D


#--------------USER INPUT----------------------

TIME_PERIOD = (datetime(2002,6,27,00), datetime(2002,6,29,0))
TIME_POINT = datetime(2002,6,28,18,30) 


SC_ID = 'C3'


PARTICLE_TYPE = "ion"  # oder "ion"


#---------------------------------------
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()


print("eflow samples:", len(flow.time_flow))
if len(flow.time_flow) > 0:
    print("eflow first:", flow.time_flow[0])
    print("eflow last :", flow.time_flow[-1])
else:
    print("eflow empty after filtering!")

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

p = PITCH_CSA(sc_id = SC_ID, time_period= TIME_PERIOD)
data_pad = p.load_pad()
pad_data = p.get_pad_var(p.data_pad)

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

t_sync = sync['time']
rapid_sync = sync['rapid']
fgm_sync = sync['fgm']
sc2gse = sync['sc2gse']

sc_data = w.spacecrafts(
    sc_id=SC_ID,
    t_sync=t_sync,
    rapid_sync=rapid_sync,
    fgm_sync=fgm_sync,
    sc2gse=sc2gse,
    time_period=None)
print("len SC time:", len(w.SC_data[SC_ID]['time']))

w.rapid_n_sc2gse(SC_ID)

pitch_angle = w.calc_pitch_angle(SC_ID)

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




def channel_label(e_idx):
    channel_phys = e_idx + 1
    energy = r.energy_table[e_idx]
    return energy, channel_phys


color1 = '#0070C0'
color2 = '#63A375'


fig = plt.figure(figsize=(16, 8))
gs = gridspec.GridSpec(
    nrows=2,
    ncols=4,
    figure=fig,
    hspace=0.4,
    wspace=0.25
)


for e_idx in range(8):

    row = e_idx // 4
    col = e_idx % 4
    ax = fig.add_subplot(gs[row, col])

    IES_energy, channel_phys = channel_label(e_idx)

    # PAD for this channel
    pitch, flux, flux_sd = w.plot_flux_of_pitch(SC_ID,
        TIME_POINT,
        e_idx
    )
    flux = np.ma.masked_where(flux <= 0, flux)
    ax.errorbar(
        pitch,
        flux,
        yerr=flux_sd,
        linestyle='none',
        marker='o',
        markersize=4,
        color=color1,
        ecolor=color2,
        capsize=2
    )

    ax.set_title(
        f'Ch {channel_phys}  ({IES_energy:.1f} keV)',
        fontsize=14
    )

    ax.grid(True)
    ax.tick_params(axis='y', labelsize=12)
    
    ax.set_xticks([45, 90, 135])
    ax.set_xticklabels(['45°', '90°', '135°'], fontsize = 12)

    # Axis labels only on outer plots
    if row == 1:
        ax.set_xlabel('Pitch Angle / °', fontsize = 14)
    if col == 0:
        ax.set_ylabel(r'Flux / $(\mathrm{cm^{2}\,s\,sr\,keV})^{-1}$', fontsize = 14)



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


# GLOBAL TITLE 
fig.suptitle(
    f'{SC_ID} – RAPID {instrument} Pitch Angle Distribution – {TIME_POINT}',
    fontsize=16, 
)
timestamp = TIME_POINT.strftime("%Y%m%d_%H%M%S")
filename = 'C:/Users/lview/OneDrive/Dokumente/Desktop/C3_RAPID_20020628/191500/Flux_Pitch/'+ f'{SC_ID}_{instrument}_'+f'{timestamp}'+'_FP_all.pdf'
print(filename)
plt.savefig(filename)
plt.show()