"""
Custom plots in 2D
==========================

Create a new 2D plot based on output data.

Download Input files

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

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

# %%
# Setting up
# ----------
# Standard Loading of input, Parsing and Initializing RoxiePlotOutputs object
#

plots = RoxiePlotOutputs(
    "../input_files/dipole_iron.post.xml", "../input_files/dipole_iron.data"
)

# %%
# Plots defined in Roxie
# ----------------------
#
# Print a list of Plots which are defined in Roxie

for pl in PlotLabels.plot2D_desc:
    lbl, desc = PlotLabels.lbl_desc_plot2D(pl)
    print(f"Coil plot {pl}, Label:  '{lbl}', Description: '{desc}'")

for pl in PlotLabels.plotMesh2D_desc:
    lbl, desc = PlotLabels.lbl_desc_mesh2D(pl)
    print(f"Mesh plot {pl}, Label:  '{lbl}', Description: '{desc}'")


# %%
# Plots defined in Roxie
# ----------------------
# Print all plots available of an output
#

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}")

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}")

# %%
# Define a new Plot and create it
# -------------------------------
# From the list of available plots, Select which fields to plot, and create the output

plot_created = Plot2D.create("My Plot2D")
plot_created.add_coilPlot("52", harm_coil=1)
# pl.add_meshPlot("35", legend=PlotLegend("e", False, 0, 1))

figure = plots.plots2d.plot_xs(plot_created)
# %%