"""绘图过程中常见的辅助性绘图

遵循 func(axes, ...) -> axes 的做法

- 辅助线, 如直线竖线等
- 箭头注释文本; 可自动计算位置
- 高亮区域
"""

import typing

import copy
from typing import (
    TYPE_CHECKING,
    Callable,
    List,
    Literal,
    Optional,
    Tuple,
    TypedDict,
    Union,
)

import numpy as np
from matplotlib.gridspec import GridSpec, SubplotSpec
from matplotlib.patches import ConnectionPatch, Rectangle
from matplotlib.ticker import EngFormatter, FuncFormatter
from typing_extensions import NotRequired, Required

if TYPE_CHECKING:
    from matplotlib.axes import Axes
    from matplotlib.figure import Figure


__all__ = [
    # 格式化工具
    'ScientificFormatter',
    'set_scientific_formatter',
    # 基础绘图组件
    'draw_reference_lines',
    'draw_highlight',
    'draw_annotate_arrow',
    'draw_text',
    # 专业绘图工具
    'draw_error_regions',
    'p_to_stars',
    'ComparisonDict',
    'draw_comparison_lines',
    # 高级布局工具
    'draw_inset_zoom',
    'create_broken_axes',
    'BrokenAxes',
    # 智能定位工具
    'calc_text_position',
    'avoid_overlap',
    'smart_position',
]


class ScientificFormatter:
    """科学记号格式化器"""

    @staticmethod
    def engineering(precision=1):
        """工程记号格式化器 (10^3, 10^6, 等)"""
        return EngFormatter(places=precision)

    @staticmethod
    def scientific(precision=2):
        """科学记号格式化器 (1.23e-4)"""
        return FuncFormatter(lambda x, p: f'{x:.{precision}e}')

    @staticmethod
    def percentage(precision=1):
        """百分比格式化器"""
        return FuncFormatter(lambda x, p: f'{x*100:.{precision}f}%')

    @staticmethod
    def pi_fraction():
        """π 分数格式化器"""

        def format_pi(x, p):
            if x == 0:
                return '0'
            coeff = x / np.pi
            if abs(coeff - round(coeff)) < 1e-10:
                coeff = int(round(coeff))
                if coeff == 1:
                    return 'π'
                elif coeff == -1:
                    return '-π'
                else:
                    return f'{coeff}π'
            else:
                return f'{coeff:.1f}π'

        return FuncFormatter(format_pi)


def set_scientific_formatter(
    ax: 'Axes',
    x_format: Literal['sci', 'eng', 'percent', 'pi'] | Callable | None = None,
    y_format: Literal['sci', 'eng', 'percent', 'pi'] | Callable | None = None,
):
    """
    为坐标轴应用科学格式化

    Parameters
    ----------
    ax : matplotlib.axes.Axes
        目标坐标轴
    x_format : str or formatter, optional
        X轴格式 ('sci', 'eng', 'percent', 'pi')
    y_format : str or formatter, optional
        Y轴格式
    """
    formatter_map = {
        'sci': ScientificFormatter.scientific(),
        'eng': ScientificFormatter.engineering(),
        'percent': ScientificFormatter.percentage(),
        'pi': ScientificFormatter.pi_fraction(),
    }

    if x_format:
        if isinstance(x_format, str):
            ax.xaxis.set_major_formatter(formatter_map[x_format])
        else:
            ax.xaxis.set_major_formatter(x_format)

    if y_format:
        if isinstance(y_format, str):
            ax.yaxis.set_major_formatter(formatter_map[y_format])
        else:
            ax.yaxis.set_major_formatter(y_format)


# ===================== 绘图辅助函数 =====================


