# coding=utf-8

import matplotlib.pyplot as plt

from nom_client.nom_client import NoMClient


nom_client = NoMClient("trapped-user-defined-spectrum", default_server_id="ext_rest_server")

sapre_model = nom_client.get_model('sapre')
sapre_model.set_params(missionDuration=100, missionDurationUnits=1)
sapre_result = nom_client.run_model(model=sapre_model)

ae8ap8_model = nom_client.get_model('ae8ap8')
ae8ap8_model.set_external_input(external_input_name="trajectory", external_input=sapre_result)
ae8ap8_results = nom_client.run_model(ae8ap8_model)
print(ae8ap8_results)

proton_res = ae8ap8_results['orbit_averaged_proton_spectrum']
proton_energies = proton_res['energy'].values
proton_integral_fluence = proton_res['integral_flux'].values
proton_differential_fluence = proton_res['differential_flux'].values

electron_res = ae8ap8_results['orbit_averaged_electron_spectrum']
electron_energies = electron_res['energy'].values
electron_integral_fluence = electron_res['integral_flux'].values
electron_differential_fluence = electron_res['differential_flux'].values

ud_trapped_model = nom_client.get_model('trapped-spectrum-ud')

ud_trapped_model.set_params(protonEnergies=proton_energies.squeeze(),
                            integralProtonData=proton_integral_fluence.squeeze(),
                            differentialProtonData=proton_differential_fluence.squeeze(),
                            electronEnergies=electron_energies.squeeze(),
                            integralElectronData=electron_integral_fluence.squeeze(),
                            differentialElectronData=electron_differential_fluence.squeeze(),
                            trajectoryDuration=100)

ud_trapped_results = nom_client.run_model(ud_trapped_model)
print(ud_trapped_results)

sd_model = nom_client.get_model('sd2')
sd_model.set_external_input(external_input_name="trappedParticleSpectrum", external_input=ae8ap8_results)
sd_result = nom_client.run_model(sd_model)
print(sd_result)

sd2_model = nom_client.get_model('sd2')
sd2_model.set_external_input(external_input_name="trappedParticleSpectrum", external_input=ud_trapped_results)
sd2_result = nom_client.run_model(sd2_model)
print(sd2_result)

eqflux_model = nom_client.get_model('eqflux')
eqflux_model.set_external_input(external_input_name="trappedParticleSpectrum",
                                external_input=ud_trapped_results)
eqflux_result = nom_client.run_model(eqflux_model)
print(eqflux_result)

depth1 = sd_result['dose_depth_curve'].get_variable_data(variable_name="thickness")
dose1 = sd_result['dose_depth_curve'].get_variable_data(variable_name="absorbed_dose")

depth2 = sd2_result['dose_depth_curve'].get_variable_data(variable_name="thickness")
dose2 = sd2_result['dose_depth_curve'].get_variable_data(variable_name="absorbed_dose")

plt.rcParams['figure.figsize'] = [15, 10]
plt.rc('font', size=18)

plt.semilogy(depth1.T, dose1[0], linewidth=2.0, markersize=12.0, label='SPENVIS models')

plt.semilogy(depth1.T, dose2[0], linewidth=2.0, markersize=12.0, label='User defined')

plt.title('Trapped Dose depth', 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("Thickness (mm Al)")
plt.ylabel("Dose (rad(Si))")
plt.show()