# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/11_matplotlib.ipynb.

# %% auto 0
__all__ = ['PlotObject', 'get_axes', 'func2axes', 'func2legend', 'easy_save', 'unify_axis_fontsize', 'unify_axes_fontsize',
           'hide_x_axes_label', 'hide_y_axes_label', 'sync_legend_colors', 'hide_legend', 'hide_legend_lines',
           'PlotOpts', 'process_figure', 'update_rcParams']

# %% ../nbs/11_matplotlib.ipynb 1
import matplotlib.pyplot as plt
from matplotlib.pyplot import Axes
from matplotlib.figure import Figure
from matplotlib.legend import Legend
from pydantic import BaseModel
from typing import Callable
import sys
from .project import plotsdir

# Define the type alias
PlotObject = Figure | Axes | list[Axes] | None

# %% ../nbs/11_matplotlib.ipynb 2
def get_axes(obj: PlotObject):
    if isinstance(obj, Axes):
        return [obj]
    elif isinstance(obj, Figure):
        return obj.axes
    elif isinstance(obj, list):
        return obj
    else:
        return plt.gcf().axes


def func2axes(obj: PlotObject, func: Callable, *args, **kwargs):
    for ax in get_axes(obj):
        func(ax, *args, **kwargs)


def func2legend(obj: PlotObject, func: Callable, *args, **kwargs):
    for ax in get_axes(obj):
        legend = ax.get_legend()
        legend and func(legend, *args, **kwargs)

# %% ../nbs/11_matplotlib.ipynb 4
def easy_save(name: str, fig: Figure = None, formats=["svg"], dir=plotsdir(), **kwargs):
    fig = fig or plt.gcf()
    path = dir / name
    return [fig.savefig(path.with_suffix(f".{fmt}"), **kwargs) for fmt in formats]

# %% ../nbs/11_matplotlib.ipynb 6
def unify_axis_fontsize(ax: Axes, fontsize: str | float):
    mainlabels = [ax.title, ax.xaxis.label, ax.yaxis.label]
    ticklabels = ax.get_xticklabels() + ax.get_yticklabels()
    for text in mainlabels + ticklabels:
        text.set_fontsize(fontsize)

# %% ../nbs/11_matplotlib.ipynb 7
def unify_axes_fontsize(
    axes: Axes,  # a single axis or a list of axes
    fontsize: str
    | float = "medium",  # string values denote sizes relative to the default font size
):
    """
    Set the fontsize of all text elements in a matplotlib axis to the same value.
    """
    func2axes(axes, unify_axis_fontsize, fontsize)

# %% ../nbs/11_matplotlib.ipynb 8
def _hide_x_axis_label(ax: Axes):
    ax.set_xlabel("")


def _hide_y_axis_label(ax: Axes):
    ax.set_ylabel(None)


def hide_x_axes_label(obj: PlotObject):
    func2axes(obj, _hide_x_axis_label)


def hide_y_axes_label(obj: PlotObject):
    func2axes(obj, _hide_y_axis_label)

# %% ../nbs/11_matplotlib.ipynb 10
def _sync_legend_colors(legend: Legend):
    """Synchronize legend text colors with their corresponding line colors."""
    for line, text in zip(legend.get_lines(), legend.get_texts()):
        text.set_color(line.get_color())


def sync_legend_colors(obj: PlotObject):
    """Synchronize legend text colors with their corresponding line colors."""
    func2legend(obj, _sync_legend_colors)


def _hide_legend(legend: Legend):
    legend.set_visible(False)


def hide_legend(obj: PlotObject):
    func2legend(obj, _hide_legend)


def _hide_legend_lines(legend: Legend):
    """Hide legend lines."""
    for line in legend.get_lines():
        line.set_visible(False)


def hide_legend_lines(obj: PlotObject):
    """Hide legend lines on all axes."""
    func2legend(obj, _hide_legend_lines)

# %% ../nbs/11_matplotlib.ipynb 12
class PlotOpts(BaseModel):
    sync_legend_colors: bool = False
    hide_legend: bool = False
    hide_legend_lines: bool = False
    hide_x_axes_label: bool = False
    hide_y_axes_label: bool = False

    def model_post_init(self, __context):
        update_rcParams(self)
        return super().model_post_init(__context)


def process_figure(obj, plot_opts: PlotOpts, **kwargs):
    # plot_opts.sync_legend_colors and sync_legend_colors(fig)
    for key, value in plot_opts.model_dump().items():
        # Get a reference to the current module function
        func = getattr(sys.modules[__name__], key)
        value and func(obj, **kwargs)

    return obj


def update_rcParams(plot_opts: PlotOpts):
    if plot_opts.hide_legend_lines:
        plt.rcParams["legend.handlelength"] = 0
        plt.rcParams["legend.handletextpad"] = -1.5
