# coding=utf-8

import matplotlib.pyplot as plt

from nom_client.nom_client import NoMClient


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


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("OMEP")

# Set the NoM server
nom_client.set_default_server("ext_rest_server")

sapre_model_gto = nom_client.get_model('sapre')
sapre_model_gto.set_params(orbitType="GEN", trajectoryFlag=1, missionDurationUnits=1,
                           trajectoryDuration=2, missionDuration=2.1)
sapre_gto_result = nom_client.run_model(sapre_model_gto)
print(sapre_gto_result)

omep_model = nom_client.get_model('omep')
omep_model.set_external_input(external_input_name="trajectory", external_input=sapre_gto_result)
omep_model.set_params(percentiles=[50,60,70,80])
omep_results = nom_client.run_model(omep_model)
print(omep_results)

summary_equivalent_fluences_pmax = omep_results.get_model_result_by_name(result_name="orbit_averaged_proton_spectrum")
energies = summary_equivalent_fluences_pmax.get_variable_data(variable_name="energy").squeeze()
percentiles = summary_equivalent_fluences_pmax.get_variable_data(variable_name="percentiles").squeeze()
diff_fluxes = summary_equivalent_fluences_pmax.get_variable_data(variable_name="differential_flux").squeeze()
diff_fluxes_units = summary_equivalent_fluences_pmax['differential_flux'].units

plot(x_data=[energies],
     y_data=[diff_fluxes.T],
     plot_types="loglog",line_types="solid",xlabel="Energy (MeV)",
     ylabel=f"Differential flux ({diff_fluxes_units})",
     labels=percentiles, title="OMEP")