# coding=utf-8
import matplotlib.pyplot as plt
import numpy as np

from nom_client.nom_client import NoMClient

conf = {
    "servers": {
        "local_server": {
            "type": "local_server",
            "models_directory": ["/home/simon/PycharmProjects/nom-models", "/home/simon/PycharmProjects/geant4-models"],
            "working_directory": "/home/simon/nom/working"
        },
        "local_rest_server": {
            "type": "rest",
            "server_addr": "http://localhost:5002",
            "api_version": "api",
            "api_key": "6b328f7a243d307a3f0803c1deff1f95c2ce74852e2110f8263911deb09a9801"
        },
        "ext_rest_server": {
            "type": "rest",
            "server_addr": "https://nom.esa.int",
            "api_version": "api",
            "api_key": "6b328f7a243d307a3f0803c1deff1f95c2ce74852e2110f8263911deb09a9801"
        }
    }
}

nom_client = NoMClient("MCICT-standalone", client_configuration=conf)

nom_client.set_default_server("local_server")

mission_dur_days = 20

sapre_model = nom_client.get_model('sapre')
sapre_model.set_params(orbitType="GEN", missionDuration=mission_dur_days, missionDurationUnits=1,
                       perigeeAltitude=18000, apogeeAltitude=18000, trajectoryFlag=0, trajectoryOrbits=30)
sapre_result = nom_client.run_model(sapre_model)
print(sapre_result)

envData = ["Mono, 10., 1.0"]

print("Running trapped model")
ap8ae8_model = nom_client.get_model('ae8ap8')

ap8ae8_model.set_external_input(external_input_name="trajectory", external_input=sapre_result)
ap8ae8_results = nom_client.run_model(ap8ae8_model)
print(ap8ae8_results)

mcict = nom_client.get_model('mcict')

mcict.set_params(particleSpecies="e-", geometryConfiguration="slab", observationStepStart=10,
                              observationStepEnd=mission_dur_days*86500, observationStepStep=3000, fluxData=envData)

mcict.new_input_group(input_group_name="input_spectra", unique_group_name="trapped_spectrum", params={
    "spectrum": ap8ae8_results,
    # "resultName": "electron_spectrum_over_trajectory",
    "resultName": "orbit_averaged_electron_spectrum",
    "agregationPeriod": 1
})

"""
L1 , 1.064,  G4_MYLAR , 300,  D , 20
L2 , 0.001,  G4_Al , 300,  G , 1 
"""

mcict.new_input_group(input_group_name="layers", unique_group_name="L1",
                                   params={
                                       "name": "layer1",
                                       "thickness": 1.064,
                                       "material": "G4_MYLAR",
                                       "temperature": 200,
                                       "type": "D",
                                       "repeat": 20
                                   })

mcict.new_input_group(input_group_name="layers", unique_group_name="L2",
                                   params={
                                       "name": "layer2",
                                       "thickness": 0.001,
                                       "material": "G4_Al",
                                       "temperature": 200,
                                       "type": "G",
                                       "repeat": 1
                                   })

mcict.new_input_group("dielectrics",
                                   params={
                                       "material": "G4_MYLAR",
                                       "permitivity": 4.7,
                                       "darkConductivity": 1.48e-18,
                                       "kp": 9.84e-14,
                                       "delta": 1.77,
                                       "ea": 0.42
                                   })

run1 = nom_client.run_model(mcict)
print(run1)

for result_name, result in run1.items():
    print(result)



#
electric_field = run1.get_model_result_by_name(result_name="electric_field")
time_vals = electric_field.get_variable_data(variable_name="time") / 86400
electric_field_vals = electric_field.get_variable_data(variable_name="electric_field", layer="layer1_0")
#
#
plt.rcParams['figure.figsize'] = [15, 10]
plt.rc('font', size=18)


plt.semilogy(time_vals.squeeze().T, electric_field_vals.squeeze().T, '.', linewidth=2.0, markersize=12.0,
             label='dose_rate_vals')


plt.title('SEP User Defined spectrum', 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.xlim(0, 10)
plt.xlabel("Energy (MeV)")
plt.ylabel("Flux (MeV)")
plt.show()
         

from typing import TypeVar, Union
import numpy as np
import numpy.typing as npt

T1 = TypeVar("T1", bound=npt.NBitBase)
T2 = TypeVar("T2", bound=npt.NBitBase)

def add(a: np.floating[T1], b: np.integer[T2]) -> np.floating[Union[T1, T2]]:
    return a + b

i8: np.int64
i4: np.int32
f8: np.float64
f4: np.float32

reveal_type(add(f8, i8))  # E: {float64}
reveal_type(add(f4, i8))  # E: {float64}
reveal_type(add(f8, i4))  # E: {float64}
reveal_type(add(f4, i4))  # E: {float32}