# coding=utf-8
import matplotlib.pyplot as plt

from nom_client.nom_client import NoMClient


def plot(x_data, y_data, plot_types, line_types, xlabel, ylabel, labels, title, xlim=None, ylim=None):
    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=14)
    plt.grid(True)
    plt.legend(loc=1)

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

    if xlim is not None:
        plt.xlim(xlim)

    if ylim is not None:
        plt.ylim(ylim)

    plt.show()


nom_client = NoMClient(project_name="Abundance model example")

trajectory_model = nom_client.get_model('trajectory')
trajectory_model.set_params(orbitSpecificationCode="NEI",missionDuration=1)
trajectory_result = nom_client.run_model(trajectory_model)
print(trajectory_result)

gcr_model = nom_client.get_model('creme-96-gcr')
gcr_model.set_external_input(external_input_name="trajectory", external_input=trajectory_result)
gcr_results = nom_client.run_model(gcr_model)
gcr_results = gcr_results['particle_spectrum']
species = list(range(1,93,1))
diff_fluxes = gcr_results.get_variable_data(variable_name="differential_flux", energy=7.1).squeeze()

he_flux = diff_fluxes[1]

diff_fluxes /= he_flux
diff_fluxes[0] = 0
diff_fluxes[1] = 0

# Get the model specification from the server.
abundance_ratio_model = nom_client.get_model('abundances')

energies = [
    7,
    8,
    12,
    18,
    26,
    38,
    55,
    80,
    115,
    170,
    240
]

abundance_ratio_model.set_params(energies=energies)
run_results = nom_client.run_model(abundance_ratio_model)


results = run_results['abundances']

print("----------------")
print(results.get_variable_data("abundance"))
print("----------------")
print(results.get_variable_data("abundance", z=3))
print("----------------")
print(results.get_variable_data("abundance", z=4))
print("----------------")

z = results.get_variable_data(variable_name="z")
abundance = results.get_variable_data(variable_name="abundance", energy=7)

plot(x_data=[z.T, species],
     y_data=[abundance.T, diff_fluxes.T],
     plot_types="semilogy",line_types="solid",xlabel="Atomic number (z)",
     ylabel=f"Abundance (%)",
     labels=['Abundance model (7 MeV)', 'GCR flux ratio (7.1 MeV)'], title="Abundance by species", xlim=[0,30])