#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Jan  5 14:05:08 2026

@author: vieweg

reads the .npy file of the merged data and plots it
"""

import numpy as np
import matplotlib.pyplot as plt
import datetime as dt
from datetime import datetime, timedelta
import matplotlib.cm as cm
from matplotlib.cm import ScalarMappable
from matplotlib.colors import BoundaryNorm
from matplotlib.colors import ListedColormap


def load_data(file):
#load merged .npy data
    all_data = np.load(file, allow_pickle=True).item()
    dates = np.array([datetime.fromisoformat(str(d)) for d in all_data['dates']])
    distances_loaded = all_data['distance'].astype(np.float64)
    position = all_data['position'].astype(np.float64)
    L_value_loaded = all_data['L_value'].astype(np.float64)
    energy_pro = all_data['energy_pro'].astype(np.float64)
    energy_el = all_data['energy_el'].astype(np.float64)
    flux_pro = all_data['mean_flux_pro'].astype(np.float64)
    flux_el = all_data['mean_flux_el'].astype(np.float64)
    x_loaded, y_loaded, z_loaded = position[:,0], position[:,1], position[:,2]
  
    return dates, distances_loaded, x_loaded, y_loaded, z_loaded, L_value_loaded, energy_pro, energy_el, flux_pro, flux_el

#Plot
def plot_data(data, start, end):
    (dates, distances_loaded, x_loaded, y_loaded, z_loaded, L_value_loaded, energy_pro, energy_el, flux_pro, flux_el) = data
    
    fig3 = plt.figure(figsize = (30,15)) #width, hight
    fig3.subplots_adjust(wspace = 0, hspace = 0)
    gs = fig3.add_gridspec(150,211,width_ratios = [1]*205 + [1]*6,height_ratios = [1]*150)
    axts = fig3.add_subplot(gs[0:50,0:205])
    axbs = fig3.add_subplot(gs[50:100,0:205], sharex = axts) # this axis shares the x axis with axds, often helpful as e.g. xlim on one panels auto adjusts the other
    axds = fig3.add_subplot(gs [100:150,0:205], sharex = axts)

    
  
    
    axts.set_ylabel(r"$J_\text{p}$ / $(\mathrm{cm²\,s\,sr\,keV})⁻¹$", fontsize = 24)
    axts.set_yscale("log")
    axts.set_xlim(start, end)
    #axts.set_ylim(1e-2, 1e6)
    axts.tick_params(axis = 'x', bottom=True, top = True, direction='in', labelsize = 20)
    axts.tick_params(axis = 'y', labelsize = 13)
    axts.tick_params(labelbottom = False)

    axbs.set_ylabel(r"$J_\text{e}$ / $(\mathrm{cm²\,s\,sr\,keV})⁻¹$", fontsize = 24)
    axbs.set_yscale("log")
    axbs.set_xlim(start, end)

    axbs.tick_params(axis = 'x', bottom=True, top = True, direction = 'in', labelsize = 20)
    axbs.tick_params(axis = 'y', labelsize = 13)
    axbs.tick_params(labelbottom = False)

   
    cmap = cm.jet   

   
    #HSPCT
    num_chanel=len(energy_pro)

    colors = cmap(np.linspace(0,1, num_chanel))
    listed = ListedColormap(colors)
    if num_chanel >1:
        last_edge = energy_pro[-1]+(energy_pro[-1]-energy_pro[-2])
 
    bounds = np.append(energy_pro, last_edge)
    norm = BoundaryNorm(boundaries=bounds, ncolors=listed.N, extend='neither')
    #legend for energy channels - no colorbar
    for k in range(num_chanel):
        axts.plot(dates, flux_pro[:,k], color = colors[k], label = f'{energy_pro[k]:g} keV')
    sm = ScalarMappable(norm=norm, cmap=listed)
    sm.set_array(None)


    axts.legend(title = 'Proton Energy', fontsize = 18, title_fontsize = 20, ncol = 4)


    #ESPCT - Electrons
    num_chanel=len(energy_el)
    colors = cmap(np.linspace(0,1, num_chanel))
    listed = ListedColormap(colors)
    if num_chanel >1:
        last_edge = energy_el[-1]+(energy_el[-1]-energy_el[-2])

    bounds = np.append(energy_el, last_edge)
    norm = BoundaryNorm(boundaries=bounds, ncolors=listed.N, extend='neither')
    for k in range(num_chanel):
        #legend for energy channels - no colorbar
        axbs.plot(dates, flux_el[:,k], color = colors[k], label = f'{energy_el[k]:g} keV')
    sm = ScalarMappable(norm=norm, cmap=listed)
    sm.set_array(None)

    
    
    
    axbs.legend(title = 'Electron Energy', fontsize = 18, title_fontsize = 20, ncol = 2)
    
    
    axds.plot(dates, x_loaded, label = 'x-coordinate (GSE)')
    axds.plot(dates, y_loaded, label = 'y-coordinate (GSE)')
    axds.plot(dates, z_loaded, label = 'z-coordinate (GSE)')   
    axds.plot(dates, distances_loaded, label = 'Radial distance') 

    axds.legend(fontsize = 18)
    axds.tick_params(axis = 'x', labelsize = 22)

    axds.set_ylabel(r"Distance $\mathrm{\,/\, R_E}$", fontsize = 24)

    axds.set_xlabel("Time", fontsize = 24)
    axds.set_xlim(start, end)
    axds.tick_params(axis = 'y', labelsize = 20)
    axts.set_title('C3 RAPID Particle Flux - SC Orbit - ' + f'{start} to '+f'{end}', fontsize = 26)
    
#please enter a .npy file name
data = load_data('/data/hiwi2/Lisa/Cluster_Data/Data_2002/C3_merged_2002.npy')
#please enter start and end date of the wanted time interval
plot = plot_data(data, datetime(2002, 1, 1, 0, 0, 0), datetime(2002, 12, 31, 23, 59, 30))
filename = '/home/asterix/vieweg/Desktop/C3_RAPID_2002.pdf'
plt.savefig(filename)
plt.show()