"""
Create animations for plots in 2D
=============================

Use the Output parser to extract info

Download Input files

* :download:`dipole_2d.data <../../../../../examples/input_files/eddy_currents.data>`
* :download:`dipole_2d.post.xml <../../../../../examples/input_files/eddy_currents.post.xml>`
"""

# %%
# Setting up
# ----------
# Standard Loading of input, Parsing and Initializing RoxiePlotOutputs object
#
from IPython.display import Image, display

from roxieapi.commons.roxie_constants import PlotLabels
from roxieapi.commons.types import Plot2D
from roxieapi.output.parser import RoxieOutputParser
from roxieapi.output.plots import RoxiePlotOutputs

xml_file_path = "../input_files/eddy_currents.post.xml"
data_file_path = "../input_files/eddy_currents.data"

plots = RoxiePlotOutputs(xml_file_path, data_file_path)
parser = RoxieOutputParser(xml_file_path)


# %%
# Show available Coil plots
# -------------------

trans = plots.output.opt[1].step[1]
print("Coil Plots:")
data_cols = trans.coilData.columns
for col in data_cols:
    if col in PlotLabels.plot2D_desc:
        lbl, desc = PlotLabels.lbl_desc_plot2D(col)
        print(f"  Plot {col}: {desc}")
for hc_idx, hc in trans.harmonicCoils.items():
    data_cols = hc.strandData
    for col in data_cols:
        if col in PlotLabels.plot2D_desc:
            lbl, desc = PlotLabels.lbl_desc_plot2D(col)
            print(f"  Harmonic coil #{hc_idx}, Plot {col}: {desc}")

# %%
# Show available Iron Mesh plots and Iron Mesh Plots for Eddy - Currents simulations
# -------------------

print("Iron mesh plots:")
data_cols = plots.output.opt[1].step[1].meshData.columns
for col in data_cols:
    if col in PlotLabels.plotMesh2D_desc:
        lbl, desc = PlotLabels.lbl_desc_mesh2D(col)
        print(f"  Plot {col}: {desc}")


# %%
# You can animate all kinds of the above plots!
# Let's plot the current densioty in the coils
# -------------------
plot_created = Plot2D.create("My 2D Animation - J")
plot_created.add_coilPlot("14", harm_coil=1)
figure = plots.plots2d.create_anim(
    plot_created,
    opt_step=1,
    start_trans_step=1,
    end_trans_step=2,
    anim_filename="current_dens_coils_animation.gif",
)
display(Image(filename="current_dens_coils_animation.gif"))


# %%
# Let's plot the Vector Potential Az in the Iron Mesh plots
# -------------------
plot_created = Plot2D.create("My 2D Animation - I")
plot_created.add_meshPlot("34")
figure = plots.plots2d.create_anim(
    plot_created,
    opt_step=1,
    start_trans_step=1,
    end_trans_step=2,
    anim_filename="az_animation.gif",
)
display(Image(filename="az_animation.gif"))


# %%
# Let's plot the current density during the excitation
# -------------------
plot_created = Plot2D.create("My 2D Animation - Jz")
plot_created.add_meshPlot("jz")
figure = plots.plots2d.create_anim(
    plot_created,
    opt_step=1,
    start_trans_step=1,
    end_trans_step=2,
    anim_filename="jz_animation.gif",
)
display(Image(filename="jz_animation.gif"))