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

### 设计哲学

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

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

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

from ._plot_utils import FigAx, _finalize_figure, _prepare_ax


def plot_group_divide(
    records: List[Dict[str, List[str]]],
    object_order: Optional[List[str]] = None,
    group_order: Optional[List[str]] = None,
    cmap: Optional[str] = None,
    title: Optional[str] = None,
    xlabel: Optional[str] = None,
    ylabel: Optional[str] = None,
    ax: Optional[plt.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)

    # 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.cm.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 = ListedColormap(colors)

    # 6. 绘制热力图
    # 颜色映射：-1→索引0(灰色), 0→索引1(第1个组色), 1→索引2(第2个组色), ...
    ax.imshow(data_matrix, cmap=cmap, 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,
                    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
