import typing
from collections import Counter
from typing import Callable, Literal, Optional, Sequence, TypedDict, Union

import numpy as np
import torch
from sklearn.model_selection import KFold
from torch.utils.data import DataLoader, Dataset, Subset
from torch.utils.data import random_split as torch_random_split

# from torch.utils.data import RandomSampler
from zyplib.utils.ensure import ensure_tensor
from zyplib.utils.print import print_debug, print_warn

# Define shuffle mode types
ShuffleMode = Union[
    bool, list[bool], Literal['shuffle-first', 'shuffle-all', 'shuffle-train']
]


def _resolve_shuffle_for_dataloaders(
    shuffle: ShuffleMode, num_dataloaders: int, mode: str = 'basic'
) -> list[bool]:
    """解析 shuffle 参数为每个 DataLoader 的具体 shuffle 设置。

    Parameters
    ----------
    - `shuffle` : `ShuffleMode`
        - 输入的 shuffle 参数
    - `num_dataloaders` : `int`
        - DataLoader 的数量
    - `mode` : `str`, optional
        - 'basic', 'split', 'kfold' 用于处理不同场景的默认行为, by default 'basic'

    Returns
    -------
    - `list[bool]`
        - 每个 DataLoader 对应的 shuffle 设置

    Raises
    ------
    - `ValueError`
        - 当 shuffle 参数无效时
    """
    if isinstance(shuffle, bool):
        return [shuffle] * num_dataloaders
    elif isinstance(shuffle, list):
        if len(shuffle) != num_dataloaders:
            raise ValueError(
                f'shuffle list length {len(shuffle)} != datasets count {num_dataloaders}'
            )
        return shuffle
    elif shuffle == 'shuffle-first':
        return [True] + [False] * (num_dataloaders - 1)
    elif shuffle == 'shuffle-all':
        return [True] * num_dataloaders
    elif shuffle == 'shuffle-train':
        if mode == 'kfold':
            # kfold 固定为 train=True, val=False
            return [True, False]
        else:
            # 其他情况等同于 shuffle-first
            return [True] + [False] * (num_dataloaders - 1)
    else:
        raise ValueError(f'Invalid shuffle mode: {shuffle}')


def _resolve_shuffle_for_kfold(
    shuffle: Union[bool, tuple[bool, bool], Literal['shuffle-train']],
) -> tuple[bool, bool]:
    """专门为 kfold 解析 shuffle 参数。

    Parameters
    ----------
    - `shuffle` : `Union[bool, tuple[bool, bool], Literal['shuffle-train']]`
        - kfold 专用的 shuffle 参数

    Returns
    -------
    - `tuple[bool, bool]`
        - (train_shuffle, val_shuffle)
    """
    if isinstance(shuffle, bool):
        return (shuffle, shuffle)
    elif isinstance(shuffle, tuple):
        return shuffle
    elif shuffle == 'shuffle-train':
        return (True, False)
    else:
        raise ValueError(f'Invalid shuffle mode for kfold: {shuffle}')


