import matplotlib.pyplot as plt

from nom_client.nom_client import NoMClient


def plot_efield(dictat_result):
    dictat_surface_voltage_over_time = dictat_result.get_model_result_by_name(result_name="surface_voltage_timeseries")
    time = dictat_surface_voltage_over_time.get_variable_data(variable_name="time")
    electric_field = dictat_surface_voltage_over_time.get_variable_data(variable_name="electric_field")

    plt.semilogy(time.T, electric_field.T, linewidth=2.0, markersize=12.0, label='Integral')

    plt.title('Electric field', fontsize=18)
    plt.grid(True)
    plt.xlabel("Time (hr)")
    plt.ylabel("Electric field (V/m)")
    plt.show()


nom_client = NoMClient("DICTAT example",
                       default_server_id="local_server",
                       debug_output=False)

sapre_model = nom_client.get_model('sapre')
sapre_model.set_params(orbitType="GEN", perigeeAltitude=600, apogeeAltitude=70000, trajectoryDuration=0.1, trajectoryFlag=1)
sapre_result = nom_client.run_model(sapre_model)
print(sapre_result)

# ---------------------------------------------------------------------------------------------------------

#print("Running DICTAT with single spectrum")

#dictat_model = nom_client.get_model('dictat')

# dictat_model.set_params(params={
#     "exposureDuration": 20,
#     "geometryType": "cylindrical",
#     "shieldMaterial": "Aluminium",
#     "shieldThickness": 0.03,
#     "dielectricMaterial": "Teflon",
#     "dielectricTemperature": 300,
#     "dielectricThickness": 0.5,
#     "numberOfGrounds": 1,
#     "groundType": "inner",
#     "coreMaterial": "Copper",
#     "coreRadius": 0.047,
#     "direction": "isotropic",
#     #"sequence": "NONE",
#     "energies": [4.00E-02, 7.00E-02, 1.00E-01, 2.50E-01, 5.00E-01, 7.50E-01, 1.00E+00, 1.50E+00, 2.00E+00, 2.50E+00,
#                  3.00E+00, 3.50E+00, 4.00E+00, 4.50E+00, 5.00E+00, 5.50E+00, 6.00E+00, 6.50E+00, 7.00E+00, 7.50E+00,
#                  8.00E+00, 8.50E+00, 1.00E+01],
#     "fluxes": [3.71E+07, 2.13E+07, 1.36E+07, 4.11E+06, 1.04E+06, 4.08E+05, 2.07E+05, 6.54E+04, 1.98E+04, 6.46E+03,
#                2.45E+03,
#                1.05E+03, 5.07E+02, 2.74E+02, 1.68E+02, 1.14E+02, 8.11E+01, 5.79E+01, 4.24E+01, 3.16E+01, 2.27E+01,
#                1.56E+01,
#                0.00E+00]
# })
#
# # sresult = nom_client.run_model(dictat_model)
# # print(sresult)
# # plot_efield(dictat_result=sresult)
#
#
# #sys.exit(0)
#
#
# # ---------------------------------------------------------------------------------------------------------
#
# print("Running DICTAT with timeseries spectrum")
#
# flux_timesseries = [
#     [0, 3.71E+07, 2.13E+07, 1.36E+07, 4.11E+06, 1.04E+06, 4.08E+05, 2.07E+05, 6.54E+04, 1.98E+04, 6.46E+03, 2.45E+03,
#      1.05E+03, 5.07E+02, 2.74E+02, 1.68E+02, 1.14E+02, 8.11E+01, 5.79E+01, 4.24E+01, 3.16E+01, 2.27E+01, 1.56E+01,
#      0.00E+00],
#     [12, 3.71E+08, 2.13E+08, 1.36E+08, 4.11E+07, 1.04E+07, 4.08E+06, 2.07E+06, 6.54E+05, 1.98E+05, 6.46E+04, 2.45E+04,
#      1.05E+04, 5.07E+03, 2.74E+03, 1.68E+03, 1.14E+03, 8.11E+02, 5.79E+02, 4.24E+02, 3.16E+02, 2.27E+02, 1.56E+02,
#      0.00E+00],
#     [24, 1.86E+08, 1.07E+08, 6.80E+07, 2.06E+07, 5.20E+06, 2.04E+06, 1.04E+06, 3.27E+05, 9.90E+04, 3.23E+04, 1.23E+04,
#      5.25E+03, 2.54E+03, 1.37E+03, 8.40E+02, 5.70E+02, 4.06E+02, 2.90E+02, 2.12E+02, 1.58E+02, 1.14E+02, 7.80E+01,
#      0.00E+00],
#     [36, 3.71E+06, 2.13E+06, 1.36E+06, 4.11E+05, 1.04E+05, 4.08E+04, 2.07E+04, 6.54E+03, 1.98E+03, 6.46E+02, 2.45E+02,
#      1.05E+02, 5.07E+01, 2.74E+01, 1.68E+01, 1.14E+01, 8.11E+00, 5.79E+00, 4.24E+00, 3.16E+00, 2.27E+00, 1.56E+00,
#      0.00E+00]]
#
# dictat_model = nom_client.get_model('dictat')
# dictat_model.set_params(params={
#     "exposureDuration": 20,
#     "geometryType": "cylindrical",
#     "shieldMaterial": "Aluminium",
#     "shieldThickness": 0.03,
#     "dielectricMaterial": "Teflon",
#     "dielectricTemperature": 300,
#     "dielectricThickness": 0.5,
#     "numberOfGrounds": 1,
#     "groundType": "inner",
#     "coreMaterial": "Copper",
#     "coreRadius": 0.047,
#     "direction": "isotropic",
#     "spectrumDefinition": "user-defined",
#     "energies": [4.00E-02, 7.00E-02, 1.00E-01, 2.50E-01, 5.00E-01, 7.50E-01, 1.00E+00, 1.50E+00, 2.00E+00, 2.50E+00,
#                  3.00E+00, 3.50E+00, 4.00E+00, 4.50E+00, 5.00E+00, 5.50E+00, 6.00E+00, 6.50E+00, 7.00E+00, 7.50E+00,
#                  8.00E+00, 8.50E+00, 1.00E+01],
#     "fluxTimeseries": flux_timesseries
# })
#
# sresult = nom_client.run_model(dictat_model)
# print(sresult)
# plot_efield(dictat_result=sresult)

