from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Callable, Literal

import matplotlib.pyplot as plt
import numpy as np

if TYPE_CHECKING:
    Figure = plt.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,
    bottom: float = None,
    left: float = None,
    right: float = 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: plt.Figure | None = None):
    """显示 matplotlib 图表"""
    if fig is not None:
        fig.tight_layout()
        plt.show()
        plt.close(fig)  # 关闭图形


def save_figure(
    fig: plt.Figure,
    path: str | Path,
    dpi: int = 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: plt.Figure = None
        self.axes: plt.Axes = 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.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:
                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 的实际结构来访问
        if self.row == 1 and self.col == 1:
            # 单个 axes 对象
            ax = self.axes
        elif self.row == 1 or self.col == 1:
            # 一维数组，使用线性索引
            ax = self.axes[self._current_index - 1]
        else:
            # 二维数组，使用二维索引
            ax = self.axes[i, j]

        return loc, ax


def make_grids(
    row: int = 1,
    col: int = 1,
    figsize: tuple = 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 = None
        self.axes = None
        self._axis_names = None
        self._current_index = 0

    def __enter__(self):
        self.fig, self.axes = plt.subplot_mosaic(
            self.layout, figsize=self.figsize, **self.mosaic_kwargs
        )
        # 获取所有 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:
                self.fn_finalize(self.fig)
        # 不自动关闭figure，让用户控制生命周期

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

    def __next__(self):
        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: list[list[str]] | str,
    figsize: tuple = 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)