class BasicXyDataset(Dataset):
    """基础的 XY 数据集，用于 PyTorch 的 DataLoader

    提供最基本的 Data, Label 数据集，可以通过 index 进行访问

    Parameters
    ----------
    - `data` : `np.ndarray | torch.Tensor`
        - 数据
    - `target` : `np.ndarray | torch.Tensor`
        - 标签
    """

    def __init__(
        self,
        data: Union[np.ndarray, torch.Tensor],
        target: Union[np.ndarray, torch.Tensor],
        data_dtype: torch.dtype = torch.float32,
        target_dtype: torch.dtype = torch.float32,
        transform: Optional[Callable] = None,
        p_trans: float = 0,
        cache_transformed: bool = False,
    ):
        """基础的 XY 数据集，用于 PyTorch 的 DataLoader

        允许传入 transform 对象进行变换或者增广等.

        当 __getitem__ 被调用时, 会按照 `p_trans` 的概率决定是否进行变换。

        参考: torchvision; zyplib.transformation

        Parameters
        ----------
        - `data` : `Union[np.ndarray, torch.Tensor]`, 数据张量
        - `labels` : `Union[np.ndarray, torch.Tensor]`, 标签张量
        - `data_dtype` : `torch.dtype`, optional
            - 数据张量的类型, 默认 torch.float32
        - `label_dtype` : `torch.dtype`, optional
            - 标签张量的类型, 默认 torch.float32
        - `transform` : `(X, y) -> (X | (X, y))`
            - 数据增广的变换函数
            - 会把 (X, y) batch 数据传入; 可以只返回 X, 也可以同时返回 (X, y)
        - `p_trans` : `float`, optional
            - 变换的概率，随机控制是否需要 transform, 默认为 0.9
        """
        self.data = ensure_tensor(data, dtype=data_dtype)
        self.targets = ensure_tensor(target, dtype=target_dtype)
        self.transform = transform
        self.p_trans = p_trans
        self.cache_transformed = cache_transformed
        self._transformed_cache = {}

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        x = self.data[idx]
        y = self.targets[idx]

        if not self.transform or self.p_trans == 0:
            # 如果没有提供变换函数，直接返回数据和标签
            return x, y

        if np.random.rand() <= self.p_trans:
            if self.cache_transformed and idx in self._transformed_cache:
                return self._transformed_cache[idx]

            result = self.transform(x, y)
            if isinstance(result, tuple):
                x, y = result
            else:
                x = result
            if self.cache_transformed:
                self._transformed_cache[idx] = (x, y)
        return x, y


def inject_data_transformation(
    dataloader: DataLoader, fn_transform: Optional[Callable] = None, p_trans: float = 1.0
):
    """动态向 DataLoader 注入数据变换逻辑

    通过修改 DataLoader 内部的 Dataset 对象，为其添加或覆盖数据变换功能。

    - 如果 `dataset` 是 `BasicXyDataset` 类型，则直接修改其 `transform` 和 `p_trans` 属性。
    - 否则，通过“猴子补丁”的方式，包装其 `__getitem__` 方法，实现动态变换。

    Parameters
    ----------
    - `dataloader` : `DataLoader`
        - 目标 DataLoader。
    - `fn_transform` : `Callable`, optional
        - 新的变换函数。它应该接受 `(X, y)` 并返回 `X` 或 `(X, y)`。
        - 如果为 `None`，则不进行任何操作。
    - `p_trans` : `float`, optional
        - 执行变换的概率, 默认为 1.0 (总是执行)。

    Notes
    -----
    猴子补丁是一种强大的技术，但也可能带来意外的副作用。此函数会打印警告，
    提醒用户这并非标准实践，最佳方式是在创建 DataLoader 前就配置好 Dataset。
    """
    if fn_transform is None:
        return

    dataset = dataloader.dataset

    # Case 1: 如果是我们自己的 Dataset，直接修改属性，最安全
    if isinstance(dataset, BasicXyDataset):
        dataset.transform = fn_transform
        if p_trans is not None:
            dataset.p_trans = p_trans
        print_debug('Injected transform into BasicXyDataset.')
        return

    # Case 2: 对通用 Dataset 进行猴子补丁
    import types

    # 保存原始的 __getitem__ 方法
    original_getitem = dataset.__getitem__

    def new_getitem(self, idx):
        # 1. 调用原始方法获取数据
        original_data = original_getitem(idx)
        # 假设原始返回是 (data, label)
        x, y = None, None
        if isinstance(original_data, tuple):
            x, y = original_data[0], original_data[1]
        else:
            x = original_data

        # 2. 以 p_trans 的概率应用新的变换
        if np.random.rand() < p_trans:
            result = fn_transform(x, y)
            if isinstance(result, tuple):
                x, y = result
            else:
                x = result
        return x, y

    # 3. 使用 types.MethodType 将新函数正确绑定为实例方法
    dataset.__getitem__ = types.MethodType(new_getitem, dataset)

    print_warn(
        '已通过猴子补丁将 transform 注入到通用 Dataset 中。'
        '这并非推荐的标准实践，建议优先使用定制的 Dataset 对象。'
    )


