"""
封装一些特殊的绘图功能

### 设计哲学

本模块中的函数专注于特殊场景的可视化需求，遵循与 plot_common.py 相同的设计规范。
"""

import typing

from typing import Dict, List, Literal, Optional, Tuple

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.axes import Axes
from matplotlib.colors import ListedColormap
from matplotlib.patches import Patch

from ._plot_utils import FigAx, _finalize_figure, _prepare_ax


def plot_group_divide(
    records: typing.Sequence[typing.Mapping[str, typing.Sequence[str]]],
    object_order: Optional[typing.Sequence[str]] = None,
    group_order: Optional[typing.Sequence[str]] = None,
    cmap: Optional[str] = None,
    title: Optional[str] = None,
    xlabel: Optional[str] = None,
    ylabel: Optional[str] = None,
    ax: Optional[Axes] = None,
    grid_label: bool = True,
    show: bool = True,
    save_path: Optional[str] = None,
    auto_close: Optional[bool] = None,
    xtick_rotation: int | Literal[45, 90] = 45,
    figsize: Optional[tuple[float, float]] = None,
    **kwargs,
) -> FigAx:
    """可视化分组划分情况。在 K fold 切分文件的时候可以用来检查文件切分是否均分。

    展示多次分组中每个对象的组别分配情况，使用网格形式显示。
    X轴为对象，Y轴为分组次数，颜色表示不同的组。

    Parameters
    ----------
    - `records` : `List[Dict[str, List[str]]]`
        - M次分组记录。每个元素是一个字典，键为组名，值为该组包含的对象列表
    - `object_order` : `Optional[List[str]]`, optional
        - 指定对象在X轴上的顺序。如果为None，则按首次出现顺序排列, by default None
    - `group_order` : `Optional[List[str]]`, optional
        - 指定组名的顺序，影响颜色分配。如果为None，则按首次出现顺序排列, by default None
    - `cmap` : `Optional[str]`, optional
        - 颜色映射方案。如果为None，则优先使用当前matplotlib配置中的cycle颜色，
        - 如果获取失败则使用Set2配色方案, by default None
    - `title` : `Optional[str]`, optional
        - 图表标题, by default None
    - `xlabel` : `Optional[str]`, optional
        - X轴标签, by default None
    - `ylabel` : `Optional[str]`, optional
        - Y轴标签, by default None
    - `ax` : `Optional[plt.Axes]`, optional
        - 绘图轴对象。如果为None则创建新图表, by default None
    - `show` : `bool`, optional
        - 是否显示图表, by default True
    - `save_path` : `Optional[str]`, optional
        - 保存路径。如果提供则保存图表, by default None
    - `auto_close` : `Optional[bool]`, optional
        - 是否自动关闭图表释放内存, by default None

    Returns
    ----------
    - `FigAx`
        - 图表对象和轴对象的元组

    Raises
    ----------
    - `ValueError`
        - 当records为空时抛出
    - `ValueError`
        - 当某个对象在同一次分组中出现在多个组时抛出

    Examples
    --------
    **基本使用**：
    ```python
    # 3次分组，每次分为train/val/test三组
    records = [
        {
            'train': ['obj1', 'obj2', 'obj3'],
            'val': ['obj4', 'obj5'],
            'test': ['obj6', 'obj7']
        },
        {
            'train': ['obj2', 'obj4', 'obj6'],
            'val': ['obj1', 'obj7'],
            'test': ['obj3', 'obj5']
        }
    ]

    fig, ax = plot_group_devide(records)
    ax.set_title('K-Fold Cross Validation Split')
    ax.set_xlabel('Objects')
    ax.set_ylabel('Fold Number')
    ```

    **指定对象和组的顺序**：
    ```python
    fig, ax = plot_group_devide(
        records,
        object_order=['obj1', 'obj2', 'obj3', 'obj4', 'obj5', 'obj6', 'obj7'],
        group_order=['train', 'val', 'test']
    )
    ```
    **可视化说明**：

    - 每个网格表示一个对象在某次分组中的组别
    - 不同颜色代表不同的组
    - 网格内的文字是组名的缩写
    - 右侧图例显示颜色与组名的对应关系
    - 如果某对象在某次分组中不存在，该网格将显示为灰色
    """
    if not records:
        raise ValueError('records 不能为空')

    fig, ax = _prepare_ax(ax, figsize=figsize)

    # 1. 提取所有对象和组名
    all_objects = set()
    all_groups = set()

    for record in records:
        for group_name, objects in record.items():
            all_groups.add(group_name)
            all_objects.update(objects)

    # 2. 确定对象和组的顺序
    if object_order is None:
        # 按首次出现顺序
        object_order = []
        for record in records:
            for group_name, objects in record.items():
                for obj in objects:
                    if obj not in object_order:
                        object_order.append(obj)
    else:
        # 检查是否包含所有对象
        missing_objects = all_objects - set(object_order)
        if missing_objects:
            object_order = list(object_order) + list(missing_objects)

    if group_order is None:
        # 按首次出现顺序
        group_order = []
        for record in records:
            for group_name in record.keys():
                if group_name not in group_order:
                    group_order.append(group_name)
    group_order = list(group_order)

    # 3. 创建对象和组的索引映射
    obj_to_idx = {obj: i for i, obj in enumerate(object_order)}
    group_to_idx = {group: i for i, group in enumerate(group_order)}

    # 4. 构建数据矩阵
    M = len(records)  # 分组次数
    N = len(object_order)  # 对象数量
    data_matrix = np.full((M, N), -1, dtype=int)  # -1 表示未分组

    for fold_idx, record in enumerate(records):
        # 检查是否有对象重复分组
        assigned_objects = set()
        for group_name, objects in record.items():
            for obj in objects:
                if obj in assigned_objects:
                    raise ValueError(
                        f"对象 '{obj}' 在第 {fold_idx+1} 次分组中出现在多个组中"
                    )
                assigned_objects.add(obj)

                if obj in obj_to_idx:
                    obj_idx = obj_to_idx[obj]
                    group_idx = group_to_idx[group_name]
                    data_matrix[fold_idx, obj_idx] = group_idx

    # 5. 设置颜色映射
    n_groups = len(group_order)

    if cmap is None:
        # 尝试从当前 matplotlib 配置中获取 cycle 颜色
        try:
            import matplotlib as mpl

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

            # 如果cycle颜色足够，使用它们；否则循环使用
            if len(cycle_colors) >= n_groups:
                group_colors = [cycle_colors[i] for i in range(n_groups)]
            else:
                group_colors = [
                    cycle_colors[i % len(cycle_colors)] for i in range(n_groups)
                ]

            # 将hex颜色转换为RGBA元组
            from matplotlib.colors import to_rgba

            group_colors = [to_rgba(color) for color in group_colors]

        except (KeyError, ImportError, Exception):
            # 如果获取cycle颜色失败，使用更好看的默认配色：Set2
            fallback_cmap = plt.get_cmap('Set2')
            group_colors = [
                fallback_cmap(i / max(1, n_groups - 1)) for i in range(n_groups)
            ]
    else:
        # 使用用户指定的颜色映射
        user_cmap = plt.cm.get_cmap(cmap)
        group_colors = [user_cmap(i / max(1, n_groups - 1)) for i in range(n_groups)]

    # 构建最终的颜色列表：先放灰色(给-1用)，再放组颜色(给0,1,2,...用)
    colors = [(0.8, 0.8, 0.8, 1.0)]  # 索引0: 灰色用于-1(未分组)
    colors.extend(group_colors)  # 索引1,2,3,...: 组颜色用于0,1,2,...

    cmap_obj = ListedColormap(colors)

    # 6. 绘制热力图
    # 颜色映射：-1→索引0(灰色), 0→索引1(第1个组色), 1→索引2(第2个组色), ...
    ax.imshow(data_matrix, cmap=cmap_obj, aspect='auto', vmin=-1, vmax=n_groups - 1)

    # 7. 添加网格内文本标注
    if grid_label:
        for i in range(M):
            for j in range(N):
                if data_matrix[i, j] == -1:
                    text = 'N/A'
                    color = 'black'
                else:
                    group_idx = data_matrix[i, j]
                    group_name = group_order[group_idx]
                    # 组名缩写：取前2个字符，或者首字母+数字
                    if len(group_name) <= 3:
                        text = group_name
                    else:
                        text = group_name[:2].upper()
                    color = 'white' if sum(colors[group_idx][:3]) < 1.5 else 'black'

                ax.text(
                    j,
                    i,
                    str(text),
                    ha='center',
                    va='center',
                    fontsize=8,
                    fontweight='bold',
                    color=color,
                )

    # 8. 设置坐标轴
    ax.set_xticks(range(N))
    ax.set_xticklabels(object_order, rotation=xtick_rotation, ha='right')
    ax.set_yticks(range(M))
    ax.set_yticklabels([f'Fold {i+1}' for i in range(M)])

    # 9. 添加图例
    legend_elements = []
    for i, group_name in enumerate(group_order):
        # 因为colors[0]是灰色(用于-1)，组颜色从colors[1]开始
        legend_elements.append(Patch(facecolor=colors[i + 1], label=group_name))

    if np.any(data_matrix == -1):
        # 灰色在colors[0]位置
        legend_elements.append(Patch(facecolor=colors[0], label='Not Assigned'))

    ax.legend(handles=legend_elements, bbox_to_anchor=(1.05, 1), loc='upper left')

    # 10. 设置标题和轴标签
    if title is not None:
        ax.set_title(title)
    if xlabel is not None:
        ax.set_xlabel(xlabel)
    if ylabel is not None:
        ax.set_ylabel(ylabel)

    # 调整布局以容纳图例
    plt.tight_layout()

    _finalize_figure(fig, show, save_path, auto_close)
    return fig, ax
