# coding=utf-8
import json

from nom_client.nom_client import NoMClient

nom_client = NoMClient(project_name="SEE tool example",
                       default_server_id="local_server")

trajectory_model = nom_client.get_model('trajectory')

trajectory_model.set_params(orbitSpecificationCode=210,
                            epoch="2020-01-01 00:00:00",
                            startDate="2020-01-01 00:00:00",
                            missionDuration=1,
                            naifCode=399,
                            circularAltitude=200000)

greet_result = nom_client.run_model(trajectory_model)
print(greet_result)

sapphire_total_fluence_model = nom_client.get_model('sapphire-total-fluence')
sapphire_total_fluence_model.set_external_input(external_input_name="trajectory", external_input=greet_result)
sapphire_total_fluence_results = nom_client.run_model(sapphire_total_fluence_model)
print(sapphire_total_fluence_results)

flux_data_example = \
    [
        [  # time step 1
            [2e11, 2e8, 2e6, 20000, 2000, 200], # particle Z 1 of len(energies)
            [1e11, 1e8, 2e6, 10000, 1000, 100]  # particle Z 2 of len(energies)
        ],
        [  # time step 2
            [2e11, 2e8, 2e6, 20000, 2000, 200], # particle Z 1 of len(energies)
            [1e11, 1e8, 2e6, 10000, 1000, 100]  # particle Z 2 of len(energies)
        ],
        [  # time step 3
            [2e11, 2e8, 2e6, 20000, 2000, 200], # particle Z 1 of len(energies)
            [1e11, 1e8, 2e6, 10000, 1000, 100]  # particle Z 2 of len(energies)
        ]
    ]

spec1_integral_fluence = sapphire_total_fluence_results["full_ion_spectrum"] \
    .get_variable_data(variable_name="integral_fluence")

spec1_differential_fluence = sapphire_total_fluence_results["full_ion_spectrum"] \
    .get_variable_data(variable_name="differential_fluence")

spec1_species = sapphire_total_fluence_results["full_ion_spectrum"]["species"].values[0]
spec1_energies = sapphire_total_fluence_results["full_ion_spectrum"]["energy"].values[0]

# full_ion_spectrum
devices_array = [
    {
        "deviceName": "dev1",
        "directIonisationMethod": -1,
        "weibull_w_direct": 28,
        "weibull_let_th_direct": 0.4,
        "weibull_s_direct": 1.5,
        "weibull_sigma_direct": 6e-8,
        "protonInducedMethod": 2,
        "weibull_w_proton": 45,
        "weibull_energy_th_proton": 5,
        "weibull_s_proton": 1.5,
        "weibull_sigma_proton": 3.6e-14
    },
    {
        "deviceName": "dev2",
        "directIonisationMethod": -1,
        "weibull_w_direct": 28,
        "weibull_let_th_direct": 0.4,
        "weibull_s_direct": 1.5,
        "weibull_sigma_direct": 6e-8,
        "protonInducedMethod": 2,
        "weibull_w_proton": 45,
        "weibull_energy_th_proton": 5,
        "weibull_s_proton": 1.5,
        "weibull_sigma_proton": 3.6e-14
    }
]

energies = [0.1, 0.11, 0.12, 0.14, 0.16, 0.18, 0.2, 0.22, 0.25, 0.28, 0.32, 0.35, 0.4, 0.45, 0.5, 0.55, 0.63, 0.71, 0.8,
            0.9, 1.0, 1.1, 1.2, 1.4, 1.6, 1.8, 2.0, 2.2, 2.5, 2.8, 3.2, 3.5, 4.0, 4.5, 5.0, 5.5, 6.3, 7.1, 8.0, 9.0,
            10.0, 11.0, 12.0, 14.0, 16.0, 18.0, 20.0, 22.0, 25.0, 28.0, 32.0, 35.0, 40.0, 45.0, 50.0, 55.0, 63.0, 71.0,
            80.0, 90.0, 100.0, 110.0, 120.0, 140.0, 160.0, 180.0, 200.0, 220.0, 250.0, 280.0, 320.0, 350.0, 400.0,
            450.0, 500.0, 550.0, 630.0, 710.0, 800.0, 900.0, 1000.0]

timesteps = ["2022-01-01T00:00", "2022-01-01T00:05"]

species = "range(1,2,3-92)"

int_flux = [spec1_integral_fluence, spec1_integral_fluence]
diff_flux = [spec1_differential_fluence, spec1_differential_fluence]

see_timeseries_new_model = nom_client.get_model('see')
see_timeseries_new_model.set_external_input(external_input_name="trajectory", external_input=greet_result)
see_timeseries_new_model.set_params(
    timesteps=timesteps,
    energies=energies,
    speciesZ=species,
    differentialData=diff_flux,
    integralData=int_flux,
    shieldThicknesses=.1,
    shieldThicknessUnit=2,
    deviceData=devices_array
)

see_timeseries_new_results = nom_client.run_model(see_timeseries_new_model)
print(see_timeseries_new_results)