def draw_reference_lines(
    ax: 'Axes',
    positions: Union[float, typing.Sequence[float], np.ndarray],
    orientation: Literal['vertical', 'horizontal'] = 'vertical',
    lim: tuple[float, float] | None = None,
    color: str = 'red',
    linestyle: str = '--',
    linewidth: float = 1.0,
    alpha: float = 0.8,
    label: Optional[str] = None,
    **kwargs,
):
    """
    绘制辅助线

    Parameters
    ----------
    - `ax` : `matplotlib.axes.Axes`
        - 目标坐标轴
    - `positions` : `float | list[float] | np.ndarray`
        - 线条位置，可以是单个值或值列表
    - `orientation` : `{'vertical', 'horizontal'}`, optional
        - 线条方向，by default 'vertical'
    - `lim` : `tuple[float, float] | None`, optional
        - 线条在交叉轴上的范围 (数据坐标)。例如，对于垂直线，lim=(ymin, ymax)。
          如果为 None, 线条将贯穿整个图表。 by default None.
        - 不同于 ax.axhline/ax.axvline 的 xmin/xmax/ymin/ymax 参数，这里是数据坐标; 不需要转换为 0-1 的相对坐标。
    - `color` : `str`, optional
        - 线条颜色, by default 'red'
    - `linestyle` : `str`, optional
        - 线条样式, by default '--'
    - `linewidth` : `float`, optional
        - 线条宽度, by default 1.0
    - `alpha` : `float`, optional
        - 透明度, by default 0.8
    - `label` : `str`, optional
        - 图例标签, by default None

    Returns
    ----------
    - `matplotlib.axes.Axes`
        - 返回修改后的坐标轴
    """
    if np.isscalar(positions):
        positions = typing.cast(typing.Sequence[float], [typing.cast(float, positions)])
    positions = np.asarray(positions)

    # 处理 lim 参数，将其从数据坐标转换为 0-1 的相对坐标
    ymin, ymax, xmin, xmax = 0, 1, 0, 1  # 默认贯穿全图
    if lim is not None:
        if len(lim) != 2:
            raise ValueError('lim must be a tuple of (min, max)')

        lim_start, lim_end = lim
        if orientation == 'vertical':
            ax_lim_start, ax_lim_end = ax.get_ylim()
            if ax_lim_end == ax_lim_start:  # Avoid division by zero
                ymin, ymax = 0, 1
            else:
                ymin = (lim_start - ax_lim_start) / (ax_lim_end - ax_lim_start)
                ymax = (lim_end - ax_lim_start) / (ax_lim_end - ax_lim_start)
        else:
            ax_lim_start, ax_lim_end = ax.get_xlim()
            if ax_lim_end == ax_lim_start:  # Avoid division by zero
                xmin, xmax = 0, 1
            else:
                xmin = (lim_start - ax_lim_start) / (ax_lim_end - ax_lim_start)
                xmax = (lim_end - ax_lim_start) / (ax_lim_end - ax_lim_start)

    style_kwargs = {
        'color': color,
        'linestyle': linestyle,
        'linewidth': linewidth,
        'alpha': alpha,
        **kwargs,
    }

    for i, pos in enumerate(positions):
        current_label = label if i == 0 else None  # 只有第一条线显示标签
        if orientation == 'vertical':
            ax.axvline(pos, ymin=ymin, ymax=ymax, label=current_label, **style_kwargs)
        else:
            ax.axhline(pos, xmin=xmin, xmax=xmax, label=current_label, **style_kwargs)

    return ax


def draw_highlight(
    ax: 'Axes',
    left: Optional[float] = None,
    right: Optional[float] = None,
    top: Optional[float] = None,
    bottom: Optional[float] = None,
    color: str = 'yellow',
    alpha: float = 0.3,
    label: Optional[str] = None,
    **kwargs,
):
    """
    高亮指定区域

    Parameters
    ----------
    - `ax` : `matplotlib.axes.Axes`
        - 目标坐标轴
    - `left` : `float`, optional
        - 左边界, by default None
    - `right` : `float`, optional
        - 右边界, by default None
    - `top` : `float`, optional
        - 上边界, by default None
    - `bottom` : `float`, optional
        - 下边界, by default None
    - `color` : `str`, optional
        - 填充颜色, by default 'yellow'
    - `alpha` : `float`, optional
        - 透明度, by default 0.3
    - `label` : `str`, optional
        - 图例标签, by default None

    Returns
    ----------
    - `matplotlib.axes.Axes`
        - 返回修改后的坐标轴
    """
    style_kwargs = {'color': color, 'alpha': alpha, 'label': label, **kwargs}

    if left is not None and right is not None and top is None and bottom is None:
        ax.axvspan(left, right, **style_kwargs)
    elif top is not None and bottom is not None and left is None and right is None:
        ax.axhspan(bottom, top, **style_kwargs)
    elif (
        left is not None and right is not None and top is not None and bottom is not None
    ):
        from matplotlib.patches import Rectangle

        rect = Rectangle((left, bottom), right - left, top - bottom, **style_kwargs)
        ax.add_patch(rect)
    else:
        xlim = ax.get_xlim()
        ylim = ax.get_ylim()
        if left is not None or right is not None:
            left_b = left if left is not None else xlim[0]
            right_b = right if right is not None else xlim[1]
            ax.axvspan(left_b, right_b, **style_kwargs)
        elif top is not None or bottom is not None:
            bottom_b = bottom if bottom is not None else ylim[0]
            top_b = top if top is not None else ylim[1]
            ax.axhspan(bottom_b, top_b, **style_kwargs)

    return ax


