# 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)

         

# coding=utf-8
import json
import math
import sys
from os import path

import numpy as np
import matplotlib.pyplot as plt

from nom_client.nom_client import NoMClient

def plot(x_data, y_data, plot_types, line_types, xlabel, ylabel, labels, title):
    for data_index in range(len(x_data)):
        if plot_types[data_index] == 'loglog':
            plt.loglog(x_data[data_index], y_data[data_index], linestyle=line_types[data_index], linewidth=2.0,
                       markersize=12.0,
                       label=labels[data_index])
        elif plot_types[data_index] == 'semilogy':
            plt.semilogy(x_data[data_index], y_data[data_index], linestyle=line_types[data_index], linewidth=2.0,
                         markersize=12.0,
                         label=labels[data_index])
        elif plot_types[data_index] == 'semilogx':
            plt.semilogx(x_data[data_index], y_data[data_index], linestyle=line_types[data_index], linewidth=2.0,
                         markersize=12.0,
                         label=labels[data_index])
        else:
            plt.plot(x_data[data_index], y_data[data_index], linestyle=line_types[data_index], linewidth=2.0,
                     markersize=12.0,
                     label=labels[data_index])

    plt.title(title, fontsize=14)
    plt.grid(True)
    plt.legend(bbox_to_anchor=(0.95, 0.95), loc=1, borderaxespad=0.,
               fancybox=True, fontsize=14, shadow=True)

    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plt.show()



a = [[1, 2, 3, 4, 5], [1, 2, 3, 4, 5]]
a = np.array(a)
print(a.shape)

#sys.exit(0)


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

import csv

gcr_energies = []
diff_data = []
int_data = []


sapre_model = nom_client.get_model('sapre')
sapre_model.set_params(orbitType="GEO", orbitStartYear=1996, orbitStartMonth=1, orbitStartDay=1)
#sapre_model.set_params(orbitSpecificationCode="GEO", epoch="1996-01-01", startDate="1996-01-01")
sapre_result = nom_client.run_model(sapre_model)
print(sapre_result)


iso_15390_model = nom_client.get_model("iso_15390")
iso_15390_model.set_params(lowestIonSpecies=1, highestIonSpecies=1)
iso_15390_model.set_external_input(external_input_name="trajectory", external_input=sapre_result)
iso_15390_res = nom_client.run_model(iso_15390_model)
print(iso_15390_res)

gcr_energies_data = iso_15390_res['particle_spectrum']["energy"].values.squeeze()
gcr_differential_flux_data = iso_15390_res['particle_spectrum']["differential_flux"].values
gcr_integral_flux_data = iso_15390_res['particle_spectrum']["integral_flux"].values



plot(x_data=[gcr_energies_data],
         y_data=[gcr_differential_flux_data[0]],
         plot_types=["loglog"],
         line_types=["solid"],
         xlabel="Energy (MeV)",
         ylabel="DFlux",
         labels=["GCR ISO"],
         title=" ")

here = path.abspath(path.dirname(__file__))
# with open(here + '/gcr_input.txt') as csv_file:
#     csv_reader = csv.reader(csv_file, delimiter=',')
#     line_count = 0
#
#     for row in csv_reader:
#         row = [float(x) for x in row]
#         gcr_energies.append(row[0])
#         diff_data.append(row[1:93])
#         int_data.append(row[93:185])
#
#         line_count += 1
#
# d= 3

# 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
    }
]



newupseto_model = nom_client.get_model('newupseto')
newupseto_model.set_external_input(external_input_name="trajectory", external_input=sapre_result)
newupseto_model.set_params(includeSEPSpectrum=0, includeTrappedSpectrum=0, includeGCRSpectrum=1)
newupseto_model.set_external_input(external_input_name="gcrSpectrum",
                                   external_input=iso_15390_res)

newupseto_model.new_input_group(input_group_name="device", unique_group_name="mydevice1",
                                params=devices_array[0]
   )

newupseto_results = nom_client.run_model(newupseto_model)


let_data_let = newupseto_results['let_spectrum']["let"].values.squeeze()
let_data_differential_flux_data = newupseto_results['let_spectrum']["differential_flux"].values
let_data_integral_flux_data = newupseto_results['let_spectrum']["integral_flux"].values



plot(x_data=[let_data_let],
         y_data=[let_data_differential_flux_data[0]],
         plot_types=["loglog"],
         line_types=["solid"],
         xlabel="Energy (MeV)",
         ylabel="DFlux",
         labels=["GCR ISO"],
         title=" ")



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

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

int_flux = [np.array(gcr_integral_flux_data.T)]
diff_flux = [np.array(gcr_differential_flux_data.T)]

see_timeseries_new_model = nom_client.get_model('see')
see_timeseries_new_model.set_external_input(external_input_name="trajectory", external_input=sapre_result)
see_timeseries_new_model.set_params(
    timesteps=timesteps,
    energies=gcr_energies_data,
    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)
seu_rates = see_timeseries_new_results['seu_rates']
direct_ionisation = see_timeseries_new_results['seu_rates'].direct_ionisation.values.squeeze()
proton_ionisation = see_timeseries_new_results['seu_rates'].proton_ionisation.values.squeeze()

print(f"{direct_ionisation} {direct_ionisation / 3.1577E-12} {3.1577E-12 / direct_ionisation}")
print(f"{proton_ionisation} {proton_ionisation / 1.4650E-13} {1.4650E-13 / proton_ionisation}")