print("Running Ae8")
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)

# ---------------------------------------------------------------------------------------------------------

# print("Running DICTAT using orbit average electron spectrum from Ae8")
#
# dictat_model = nom_client.get_model('dictat')
# dictat_model.set_external_input(external_input_name="electronSpectrum", external_input=ap8ae8_results,
#                                 result_name='orbit_averaged_electron_spectrum')
# dictat_model.set_params(params={
#     "geometryType": "cylindrical",
#     "shieldMaterial": "Aluminium",
#     "shieldThickness": 0.03,
#     "dielectricMaterial": "Teflon",
#     "dielectricTemperature": 300,
#     "dielectricThickness": 0.5,
#     "numberOfGrounds": 1,
#     "groundType": "inner",
#     "coreMaterial": "Copper",
#     "coreRadius": 0.047,
#     "direction": "isotropic"
# })
#
# sresult = nom_client.run_model(dictat_model)
# print(sresult)
# plot_efield(dictat_result=sresult)

# ---------------------------------------------------------------------------------------------------------

print("Running DICTAT using orbit electron spectrum from Ae8 (takes a while)")

dictat_model = nom_client.get_model('dictat')
dictat_model.set_external_input(external_input_name="electronSpectrum", external_input=ap8ae8_results['electron_spectrum_over_trajectory'],
                                result_name='electron_spectrum_over_trajectory')
dictat_model.set_params(params={
    "geometryType": "cylindrical",
    "shieldMaterial": "Aluminium",
    "shieldThickness": 0.1,
    "dielectricMaterial": "Teflon",
    "dielectricTemperature": 300,
    "dielectricThickness": 0.1,
    "numberOfGrounds": 1,
    "groundType": "inner",
    "coreMaterial": "Copper",
    "coreRadius": 0.1,
    "direction": "isotropic",
    "spectrumDefinition": "user-defined"
})

sresult = nom_client.run_model(dictat_model)
print(sresult)
plot_efield(dictat_result=sresult)