import typing
from dataclasses import dataclass
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


__all__ = [
    'show_figure',
    'save_figure',
    'make_grids',
    'GridContext',
    'make_figure_layout',
    'LayoutContext',
    'adjust_subplot_spacing',
]


def adjust_subplot_spacing(
    fig: 'Figure', wspace_delta: float = 0, hspace_delta: float = 0
):
    """
    相对于当前状态，微调 Figure 中子图的水平和垂直间距。

    这个函数非常适合在 `tight_layout` 自动调整后进行手动修正。

    Parameters
    ----------
    - `fig` : `plt.Figure`
        - 需要调整的 Matplotlib Figure 对象。
    - `wspace_delta` : `float`, optional
        - 水平间距 (width space) 的增量。正值增大间距，负值减小间距。默认为 0。
        - 这个值是相对于子图平均宽度的比例。
    - `hspace_delta` : `float`, optional
        - 垂直间距 (height space) 的增量。正值增大间距，负值减小间距。默认为 0。
        - 这个值是相对于子图平均高度的比例。

    Returns
    -------
    - `None`
        - 函数直接修改传入的 Figure 对象。
    """
    # 获取当前的间距值
    # 如果之前没有设置过，这些值会有默认值（例如，从 rcParams 读取）
    current_wspace = fig.subplotpars.wspace
    current_hspace = fig.subplotpars.hspace

    # 计算新的间距值
    # 我们使用 max(0, ...) 来防止间距变为负数，虽然 Matplotlib 也能处理，但这更稳健
    new_wspace = max(0, current_wspace + wspace_delta)
    new_hspace = max(0, current_hspace + hspace_delta)

    # 应用新的间距值
    fig.subplots_adjust(wspace=new_wspace, hspace=new_hspace)

    # 在交互式环境中，需要重绘 canvas 才能看到变化
    if fig.canvas and hasattr(fig.canvas, 'draw_idle'):
        fig.canvas.draw_idle()


def adjust_figure_margins(
    fig: 'Figure',
    top: float | None = None,
    bottom: float | None = None,
    left: float | None = None,
    right: float | None = None,
):
    """调整 Figure 的整体边距，用于控制 suptitle、图例或轴标签与边缘的距离。"""
    adjust_kwargs = {'top': top, 'bottom': bottom, 'left': left, 'right': right}
    filtered_kwargs = {k: v for k, v in adjust_kwargs.items() if v is not None}
    if filtered_kwargs:
        fig.subplots_adjust(**filtered_kwargs)
    if fig.canvas and hasattr(fig.canvas, 'draw_idle'):
        fig.canvas.draw_idle()


def show_figure(fig: Figure | None = None):
    """显示 matplotlib 图表"""
    if fig is not None:
        fig.tight_layout()
        plt.show()
        plt.close(fig)  # 关闭图形


def save_figure(
    fig: Figure,
    path: str | Path,
    dpi: int | None = None,
    transparent: bool = True,
    bbox_inches: str = 'tight',
    **savefig_kwargs,
):
    """
    以标准化的参数保存 Matplotlib 图表，并自动创建目录。
    """
    filepath = Path(path)
    filepath.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(
        filepath,
        dpi=dpi,
        transparent=transparent,
        bbox_inches=bbox_inches,
        **savefig_kwargs,
    )
    plt.close(fig)  # 保存后关闭图像，释放内存
    print(f'Figure saved to: {filepath}')


@dataclass
class GridLoc:
    i: int
    j: int
    index: int


class GridContext:
    """可迭代的子图网格上下文管理器"""

    def __init__(self, row, col, figsize, tight_layout, fn_finalize, **subplots_kwargs):
        self.row = row
        self.col = col
        self.figsize = figsize
        self.tight_layout = tight_layout
        self.fn_finalize = fn_finalize
        self.subplots_kwargs = subplots_kwargs
        self.fig: Figure | None = None
        self.axes: Axes | np.ndarray | None = None
        self._current_index = 0

    def __enter__(self):
        # 创建图形和子图
        self.fig, self.axes = plt.subplots(
            self.row, self.col, figsize=self.figsize, **self.subplots_kwargs
        )
        if self.fig is None:
            raise AssertionError('Expected matplotlib to create a figure')

        if self.tight_layout:
            self.fig.tight_layout()

        # 保持 matplotlib 原生的 axes 结构，符合用户直觉
        # 1x1: self.axes 是单个 Axes 对象
        # 1xN 或 Nx1: self.axes 是一维 ndarray
        # NxM: self.axes 是二维 ndarray

        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is None:  # 仅在没有异常时执行
            if self.fn_finalize and self.fig is not None:
                self.fn_finalize(self.fig)
        # 注意：不自动关闭figure，让用户控制生命周期

    def __iter__(self):
        self._current_index = 0
        return self

    def __next__(self):
        if self._current_index >= self.row * self.col:
            raise StopIteration

        i = self._current_index // self.col
        j = self._current_index % self.col
        loc = GridLoc(i, j, self._current_index)

        self._current_index += 1
        axes = self.axes
        if axes is None:
            raise RuntimeError('GridContext must be entered before iteration')

        # 内部处理兼容性：根据 axes 的实际结构来访问
        if self.row == 1 and self.col == 1:
            # 单个 axes 对象
            ax = typing.cast(Axes, axes)
        elif self.row == 1 or self.col == 1:
            # 一维数组，使用线性索引
            ax = typing.cast(np.ndarray, axes)[self._current_index - 1]
        else:
            # 二维数组，使用二维索引
            ax = typing.cast(np.ndarray, axes)[i, j]

        return loc, ax


