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

# %% auto 0
__all__ = ['PlotObject', 'module_path', 'styles_path', 'stylesheets', '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
import os
import beforerr.matplotlib as bmpl
from ..project import plotsdir

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

# %% ../../nbs/11_matplotlib.ipynb 2
# | code-summary: Register the bundled stylesheets in the matplotlib style library
# https://github.com/garrettj403/SciencePlots/blob/master/scienceplots/__init__.py
module_path = bmpl.__path__[0]
styles_path = os.path.join(module_path, "styles")

# Reads styles in /styles folder and all subfolders
stylesheets = {}  # plt.style.library is a dictionary
for folder, _, _ in os.walk(styles_path):
    new_stylesheets = plt.style.core.read_style_directory(folder)
    stylesheets.update(new_stylesheets)

# Update dictionary of styles - plt.style.library
plt.style.core.update_nested_dict(plt.style.library, stylesheets)
# Update `plt.style.available`, copy-paste from:
# https://github.com/matplotlib/matplotlib/blob/a170539a421623bb2967a45a24bb7926e2feb542/lib/matplotlib/style/core.py#L266  # noqa: E501
plt.style.core.available[:] = sorted(plt.style.library.keys())

# %% ../../nbs/11_matplotlib.ipynb 3
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 5
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 7
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 8
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 9
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 11
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 13
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
