from contextlib import ExitStack
from pathlib import Path
from typing import TYPE_CHECKING, Callable, Literal

import matplotlib.pyplot as plt
import numpy as np
from typing_extensions import deprecated

from zyplib.vis.mpl_utils.draw import set_scientific_formatter

if TYPE_CHECKING:
    Figure = plt.Figure

__all__ = [
    'MplStyleContext',
    'register_mpl_styles',
    'style_ax',
    'style_figure',
    'sample_color_list',
    'use_color_cycle',
    'use_font_size',
    'use_font_family',
    'color_cycle_context',
    'get_current_color_cycle',
]


class MplStyleContext:
    """
    一个结合了样式表 (style) 和精细参数 (rc) 的 Matplotlib 上下文管理器。

    允许你临时应用一个基础样式，并在此之上覆盖一些特定的 rcParams。

    Parameters
    ----------
    style : str or list, optional
        要应用的 Matplotlib 样式名称或名称列表。
    rc : dict, optional
        一个包含 rc 参数的字典，用于覆盖样式中的设置。
    """

    def __init__(self, style: str = None, rc: dict = None):
        self.style = style
        self.rc = rc
        self._stack = ExitStack()

    def __enter__(self):
        # ExitStack 会处理所有进入的上下文
        # 1. 首先进入 style 的上下文
        if self.style:
            self._stack.enter_context(plt.style.context(self.style))

        # 2. 接着进入 rc_params 的上下文，这样 rc 的设置会覆盖 style 的设置
        if self.rc:
            self._stack.enter_context(plt.rc_context(self.rc))

        return self

    def __exit__(self, exc_type, exc_value, traceback):
        # ExitStack 会自动以后进先出(LIFO)的顺序调用所有上下文的 __exit__ 方法
        self._stack.close()


def register_mpl_styles():
    """
    注册 matplotlib 内置的样式

    源码
    ----------
    本代码来自: https://github.com/garrettj403/SciencePlots/blob/master/scienceplots/__init__.py
    """
    import os

    current_dir_path = os.path.dirname(os.path.abspath(__file__))
    styles_path = os.path.join(current_dir_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())

    return set(plt.style.library.keys())  # 返回所有可用样式的名称


def style_ax(
    ax: plt.Axes,
    title: str = None,
    xlabel: str = None,
    ylabel: str = None,
    xlim: tuple[float, float] = None,
    ylim: tuple[float, float] = None,
    xscale: Literal['linear', 'log']
    | None = None,  # 注意如果设置了 xticklabel 等一定不要设置这个参数
    yscale: Literal['linear', 'log'] | None = None,
    title_weight: Literal['normal', 'bold'] = 'normal',
    hide_spines: list[Literal['top', 'right', 'bottom', 'left']]
    | Literal['all', 'none'] = ['top', 'right'],
    grid: bool | None = None,
    grid_style: str = '--',
    grid_alpha: float = 0.5,
    fontsize_title: int = None,
    fontsize_label: int = None,
    fontsize_tick: int = None,
    tick_direction: Literal['in', 'out', 'inout'] | None = None,
    tick_label_rotation_x: int = 0,
    xtick_format: Literal['sci', 'eng', 'percent', 'pi'] | Callable[[str], str] = None,
    ytick_format: Literal['sci', 'eng', 'percent', 'pi'] | Callable[[str], str] = None,
    aspect: float | Literal['auto', 'equal'] = 'auto',
    legend: bool = True,
    legend_loc: str = 'best',
    legend_frameon: bool = True,
):
    """
    对指定的 Axes 对象应用一套统一的样式。
    """
    if title:
        ax.set_title(title, fontsize=fontsize_title, weight=title_weight)  # 加粗标题
    if xlabel:
        ax.set_xlabel(xlabel, fontsize=fontsize_label)
    if ylabel:
        ax.set_ylabel(ylabel, fontsize=fontsize_label)

    if xlim:
        ax.set_xlim(xlim)
    if ylim:
        ax.set_ylim(ylim)
    if xscale is not None:
        ax.set_xscale(xscale)
    if yscale is not None:
        ax.set_yscale(yscale)

    # 边框
    if hide_spines == 'all':
        spines_to_hide = ['top', 'right', 'bottom', 'left']
    elif hide_spines == 'none' or hide_spines is None:
        spines_to_hide = []
    else:
        spines_to_hide = hide_spines
    for spine in spines_to_hide:
        if spine in ax.spines:
            ax.spines[spine].set_visible(False)

    # 网格
    if grid:
        ax.grid(True, which='major', linestyle=grid_style, alpha=grid_alpha)

    # 新增：刻度样式
    ax.tick_params(
        axis='both', which='major', direction=tick_direction, labelsize=fontsize_tick
    )
    # 只有当需要设置旋转角度时才获取和设置刻度标签，避免重置用户自定义的标签
    if tick_label_rotation_x != 0:
        plt.setp(ax.get_xticklabels(), rotation=tick_label_rotation_x)

    # 刻度标签格式化
    if xtick_format or ytick_format:
        set_scientific_formatter(ax, x_format=xtick_format, y_format=ytick_format)

    # 宽高比
    ax.set_aspect(aspect)

    # 新增：图例处理 (仅当 ax 中有可生成图例的 artist 时生效)
    if legend and ax.get_legend_handles_labels()[0]:
        ax.legend(loc=legend_loc, frameon=legend_frameon, fontsize=fontsize_tick)

    return ax


