import os
import numpy as np
import matplotlib.pyplot as plt

from duat import config, plot, run

"""
Explanation: Introductory demo
This document is a functional Jupyter notebook showing how to work with OSIRIS using the duat Python interface. You can also use it with python scripts or from the interpreter, but the notebook is a handy way to use the package I suggest you to consider.
The updated version of this notebook can be download here.
First, import some modules.
End of explanation
"""


# Create a config file with the defaults
sim = config.ConfigFile(1)  # Argument -> dimension
# Check the generated code
print(sim)

# Let's change some parameters:
# Parameters can be edited using the python item access notation
sim["time"]["tmax"] = 30.0
# Beware the python indexes starting at zero. E.g., "1" is the second particle species
sim["species_list"][0]["species"]["vfl"] = [0.0, 0.0, 0.6]
sim["species_list"][0]["species"]["num_par_x"] = [200]
sim["species_list"][1]["species"]["vfl"] = [0.0, 0.0, -0.6]
sim["species_list"][1]["species"]["num_par_x"] = [200]

# Order of output is handled by duat. We can now change parameters that will appear before in the generated file
sim["species_list"][0]["diag_species"].set_pars(ndump_fac=1, reports="ene")

# We can benefit from python preprocessing power
ene_bins = np.arange(0, 0.5, 0.02)
sim["species_list"][1]["diag_species"].set_pars(ndump_fac=1, ndump_fac_pha=1, pha_ene_bin="x1_|charge|",
                                                ene_bins=ene_bins, n_ene_bins=len(ene_bins))

# Even if a section was not created before, accessing it with the index notation will create it
sim["diag_emf"]["reports"] = ["e1", "e2", "e3"]
sim["diag_emf"]["ndump_fac"] = 5

# This also works with lists
sim["zpulse_list"][0].set_pars(a0=1.0, omega0=1.0, phase=0.0, pol_type=1, pol=0, propagation="forward",
                               lon_type="gaussian", lon_x0=40, lon_range=20, lon_duration=50, per_type="plane")

# Check the generated code again
print(sim)

"""
Explanation: If a warning on not found executables was raised for you, add some code setting the path to the folder with the executables:
run.set_osiris_path(path.join("path", "to", "osiris", "folder"))
Creating a simulation
The config module offers functionality to create a simulation.
End of explanation
"""


# A directory for the run
run_dir = os.path.join(os.path.expanduser("~"), "test-run")

# Run the simulation. This also checks for errors, but note they are not always detected
myrun = run.run_config(sim, run_dir, clean_dir=True)

# The created run instance offers some information of the run in progress.
myrun

# Other methods are available in the Run object. See the documentation for more information.
print("Size in disk of the run: %.2f MiB" % (myrun.get_size()/1024/1024))

"""
Explanation: Running a simulation
The config module offers functionality to create a simulation
End of explanation
"""


# Diagnostic objects can be used even if the run is in progress
diagnostics = myrun.get_diagnostic_list()
for d in diagnostics:
    print(d)

# Methods like time_1d_colormap allow to represent a map
fig, ax = diagnostics[0].time_1d_colormap(dataset_selector=np.sum)
# The returned Figure and Axes instances allow for customization and export.
# Note the method itself takes a parameter for automatic exportation

# The time_1d_animation method allows to visualize a function in time
fig, ax, anim = diagnostics[2].time_1d_animation()
# This can be exported using the anim object or given a parameter to the function.

# In Jupyter you can plot this with:
from IPython.display import HTML
HTML(anim.to_html5_video())
# mpeg must be installed for this to work!
# The video can be downloaded from the output.

# Do the same as a color map
fig, ax = diagnostics[0].time_1d_colormap(axes_selector=(np.sum,))

# For manual manipulation use the get_generator method
gen = diagnostics[1].get_generator()
# This returns a generator that provides data when iterated
for snapshot in gen:
    print(snapshot)

"""
Explanation: Plotting the results
The plot module offers functionality to create a simulation
End of explanation
"""