def make_basic_dataloader(
    *datasets: tuple[np.ndarray, np.ndarray] | Dataset[typing.Any],
    batch_size: int = 64,
    num_workers: int = 0,
    pin_memory: bool = False,
    shuffle: ShuffleMode = 'shuffle-first',
    dtype: Union[torch.dtype, tuple[torch.dtype, torch.dtype]] = torch.float32,
    **kwargs,
) -> tuple[DataLoader, ...] | DataLoader:
    """为每个输入的数据集创建一个 DataLoader。

    本函数可以接收一个或多个数据集，每个数据集可以是 (X, y) 元组或 Dataset 对象。
    它会为每个输入的数据集创建一个对应的 DataLoader。

    Parameters
    ----------
    - `*datasets` : `Union[tuple[np.ndarray, np.ndarray], Dataset]`
        - 一个或多个数据集。每个数据集可以是 `(X, y)` 元组或 `torch.utils.data.Dataset` 对象。
    - `batch_size` : `int`, optional
        - 每个批次的样本数量，默认为 64。
    - `shuffle` : `ShuffleMode`, optional
        - 数据打乱模式，支持以下选项：
          - `bool`: 所有 DataLoader 使用相同的 shuffle 设置
          - `list[bool]`: 为每个 DataLoader 单独指定 shuffle 设置
          - `'shuffle-first'`: 只有第一个 DataLoader shuffle（默认，适用于 train, val, test 顺序）
          - `'shuffle-all'`: 所有 DataLoader 都 shuffle
          - `'shuffle-train'`: 等同于 'shuffle-first'
    - `num_workers` : `int`, optional
        - 用于数据加载的子进程数量，默认为 0。
    - `pin_memory` : `bool`, optional
        - 如果为 True，数据加载器将把张量复制到 CUDA 固定内存中，默认为 False。
    - `dtype` : `Union[torch.dtype, tuple[torch.dtype, torch.dtype]]`, optional
        - 数据类型。可以是单个 dtype (应用于数据和标签)，也可以是 (data_dtype, label_dtype) 的元组。默认为 torch.float32。
    - `**kwargs` : `dict`, optional
        - DataLoader 的其他参数。

    Returns
    -------
    - `tuple[DataLoader, ...]`
        - 返回一个包含多个 DataLoader 的元组，顺序与输入数据集一致。

    Examples
    --------
    ```python
    >>> # 默认行为：只有第一个（训练集）shuffle
    >>> train_loader, val_loader, test_loader = make_basic_dataloader(
    ...     (X_train, y_train),
    ...     (X_val, y_val),
    ...     (X_test, y_test),
    ...     batch_size=32
    ... )
    >>>
    >>> # 精确控制每个 DataLoader 的 shuffle
    >>> train_loader, val_loader, test_loader = make_basic_dataloader(
    ...     (X_train, y_train),
    ...     (X_val, y_val),
    ...     (X_test, y_test),
    ...     shuffle=[True, False, False],
    ...     batch_size=32
    ... )
    ```
    """
    dataloaders = []

    # 计算有效数据集数量用于 shuffle 解析
    num_valid_datasets = 0
    for ds_input in datasets:
        if ds_input is None:
            continue
        if isinstance(ds_input, tuple):
            if len(ds_input[0]) == 0:
                continue
        elif isinstance(ds_input, Dataset):
            dataset_obj = typing.cast(typing.Sized, ds_input)
            if len(dataset_obj) == 0:
                continue
        num_valid_datasets += 1

    # 解析 shuffle 模式
    shuffle_settings = _resolve_shuffle_for_dataloaders(
        shuffle, num_valid_datasets, 'basic'
    )

    data_dtype: torch.dtype
    label_dtype: torch.dtype
    if isinstance(dtype, tuple):
        data_dtype, label_dtype = dtype
    else:
        data_dtype = dtype
        label_dtype = dtype

    shuffle_idx = 0
    for i, ds_input in enumerate(datasets):
        if ds_input is None:
            continue

        dataset: Dataset[typing.Any]
        if isinstance(ds_input, tuple):
            if len(ds_input[0]) == 0:
                continue
            dataset = BasicXyDataset(
                *ds_input, data_dtype=data_dtype, target_dtype=label_dtype
            )
        elif isinstance(ds_input, Dataset):
            dataset_obj = typing.cast(typing.Sized, ds_input)
            if len(dataset_obj) == 0:
                continue
            dataset = typing.cast(Dataset[typing.Any], ds_input)
        else:
            raise TypeError(
                f'Unsupported dataset type: {type(ds_input)}. Must be a tuple (X, y) or a torch.utils.data.Dataset.'
            )

        # 使用对应的 shuffle 设置
        current_shuffle = shuffle_settings[shuffle_idx]
        shuffle_idx += 1

        loader = DataLoader(
            dataset,
            batch_size=batch_size,
            shuffle=current_shuffle,
            num_workers=num_workers,
            pin_memory=pin_memory,
            **kwargs,
        )
        dataloaders.append(loader)

    return tuple(dataloaders) if len(dataloaders) > 1 else dataloaders[0]