def style_figure(
    fig: plt.Figure,
    suptitle: str = None,
    suptitle_fontsize: int = 16,
    tight_layout: bool = True,
    style_each_ax: bool = True,
    # 使用 **kwargs 捕获所有要传递给 style_ax 的参数
    **style_ax_kwargs,
):
    """
    对整个 Figure 对象应用样式，包括 Figure 级别的属性（如总标题、布局）
    以及对其包含的所有 Axes 应用 style_ax。

    Examples
    ----------

    ```py
    fig, axes = plt.subplots(2, 2)
    axes_flat = axes.flatten()
    for i in range(4):
        axes_flat[i].plot(np.random.randn(100).cumsum())

    # 所有绘图完成后，对整个 figure 进行一次性样式化
    style_figure(
        fig,
        suptitle='Main Title',
        xlabel='Time',
        ylabel='Value'
    )
    plt.show()
    ```

    Parameters
    ----------
    - `fig` : plt.Figure
        要设置样式的 Figure 对象。
    - `suptitle` : str, optional
        Figure 的总标题。
    - `suptitle_fontsize` : int, optional
        总标题的字体大小。
    - `tight_layout` : bool, optional
        是否应用 tight_layout 来自动调整子图布局, 默认为 True。
    - **style_ax_kwargs :
        所有其他关键字参数都将传递给 style_ax 函数，
        应用于此 Figure 中的每一个 Axes。
    """
    if suptitle:
        fig.suptitle(suptitle, fontsize=suptitle_fontsize)

    # 遍历 Figure 中的所有 Axes 并应用 style_ax
    if style_each_ax:
        for ax in fig.get_axes():
            style_ax(ax, **style_ax_kwargs)

    # 最后应用布局调整，确保所有元素（包括新加的标签）位置正确
    if tight_layout:
        # 在 suptitle 之后调用，以正确处理其空间
        fig.tight_layout(rect=[0, 0, 1, 0.96] if suptitle else None)

    return fig


def sample_color_list(
    n: int, cmap: str = 'viridis', start: float = 0.0, stop: float = 1.0
) -> list[tuple]:
    """
    从指定的 colormap 中获取 n 个均匀分布的颜色。

    Parameters
    ----------
    - `n` : `int`
        需要获取的颜色数量。
    - `cmap` : `str`, optional
        Colormap 的名称, 默认为 'viridis'。
    - `start` : `float`, optional
        从 colormap 的哪个位置开始采样 (0.0 到 1.0), 默认为 0.0。
    - `stop` : `float`, optional
        在 colormap 的哪个位置结束采样 (0.0 到 1.0), 默认为 1.0。

    Returns
    -------
    - list[tuple]
        一个包含 n 个 RGBA 元组的列表。
    """
    colormap = plt.get_cmap(cmap)
    return [colormap(i) for i in np.linspace(start, stop, n)]


def use_font_family(
    type: Literal['serif', 'sans-serif', 'monospace'] | str = 'sans-serif',
    policy: Literal['apply', 'context'] = 'apply',
):
    import matplotlib.pyplot as plt

    if policy == 'context':
        return MplStyleContext(rc={'font.family': type})
    elif policy == 'apply':
        plt.rcParams['font.family'] = type
    return