def draw_annotate_arrow(
    ax: 'Axes',
    xy: Tuple[float, float],
    text: str,
    xytext: Optional[Tuple[float, float]] = None,
    strategy: Literal['smart', 'simple', 'manual'] = 'simple',
    arrow_color: str = 'black',
    text_color: str = 'black',
    fontsize: int = 10,
    arrowstyle: str = '->',
    shrink: float = 5,
    ha: str = 'center',
    va: str = 'bottom',
    **kwargs,
):
    """
    智能箭头注释

    Parameters
    ----------
    - `ax` : `matplotlib.axes.Axes`
        - 目标坐标轴
    - `xy` : `tuple[float, float]`
        - 注释点的坐标 (x, y)
    - `text` : `str`
        - 注释文本
    - `xytext` : `tuple[float, float]`, optional
        - 文本位置，None时自动计算, by default None
    - `strategy` : `{'smart', 'simple', 'manual'}`, optional
        - 定位策略。如果提供了 `xytext`，则忽略此参数。
        - 'manual' 策略下，若 `xytext` 未提供，则使用一个简单的默认偏移。
        - by default 'simple'.
    - `arrow_color` : `str`, optional
        - 箭头颜色, by default 'black'
    - `text_color` : `str`, optional
        - 文本颜色, by default 'black'
    - `fontsize` : `int`, optional
        - 字体大小, by default 10
    - `arrowstyle` : `str`, optional
        - 箭头样式, by default '->'
    - `shrink` : `float`, optional
        - 箭头从两端收缩的点数 (points), by default 5.0

    Returns
    ----------
    - `matplotlib.axes.Axes`
        - 返回修改后的坐标轴
    """
    # xytext 的计算逻辑
    if xytext is None:
        if strategy in ['simple', 'smart']:
            xytext = calc_text_position(
                ax, xy, text, strategy=typing.cast(Literal['smart', 'simple'], strategy)
            )
        else:  # 'manual' 策略或未知策略的回退情况
            x, y = xy
            xlim = ax.get_xlim()
            ylim = ax.get_ylim()
            offset_x = (xlim[1] - xlim[0]) * 0.1
            offset_y = (ylim[1] - ylim[0]) * 0.1
            xytext = (x + offset_x, y + offset_y)

    arrowprops = {
        'arrowstyle': arrowstyle,
        'color': arrow_color,
        'shrinkA': shrink,
        'shrinkB': shrink,
    }
    for k, v in kwargs.items():
        if k.startswith('arrow'):
            prop_name = k.replace('arrow_', '')
            arrowprops[prop_name] = v

    text_kwargs = {'fontsize': fontsize, 'color': text_color}
    for k, v in kwargs.items():
        if not k.startswith('arrow'):
            text_kwargs[k] = v

    ax.annotate(
        text, xy=xy, xytext=xytext, arrowprops=arrowprops, ha=ha, va=va, **text_kwargs
    )
    return ax


def draw_text(
    ax: 'Axes',
    text: str,
    xy: Tuple[float, float],
    anchor: Literal[
        'center',
        'top',
        'bottom',
        'left',
        'right',
        'top_left',
        'top_right',
        'bottom_left',
        'bottom_right',
    ] = 'center',
    color: str = 'black',
    fontsize: int = 10,
    **kwargs,
):
    """
    在指定坐标点绘制文本，并精确控制该坐标点对应于文本框的位置。

    Parameters
    ----------
    - `ax` : `matplotlib.axes.Axes`
        - 目标坐标轴。
    - `text` : `str`
        - 要显示的文本内容。
    - `xy` : `tuple[float, float]`
        - 文本的锚点坐标 (x, y)，使用数据坐标系。
    - `anchor` : str, optional
        - 定义 `xy` 坐标点是文本框的哪个位置。
        - by default 'center'.
    - `color` : `str`, optional
        - 文本颜色, by default 'black'.
    - `fontsize` : `int`, optional
        - 字体大小, by default 10.
    - `**kwargs`
        - 其他所有传递给 `ax.text` 的关键字参数。

    Returns
    ----------
    - `matplotlib.axes.Axes`
        - 返回修改后的坐标轴。
    """
    anchor_map = {
        'center': {'ha': 'center', 'va': 'center'},
        'top': {'ha': 'center', 'va': 'top'},
        'bottom': {'ha': 'center', 'va': 'bottom'},
        'left': {'ha': 'left', 'va': 'center'},
        'right': {'ha': 'right', 'va': 'center'},
        'top_left': {'ha': 'left', 'va': 'top'},
        'top_right': {'ha': 'right', 'va': 'top'},
        'bottom_left': {'ha': 'left', 'va': 'bottom'},
        'bottom_right': {'ha': 'right', 'va': 'bottom'},
    }
    if anchor not in anchor_map:
        raise ValueError(
            f"Invalid anchor: '{anchor}'. Must be one of {list(anchor_map.keys())}"
        )

    alignment_props = anchor_map[anchor]
    text_props = {'color': color, 'fontsize': fontsize, **alignment_props, **kwargs}
    ax.text(x=xy[0], y=xy[1], s=text, **text_props)
    return ax


# ===================== 专业绘图工具 =====================


def draw_error_regions(
    ax: 'Axes',
    x: np.ndarray,
    y: np.ndarray,
    yerr: Optional[np.ndarray] = None,
    xerr: Optional[np.ndarray] = None,
    color: str = 'blue',
    alpha: float = 0.3,
    label: Optional[str] = None,
    **kwargs,
):
    """
    绘制误差区间

    Parameters
    ----------
    - `ax` : `matplotlib.axes.Axes`
        - 目标坐标轴
    - `x` : `np.ndarray`
        - X轴数据
    - `y` : `np.ndarray`
        - Y轴数据
    - `yerr` : `np.ndarray`, optional
        - Y方向误差, by default None
    - `xerr` : `np.ndarray`, optional
        - X方向误差, by default None
    - `color` : `str`, optional
        - 填充颜色, by default 'blue'
    - `alpha` : `float`, optional
        - 透明度, by default 0.3
    - `label` : `str`, optional
        - 图例标签, by default None

    Returns
    ----------
    - `matplotlib.axes.Axes`
        - 返回修改后的坐标轴
    """
    style_kwargs = {'color': color, 'alpha': alpha, **kwargs}

    # 仅为第一个绘制的区域添加标签，避免图例重复
    has_labeled = False
    if yerr is not None:
        ax.fill_between(x, y - yerr, y + yerr, label=label, **style_kwargs)
        has_labeled = True

    if xerr is not None:
        current_label = label if not has_labeled else None
        ax.fill_betweenx(y, x - xerr, x + xerr, label=current_label, **style_kwargs)

    return ax