# alias
make_dataloader_from_xy = make_basic_dataloader


def make_split_dataloader(
    xy: tuple[np.ndarray, np.ndarray] | Dataset[typing.Any],
    split_size: Union[int, Sequence[Union[int, float]]],
    batch_size: int = 64,
    num_workers: int = 0,
    pin_memory: bool = False,
    shuffle: ShuffleMode = 'shuffle-first',
    dtype: Union[torch.dtype, tuple[torch.dtype, torch.dtype]] = torch.float32,
    **kwargs,
) -> tuple[DataLoader, ...]:
    """从单个 (X, y) 数据集创建多个 DataLoader，按指定大小或数量进行划分。

    本函数是对 `torch.utils.data.random_split` 和 `DataLoader` 的封装，
    可以方便地将一个完整的数据集划分为多个部分，并为每个部分创建 DataLoader。

    Parameters
    ----------
    - `xy` : `tuple[np.ndarray, np.ndarray] | Dataset`
        - 包含数据和标签的元组或数据集对象。
    - `split_size` : `int | Sequence[Union[int, float]]`
        - 划分方式。
        - 如果为 `int`，则将数据集平均划分为 N 份。
        - 如果为 `Sequence`，则按指定的长度或比例进行划分。
    - `batch_size` : `int`, optional
        - 每个批次的样本数量，默认为 64。
    - `shuffle` : `ShuffleMode`, optional
        - 数据打乱模式，支持以下选项：
          - `bool`: 所有 DataLoader 使用相同的 shuffle 设置
          - `list[bool]`: 为每个 DataLoader 单独指定 shuffle 设置
          - `'shuffle-first'`: 只有第一个 DataLoader shuffle（默认，通常是训练集）
          - `'shuffle-all'`: 所有 DataLoader 都 shuffle
          - `'shuffle-train'`: 等同于 'shuffle-first'
    - `num_workers` : `int`, optional
        - 用于数据加载的子进程数量，默认为 0。
    - `pin_memory` : `bool`, optional
        - 如果为 True，数据加载器将把张量复制到 CUDA 固定内存中，默认为 False。
    - `dtype` : `Union[torch.dtype, tuple[torch.dtype, torch.dtype]]`, optional
        - 数据类型。可以是单个 dtype (应用于数据和标签)，也可以是 (data_dtype, label_dtype) 的元组。默认为 torch.float32。
    - `**kwargs` : `dict`, optional
        - DataLoader 的其他参数。

    Returns
    -------
    - `tuple[DataLoader, ...]`
        - 返回一个包含多个 DataLoader 的元组。
    """
    if isinstance(xy, tuple):
        data_dtype: torch.dtype
        label_dtype: torch.dtype
        if isinstance(dtype, tuple):
            data_dtype, label_dtype = dtype
        else:
            data_dtype = dtype
            label_dtype = dtype
        dataset = BasicXyDataset(*xy, data_dtype=data_dtype, target_dtype=label_dtype)
    else:
        dataset = xy

    lengths: Sequence[Union[int, float]]
    if isinstance(split_size, int):
        k = split_size
        if k <= 0:
            raise ValueError('If split_size is an integer, it must be positive.')
        total_len = len(typing.cast(typing.Sized, dataset))
        base_len = total_len // k
        remainder = total_len % k
        lengths = [base_len + 1 if i < remainder else base_len for i in range(k)]
    else:
        lengths = split_size

    split_datasets = torch_random_split(dataset, lengths)
    num_splits = len(split_datasets)

    # 解析 shuffle 模式
    shuffle_settings = _resolve_shuffle_for_dataloaders(shuffle, num_splits, 'split')

    dataloaders = []
    for i, subset in enumerate(split_datasets):
        loader = DataLoader(
            subset,
            batch_size=batch_size,
            shuffle=shuffle_settings[i],
            num_workers=num_workers,
            pin_memory=pin_memory,
            **kwargs,
        )
        dataloaders.append(loader)

    return tuple(dataloaders)