def use_color_cycle(colors: list, policy: Literal['apply', 'context'] = 'apply'):
    """设置 Matplotlib 的颜色循环。

    默认 policy 是 'context'，这意味着颜色循环只在上下文管理器内有效。
    如果 policy 是 'apply'，则颜色循环会被永久应用到当前的 Matplotlib 会话。

    **重要**: 必须在创建 Figure/Axes 之前使用此上下文管理器！

    ```py
    # ✓ 正确的使用方式：在创建 subplots 之前设置颜色循环
    my_colors = ['#ff6666', '#66ff66', '#6666ff']
    with use_color_cycle(my_colors):
        fig, ax = plt.subplots()  # 在上下文内创建图表
        for i in range(3):
            ax.plot(np.arange(10) + i, label=f'Line {i+1}')

    # ✓ 使用 colormap 但是 apply 到全局
    use_color_cycle('viridis', policy='apply')
    fig, axes = plt.subplots(2, 2)
    ...

    # ✗ 错误的使用方式：在创建 subplots 之后设置颜色循环（不会生效）
    fig, ax = plt.subplots()
    with use_color_cycle(my_colors):  # 这样不会生效！
        ax.plot(...)
    ```

    Parameters
    ----------
    - `colors` : `list`
        - 颜色列表 (例如 ['#ff0000', ...])。
    - `policy` : `Literal['apply', 'context']`, optional
        - 上下文管理器的作用范围，默认为 'context'。
    """
    from cycler import cycler

    custom_cycler = cycler(color=colors)
    if policy == 'context':
        return MplStyleContext(rc={'axes.prop_cycle': custom_cycler})
    elif policy == 'apply':
        plt.rcParams['axes.prop_cycle'] = custom_cycler
    return