def p_to_stars(p_value: float) -> str:
    """根据 p-value 返回显著性星号。"""
    if p_value < 0.001:
        return '***'
    elif p_value < 0.01:
        return '**'
    elif p_value < 0.05:
        return '*'
    else:
        return 'ns'  # Not Significant


class ComparisonDict(TypedDict, total=False):
    """
    使用 TypedDict 描述一次组间比较所需的数据。

    Keys:
    - x1: 第一个比较组的 x 轴坐标。
    - y1: 第一个比较组的 y 轴值 (例如条形图的高度)。
    - x2: 第二个比较组的 x 轴坐标。
    - y2: 第二个比较组的 y 轴值。
    - text (optional): 显示在比较线上方的文本。
    - p_value (optional): 如果提供了 p 值，并且 text 未定义，会自动生成星号。
    """

    x1: Required[float]
    y1: Required[float]
    x2: Required[float]
    y2: Required[float]
    text: NotRequired[str]
    p_value: NotRequired[float]


def _as_2d_float_array(values: object) -> np.ndarray:
    array = np.asarray(values, dtype=float)
    if array.ndim == 1:
        array = np.atleast_2d(array)
    return typing.cast(np.ndarray, array)


def _as_axes_list(values: object) -> list['Axes']:
    if isinstance(values, np.ndarray):
        return typing.cast(list['Axes'], values.ravel().tolist())
    return [typing.cast('Axes', values)]


def draw_comparison_lines(
    ax: 'Axes',
    comparisons: typing.Sequence[ComparisonDict],
    stack_gap: float = 0.05,
    abs_stack_gap: Optional[float] = None,
    color: str = 'black',
    linewidth: float = 1.0,
    **line_kwargs,
):
    """
    绘制分组数据对比线，垂直线能精确连接到数据点，并具有自动堆叠功能。

    Parameters
    ----------
    - `ax` : `matplotlib.axes.Axes`
        - 目标坐标轴
    - `comparisons` : `List[ComparisonDict]`
        - 比较对象字典的列表。
    - `stack_gap` : `float`, optional
        - 比较线之间的垂直间隙（占 Y 轴范围的比例）, by default 0.05
    - `abs_stack_gap` : `float`, optional
        - 如果提供，将覆盖 stack_gap 的计算。表示绝对的 Y 轴数据单位间隙。
    - `color` : `str`, optional
        - 线条颜色, by default 'black'
    - `linewidth` : `float`, optional
        - 线条宽度, by default 1.0

    Returns
    ----------
    - `matplotlib.axes.Axes`
        - 返回修改后的坐标轴
    """
    if not comparisons:
        return ax

    sorted_comparisons = sorted(comparisons, key=lambda comp: max(comp['y1'], comp['y2']))

    y_min, y_max_orig = ax.get_ylim()
    y_range = y_max_orig - y_min
    if abs_stack_gap is None:
        abs_stack_gap = stack_gap * y_range

    level_manager = []
    max_y_drawn = y_max_orig

    for comp in sorted_comparisons:
        x1, y1, x2, y2 = comp['x1'], comp['y1'], comp['x2'], comp['y2']
        x_start, x_end = min(x1, x2), max(x1, x2)
        data_y_max = max(y1, y2)
        horizontal_y = data_y_max + abs_stack_gap

        # 碰撞检测和堆叠
        is_colliding = True
        while is_colliding:
            is_colliding = False
            for drawn_x_start, drawn_x_end, drawn_y in level_manager:
                # 检查 x 轴区间是否重叠
                if max(x_start, drawn_x_start) < min(x_end, drawn_x_end):
                    # 检查 y 轴位置是否过近
                    if abs(horizontal_y - drawn_y) < abs_stack_gap:
                        horizontal_y = drawn_y + abs_stack_gap
                        is_colliding = True
                        break  # 重新开始检查新高度
            if is_colliding:
                continue  # 如果发生了碰撞并调整了高度，则重新检查所有已绘制的线
            else:
                break  # 没有碰撞，当前高度可用

        level_manager.append((x_start, x_end, horizontal_y))
        max_y_drawn = max(max_y_drawn, horizontal_y)

        # 绘制线条
        ax.plot(
            [x1, x1], [y1, horizontal_y], color=color, linewidth=linewidth, **line_kwargs
        )
        ax.plot(
            [x2, x2], [y2, horizontal_y], color=color, linewidth=linewidth, **line_kwargs
        )
        ax.plot(
            [x1, x2],
            [horizontal_y, horizontal_y],
            color=color,
            linewidth=linewidth,
            **line_kwargs,
        )

        # 绘制文本
        display_text = comp.get('text')
        if display_text is None and 'p_value' in comp:
            display_text = p_to_stars(comp['p_value'])

        if display_text:
            ax.annotate(
                text=display_text,
                xy=((x1 + x2) / 2, horizontal_y),
                xytext=(0, 2),
                textcoords='offset points',
                ha='center',
                va='bottom',
                fontsize=line_kwargs.get('fontsize', 10),
            )

    ax.set_ylim(y_min, max_y_drawn + abs_stack_gap)
    return ax