def make_grids(
    row: int = 1,
    col: int = 1,
    figsize: tuple | None = None,
    tight_layout: bool = True,
    fn_finalize: Callable[[Figure], None] = show_figure,
    **subplots_kwargs,
):
    """创建 matplotlib 的子图网格，返回可迭代的上下文管理器

    ```py
    # 上下文管理器用法 - 迭代访问
    with make_grids(2, 3) as grid_ctx:
        for loc, ax in grid_ctx:
            # loc.i, loc.j, loc.index 分别为子图的行列索引和索引
            ax.plot(np.random.randn(100))

    # 直接访问 axes（符合 matplotlib 原生行为）
    with make_grids(2, 3) as grid_ctx:
        # 1x1: grid_ctx.axes 是单个 Axes 对象
        # 1xN 或 Nx1: grid_ctx.axes 是一维 ndarray
        # NxM: grid_ctx.axes 是二维 ndarray
        grid_ctx.axes[0, 1].plot(data)  # 对于 2x3 网格
    ```

    Parameters
    ----------
    - `row` : `int`, optional
        - 行数, 默认为 1
    - `col` : `int`, optional
        - 列数, 默认为 1
    - `figsize` : `tuple`, optional
        - 图形大小 (宽, 高), 单位为英寸, 默认为 None (使用 matplotlib 默认大小)
    - `tight_layout` : `bool`, optional
        - 是否使用 tight_layout, 默认为 True
    - `fn_finalize` : `Callable[['Figure'], None]`, optional
        - 图表的最终处理函数, 默认为 None
    - `subplots_kwargs` : `dict`, optional
        - 其他参数, 传递给 `plt.subplots`

    Returns
    ----------
    - `GridContext`
        - 可迭代的上下文管理器，支持 with 语句和 for 循环
        - `grid_ctx.axes` 的结构完全符合 matplotlib 原生的 `plt.subplots()` 行为
    """
    return GridContext(row, col, figsize, tight_layout, fn_finalize, **subplots_kwargs)


class LayoutContext:
    """可迭代的 mosaic 布局上下文管理器"""

    def __init__(self, layout, figsize, fn_finalize, **mosaic_kwargs):
        self.layout = layout
        self.figsize = figsize
        self.fn_finalize = fn_finalize
        self.mosaic_kwargs = mosaic_kwargs
        self.fig: Figure | None = None
        self.axes: dict[str, Axes] | None = None
        self._axis_names: list[str] | None = None
        self._current_index = 0

    def __enter__(self):
        self.fig, self.axes = plt.subplot_mosaic(
            self.layout, figsize=self.figsize, **self.mosaic_kwargs
        )
        if self.axes is None:
            raise AssertionError('Expected subplot_mosaic to create axes')
        # 获取所有 axis 的名称，保持顺序
        self._axis_names = list(self.axes.keys())
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is None:  # 仅在没有异常时执行
            if self.fn_finalize and self.fig is not None:
                self.fn_finalize(self.fig)
        # 不自动关闭figure，让用户控制生命周期

    def __iter__(self):
        self._current_index = 0
        return self

    def __next__(self):
        if self._axis_names is None or self.axes is None:
            raise RuntimeError('LayoutContext must be entered before iteration')
        if self._current_index >= len(self._axis_names):
            raise StopIteration

        name = self._axis_names[self._current_index]
        ax = self.axes[name]
        self._current_index += 1

        return name, ax


def make_figure_layout(
    layout: typing.Sequence[typing.Sequence[str]] | str,
    figsize: tuple | None = None,
    fn_finalize: Callable[['Figure'], None] = show_figure,
    **mosaic_kwargs,
):
    """
    使用 subplot_mosaic 创建一个复杂的、命名的子图布局，返回可迭代的上下文管理器。

    ```py
    # 上下文管理器用法 - 可以直接访问 fig 和 axes
    layout = [['A', 'B'],
              ['A', 'C']]
    with make_figure_layout(layout) as layout_ctx:
        layout_ctx.axes['A'].plot(...)
        layout_ctx.axes['B'].hist(...)
        layout_ctx.axes['C'].scatter(...)

    # 可迭代用法 - 遍历所有子图
    with make_figure_layout(layout) as layout_ctx:
        for name, ax in layout_ctx:
            ax.plot(...)
            ax.set_title(f'Subplot {name}')
    ```

    Parameters
    ----------
    - `layout` : `list[list[str]]` or `str`
        - 传递给 plt.subplot_mosaic 的布局定义
    - `figsize` : `tuple`, optional
        - 图形大小, 默认为 None
    - `fn_finalize` : `Callable[['Figure'], None]`, optional
        - 图表的最终处理函数, 默认为 show_figure
    - `mosaic_kwargs` : `dict`, optional
        - 传递给 plt.subplot_mosaic 的其他参数

    Returns
    ----------
    - `LayoutContext`
        - 可迭代的上下文管理器，支持 with 语句和 for 循环
    """
    return LayoutContext(layout, figsize, fn_finalize, **mosaic_kwargs)
