import datetime as dt

import matplotlib.pyplot as plt

from nom_client.nom_client import NoMClient

plt.rcParams['figure.figsize'] = [15, 10]
plt.rc('font', size=18)


def mjd50_to_datetime(mjd):
    dt1950 = dt.datetime(1950, 1, 1, 0, 0, 0, 0)

    if hasattr(mjd, "__len__"):
        rval = []
        for m in mjd:
            t = dt1950 + dt.timedelta(m)
            rval.append(t)
    else:
        rval = dt1950 + dt.timedelta(mjd)
    return rval


def plot(x_data, y_data, plot_types, line_types, xlabel, ylabel, labels, title):
    if not isinstance(plot_types, list):
        plot_types = [plot_types] * len(x_data)

    if not isinstance(line_types, list):
        line_types = [line_types] * len(x_data)

    if not isinstance(labels, list):
        labels = [labels] * len(x_data)

    for data_index in range(len(x_data)):
        if plot_types[data_index] == 'loglog':
            plt.loglog(x_data[data_index], y_data[data_index], linestyle=line_types[data_index], linewidth=2.0,
                       markersize=12.0,
                       label=labels[data_index])
        elif plot_types[data_index] == 'semilogy':
            plt.semilogy(x_data[data_index], y_data[data_index], linestyle=line_types[data_index], linewidth=2.0,
                         markersize=12.0,
                         label=labels[data_index])
        elif plot_types[data_index] == 'semilogx':
            plt.semilogx(x_data[data_index], y_data[data_index], linestyle=line_types[data_index], linewidth=2.0,
                         markersize=12.0,
                         label=labels[data_index])
        else:
            plt.plot(x_data[data_index], y_data[data_index], linestyle=line_types[data_index], linewidth=2.0,
                     markersize=12.0,
                     label=labels[data_index])

    plt.title(title, fontsize=18)
    plt.grid(True)
    plt.legend(bbox_to_anchor=(0.95, 0.95), loc=1, borderaxespad=0.,
               fancybox=True, fontsize=18, shadow=True)

    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plt.show()


nom_client = NoMClient("IRENE example")

trajectory_model = nom_client.get_model('trajectory')
trajectory_model.set_params(orbitSpecificationCode=205,
                            naifCode=399, apogeeAltitude=20000, perigeeAltitude=600,
                            missionDuration=0.25)
trajectory_result = nom_client.run_model(trajectory_model)
print(trajectory_result)

mjd = trajectory_result['trajectory'].get_variable_data(variable_name="MJD")
alt = trajectory_result['trajectory'].get_variable_data(variable_name="altitude")

plot(x_data=[mjd.T],
     y_data=[alt.T],
     plot_types="plot",line_types="solid",xlabel="MJD",
     ylabel=f"Altitude (m)",
     labels=[''], title="IRENE by percentile")


print("Running Ax8 model")
ap8ae8_model = nom_client.get_model('ae8ap8')
ap8ae8_model.set_external_input(external_input_name="trajectory", external_input=trajectory_result)
ap8ae8_results = nom_client.run_model(ap8ae8_model)
print(ap8ae8_results)

print("Running IRENE model")
irene_model = nom_client.get_model('irene')
irene_model.set_params(mode="mean_percentiles", percentiles=[75, 90])

irene_model.set_external_input(external_input_name="trajectory", external_input=trajectory_result)
irene_results = nom_client.run_model(irene_model)
print(irene_results)

# Get Ax8 results
energy_ax8 = ap8ae8_results['orbit_averaged_proton_spectrum'].get_variable_data(variable_name="energy")
int_flux_ax8 = ap8ae8_results['orbit_averaged_proton_spectrum'].get_variable_data(variable_name="integral_flux")
diff_flux_ax8 = ap8ae8_results['orbit_averaged_proton_spectrum'].get_variable_data(variable_name="differential_flux")

# Get Ax9 results
energy_ax9 = irene_results['orbit_averaged_proton_spectrum'].get_variable_data(variable_name="energy")
energy_units = irene_results['orbit_averaged_proton_spectrum'].get_variable_metadata(
    variable_name="energy")['units']

diff_flux_ax9 = irene_results['orbit_averaged_proton_spectrum'].get_variable_data(variable_name="differential_flux",
                                                                                  percentile=50).squeeze()
diff_units = irene_results['orbit_averaged_proton_spectrum'].get_variable_metadata(
    variable_name="differential_flux")['units']

plot(x_data=[energy_ax8.T, energy_ax9.T],
     y_data=[diff_flux_ax8.T, diff_flux_ax9.T],
     plot_types="loglog",line_types="solid",xlabel=f"Energy ({energy_units})",
     ylabel=f"Differential flux ({diff_units})",
     labels=['AP8', 'IRENE 50%'], title="AP8 vs. IRENE")

# Plot all percentiles

diff_flux_ax9_all_pc = irene_results['orbit_averaged_proton_spectrum'].get_variable_data(variable_name="differential_flux")

plot(x_data=[energy_ax9.T],
     y_data=[diff_flux_ax9_all_pc.T],
     plot_types="loglog",line_types="solid", xlabel=f"Energy ({energy_units})",
     ylabel=f"Differential flux ({diff_units})",
     labels=['50%, 75%, 90%'], title="IRENE by percentile")


mjd_over_traj_ax9 = irene_results['proton_spectrum_over_trajectory'].get_variable_data(variable_name="mjd").squeeze()
irene_diff_flux_timeseries = irene_results['proton_spectrum_over_trajectory']

diff_flux_over_traj_ax9 = irene_diff_flux_timeseries.get_variable_data(variable_name="differential_flux")


datetimes = mjd50_to_datetime(mjd.squeeze()[1::40])

diff_flux_over_traj_ax9_50pc = diff_flux_over_traj_ax9[0]

plot(x_data=[mjd_over_traj_ax9.T],
     y_data=[diff_flux_over_traj_ax9_50pc],
     plot_types="semilogy",line_types="solid",xlabel="MJD (hr)",
     ylabel=f"Differential flux ({diff_units})",
     labels=energy_ax9.squeeze().T, title="IRENE")

plot(x_data=[energy_ax9.T],
     y_data=[diff_flux_over_traj_ax9_50pc[1::40].T],
     plot_types="loglog",line_types="solid",xlabel=f"Energy ({energy_units})",
     ylabel=f"Differential flux ({diff_units})",
     labels=[datetimes], title="IRENE")