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("TREP model example", default_server_id="ext_rest_server")

sapre_model = nom_client.get_model('sapre')
sapre_model.set_params(orbitType="GEN")
sapre_result = nom_client.run_model(sapre_model)


print("Running trapped model")
trep_model = nom_client.get_model('trep')
trep_model.set_params(trappedProtonModel=1, trappedElectronModel=1)
trep_model.set_external_input(external_input_name="trajectory", external_input=sapre_result)
trep_results = nom_client.run_model(trep_model)
print(trep_results)


energy1_ax8 = trep_results['orbit_averaged_proton_spectrum'].get_variable_data(variable_name="energy")
int_flux1_ax8 = trep_results['orbit_averaged_proton_spectrum'].get_variable_data(variable_name="integral_flux")
diff_flux1_ax8 = trep_results['orbit_averaged_proton_spectrum'].get_variable_data(variable_name="differential_flux")



print("Running trapped model")
trep_model = nom_client.get_model('trep')
trep_model.set_params(trappedProtonModel=2, trappedElectronModel=1)
trep_model.set_external_input(external_input_name="trajectory", external_input=sapre_result)
trep_results = nom_client.run_model(trep_model)
print(trep_results)


energy2_ax8 = trep_results['orbit_averaged_proton_spectrum'].get_variable_data(variable_name="energy")
int_flux2_ax8 = trep_results['orbit_averaged_proton_spectrum'].get_variable_data(variable_name="integral_flux")
diff_flux2_ax8 = trep_results['orbit_averaged_proton_spectrum'].get_variable_data(variable_name="differential_flux")





plot(x_data=[energy1_ax8.T, energy2_ax8.T],
     y_data=[diff_flux1_ax8.T, diff_flux2_ax8.T],
     plot_types="loglog",line_types="solid",xlabel="Energy (MeV)",
     ylabel=f"Differential flux ()",
     labels=['AP8', 's'], title="AP8")