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

# %% auto 0
__all__ = ['PlotObject', 'module_path', 'styles_path', 'stylesheets', 'get_axes', 'func2axes', 'func2lines', 'func2legend',
           'figsize', 'easy_save', 'unify_axis_fontsize', 'unify_axes_fontsize', 'hide_xlabels', 'hide_ylabels',
           'create_legends', 'set_linewidth', 'sync_legend_colors', 'hide_legends', 'hide_legend_lines', 'PlotOpts',
           'process_figure', 'update_rcParams']

# %% ../../../nbs/11_matplotlib.ipynb 1
import matplotlib.pyplot as plt
from matplotlib.pyplot import Axes, Line2D
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 func2lines(obj: PlotObject, func: Callable, *args, **kwargs):
    for ax in get_axes(obj):
        for line in ax.get_lines():
            func(line, *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 figsize(col=1, width: float = None, height: float = None, ratio: float = None):
    """Return figure size for a given height and column width for publication.

    Parameters
    ----------
    col : int, optional
        Number of columns the figure spans.
    ratio : float, optional
        Aspect ratio of the figure (height/width).
    """
    DEFAULT_WIDTH = 3.4
    DEFAULT_RATIO = 0.68

    width = width or DEFAULT_WIDTH * col
    ratio = ratio or DEFAULT_RATIO
    height = height or ratio * width

    return width, height

# %% ../../../nbs/11_matplotlib.ipynb 7
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 9
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 10
def unify_axes_fontsize(
    obj: PlotObject = None,  # 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(obj, unify_axis_fontsize, fontsize)

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


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


def hide_xlabels(obj: PlotObject = None):
    func2axes(obj, _hide_x_axis_label)


def hide_ylabels(obj: PlotObject = None):
    func2axes(obj, _hide_y_axis_label)


def _create_legend(ax: Axes, handletextpad=0, handlelength=0, **kwargs):
    if ax.get_legend_handles_labels()[0]:
        ax.legend(handletextpad=handletextpad, handlelength=handlelength, **kwargs)


def create_legends(obj: PlotObject = None, **kwargs):
    func2axes(obj, _create_legend, **kwargs)

# %% ../../../nbs/11_matplotlib.ipynb 13
# change line width in all axes
def _set_linewidth(obj: Line2D, w):
    obj.set_linewidth(w)


def set_linewidth(w, obj: PlotObject = None):
    func2lines(obj, _set_linewidth, w)

# %% ../../../nbs/11_matplotlib.ipynb 15
def _sync_legend_colors(legend: Legend):
    for line, text in zip(legend.get_lines(), legend.get_texts()):
        text.set_color(line.get_color())


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


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


def hide_legends(obj: PlotObject = None):
    """Hide legends on all axes."""
    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 = None):
    """Hide legend lines on all axes."""
    func2legend(obj, _hide_legend_lines)

# %% ../../../nbs/11_matplotlib.ipynb 17
class PlotOpts(BaseModel):
    """Options to apply to a matplotlib object."""

    sync_legend_colors: bool = False
    hide_legends: bool = False
    hide_legend_lines: bool = False
    hide_xlabels: bool = False
    hide_ylabels: bool = False

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

    def process_figure(self, obj, **kwargs):
        return process_figure(obj, self, **kwargs)


def process_figure(obj, plot_opts: BaseModel, **kwargs):
    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