def make_kfold_dataloader(
    *,
    xy: tuple[np.ndarray, np.ndarray] | None = None,
    dataset: Dataset[typing.Any] | None = None,
    k: int = 5,
    random_seed: int | None = None,
    batch_size: int = 64,
    num_workers: int = 0,
    pin_memory: bool = False,
    shuffle: bool | tuple[bool, bool] | Literal['shuffle-train'] = 'shuffle-train',
    dtype: Union[torch.dtype, tuple[torch.dtype, torch.dtype]] = torch.float32,
    **kwargs,
):
    """根据输入的 xy 或者 dataset 创建 KFold 的 DataLoader(s)。

    > 注意，这里是导入 XY 之后再进行 KFold；一般而言，应该是导入文件的时候就直接进行划分，推荐使用 `make_xy_pipeline` 的 `kfold_pipeline` 函数

    从输入的数据中创建 KFold 的 DataLoader(s)。

    xy 和 dataset 二选一。

    Parameters
    ----------
    - `xy` : `tuple[np.ndarray, np.ndarray]`
        - 包含数据和标签的元组。
    - `dataset` : `Dataset`, 可选
        - 数据集。
    - `k` : `int`, 可选
        - 折数，默认为 5。
    - `random_seed` : `int | None`, 可选
        - 随机种子，默认为 None。
    - `batch_size` : `int`, 可选
        - 每个批次的样本数量，默认为 64。
    - `num_workers` : `int`, 可选
        - 用于数据加载的子进程数量，默认为 0。
    - `pin_memory` : `bool`, 可选
        - 如果为 True，数据加载器将把张量复制到 CUDA 固定内存中，默认为 False。
    - `shuffle` : `Union[bool, tuple[bool, bool], Literal['shuffle-train']]`, optional
        - K折交叉验证的 shuffle 模式：
          - `bool`: 训练集和验证集使用相同的 shuffle 设置
          - `tuple[bool, bool]`: 分别指定 (train_shuffle, val_shuffle)
          - `'shuffle-train'`: 训练集 shuffle，验证集不 shuffle（默认，推荐）
    - `dtype` : `Union[torch.dtype, tuple[torch.dtype, torch.dtype]]`, optional
        - 数据类型。可以是单个 dtype (应用于数据和标签)，也可以是 (data_dtype, label_dtype) 的元组。默认为 torch.float32。
    - `**kwargs` : `dict`, 可选
        - DataLoader 的其他参数，默认为 None。

    Yields
    ----------
    - `tuple[DataLoader, DataLoader]`
        - 训练和验证数据集的 DataLoader。
    """
    # xy 和 dataset 只能有一个是 None，一个不是 None
    assert (
        xy is not None or dataset is not None
    ), 'xy 和 dataset 只能有一个是 None，一个不是 None'
    if xy is not None:
        data_dtype: torch.dtype
        label_dtype: torch.dtype
        if isinstance(dtype, tuple):
            data_dtype, label_dtype = dtype
        else:
            data_dtype = dtype
            label_dtype = dtype
        dataset = BasicXyDataset(*xy, data_dtype=data_dtype, target_dtype=label_dtype)
    else:
        dataset = typing.cast(Dataset[typing.Any], dataset)

    # 解析 kfold 的 shuffle 设置
    train_shuffle, val_shuffle = _resolve_shuffle_for_kfold(shuffle)

    kf = KFold(n_splits=k, shuffle=True, random_state=random_seed)

    sized_dataset = typing.cast(Dataset[typing.Any], dataset)
    for train_idx, val_idx in kf.split(sized_dataset):  # pyright: ignore[reportArgumentType]
        train_dataset = Subset(sized_dataset, train_idx)  # pyright: ignore[reportArgumentType]
        val_dataset = Subset(sized_dataset, val_idx)  # pyright: ignore[reportArgumentType]
        train_loader = DataLoader(
            train_dataset,
            batch_size=batch_size,
            shuffle=train_shuffle,
            num_workers=num_workers,
            pin_memory=pin_memory,
            **kwargs,
        )
        val_loader = DataLoader(
            val_dataset,
            batch_size=batch_size,
            shuffle=val_shuffle,
            num_workers=num_workers,
            pin_memory=pin_memory,
            **kwargs,
        )
        yield train_loader, val_loader


