# coding=utf-8
"""
title: Simple GCR model example
description: Comparing two GCR models, with and without geomagnetic shielding
keywords: GCR, greet, iso_15390, dlr-gcr, to_units
"""

import matplotlib.pyplot as plt

from nom_client.nom_client import NoMClient

conf = {
    "servers": {
        "ext_rest_server": {
            "type": "rest",
            "server_addr": "https://nom.esa.int",
            "api_version": "api",
            "api_key": ""
        }
    }
}

nom_client = NoMClient("GCR models examples", client_configuration=conf, default_server_id="ext_rest_server")

greet_model = nom_client.get_model('greet')
greet_model.set_params(apogeeAltitude=200, perigeeAltitude=200)
greet_result = nom_client.run_model(model=greet_model)
print(greet_result)

""" ISO GCR Model """

iso_model_no_geo_shielding = nom_client.get_model('iso_15390')
iso_model_no_geo_shielding.set_params(useGeomagneticShielding=0)
iso_model_no_geo_shielding.set_external_input(external_input_name="trajectory", external_input=greet_result)
iso_no_geo_shielding_result = nom_client.run_model(model=iso_model_no_geo_shielding)
print(iso_no_geo_shielding_result)

iso_model_geo_shielding = nom_client.get_model('iso_15390')
iso_model_geo_shielding.set_params(useGeomagneticShielding=1)
iso_model_geo_shielding.set_external_input(external_input_name="trajectory", external_input=greet_result)
iso_geo_shielding_result = nom_client.run_model(model=iso_model_geo_shielding)
print(iso_geo_shielding_result)

""" DLR GCR Model """

dlr_no_geo_shielding = nom_client.get_model('dlr-gcr')
dlr_no_geo_shielding.set_params(msm_enable=False)
dlr_no_geo_shielding.set_external_input(external_input_name="trajectory", external_input=greet_result)
dlr_no_geo_shielding_result = nom_client.run_model(model=dlr_no_geo_shielding)
print(dlr_no_geo_shielding_result)

dlr_geo_shielding = nom_client.get_model('dlr-gcr')
dlr_geo_shielding.set_params(msm_enable=True)
dlr_geo_shielding.set_external_input(external_input_name="trajectory", external_input=greet_result)
dlr_geo_shielding_result = nom_client.run_model(model=dlr_geo_shielding)
print(dlr_geo_shielding_result)

iso_gcr_spectrum_no_geo_shielding = iso_no_geo_shielding_result.get_model_result_by_name(
    result_name="particle_spectrum")
iso_gcr_spectrum_geo_shielding = iso_geo_shielding_result.get_model_result_by_name(result_name="particle_spectrum")

dlr_spectrum_no_geo_shielding = dlr_no_geo_shielding_result.get_model_result_by_name(result_name="particle_spectrum")
dlr_spectrum_geo_shielding = dlr_geo_shielding_result.get_model_result_by_name(result_name="particle_spectrum")

energy_iso = iso_gcr_spectrum_no_geo_shielding.get_variable_data(variable_name="energy")
energy_dlr = dlr_spectrum_no_geo_shielding.get_variable_data(variable_name="energy")

int_flux_iso_no_geo_shielding = iso_gcr_spectrum_no_geo_shielding.get_variable_data(variable_name="integral_flux",
                                                                                    species="H")
int_flux_iso_geo_shielding = iso_gcr_spectrum_geo_shielding.get_variable_data(variable_name="integral_flux",
                                                                              species="H")

int_flux_dlr_no_geo_shielding_10 = dlr_spectrum_no_geo_shielding.get_variable_data(variable_name="integral_flux",
                                                                                   to_units="m^-2 s^-1 sr^-1",
                                                                                   quantile=5)
int_flux_dlr_no_geo_shielding_90 = dlr_spectrum_no_geo_shielding.get_variable_data(variable_name="integral_flux",
                                                                                   to_units="m^-2 s^-1 sr^-1",
                                                                                   quantile=5)

int_flux_dlr_geo_shielding_10 = dlr_spectrum_geo_shielding.get_variable_data(variable_name="integral_flux",
                                                                             to_units="m^-2 s^-1 sr^-1", quantile=5)
int_flux_dlr_geo_shielding_90 = dlr_spectrum_geo_shielding.get_variable_data(variable_name="integral_flux",
                                                                             to_units="m^-2 s^-1 sr^-1", quantile=5)

plt.rcParams['figure.figsize'] = [15, 10]
plt.rc('font', size=18)
plt.semilogy(energy_iso.T, int_flux_iso_no_geo_shielding.T, '-', linewidth=2.0, markersize=12.0,
             label='ISO (no GEO shielding)')
plt.semilogy(energy_iso.T, int_flux_iso_geo_shielding.T, '-', linewidth=2.0, markersize=12.0,
             label='ISO (GEO shielding)')

plt.semilogy(energy_dlr.T, int_flux_dlr_no_geo_shielding_90.T, '--', linewidth=2.0, markersize=12.0,
             label='DLR  (no GEO shielding)')
plt.semilogy(energy_dlr.T, int_flux_dlr_geo_shielding_90.T, '--', linewidth=2.0, markersize=12.0,
             label='DLR  (GEO shielding)')

plt.title('GCR integral spectrum', 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.xlim(0, 40000)
plt.show()