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

from nom_client.nom_client import NoMClient


nom_client = NoMClient("SAPPHIRE total fluence example")

sapre_model_gto = nom_client.get_model('sapre')
sapre_model_gto.set_params(orbitType="GEN")
sapre_gto_result = nom_client.run_model(sapre_model_gto)
print(sapre_gto_result)

# Get the model specification from the server.
sapphire_model = nom_client.get_model('sapphire-total-fluence')
sapphire_model.set_external_input(external_input_name="trajectory", external_input=sapre_gto_result)

# Geomagnetic shielding
sapphire_results_no_geo_sheilding = nom_client.run_model(sapphire_model)
sapphire_model.set_params(useGeomagneticShielding=1)
sapphire_results_geo_sheilding = nom_client.run_model(sapphire_model)

sep_proton_spectrum_no_geo_sheilding = sapphire_results_no_geo_sheilding.get_model_result_by_name(result_name="sep_proton_spectrum")
differential_fluence_no_geo_sheilding = sep_proton_spectrum_no_geo_sheilding.get_variable_data(variable_name="differential_fluence")

sep_proton_spectrum_geo_sheilding = sapphire_results_geo_sheilding.get_model_result_by_name(result_name="sep_proton_spectrum")
differential_fluence_geo_sheilding = sep_proton_spectrum_geo_sheilding.get_variable_data(variable_name="differential_fluence")

energy = sep_proton_spectrum_no_geo_sheilding.get_variable_data(variable_name="energy")[0] # same energies with or without geo shielding
energy = np.array(energy)


fig, ax = plt.subplots()
flux_units = sep_proton_spectrum_no_geo_sheilding['differential_fluence'].units
for fl_data in [differential_fluence_no_geo_sheilding, differential_fluence_geo_sheilding]:
    plt.loglog(energy, fl_data.T, '-', linewidth=2.0, markersize=12.0)

plt.rc('font', size=12)
plt.grid(True)
plt.legend(labels=["No shielding (default)", "Shielding"])
plt.xlabel("Energy")
plt.ylabel(f"Flux ({flux_units})")
plt.show()


# Datasets

datasets = ["sapphire", "sapphire-hex"]
dataset_labels = ["sapphire (default)", "sapphire-hex"]

flux_results_per_dataset = []
for dataset in datasets:
    sapphire_model.set_params(dataSet=dataset)
    results = nom_client.run_model(sapphire_model)
    sep_proton_spectrum = results.get_model_result_by_name(result_name="sep_proton_spectrum")
    flux = sep_proton_spectrum.get_variable_data(variable_name="differential_fluence")
    flux_results_per_dataset.append(flux)


for fl_data in flux_results_per_dataset:
    plt.loglog(energy, fl_data.T, '-', linewidth=2.0, markersize=12.0)

plt.rc('font', size=12)
plt.grid(True)
plt.legend(labels=dataset_labels)
plt.xlabel("Energy")
plt.ylabel(f"Flux ({flux_units})")
plt.show()

# Confience levels

confidence_levels = [50, 70, 95, 99]
confidence_levels_labels = ["50%", "70%", "95% (default)", "99%"]

flux_results_per_conflevel = []
for confidence_level in confidence_levels:
    sapphire_model.set_params(confidenceLevel=confidence_level)
    results = nom_client.run_model(sapphire_model)
    sep_proton_spectrum = results.get_model_result_by_name(result_name="sep_proton_spectrum")
    flux = sep_proton_spectrum.get_variable_data(variable_name="differential_fluence")
    flux_results_per_conflevel.append(flux)


for fl_data in flux_results_per_conflevel:
    plt.loglog(energy, fl_data.T, '-', linewidth=2.0, markersize=12.0)

plt.rc('font', size=12)
plt.grid(True)
plt.legend(labels=confidence_levels_labels)
plt.xlabel("Energy")
plt.ylabel(f"Flux ({flux_units})")
plt.show()