def use_font_size(
    fontsize: int | None = None,
    titlefontsize: int | None = None,
    xlabelsize: int | None = None,
    ylabelsize: int | None = None,
    xticklabelsize: int | None = None,
    yticklabelsize: int | None = None,
    legendfontsize: int | None = None,
    legend_title_fontsize: int | None = None,
    axes_titlesize: int | None = None,
    figure_titlesize: int | None = None,
    figure_labelsize: int | None = None,
    mathtext_fontsize: int | None = None,
    howto_change: Literal['increasement', 'quantization'] = 'increasement',
    policy: Literal['apply', 'context'] = 'apply',
    increase_all: int | None = None,
):
    """一个适应实际需求的对 Matplotlib 字体大小进行调整的函数

    在科研论文绘图中，字体大小的统一调整是常见需求。此函数提供两种调整模式：
    1. 相对调整 (increasement): 在当前字体基础上增减
    2. 绝对设置 (quantization): 直接设置绝对值

    **重要说明**：matplotlib 中的字体大小字符串值（如 'medium', 'large'）是相对值，
    会根据 `font.size` 动态计算。例如当 `font.size=10` 时，'medium'=10, 'large'=12；
    当 `font.size=14` 时，'medium'=14, 'large'=16.8。这种设计允许通过调整基础字体
    大小来统一缩放所有字体元素。

    Parameters
    ----------
    - `fontsize` : `int`, optional
        基础字体大小，对应 'font.size'
    - `titlefontsize` : `int`, optional
        标题字体大小，对应 'axes.titlesize'
    - `xlabelsize` : `int`, optional
        X轴标签字体大小，对应 'axes.labelsize'
    - `ylabelsize` : `int`, optional
        Y轴标签字体大小，对应 'axes.labelsize'
    - `xticklabelsize` : `int`, optional
        X轴刻度标签字体大小，对应 'xtick.labelsize'
    - `yticklabelsize` : `int`, optional
        Y轴刻度标签字体大小，对应 'ytick.labelsize'
    - `legendfontsize` : `int`, optional
        图例字体大小，对应 'legend.fontsize'
    - `legend_title_fontsize` : `int`, optional
        图例标题字体大小，对应 'legend.title_fontsize'
    - `axes_titlesize` : `int`, optional
        子图标题字体大小，对应 'axes.titlesize'
    - `figure_titlesize` : `int`, optional
        图片总标题字体大小，对应 'figure.titlesize'
    - `figure_labelsize` : `int`, optional
        图片标签字体大小，对应 'figure.labelsize'
    - `mathtext_fontsize` : `int`, optional
        数学文本字体大小，需要结合 'font.size' 使用
    - `howto_change` : `Literal['increasement', 'quantization']`, optional
        调整模式, 默认为 'increasement'
        - 'increasement': 相对调整，在当前值基础上增减
        - 'quantization': 绝对调整，直接设置为指定值
    - `policy` : `Literal['apply', 'context']`, optional
        应用策略, 默认为 'apply'
        - 'apply': 永久应用到当前会话
        - 'context': 作为上下文管理器临时应用
    - `increase_all` : `int`, optional
        一次性调整所有字体大小（仅在 howto_change='increasement' 时有效）

    Returns
    -------
    - `MplStyleContext` or `None`
        当 policy='context' 时返回上下文管理器，否则返回 None

    Examples
    --------
    ```python
    # 全局增加所有字体 2pt
    use_font_size(increase_all=2)

    # 设置论文标准字体大小
    use_font_size(
        fontsize=10,
        titlefontsize=12,
        legendfontsize=9,
        howto_change='quantization'
    )

    # 临时调整字体
    with use_font_size(increase_all=3, policy='context'):
        fig, ax = plt.subplots()
        ax.plot([1, 2, 3])
    ```
    """
    import matplotlib.pyplot as plt

    # 收集需要修改的 rc 参数
    rc_dict = {}

    # 字体大小参数映射 (参数名 -> rcParams 键名)
    font_params = {
        'fontsize': 'font.size',
        'titlefontsize': 'axes.titlesize',
        'axes_titlesize': 'axes.titlesize',
        'xlabelsize': 'axes.labelsize',
        'ylabelsize': 'axes.labelsize',
        'xticklabelsize': 'xtick.labelsize',
        'yticklabelsize': 'ytick.labelsize',
        'legendfontsize': 'legend.fontsize',
        'legend_title_fontsize': 'legend.title_fontsize',
        'figure_titlesize': 'figure.titlesize',
        'figure_labelsize': 'figure.labelsize',
    }

    # 处理 increase_all 参数 - 一次性调整所有字体
    if increase_all is not None and howto_change == 'increasement':
        for rc_key in font_params.values():
            if rc_key == 'legend.title_fontsize':
                # legend.title_fontsize 默认为 None，需要特殊处理
                current_val = plt.rcParams.get('legend.fontsize', 'medium')
            else:
                current_val = plt.rcParams[rc_key]

            # 处理字符串类型的大小值 (如 'medium', 'large' 等)
            # 注意：这些字符串值是相对于 font.size 的倍数关系
            if isinstance(current_val, str):
                size_map = {
                    'xx-small': 0.58,  # ≈ 0.58 × font.size
                    'x-small': 0.69,  # ≈ 0.69 × font.size
                    'small': 0.83,  # ≈ 0.83 × font.size
                    'smaller': 0.83,  # 同 small
                    'medium': 1.0,  # = font.size
                    'large': 1.2,  # = 1.2 × font.size
                    'larger': 1.2,  # 同 large
                    'x-large': 1.44,  # ≈ 1.44 × font.size
                    'xx-large': 1.73,  # ≈ 1.73 × font.size
                }
                base_size = plt.rcParams['font.size']
                current_val = size_map.get(current_val, 1.0) * base_size
            elif current_val is None:
                current_val = 10  # 默认大小

            rc_dict[rc_key] = float(current_val) + increase_all

    # 处理单独设置的参数
    params = locals()
    for param_name, rc_key in font_params.items():
        value = params.get(param_name)
        if value is not None:
            if howto_change == 'increasement':
                # 获取当前值
                if rc_key == 'legend.title_fontsize':
                    current_val = plt.rcParams.get('legend.fontsize', 'medium')
                else:
                    current_val = plt.rcParams[rc_key]

                # 处理字符串类型
                if isinstance(current_val, str):
                    size_map = {
                        'xx-small': 0.58,
                        'x-small': 0.69,
                        'small': 0.83,
                        'smaller': 0.83,
                        'medium': 1.0,
                        'large': 1.2,
                        'larger': 1.2,
                        'x-large': 1.44,
                        'xx-large': 1.73,
                    }
                    base_size = plt.rcParams['font.size']
                    current_val = size_map.get(current_val, 1.0) * base_size
                elif current_val is None:
                    current_val = 10

                rc_dict[rc_key] = float(current_val) + value
            else:  # quantization
                rc_dict[rc_key] = float(value)

    # 处理数学文本字体大小
    if mathtext_fontsize is not None:
        if howto_change == 'increasement':
            current_size = plt.rcParams['font.size']
            if isinstance(current_size, str):
                current_size = 10  # 默认值
            rc_dict['font.size'] = float(current_size) + mathtext_fontsize
        else:
            rc_dict['font.size'] = float(mathtext_fontsize)

    # 应用设置
    if policy == 'context':
        return MplStyleContext(rc=rc_dict)
    elif policy == 'apply':
        for key, value in rc_dict.items():
            plt.rcParams[key] = value

    return None


@deprecated('请使用 use_color_cycle 函数替代')
def color_cycle_context(colors: list | str):
    """
    等价于  `use_color_cycle(colors, policy='context')`。

    Returns
    -------
    - contextlib.AbstractContextManager
        一个 Matplotlib 的 rc_context 上下文管理器。
    """
    return use_color_cycle(colors, policy='context')


def get_current_color_cycle() -> list[str]:
    """获取当前的颜色循环

    Returns
    -------
    - `list[str]`: 当前的颜色循环, color hex
    """
    import matplotlib as mpl

    prop_cycle = mpl.rcParams['axes.prop_cycle']
    cycle_colors = prop_cycle.by_key()['color']
    return cycle_colors
