# 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("IEM", debug_output=False)

# Set the NoM server
nom_client.set_default_server("local_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)

iem_model = nom_client.get_model('iem')
iem_model.set_params(mode="pk_wc_fluence")
iem_results = nom_client.run_model(iem_model)
print(iem_results)

summary_equivalent_fluences_pmax = iem_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")

print("Running SD2 model")
sd2_model = nom_client.get_model('sd2')
sd2_model.set_external_input(external_input_name="trappedParticleSpectrum", external_input=iem_results)
sd2_model.set_params(detectorMaterial=1,
                     shieldDepthValues=1,
                     userShieldDepths=[0.05, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.8, 1,
                                       1.5, 2, 2.5, 3, 4, 5, 6, 7, 8, 9, 10, 12, 14, 16, 18, 20])
sd2_result = nom_client.run_model(sd2_model)
print(sd2_result)