from os import path
import matplotlib.pyplot as plt

plt.rcParams['figure.figsize'] = [15, 10]
plt.rc('font', size=18)


def plot(x_data, y_data, plot_types, line_types, xlabel, ylabel, labels, title):
    if not isinstance(plot_types, list):
        plot_types = [plot_types] * len(x_data)

    if not isinstance(line_types, list):
        line_types = [line_types] * len(x_data)

    if not isinstance(labels, list):
        labels = [labels] * len(x_data)

    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=18)
    plt.grid(True)
    plt.legend(bbox_to_anchor=(0.95, 0.95), loc=1, borderaxespad=0.,
               fancybox=True, fontsize=18, shadow=True)

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


from nom_client.nom_client import NoMClient
from nom_client.results.numeric_model_result import NumericModelResult

nom_client = NoMClient("AE8AP8 model example", default_server_id="local_server", debug_output=True)

trajectory_model = nom_client.get_model('trajectory')
trajectory_model.set_params(orbitSpecificationCode=205,
                            epoch="2020-01-01 00:00:00",
                            missionDuration=30,
                            naifCode=399,
                            perigeeAltitude=500, apogeeAltitude=18000)
trajectory_result = nom_client.run_model(trajectory_model)

print("Running SEP model")
sapphire_total_fluence_model = nom_client.get_model('sapphire-total-fluence')
sapphire_total_fluence_model.set_params(solarCycleOffsetFlag=1, solarCycleOffset=0)
sapphire_total_fluence_model.set_external_input(external_input_name="trajectory", external_input=trajectory_result)
sapphire_total_fluence_results = nom_client.run_model(sapphire_total_fluence_model)
print(sapphire_total_fluence_results)

ae8ap8_model = nom_client.get_model('ae8ap8')
ae8ap8_model.set_external_input(external_input_name="trajectory", external_input=trajectory_result)
ae8ap8_results = nom_client.run_model(ae8ap8_model)
print(ae8ap8_results)

proton_res = ae8ap8_results['orbit_averaged_proton_spectrum']
proton_energies = proton_res['energy'].values
proton_integral_fluence = proton_res['integral_flux'].values
proton_differential_fluence = proton_res['differential_flux'].values

electron_res = ae8ap8_results['orbit_averaged_electron_spectrum']
electron_energies = electron_res['energy'].values
electron_integral_fluence = electron_res['integral_flux'].values
electron_differential_fluence = electron_res['differential_flux'].values

ud_trapped_model = nom_client.get_model('ud-spectrum')

ud_trapped_model.set_params(species="1",
                            integralData=proton_integral_fluence,
                            differentialData=proton_differential_fluence,
                            energies=proton_energies.squeeze())

ud_trapped_results = nom_client.run_model(ud_trapped_model)
print(ud_trapped_results)

here = path.abspath(path.dirname(__file__))

sapphire_total_fluence_results['full_ion_spectrum'].to_csv(here + "/" + "eee.csv")

new_res = NumericModelResult.fromCSV(file_path=here + "/" + "eee.csv")

print(sapphire_total_fluence_results['full_ion_spectrum'])
print(new_res)

print("Running SD2 model with actual result")
sd2_model = nom_client.get_model('sd2')
#sd2_model.set_external_input(external_input_name="trappedParticleSpectrum", external_input=ap8ae8_results)
sd2_model.set_external_input(external_input_name="solarParticleSpectrum", external_input=sapphire_total_fluence_results)
#sd2_model.set_external_input(external_input_name="gcrSpectrum", external_input=creme_96_gcr_results)
sd2_model.set_params(detectorMaterial=1,
                     shieldDepthValues=1,
                     userShieldDepths=[0.05, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.8, 1, 1.5, 2, 2.5, 3, 4, 5, 6, 7, 8, 9, 10, 12, 14, 16, 18, 20])
sd2_result = nom_client.run_model(sd2_model)
print(sd2_result)




print("Running SD2 model with upload result")
sd2_model2 = nom_client.get_model('sd2')
#sd2_model.set_external_input(external_input_name="trappedParticleSpectrum", external_input=ap8ae8_results)
sd2_model2.set_external_input(external_input_name="solarParticleSpectrum", external_input=new_res)
#sd2_model.set_external_input(external_input_name="gcrSpectrum", external_input=creme_96_gcr_results)
sd2_model2.set_params(detectorMaterial=1,
                     shieldDepthValues=1,
                     userShieldDepths=[0.05, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.8, 1, 1.5, 2, 2.5, 3, 4, 5, 6, 7, 8, 9, 10, 12, 14, 16, 18, 20])
sd2_result2 = nom_client.run_model(sd2_model2)
print(sd2_result2)




sd2_thicknesses = sd2_result['dose_depth_curve'].get_variable_data(variable_name="thickness")
sd2_dose = sd2_result['dose_depth_curve'].get_variable_data(variable_name="absorbed_dose")

sd2_thicknesses2 = sd2_result2['dose_depth_curve'].get_variable_data(variable_name="thickness")
sd2_dose2 = sd2_result2['dose_depth_curve'].get_variable_data(variable_name="absorbed_dose")

plot(x_data=[sd2_thicknesses.T, sd2_thicknesses2.T],
     y_data=[sd2_dose.T, sd2_dose2.T],
     plot_types="semilogy", line_types="solid", xlabel="Energy (MeV)",
     ylabel=f"Differential flux",
     labels=['Actual', 'Upload'], title="Ax8 differential fluxes")