# coding=utf-8
"""
category: tutorial
title: Model discovery and model exploration
description: Explore the models available, check what inputs it needs and get results from a model
keywords: niel, ae8ap8, trajectory, describe_inputs, get_external_inputs, get_result_names, get_model_result_by_name, get_variable_data, slicing data
"""
from pprint import pprint

from nom_client.nom_client import NoMClient


"""
Initialise the NoM Client with a project name (optional) and the configuration data.
"""
nom_client = NoMClient(project_name="Model discovery and model exploration")


"""
Determine the status of the NoM Server(s) specified in the configuration.
"""
print(nom_client.status())


"""
Find models with category 'radiation'
A dictionary is returned with model name as keys.
"""

print("Radiation related models")
print("---------------------------")

with_radiation = nom_client.list_models(categories=["radiation"])

for model_name, model_description in with_radiation.items():
    print(f"{model_name:30}: {model_description}")

print("========================")
print("Effects models")
print("---------------------------")
"""
Find models with category 'effects'
"""
with_effects = nom_client.list_models(categories=["effects"])

for model_name, model_description in with_effects.items():
    print(f"{model_name:30}: {model_description}")

"""
Find models with category 'radiation' AND 'effects'
"""

print("========================")
print("Both radiation and effects models")
print("---------------------------")

both = {k:with_radiation.get(k,with_effects.get(k)) for k in set(with_radiation) & set(with_effects)}
for model_name, model_description in both.items():
    print(f"{model_name:30}: {model_description}")

"""
Now lets explore an individual model, the NIEL model

You can list the model inputs and their defaults ...
"""
niel_model = nom_client.get_model(model_name='niel')

# To describe all inputs
print(niel_model.describe_inputs())

# To describe a single input
print(niel_model.describe_inputs(input_name="damageCurve"))


# Lets set the damageCurve to Huhtinen Si
niel_model.set_params(damageCurve=4)

# The NIEL model requires some external inputs to be satisfied before we can run it

pprint(niel_model.get_external_inputs())

"""
Lets give the NIEL model a trapped particle spectrum using the ae8ap8 model. Similar investigation of the ae8ap8 model
requires us to give the ae8ap8 model a trajectory
"""
trajectory_model = nom_client.get_model(model_name="trajectory")
trajectory_result = nom_client.run_model(model=trajectory_model)

ae8ap8_model = nom_client.get_model(model_name="ae8ap8")
ae8ap8_model.set_external_input(external_input_name="trajectory", external_input=trajectory_result)
ae8ap8_result = nom_client.run_model(model=ae8ap8_model)


""" 
OK, now we can run the NIEL model after we add the trapped particle spectrum external input
Run the model
"""
niel_model.set_external_input(external_input_name="trappedParticleSpectrum", external_input=ae8ap8_result)
niel_results = nom_client.run_model(model=niel_model)
print(niel_results)


"""
List all of the model results from this model. This can be used to get the approriate model result 
from the results returned by the model
"""
print(niel_model.get_result_names())


# Lets get the 'shielded_trapped_proton_spectrum'
niel_result_data = niel_results.get_model_result_by_name(result_name="shielded_trapped_proton_spectrum")


# List all the data variable names for the 'shielded_trapped_proton_spectrum' result
pprint(niel_model.get_result_data_names(result_name="shielded_trapped_proton_spectrum"))

# Lets get the 'integral_fluence' data
integral_flux = niel_result_data.get_variable_data(
    variable_name="integral_fluence")

print(integral_flux)

# The integral fluence data is a 2-dimensional, as can be confirmed by:
print(niel_result_data.get_variable_dimension(variable_name="integral_fluence"))

# We can explore this variable further using:
print(niel_result_data.get_variable_axes(variable_name="integral_fluence"))

"""
>>> {'integral_fluence': {'row_varying': 'energy', 'non_row_varying': ['thickness']}}
The output indicates there are two dimensions to the integral fluence data, energy and thickness.
The energy axis is a row-varying variable and represents the independent variable.
The thickness axis is non-row varying and there will be a row of integral fluence data (over energy) for each thickness value
within the thickness variable
"""

# If we look at the shape of the integral fluence data (we can use the shape property as it is a numpy array
int_fluence_shape = integral_flux.shape
print(int_fluence_shape)
"""
>>> (26, 29)
So we have 29 (energies) and 26 (thicknesses) of data
"""

# We can grab the energy and thicnkess variables
energies = niel_result_data.get_variable_data(variable_name="energy").squeeze()
print(f"# of energies: {len(energies)}")
print(energies)
thicknesses = niel_result_data.get_variable_data(variable_name="thickness").squeeze()
print(f"# of thicknesses: {len(thicknesses)}")
print(thicknesses)

# You can also slice the data using the values of the axis data, e.g.
print(niel_result_data.get_variable_data(
      variable_name="integral_fluence", thickness="0.050 mm"))

# integral flu value for energy 1e-1 MeV and thickness 0.050 mm
print(niel_result_data.get_variable_data(
      variable_name="integral_fluence", thickness="0.050 mm",
                                        energy=1.0e-01))