import math
import matplotlib.pyplot as plt

from nom_client.nom_client import NoMClient

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

m_to_cm = 1e-4
to_omni = 4 * math.pi
mis_dur_in_days = 1
mis_dur_in_s = mis_dur_in_days * 86400

shieldThickness = 3
species_name_to_plot = "Fe"
atomic_number_to_plot = 26


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=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()


nom_client = NoMClient(project_name="Shielded spectrum example", default_server_id="ext_rest_server")

sapre_model_gto = nom_client.get_model('sapre')
sapre_model_gto.set_params(orbitType="NEI", missionDuration=mis_dur_in_days, missionDurationUnits=1)
sapre_gto_result = nom_client.run_model(sapre_model_gto)
print(sapre_gto_result)

lowestIonSpecies = 1
highestIonSpecies = 92

iso_15390_model = nom_client.get_model('iso_15390')
iso_15390_model.set_params(lowestIonSpecies=lowestIonSpecies, highestIonSpecies=highestIonSpecies)
iso_15390_model.set_external_input(external_input_name="trajectory", external_input=sapre_gto_result)
iso_15390_results = nom_client.run_model(iso_15390_model)
print(iso_15390_results)

gcr_spectrum = iso_15390_results.get_model_result_by_name(result_name="particle_spectrum")
gcr_spectrum_energies = gcr_spectrum.get_variable_data(variable_name="energy").squeeze()
gcr_spectrum_diff_fluxes = gcr_spectrum.get_variable_data(variable_name="differential_flux",
                                                          species=species_name_to_plot).squeeze()
gcr_spectrum_int_fluxes = gcr_spectrum.get_variable_data(variable_name="integral_flux",
                                                         species=species_name_to_plot).squeeze()

sapphire_model = nom_client.get_model('sapphire-total-fluence')
sapphire_model.set_params(lowestIonSpecies=lowestIonSpecies, highestIonSpecies=highestIonSpecies)
sapphire_model.set_external_input(external_input_name="trajectory", external_input=sapre_gto_result)
sapphire_results = nom_client.run_model(sapphire_model)
print(sapphire_results)

sep_spectrum = sapphire_results.get_model_result_by_name(result_name="full_ion_spectrum")
sep_spectrum_energies = sep_spectrum.get_variable_data(variable_name="energy").squeeze()
sep_spectrum_diff_fluxes = sep_spectrum.get_variable_data(variable_name="differential_fluence",
                                                          species=species_name_to_plot).squeeze()
sep_spectrum_int_fluxes = sep_spectrum.get_variable_data(variable_name="integral_fluence",
                                                         species=species_name_to_plot).squeeze()

plot(x_data=[gcr_spectrum_energies, sep_spectrum_energies],
     y_data=[m_to_cm * to_omni * gcr_spectrum_int_fluxes, sep_spectrum_int_fluxes / mis_dur_in_s],
     plot_types=["loglog", "loglog"],
     line_types=["solid", "solid"],
     xlabel="Energy (MeV)",
     ylabel="DIfferential flux (cm^-2 s^-1 MeV^-1)",
     labels=["GCR", "SAPPHIRE total fluence"],
     title="Unshielded GCR / SEP spectra")

ss_model = nom_client.get_model('shielded-spectrum')
ss_model.set_params(shieldThickness=shieldThickness)
ss_model.set_external_input(external_input_name="spectra", external_input=[sep_spectrum, gcr_spectrum])
ss_results = nom_client.run_model(ss_model)
print(ss_results)

shielded_spectrum = ss_results.get_model_result_by_name(result_name="shielded_spectrum")
calc_shielded_energy = shielded_spectrum.get_variable_data(variable_name="energy").squeeze()
calc_shielded_diff_flux = shielded_spectrum.get_variable_data(variable_name="differential_fluence",
                                                              z=atomic_number_to_plot).squeeze()
calc_shielded_int_flux = shielded_spectrum.get_variable_data(variable_name="integral_fluence",
                                                             z=atomic_number_to_plot).squeeze()


# Calculate the stopping powers

spectrum_energy_min = min(calc_shielded_energy)
spectrum_energy_max = max(calc_shielded_energy)

stopping_powers_model = nom_client.get_model('stopping-powers')
stopping_powers_model.set_params(energyLow=spectrum_energy_min, energyHigh=spectrum_energy_max, numberEnergies=66,
                      material="G4_Si", particle="all", data_source="c86")
stopping_powers_result = nom_client.run_model(stopping_powers_model)
stopping_powers = stopping_powers_result.get_model_result_by_name(result_name="stopping_powers")
stopping_powers_energy_vals = stopping_powers.get_variable_data(variable_name="energy", scenario="G4_Si_Fe56")
stopping_powers_S_vals = stopping_powers.get_variable_data(variable_name="stopping_power", scenario="G4_Si_Fe56")
stopping_powers_scenario = stopping_powers.get_variable_data(variable_name="scenario")