class LabelCounter(TypedDict):
    number: int
    ratio: float


def _label_to_hashable(
    label: Union[np.ndarray, torch.Tensor, int, float],
) -> Union[int, float, str]:
    """将标签转换为可哈希的形式用于计数。

    Parameters
    ----------
    - `label` : `Union[np.ndarray, torch.Tensor, int, float]`
        - 输入的标签，可能是标量、一维数组或张量

    Returns
    -------
    - `Union[int, float, str]`
        - 可哈希的标签表示
    """
    # 如果已经是标量，直接返回
    if np.isscalar(label):
        return typing.cast(Union[int, float, str], label)

    # 转换为 numpy 数组进行处理
    if isinstance(label, torch.Tensor):
        label = label.cpu().numpy()

    # 压缩维度
    label = np.squeeze(label)

    # 处理 0 维数组（标量）
    if label.ndim == 0:
        return label.item()

    # 处理一维数组
    if label.ndim == 1:
        # 检查是否为 one-hot 编码
        if len(label) > 1 and all(x in [0, 1] for x in label) and sum(label) == 1:
            return int(np.argmax(label))
        # 如果只有一个元素，返回该元素
        elif len(label) == 1:
            return typing.cast(Union[int, float, str], label[0])
        # 多元素非 one-hot，转换为字符串表示
        else:
            return str(label.tolist())

    # 多维数组，展平后转换为字符串
    return str(label.flatten().tolist())


def count_dataloader_labels(
    loader: DataLoader,
) -> dict[Union[int, float, str], LabelCounter]:
    """检查数据加载器是否平衡"""
    all_labels = []
    for _, labels in loader:
        # 如果是批次标签，逐个处理
        if isinstance(labels, (torch.Tensor, np.ndarray)):
            # 确保是 numpy 数组
            if isinstance(labels, torch.Tensor):
                labels = labels.cpu().numpy()

            # 处理批次中的每个标签
            for single_label in labels:
                hashable_label = _label_to_hashable(single_label)
                all_labels.append(hashable_label)
        else:
            # 单个标签的情况
            hashable_label = _label_to_hashable(labels)
            all_labels.append(hashable_label)

    counter = Counter(all_labels)
    total_samples = sum(counter.values())

    if total_samples == 0:
        return {}

    return {
        label: {'number': number, 'ratio': number / total_samples}
        for label, number in counter.items()  # 移除排序避免类型混合问题
    }