# ===================== 高级布局工具 =====================


def draw_inset_zoom(
    ax: 'Axes',
    bounds: Tuple[float, float, float, float],
    inset_position: Tuple[float, float, float, float],
    connect_corners: Optional[Tuple[int, int]] = (3, 4),
    mark_kwargs: Optional[typing.Mapping] = None,
    connector_kwargs: Optional[typing.Mapping] = None,
) -> 'Axes':
    """
    在现有坐标轴上创建一个局部放大的内嵌图。

    此函数会在 `ax` 上绘制一个矩形标记放大的区域，并创建一个新的内嵌坐标轴
    `ax_inset` 来显示放大后的内容，同时绘制连接线。

    Parameters
    ----------
    - `ax` : `matplotlib.axes.Axes`
        - 父坐标轴。
    - `bounds` : `tuple[float, float, float, float]`
        - 要放大区域的边界，格式为 `(x, y, width, height)`，使用数据坐标。
    - `inset_position` : `tuple[float, float, float, float]`
        - 内嵌图在父坐标轴中的位置和大小，格式为 `(x, y, width, height)`，
          使用父坐标轴的相对坐标 (0-1 范围)。
    - `connect_corners` : `tuple[int, int] | None`, optional
        - 指定连接线的端点。元组中的两个数字分别代表标记框和内嵌图的角点。
        - 角点编号: 1 (右上), 2 (左上), 3 (左下), 4 (右下)。
        - `(3, 4)` 表示连接标记框的左下角(3)和内嵌图的右下角(4)。
        - 如果为 `None`，则不绘制连接线。 by default (3, 4).
    - `mark_kwargs` : `dict`, optional
        - 传递给 `matplotlib.patches.Rectangle` 的参数，用于样式化标记框。
        - 默认: `{'ec': 'black', 'fc': 'none', 'lw': 1.0}`.
    - `connector_kwargs` : `dict`, optional
        - 传递给 `matplotlib.patches.ConnectionPatch` 的参数，用于样式化连接线。
        - 默认: `{'ec': 'black', 'lw': 0.8, 'arrowstyle': '-'}`.

    Returns
    -------
    - `matplotlib.axes.Axes`
        - 返回新创建的内嵌坐标轴 `ax_inset`，可继续对其进行操作。
    """
    # 1. 创建内嵌坐标轴
    ax_inset = ax.inset_axes(inset_position)

    # 2. [最终修正] 提取安全样式属性并重新绘制
    # 复制线图 (plot)
    for line in ax.lines:
        xy_data = _as_2d_float_array(line.get_xydata())
        if xy_data.shape[1] < 2:
            continue
        x_data = xy_data[:, 0]
        y_data = xy_data[:, 1]
        # 提取安全的视觉样式属性
        style_props = {
            'color': line.get_color(),
            'linestyle': line.get_linestyle(),
            'linewidth': line.get_linewidth(),
            'marker': line.get_marker(),
            'markersize': line.get_markersize(),
            'markerfacecolor': line.get_markerfacecolor(),
            'markeredgecolor': line.get_markeredgecolor(),
            'alpha': line.get_alpha(),
            'label': line.get_label(),
        }
        ax_inset.plot(x_data, y_data, **style_props)

    # 复制散点图 (scatter)
    for collection in ax.collections:
        offsets = _as_2d_float_array(collection.get_offsets())
        if offsets.shape[1] < 2 or offsets.size == 0:
            continue
        # 提取安全的视觉样式属性
        style_props = {
            'c': collection.get_facecolor(),
            'marker': collection.get_paths()[0] if collection.get_paths() else 'o',
            'edgecolor': collection.get_edgecolor(),
            'alpha': collection.get_alpha(),
            'label': collection.get_label(),
        }
        get_sizes = getattr(collection, 'get_sizes', None)
        if callable(get_sizes):
            sizes = np.asarray(get_sizes(), dtype=float)
            if sizes.size > 0:
                style_props['s'] = sizes
        ax_inset.scatter(offsets[:, 0], offsets[:, 1], **style_props)

    # 3. 设置内嵌图的坐标范围以实现“放大”
    x, y, width, height = bounds
    ax_inset.set_xlim(x, x + width)
    ax_inset.set_ylim(y, y + height)
    ax_inset.tick_params(axis='both', which='major', labelsize=8)

    # 4. 在父坐标轴上绘制标记框
    mark_style = {'ec': 'black', 'fc': 'none', 'lw': 1.0, 'transform': ax.transData}
    if mark_kwargs:
        mark_style.update(mark_kwargs)
    rect = Rectangle((x, y), width, height, **mark_style)
    ax.add_patch(rect)

    # 5. 绘制连接线
    if connect_corners:
        connector_style = {'ec': 'black', 'lw': 0.8, 'arrowstyle': '-'}
        if connector_kwargs:
            connector_style.update(connector_kwargs)

        rect_corners = [(x + width, y + height), (x, y + height), (x, y), (x + width, y)]
        inset_corners = [(1, 1), (0, 1), (0, 0), (1, 0)]

        c1_idx, c2_idx = connect_corners[0] - 1, connect_corners[1] - 1

        con = ConnectionPatch(
            xyA=rect_corners[c1_idx],
            coordsA=ax.transData,
            xyB=inset_corners[c2_idx],
            coordsB=ax_inset.transAxes,
            **connector_style,
        )
        ax.figure.add_artist(con)

    return ax_inset