plot(x_data=[stopping_powers_energy_vals.squeeze()],
     y_data=[stopping_powers_S_vals.squeeze()],
     plot_types=["loglog"],
     line_types=["solid"],
     xlabel="Energy (MeV)",
     ylabel="Stopping power (MeV cm g-1)",
     labels=["Stopping power", ],
     title=f"Stopping power ({species_name_to_plot})")


# Calculate the LET spectra

let_model = nom_client.get_model('let')
let_model.set_external_input(external_input_name="spectra", external_input=[shielded_spectrum])
let_model.set_external_input(external_input_name="stopping_power", external_input=stopping_powers_result)
let_results = nom_client.run_model(let_model)

calc_let_spectrum = let_results.get_model_result_by_name(result_name="let_spectrum")
calc_let = calc_let_spectrum.get_variable_data(variable_name="let").squeeze()
calc_diff_flux = calc_let_spectrum.get_variable_data(variable_name="differential_fluence").squeeze()
calc_int_flux = calc_let_spectrum.get_variable_data(variable_name="integral_fluence").squeeze()

# Run longupset as a comparison

longupset_model = nom_client.get_model('longupset')
longupset_model.set_external_input(external_input_name="trajectory", external_input=sapre_gto_result)
longupset_model.set_params(includeSEPSpectrum=1, includeTrappedSpectrum=0, includeGCRSpectrum=1,
                           shieldThickness=shieldThickness, shieldThicknessUnit=2)
longupset_model.set_external_input(external_input_name="solarParticleSpectrum",
                                   external_input=sapphire_results)
longupset_model.set_external_input(external_input_name="gcrSpectrum",
                                   external_input=iso_15390_results)
longupset_model.new_input_group(input_group_name="device", unique_group_name="mydevice1", params={
    "deviceMaterial": 0,
    "deviceName": "DEFAULT1"
})

longupset_result = nom_client.run_model(longupset_model)
print(longupset_result)

longupset_energy_spectrum = longupset_result.get_model_result_by_name(result_name="energy_spectrum_all_species")
longupset_shielded_energies = longupset_energy_spectrum.get_variable_data(variable_name="energy").squeeze()
longupset_shielded_differential_flux = longupset_energy_spectrum.get_variable_data(variable_name="differential_fluence",
                                                                                   species="H").squeeze()
longupset_shielded_integral_flux = longupset_energy_spectrum.get_variable_data(variable_name="integral_fluence",
                                                                               species=species_name_to_plot).squeeze()

longupset_shielded_energies = longupset_shielded_energies.squeeze()
longupset_shielded_differential_flux = longupset_shielded_differential_flux.squeeze()

let_spectrum = longupset_result.get_model_result_by_name(result_name="let_spectrum")
let_let = let_spectrum.get_variable_data(variable_name="let").squeeze()
let_differential_flux = let_spectrum.get_variable_data(variable_name="differential_fluence").squeeze()
let_integral_flux = let_spectrum.get_variable_data(variable_name="integral_fluence").squeeze()

plot(x_data=[gcr_spectrum_energies, sep_spectrum_energies, longupset_shielded_energies, calc_shielded_energy],
     y_data=[m_to_cm * to_omni * gcr_spectrum_int_fluxes, sep_spectrum_int_fluxes / mis_dur_in_s,
             to_omni * m_to_cm * longupset_shielded_integral_flux / mis_dur_in_s, calc_shielded_int_flux],
     plot_types=['loglog', 'loglog', 'loglog', 'loglog'],
     line_types=["dashed", "dashed", "solid", "solid"],
     xlabel="Energy (MeV)",
     ylabel="Differential flux (cm^-2 s^-1 MeV^-1)",
     labels=["GCR / Unshielded", "SAPPHIRE / Unshielded", "LONGUPSET / Shielded", "Shielded-spectrum model / Shielded"],
     title="Shielded spectra")

if shieldThickness == 0:
    plot_title = "Unshielded LET spectra"
else:
    plot_title = f"Shielded LET spectra ({shieldThickness} g/cm2)"

plot(x_data=[let_let, calc_let],
     y_data=[to_omni * let_differential_flux / (1e4 * 86400), calc_diff_flux],
     plot_types=['loglog', 'loglog'],
     line_types=["solid", "dashed"],
     xlabel="LET (MeV cm^2 g^-1)",
     ylabel="Differential flux cm^-2 s^-1 (MeV cm^2 g^-1)^-1",
     labels=["LONGUPSET", "LET model"],
     title=plot_title)