import sys

import matplotlib.pyplot as plt

from nom_client.nom_client import NoMClient


def plot_efield(dictat_result, legend=None):
    if type(dictat_result) != list:
        dictat_result = [dictat_result]

    for dic_res in dictat_result:
        dictat_surface_voltage_over_time = dic_res.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')

    if legend is not None:
        plt.legend(labels=legend)

    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="ext_rest_server",
                       debug_output=False)

sapre_model = nom_client.get_model('sapre')
sapre_model.set_params(orbitType="GEN", perigeeAltitude=300, apogeeAltitude=40000)
sapre_result = nom_client.run_model(sapre_model)
print(sapre_result)

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

print("Running DICTAT with single spectrum")

dictat_model = nom_client.get_model('dictat')

base_dictat_settings = {
    "exposureDuration": 20,
    "geometryType": "planar",
    "shieldMaterial": "Aluminium",
    "shieldThickness": 0.1,
    "dielectricMaterial": "Teflon",
    "dielectricTemperature": 298,
    "dielectricThickness": 0.1,
    "numberOfGrounds": 1,
    "groundType": "inner",
    "coreMaterial": "Copper",
    "coreRadius": 0.047,
    "direction": "isotropic",
    "dielectricRicDoseRateKp": 1e-13
}


single_spec_settings = base_dictat_settings

single_spec_settings["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]

single_spec_settings["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]

dictat_model.set_params(params=single_spec_settings)
sresult = nom_client.run_model(dictat_model)
print(sresult)
plot_efield(dictat_result=sresult)

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

print("Running DICTAT with timeseries spectrum")

timmeseries_spec_settings = base_dictat_settings

timmeseries_spec_settings["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]

timmeseries_spec_settings["fluxTimeseries"] = [
    [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],
    [6, 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],
    [12, 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],
    [18, 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=timmeseries_spec_settings)
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")

orbit_averaged_electron_spectrum = ap8ae8_results["orbit_averaged_electron_spectrum"]

dictat_model = nom_client.get_model('dictat')
dictat_model.set_external_input(external_input_name="electronSpectrum", external_input=orbit_averaged_electron_spectrum)
dictat_model.set_params(params=base_dictat_settings)
sresult1 = nom_client.run_model(dictat_model)
print(sresult1)

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

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

electron_spectrum_over_trajectory = ap8ae8_results["electron_spectrum_over_trajectory"]

dictat_model = nom_client.get_model('dictat')
dictat_model.set_external_input(external_input_name="electronSpectrum", external_input=electron_spectrum_over_trajectory)
dictat_model.set_params(params=base_dictat_settings)
sresult2 = nom_client.run_model(dictat_model)
print(sresult2)

plot_efield(dictat_result=[sresult1, sresult2], legend=["Orbit ave. e- spectrum", "e- spectrum over orbit"])