# coding=utf-8
import os 
dir_path = os.path.dirname(os.path.realpath(__file__))

import nom_client as nc
from nom_client.nom_client import NoMClient
print(nc.__file__)

print('NoM Client Version: ' + nc.__version__)

server_id = 'ext_rest_server'
server_id = 'local_server'
nom_client = NoMClient("LARB-example", default_server_id=server_id)

if False:
    trajectory_model = nom_client.get_model('trajectory')
    trajectory_model.set_params(orbitSpecificationCode=205,
                                epoch="2020-01-01 00:00:00",
                                missionDuration=10,
                                naifCode=399,
                                inclination=60,
                                perigeeAltitude=500,
                                apogeeAltitude=600)
    trajectory_result = nom_client.run_model(trajectory_model)
else:
    trajectory_model = nom_client.get_model('sapre')
    trajectory_model.set_params(
        planet=3,
        missionDurationUnits=1,
        missionDuration=11,
        orbitType='GEN', altitudeSpecification=2, altitude=500, inclination=60.0,
        trajectoryDuration=1,
        trajectoryFlag=1
    )
    trajectory_result = nom_client.run_model(model=trajectory_model)

print(trajectory_result)

print("")
print(".... getting LARB ....")

local_larb_model = nom_client.get_model('trep_larb')  # larb_model_spec )

local_larb_model.set_params(F107=210, F107a=230)
local_larb_model.set_external_input(
    external_input_name='trajectory', external_input=trajectory_result)
#                        
resLarb = nom_client.run_model(model=local_larb_model)
nom_client.write_output_files( 'trep_larb', location=dir_path)
print("")

print(resLarb)
outLarb = resLarb['orbit_averaged_proton_spectrum']

energy = outLarb.get_variable_data('energy').squeeze()
iflux = outLarb.get_variable_data('integral_flux').squeeze()
dflux = outLarb.get_variable_data('differential_flux').squeeze()

for i, e in enumerate(energy):
    print("{0:8.2f}, {1:10.3e}, {2:10.3e}".format(e, iflux[i], dflux[i]))

print("Running SD2 model")
sd2_model = nom_client.get_model('sd2')
sd2_model.set_external_input(external_input_name="trappedParticleSpectrum", external_input=resLarb)
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)