class BrokenAxes:
    """
    一个代理类，用于模拟一个单一的坐标轴，但内部将绘图命令
    分发到一组真实的、具有断裂效果的子坐标轴上。
    此实现参考了 Matplotlib 官方的 Broken Axis 示例。
    """

    def __init__(
        self,
        fig: 'Figure',
        subplot_spec: SubplotSpec | tuple[int, int, int],
        breaks: typing.Sequence[Tuple[float, float]],
        orientation: Literal['vertical', 'horizontal'] = 'vertical',
        d: float = 0.015,  # 标记尺寸
    ):
        self.orientation = orientation
        self.breaks = breaks
        self.d = d
        self.axes: list[Axes]

        if isinstance(subplot_spec, tuple):
            gs_main = GridSpec(*subplot_spec[:2], figure=fig)[subplot_spec[2] - 1]
        else:
            gs_main = subplot_spec

        n_breaks = len(breaks)
        # 增加 hspace/wspace 以提供视觉间距
        if orientation == 'vertical':
            height_ratios = [b[1] - b[0] for b in breaks]
            gs_sub = gs_main.subgridspec(
                n_breaks, 1, hspace=0.1, height_ratios=height_ratios
            )
            created_axes = gs_sub.subplots(sharex=True)
            self.axes = _as_axes_list(created_axes)[::-1]
        else:  # horizontal
            width_ratios = [b[1] - b[0] for b in breaks]
            gs_sub = gs_main.subgridspec(
                1, n_breaks, wspace=0.1, width_ratios=width_ratios
            )
            created_axes = gs_sub.subplots(sharey=True)
            self.axes = _as_axes_list(created_axes)

        for i, ax in enumerate(self.axes):
            if orientation == 'vertical':
                ax.set_ylim(breaks[i])
                ax.spines[['right']].set_visible(False)
                if i != n_breaks - 1:
                    ax.spines[['top']].set_visible(False)
                if i != 0:
                    ax.spines[['bottom']].set_visible(False)
            else:  # horizontal
                ax.set_xlim(breaks[i])
                ax.spines[['top', 'bottom']].set_visible(
                    False
                )  # [修正] 水平模式也隐藏上下 spine
                if i != n_breaks - 1:
                    ax.spines[['right']].set_visible(False)
                if i != 0:
                    ax.spines[['left']].set_visible(False)

            # 绘制断裂标记
            if i > 0:
                pos = 'bottom' if orientation == 'vertical' else 'left'
                self._draw_break_marks(ax, pos)
            if i < n_breaks - 1:
                pos = 'top' if orientation == 'vertical' else 'right'
                self._draw_break_marks(ax, pos)

    def _draw_break_marks(self, ax, position):
        kwargs = dict(transform=ax.transAxes, color='k', clip_on=False, lw=1)
        if position == 'bottom':
            ax.plot((-self.d, +self.d), (-self.d, +self.d), **kwargs)
            ax.plot((1 - self.d, 1 + self.d), (-self.d, +self.d), **kwargs)
        elif position == 'top':
            ax.plot((-self.d, +self.d), (1 - self.d, 1 + self.d), **kwargs)
            ax.plot((1 - self.d, 1 + self.d), (1 - self.d, 1 + self.d), **kwargs)
        elif position == 'left':
            ax.plot((-self.d, +self.d), (-self.d, +self.d), **kwargs)
            ax.plot((-self.d, +self.d), (1 - self.d, 1 + self.d), **kwargs)
        elif position == 'right':
            ax.plot((1 - self.d, 1 + self.d), (-self.d, +self.d), **kwargs)
            ax.plot((1 - self.d, 1 + self.d), (1 - self.d, 1 + self.d), **kwargs)

    def __getattr__(self, name):
        # 绘图命令需要分发
        if name in ['plot', 'bar', 'scatter', 'errorbar', 'axhline', 'axvline', 'grid']:

            def method_proxy(*args, **kwargs):
                results = []
                for ax in self.axes:
                    results.append(getattr(ax, name)(*args, **kwargs))
                return results

            return method_proxy

        # [核心修正] text 方法需要特殊处理
        if name == 'text':

            def text_proxy(x, y, s, **kwargs):
                for ax in self.axes:
                    if self.orientation == 'vertical':
                        ymin, ymax = ax.get_ylim()
                        if ymin <= y <= ymax:
                            ax.text(x, y, s, **kwargs)
                    else:  # horizontal
                        xmin, xmax = ax.get_xlim()
                        if xmin <= x <= xmax:
                            ax.text(x, y, s, **kwargs)

            return text_proxy

        # [核心修正] 标签和图例设置，区分方向
        if name == 'set_xlabel':
            return getattr(self.axes[0], 'set_xlabel')
        if name == 'set_ylabel':
            if self.orientation == 'vertical':

                def set_ylabel_proxy(label, **kwargs):
                    self.axes[0].figure.text(
                        0.04,
                        0.5,
                        label,
                        ha='center',
                        va='center',
                        rotation='vertical',
                        **kwargs,
                    )

                return set_ylabel_proxy
            else:  # horizontal
                return getattr(self.axes[0], 'set_ylabel')

        if name == 'set_title':
            ax_to_use = (
                self.axes[-1]
                if self.orientation == 'vertical'
                else self.axes[len(self.axes) // 2]
            )
            return getattr(ax_to_use, 'set_title')

        if name == 'legend':
            return getattr(self.axes[-1], 'legend')

        # 其他属性默认返回第一个（最底层/最左侧）轴的属性
        return getattr(self.axes[0], name)


def create_broken_axes(
    fig: 'Figure',
    subplot_spec: SubplotSpec | tuple[int, int, int],
    breaks: typing.Sequence[Tuple[float, float]],
    orientation: Literal['vertical', 'horizontal'] = 'vertical',
    d: float = 0.015,
) -> 'BrokenAxes':
    """
    创建一个断裂坐标轴，并返回一个可直接用于绘图的代理对象。

    参考: [Broken axis — Matplotlib 3.10.3 documentation](https://matplotlib.org/stable/gallery/subplots_axes_and_figures/broken_axis.html)
    """
    return BrokenAxes(fig, subplot_spec, breaks, orientation, d=d)


# ===================== 智能辅助工具函数 =====================


def calc_text_position(
    ax: 'Axes',
    xy: Tuple[float, float],
    text: str,
    avoid_data: bool = True,
    strategy: Literal['smart', 'simple'] = 'smart',
) -> Tuple[float, float]:
    """
    计算文本标注的最佳位置

    Parameters
    ----------
    - `ax` : `matplotlib.axes.Axes`
        - 目标坐标轴
    - `xy` : `tuple[float, float]`
        - 注释点的坐标 (x, y)
    - `text` : `str`
        - 注释文本
    - `avoid_data` : `bool`, optional
        - 是否避免与数据点重叠, by default True
    - `strategy` : `{'smart', 'simple'}`, optional
        - 定位策略, by default 'smart'

    Returns
    ----------
    - `tuple[float, float]`
        - 计算得到的文本位置 (x, y)
    """
    x, y = xy
    xlim = ax.get_xlim()
    ylim = ax.get_ylim()
    x_range = xlim[1] - xlim[0]
    y_range = ylim[1] - ylim[0]

    # 回退策略
    fallback_pos = (x + x_range * 0.05, y + y_range * 0.05)

    if strategy == 'simple':
        return fallback_pos

    elif strategy == 'smart':
        candidates = [
            (x + x_range * 0.1, y + y_range * 0.1),  # 右上
            (x - x_range * 0.1, y + y_range * 0.1),  # 左上
            (x + x_range * 0.1, y - y_range * 0.1),  # 右下
            (x - x_range * 0.1, y - y_range * 0.1),  # 左下
        ]
        valid_candidates = [
            c
            for c in candidates
            if xlim[0] <= c[0] <= xlim[1] and ylim[0] <= c[1] <= ylim[1]
        ]

        if not valid_candidates:
            return fallback_pos

        # [增强] 不仅考虑散点图，还考虑线图数据
        if avoid_data:
            all_points = []
            for collection in ax.collections:
                if hasattr(collection, 'get_offsets'):
                    offsets = _as_2d_float_array(collection.get_offsets())
                    if offsets.size > 0:
                        all_points.append(offsets)
            for line in ax.lines:
                if hasattr(line, 'get_xydata'):
                    xy_data = _as_2d_float_array(line.get_xydata())
                    if xy_data.size > 0:
                        all_points.append(xy_data)

            if all_points:
                all_points_np = np.vstack(all_points)
                best_pos = valid_candidates[0]
                max_min_dist = -1

                for candidate in valid_candidates:
                    dists = np.sqrt(np.sum((all_points_np - candidate) ** 2, axis=1))
                    min_dist = np.min(dists)
                    if min_dist > max_min_dist:
                        max_min_dist = min_dist
                        best_pos = candidate
                return best_pos

        return valid_candidates[0]

    return fallback_pos


def avoid_overlap(
    ax,
    annotations: typing.Sequence[Tuple[Tuple[float, float], str]],
    initial_positions: Optional[typing.Sequence[Tuple[float, float]]] = None,
    method: Literal['force_directed', 'greedy', 'simple'] = 'force_directed',
    min_distance: float = 0.02,
) -> List[Tuple[float, float]]:
    """
    自动调整多个注释的位置避免重叠

    Parameters
    ----------
    - `ax` : `matplotlib.axes.Axes`
        - 目标坐标轴
    - `annotations` : `list[tuple[tuple[float, float], str]]`
        - 注释列表，每个元素为 ((x, y), text)
    - `initial_positions` : `list[tuple[float, float]]`, optional
        - 提供优化的起始位置。如果为 None, 则自动计算。
    - `method` : `{'force_directed', 'greedy', 'simple'}`, optional
        - 避免重叠的方法, by default 'force_directed'
    - `min_distance` : `float`, optional
        - 最小间距（相对于轴范围的比例）, by default 0.02

    Returns
    ----------
    - `list[tuple[float, float]]`
        - 调整后的文本位置列表
    """
    if len(annotations) <= 1:
        return (
            list(initial_positions)
            if initial_positions
            else [calc_text_position(ax, pos, text) for pos, text in annotations]
        )

    xlim, ylim = ax.get_xlim(), ax.get_ylim()
    x_range, y_range = xlim[1] - xlim[0], ylim[1] - ylim[0]
    min_dist = min_distance * min(x_range, y_range)

    # [修正] 使用提供的初始位置或自动计算
    if initial_positions:
        positions = list(initial_positions)
    else:
        positions = [
            calc_text_position(ax, pos, text, strategy='simple')
            for pos, text in annotations
        ]

    if method == 'simple' or len(positions) <= 1:
        return positions

    if method == 'greedy':
        for i in range(len(positions)):
            for j in range(i):
                p1, p2 = np.array(positions[i]), np.array(positions[j])
                dist = np.linalg.norm(p1 - p2)
                if dist < min_dist:
                    direction = (p1 - p2) / dist
                    positions[i] = tuple(p2 + direction * min_dist)
        return positions

    elif method == 'force_directed':
        positions_np = np.array(positions)
        for _ in range(50):
            forces = np.zeros_like(positions_np)
            for i in range(len(positions_np)):
                for j in range(i + 1, len(positions_np)):
                    diff = positions_np[i] - positions_np[j]
                    dist = np.linalg.norm(diff)
                    if 0 < dist < min_dist:
                        force = diff / dist * (min_dist - dist) * 0.1
                        forces[i] += force
                        forces[j] -= force
            positions_np += forces
            if np.max(np.linalg.norm(forces, axis=1)) < 1e-6:
                break
        return [(float(x), float(y)) for x, y in positions_np]

    return positions


def smart_position(
    ax,
    points: typing.Sequence[Tuple[float, float]],
    texts: typing.Sequence[str],
    spacing: Union[float, Literal['auto']] = 'auto',
    prefer_side: Literal['auto', 'left', 'right', 'top', 'bottom'] = 'auto',
) -> List[Tuple[float, float]]:
    """
    为一组点自动分配最佳的注释位置

    Parameters
    ----------
    - `ax` : `matplotlib.axes.Axes`
        - 目标坐标轴
    - `points` : `list[tuple[float, float]]`
        - 要注释的点列表
    - `texts` : `list[str]`
        - 对应的文本列表
    - `spacing` : `float | 'auto'`, optional
        - 注释间距，'auto'时自动计算, by default 'auto'
    - `prefer_side` : `{'auto', 'left', 'right', 'top', 'bottom'}`, optional
        - 偏好的注释方向, by default 'auto'

    Returns
    ----------
    - `list[tuple[float, float]]`
        - 计算得到的文本位置列表
    """
    if len(points) != len(texts):
        raise ValueError('points 和 texts 长度必须相同')
    if not points:
        return []

    xlim, ylim = ax.get_xlim(), ax.get_ylim()
    x_range, y_range = xlim[1] - xlim[0], ylim[1] - ylim[0]

    if spacing == 'auto':
        spacing = min(x_range, y_range) * 0.05

    if prefer_side == 'auto':
        points_array = np.array(points)
        center_x, center_y = np.mean(points_array[:, 0]), np.mean(points_array[:, 1])
        space_right = xlim[1] - center_x
        space_left = center_x - xlim[0]
        space_top = ylim[1] - center_y
        space_bottom = center_y - ylim[0]
        spaces = {
            'right': space_right,
            'left': space_left,
            'top': space_top,
            'bottom': space_bottom,
        }
        prefer_side = typing.cast(
            Literal['left', 'right', 'top', 'bottom'],
            max(spaces, key=lambda side: spaces[side]),
        )

    resolved_spacing = float(spacing)

    initial_positions = []
    for i, (x, y) in enumerate(points):
        offset_factor = i - (len(points) - 1) / 2
        if prefer_side == 'right':
            offset = (resolved_spacing, resolved_spacing * offset_factor)
        elif prefer_side == 'left':
            offset = (-resolved_spacing, resolved_spacing * offset_factor)
        elif prefer_side == 'top':
            offset = (resolved_spacing * offset_factor, resolved_spacing)
        else:  # bottom
            offset = (resolved_spacing * offset_factor, -resolved_spacing)
        initial_positions.append((x + offset[0], y + offset[1]))

    # [核心修正] 将计算出的初始位置传递给 avoid_overlap 函数进行优化
    annotations = list(zip(points, texts))
    # 使用 'force_directed' 以获得更优的布局效果
    return avoid_overlap(
        ax, annotations, initial_positions=initial_positions, method='force_directed'
    )
