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

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.axes import Axes
from matplotlib.figure import Figure
from typing_extensions import deprecated

from zyplib.vis.mpl_utils.draw import set_scientific_formatter

__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 = None, rc: typing.Mapping[str, typing.Any] | None = 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(dict(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')
    style_core = typing.cast(typing.Any, plt.style.core)

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

    # Update dictionary of styles - plt.style.library
    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
    style_core.available[:] = sorted(plt.style.library.keys())

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


def style_ax(
    ax: Axes,
    title: str | None = None,
    xlabel: str | None = None,
    ylabel: str | None = None,
    xlim: tuple[float, float] | None = None,
    ylim: tuple[float, float] | None = None,
    xscale: Literal['linear', 'log']
    | None = None,  # 注意如果设置了 xticklabel 等一定不要设置这个参数
    yscale: Literal['linear', 'log'] | None = None,
    title_weight: Literal['normal', 'bold'] = 'normal',
    hide_spines: typing.Sequence[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 = None,
    fontsize_label: int | None = None,
    fontsize_tick: int | None = 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 = None,
    ytick_format: Literal['sci', 'eng', 'percent', 'pi']
    | Callable[[str], str]
    | None = 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: Figure,
    suptitle: str | None = 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 之后调用，以正确处理其空间
        rect = (0.0, 0.0, 1.0, 0.96) if suptitle else None
        fig.tight_layout(rect=rect)

    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: typing.Sequence, 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: typing.Sequence | 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