def _dataloader_is_shuffle(loader):
    from torch.utils.data import RandomSampler

    # 1. 标准 DataLoader
    if hasattr(loader, 'sampler') and isinstance(loader.sampler, RandomSampler):
        return True

    # 2. 某些自定义 DataLoader 可能有 `shuffle` 属性
    if hasattr(loader, 'shuffle') and loader.shuffle:
        return True

    # 3. 如果 sampler 不是 RandomSampler，则认为没有 shuffle
    if hasattr(loader, 'sampler'):
        return False

    # 4. 无法判断，返回 None
    return None


def analyze_dataloader(
    loader: DataLoader,
    prefix_name: str = 'DataLoader',
    balance_range: tuple[float, float] = (0.4, 0.6),
    verbose: bool = True,
):
    """分析 DataLoader 的内部数据组成，检查是否存在数据不平衡的情况。

    Parameters
    ----------
    - `loader` : `DataLoader`
        - DataLoader 对象
    - `prefix_name` : `str`, optional
        - 前缀名称，默认为 'DataLoader'
    - `balance_range` : `tuple[float, float]`, optional
        - 被认为平衡的标签比例范围，默认为 (0.4, 0.6)，对应了二分类问题中标签比例在 40% 到 60% 之间
    - `verbose` : `bool`, optional
        - 是否打印详细信息，默认为 True

    Returns
    ----------
    - `dict`
    """
    # --- 1. 执行一次数据遍历和计数 (修复了重复调用的效率问题) ---
    label_distribution = count_dataloader_labels(loader)

    # --- 2. 收集基本信息 ---
    # loader_num_samples = loader_len * loader_batch_size 这个计算只有在 drop_last=True 时才是精确的。如果 drop_last=False，最后一个 batch 的大小可能小于 `batch_size`
    # loader_batch_size = loader.batch_size
    # loader_num_samples = loader_len * loader_batch_size
    num_samples = sum(c['number'] for c in label_distribution.values())
    num_batches = len(loader)
    is_shuffle = _dataloader_is_shuffle(loader)

    # --- 3. 分析平衡性并收集问题 ---
    balance_issues = []
    for label, counter in label_distribution.items():
        ratio = counter['ratio']
        if not (balance_range[0] <= ratio <= balance_range[1]):
            balance_issues.append(label)

    # --- 4. 按需打印报告 (所有打印信息均为中文) ---
    if verbose:
        print_debug(f'{prefix_name} 分析报告'.center(30, '-'))
        print_debug(f'样本总数: {num_samples}')
        print_debug(f'批次总数: {num_batches}')
        print_debug(f'是否打乱: {is_shuffle}')
        print_debug('标签分布:')
        if not label_distribution:
            print_debug('  未找到任何标签。')
        else:
            for label, counter in label_distribution.items():
                status_msg = ''
                if label in balance_issues:
                    status_msg = f' (比例异常, 不在 [{balance_range[0]:.0%}, {balance_range[1]:.0%}] 范围内)'
                    # 使用 print_warn 突出显示问题
                    print_warn(
                        f"  - 标签 {label}: {counter['ratio']:.2%} "
                        f"({counter['number']} / {num_samples}){status_msg}"
                    )
                else:
                    print_debug(
                        f"  - 标签 {label}: {counter['ratio']:.2%} "
                        f"({counter['number']} / {num_samples})"
                    )

        if not balance_issues and label_distribution:
            print_debug('根据给定范围，数据集标签比例被认为是平衡的。')
        print_debug('--- 报告结束 ---')

    # --- 5. 返回结构化的中文键名结果 ---
    return {
        'num_samples': num_samples,
        'num_batches': num_batches,
        'label_ratios': label_distribution,
        'is_shuffle': is_shuffle,
        'balance_issues': balance_issues,
    }
