"""
本模块主要提供了一些数据集导入的 pipeline 相关的函数，是对脑电信号数据导入、预处理流程的高级封装

所有函数均满足:

- 输入: 一些参数，以及一些回调函数
- 输出: 返回数据集 (X, y)

"""

import typing

from functools import reduce
from typing import Callable, Literal, Protocol, Sequence, cast, runtime_checkable

import numpy as np
from sklearn.model_selection import KFold

from zyplib.dataset.split import random_split
from zyplib.preprocessing.labels import to_onehot
from zyplib.utils.print import print_info
from zyplib.utils.progress_bar import rich_progress

__all__ = [
    'fileseg_pipeline',
    'fileloader_pipeline',
    'basic_pipeline',
    'simple_pipeline',
    'train_test_pipeline',
    'kfold_pipeline',
]


@runtime_checkable
class FuncPipeline(Protocol):
    """函数类型，输入文件路径列表，输出数据集 (X, y)"""

    def __call__(self, files: typing.Sequence[str]) -> tuple[np.ndarray, np.ndarray]: ...


@runtime_checkable
class ReadFunction(Protocol):
    def __call__(self, file_path: str) -> np.ndarray: ...


@runtime_checkable
class SegmentFunction(Protocol):
    def __call__(self, signal: np.ndarray) -> np.ndarray: ...


@runtime_checkable
class LabelFunction(Protocol):
    def __call__(self, file_path: str) -> int | list[int]: ...


@runtime_checkable
class SignalProcessFunction(Protocol):
    def __call__(self, signal: np.ndarray) -> np.ndarray: ...


@runtime_checkable
class EncodeLabelFunction(Protocol):
    def __call__(self, label: typing.Sequence[int] | np.ndarray) -> np.ndarray: ...


@runtime_checkable
class TransformFunction(Protocol):
    def __call__(self, file_path: str) -> tuple[np.ndarray, np.ndarray]: ...


def _default_encode_label(label: typing.Sequence[int] | np.ndarray) -> np.ndarray:
    return np.asarray(to_onehot(label, squeeze_2_class=True))


def _identity_signal(signal: np.ndarray) -> np.ndarray:
    return signal


def _normalize_file_groups(
    all_files: typing.Sequence[str] | Sequence[typing.Sequence[str]],
) -> list[list[str]]:
    if not all_files:
        raise ValueError("The 'all_files' list cannot be empty.")

    first_group = all_files[0]
    if not first_group:
        raise ValueError("The 'all_files' list cannot be empty.")

    if isinstance(first_group, str):
        flat_files = cast(typing.Sequence[str], all_files)
        return [sorted(flat_files)]

    grouped_files = cast(Sequence[typing.Sequence[str]], all_files)
    return [sorted(group) for group in grouped_files]


def fileseg_pipeline(
    files: typing.Sequence[str],
    fn_read: ReadFunction,
    label: int | LabelFunction,
    fn_segment: SegmentFunction | None = None,
    fn_before_segment: SignalProcessFunction | None = None,
    fn_after_segment: SignalProcessFunction | None = None,
    fn_encode_label: EncodeLabelFunction | None = _default_encode_label,
    show_progress: bool = True,
    random_shuffle: bool | int = False,
) -> tuple[np.ndarray, np.ndarray]:
    """从文件中读取数据，并自动切片、制作用于训练的数据集

    本 Pipeline 适用于:
    - 从文件中读取完整的脑电
    - 需要对完整脑电进行切分
    - 无特殊标签，只需根据文件进行标记即可
    - 输出: 数据集 (X, y)
        - X: 数据，shape = (N, C, T)
        - y: 标签，shape = (N,)

    Pipeline 流程:
    ----------

    1. 从文件中读取完整的数据 X
    2. 对数据 X 进行预处理
    3. 对数据 X 进行切片
    4. 对数据 X 进行后处理
    5. 制作标签 y
    6. 返回数据集 (X, y)

    Parameters
    ----------
    - `files` : `list[str]`
        - 文件路径列表
    - `fn_read` : `(file_path: str) -> np.ndarray`
        - 文件读取函数
    - `fn_segment` : `(signal: np.ndarray) -> np.ndarray`
        - 数据切片函数
        - 输入一个 `[C, T]` 的信号
        - 输出一个 `[N, C, T]` 的信号
    - `label` : `int | (file_path: str) -> int`
        - 标签，如果为 int，则所有数据标签相同；如果为函数，则根据文件名生成标签
    - `fn_before_segment` : `(signal: [C, T]) -> [C, T]`, optional
        - 数据切片前的处理函数
    - `fn_after_segment` : `(signal: [N, C, T]) -> [N, C, T]`, optional
        - 数据切片后的处理函数
    - `fn_encode_label` : `(label: list | np.ndarray) -> np.ndarray`, optional
        - 标签编码函数, 默认使用 `to_onehot`
    - `show_progress` : `bool`, optional
        - 是否显示处理进度, by default `True`
    - `random_shuffle` : `bool | int`, optional
        - 是否在返回前随机打乱数据
        - `False`: 不打乱数据
        - `True`: 使用默认随机种子打乱数据
        - `int`: 使用指定的整数作为随机种子打乱数据
        - 注意: 此参数不会影响全局随机状态

    Returns
    ----------
    - `X` : `np.ndarray`, shape = (N, C, T)
        - 数据
    - `y` : `np.ndarray`, shape = (N, n_classes)
        - 标签
    """
    if not files:
        raise ValueError('The files list is empty. Please provide at least one file.')

    segment_fn: SegmentFunction = (
        fn_segment if fn_segment is not None else _identity_signal
    )
    before_segment_fn: SignalProcessFunction = (
        fn_before_segment if fn_before_segment is not None else _identity_signal
    )
    after_segment_fn: SignalProcessFunction = (
        fn_after_segment if fn_after_segment is not None else _identity_signal
    )
    fn_label = label if callable(label) else (lambda _: label)

    # 处理单个文件
    def load_single_file(fpath: str):
        signal = fn_read(fpath)  # 读取数据
        signal = before_segment_fn(signal)  # 数据预处理
        segments = segment_fn(signal)  # 数据切片
        segments = after_segment_fn(segments)
        N = len(segments)
        label_ = fn_label(fpath)  # 制作标签
        labels = []
        if isinstance(label_, list):
            labels = label_
        else:
            labels = [label_] * N
        return segments, labels

    # 批量处理所有文件
    if show_progress:
        import os

        iterable = rich_progress(
            files, description='处理中', fn_get_status=lambda f: os.path.basename(f)
        )
    else:
        iterable = files
    results = map(load_single_file, iterable)

    # 合并数据
    def merge_data(
        Xy: tuple[list[np.ndarray], list[int]],
        pair: tuple[np.ndarray, typing.Sequence[int]],
    ):
        segments, labels = pair
        if len(segments) > 0:
            Xy[0].append(segments)
            Xy[1].extend(labels)
        return Xy

    Xy = reduce(merge_data, results, ([], []))

    if not Xy[0]:
        X = np.array([])
        y = np.array([], dtype=int)
    else:
        X = np.concatenate(Xy[0], axis=0)
        y = np.array(Xy[1], dtype=int)

    encode_label_fn = fn_encode_label or _default_encode_label
    y = encode_label_fn(y)

    # 随机打乱数据
    if random_shuffle is not False:
        # 创建独立的随机数生成器实例
        if isinstance(random_shuffle, int):
            # 使用指定的种子
            rand_gen = np.random.RandomState(seed=random_shuffle)
        elif random_shuffle is True:
            # 使用默认种子
            rand_gen = np.random.RandomState()
        else:
            raise ValueError('random_shuffle must be a boolean or an integer')

        # 生成随机索引并打乱数据
        if len(X) > 0:  # 确保数据不为空
            indices = rand_gen.permutation(len(X))
            X = X[indices]
            y = y[indices]

    return X, y


basic_pipeline = fileseg_pipeline


def fileloader_pipeline(
    files: typing.Sequence[str],
    fn_fileloader: TransformFunction,
    fn_encode_label: EncodeLabelFunction | None = _default_encode_label,
    show_progress: bool = True,
    random_shuffle: bool | int = False,
) -> tuple[np.ndarray, np.ndarray]:
    """从文件中读取，制作用于训练的数据集

    本函数相对 fileseg_pipeline 更加灵活，核心是 fn_fileloader 函数，可以用于处理数据并制作数据集

    ```py
    fn_fileloader(file_path) -> (Data, Label)
    ```

    Pipeline 流程:
    ----------

    1. 从每个文件中读取 (X, y) 数据
    2. 合并所有数据
    3. 对标签进行编码
    4. 返回数据集 (X, y)

    Parameters
    ----------
    - `files` : `list[str]`
        - 文件路径列表
    - `fn_fileloader` : `(file_path: str) -> (np.ndarray, np.ndarray)`
        - 文件变换函数，输入文件路径，输出 (X, y) 数据
        - X: shape = (N, C, T)，数据
        - y: shape = (N,)，标签
    - `fn_encode_label` : `(label: list | np.ndarray) -> np.ndarray`, optional
        - 标签编码函数, 默认使用 `to_onehot`
    - `show_progress` : `bool`, optional
        - 是否显示处理进度, by default `True`
    - `random_shuffle` : `bool | int`, optional
        - 是否在返回前随机打乱数据
        - `False`: 不打乱数据
        - `True`: 使用默认随机种子打乱数据
        - `int`: 使用指定的整数作为随机种子打乱数据
        - 注意: 此参数不会影响全局随机状态

    Returns
    ----------
    - `X` : `np.ndarray`, shape = (N, C, T)
        - 数据
    - `y` : `np.ndarray`, shape = (N, n_classes)
        - 标签
    """
    if not files:
        raise ValueError('The files list is empty. Please provide at least one file.')

    # 处理单个文件
    def load_single_file(fpath: str):
        X, y = fn_fileloader(fpath)  # 读取数据
        return X, y

    # 批量处理所有文件
    if show_progress:
        import os

        iterable = rich_progress(
            files, description='处理中', fn_get_status=lambda f: os.path.basename(f)
        )
    else:
        iterable = files
    results = map(load_single_file, iterable)

    # 合并数据
    def merge_data(
        Xy: tuple[list[np.ndarray], list[np.ndarray]], pair: tuple[np.ndarray, np.ndarray]
    ):
        X_batch, y_batch = pair
        if len(X_batch) > 0:
            Xy[0].append(X_batch)
            Xy[1].append(y_batch)
        return Xy

    Xy = reduce(merge_data, results, ([], []))

    if not Xy[0]:
        X = np.array([])
        y = np.array([], dtype=int)
    else:
        X = np.concatenate(Xy[0], axis=0)
        y = np.concatenate(Xy[1], axis=0)

    encode_label_fn = fn_encode_label or _default_encode_label
    y = encode_label_fn(y)

    # 随机打乱数据
    if random_shuffle is not False:
        # 创建独立的随机数生成器实例
        if isinstance(random_shuffle, int):
            # 使用指定的种子
            rand_gen = np.random.RandomState(seed=random_shuffle)
        elif random_shuffle is True:
            # 使用默认种子
            rand_gen = np.random.RandomState()
        else:
            raise ValueError('random_shuffle must be a boolean or an integer')

        # 生成随机索引并打乱数据
        if len(X) > 0:  # 确保数据不为空
            indices = rand_gen.permutation(len(X))
            X = X[indices]
            y = y[indices]

    return X, y


simple_pipeline = fileloader_pipeline  # Alias


def train_test_pipeline(
    split_size: float | int | tuple,
    all_files: typing.Sequence[str] | Sequence[typing.Sequence[str]],
    random_seed: int | None = None,
    return_split_files: bool = False,
    *,
    fn_pipeline: FuncPipeline
    | Literal['fileseg', 'fileloader', 'basic', 'simple'] = 'fileseg',
    pipeline_kwargs: typing.Mapping | None = None,
):
    """在导入文件的时候，自动将不同文件划分为训练集、验证集和测试集

    `split_size` 控制划分方式:
    - `number`: 视为训练集大小，剩余为测试集 (2-way split)。
    - `(train, test)`: 指定训练集和测试集大小 (2-way split)。
    - `(train, val, test)`: 指定训练、验证和测试集大小 (3-way split)。
    大小可以是 `int` (绝对数量) 或 `float` (比例)。

    `fn_pipeline` 回调函数: `(list[str]) -> (np.ndarray, np.ndarray)`
    - 输入: 文件路径列表
    - 输出: 数据集 (X, y)
    当 `fn_pipeline` 为字符串时，使用对应的预设 pipeline。

    #### 使用示例

    ```py
    # 2-way split with ratio
    X_train, X_test, y_train, y_test = train_test_pipeline(
        0.8, all_files, fn_pipeline=my_load_xy
    )

    # 3-way split with ratio
    X_train, X_val, X_test, y_train, y_val, y_test = train_test_pipeline(
        (0.7, 0.1, 0.2), all_files, fn_pipeline=my_load_xy
    )
    ```

    Parameters
    ----------
    - `split_size` : `float | int | tuple`
        - 数据集划分大小或比例。
        - `number`: 训练集大小，剩余为测试集。
        - `(train, test)`: 训练集和测试集大小。
        - `(train, val, test)`: 训练集、验证集和测试集大小。
    - `all_files` : `list[str] | Sequence[list[str]]`
        - 所有文件路径列表。
        - `list[str]`: 直接对整个文件列表进行划分。
        - `Sequence[list[str]]`: 按组进行分层划分。每个子列表被视为一个独立的组（例如，一个类别），
          切分比例将独立应用于每个组，以确保在训练/测试集中保持类别平衡。
    - `random_seed` : `int | None`, optional
        - 随机种子; 当 `random_seed` 为 `None` 时，不设置随机种子。
    - `return_split_files` : `bool`, optional
        - 是否返回划分后的文件路径列表。
    - `fn_pipeline` : `FuncPipeline | Literal['fileseg', 'fileloader', 'basic', 'simple']`, optional
        - 数据处理 pipeline; 默认为 `'fileseg'`，使用 `fileseg_pipeline`。
        - `'fileseg'` / `'basic'`: 使用 `fileseg_pipeline`（信号切片流程）
        - `'fileloader'` / `'simple'`: 使用 `fileloader_pipeline`（灵活加载流程）
    - `**pipeline_kwargs` :
        - pipeline 函数的参数
        - 当使用 `'fileseg'` / `'basic'` 时，包括: `['fn_read', 'fn_segment', 'label', 'fn_before_segment', 'fn_after_segment', 'show_progress', 'fn_encode_label', 'random_shuffle']`
        - 当使用 `'fileloader'` / `'simple'` 时，包括: `['fn_fileloader', 'fn_encode_label', 'show_progress', 'random_shuffle']`

    Returns
    ----------
    - `*data`: `np.ndarray`
        - 划分后的数据集。根据 `split_size` 的不同，返回 `(X_train, X_test, y_train, y_test)` 或 `(X_train, X_val, X_test, y_train, y_val, y_test)`。
    - `*files`: `list[str]`, optional
        - 划分后的文件列表。当 `return_split_files` 为 `True` 时返回。
    """
    all_files = _normalize_file_groups(all_files)

    train_files, val_files, test_files = [], [], []
    has_val = isinstance(split_size, tuple) and len(split_size) == 3

    for group in all_files:
        N = len(group)
        sizes = []
        if isinstance(split_size, (int, float)):
            train_s = int(N * split_size) if isinstance(split_size, float) else split_size
            test_s = N - train_s
            sizes = (train_s, test_s)
        elif isinstance(split_size, tuple):
            raw_sizes = [int(N * s) if isinstance(s, float) else s for s in split_size]
            if isinstance(split_size[-1], float):
                raw_sizes[-1] = N - sum(raw_sizes[:-1])
            sizes = tuple(raw_sizes)
        else:
            raise TypeError(f'Unsupported type for split_size: {type(split_size)}')

        if sum(sizes) > N:
            raise ValueError(
                f'Sum of split sizes ({sum(sizes)}) cannot exceed the number of files ({N}).'
            )

        splits = random_split(group, sizes, seed=random_seed)

        train_files.extend(splits[0])
        if has_val:
            val_files.extend(splits[1])
            test_files.extend(splits[2])
        else:
            test_files.extend(splits[1])

    pipeline_kwargs = pipeline_kwargs or {}
    if fn_pipeline in ('fileseg', 'basic'):
        fn_pipeline = lambda files: fileseg_pipeline(files, **pipeline_kwargs)
    elif fn_pipeline in ('fileloader', 'simple'):
        fn_pipeline = lambda files: fileloader_pipeline(files, **pipeline_kwargs)

    def process_files(file_list: list[str], name: str) -> tuple[np.ndarray, np.ndarray]:
        if not file_list:
            return np.array([]), np.array([])
        print_info(f'Processing {name} files...')
        return fn_pipeline(file_list)

    X_train, y_train = process_files(train_files, 'train')
    X_test, y_test = process_files(test_files, 'test')

    if has_val:
        X_val, y_val = process_files(val_files, 'val')
        data_to_return = (X_train, X_val, X_test, y_train, y_val, y_test)
        files_to_return = (train_files, val_files, test_files)
    else:
        data_to_return = (X_train, X_test, y_train, y_test)
        files_to_return = (train_files, test_files)

    if return_split_files:
        return (*data_to_return, *files_to_return)
    else:
        return data_to_return


def kfold_pipeline(
    k: int,
    all_files: typing.Sequence[str] | Sequence[typing.Sequence[str]],
    random_seed: int | None = None,
    yield_split_files: bool = False,
    *,
    fn_pipeline: FuncPipeline
    | Literal['fileseg', 'fileloader', 'basic', 'simple'] = 'fileseg',
    pipeline_kwargs: typing.Mapping | None = None,
):
    """在导入文件的时候，自动将不同文件划分为 k 折, 通过 `yield` 返回每次的 XY 数据

    - 需要一个 `fn_pipeline: (list[str]) -> (np.ndarray, np.ndarray)` 的回调函数
        - 输入: 文件路径列表
        - 输出: 数据集 (X, y)
    - 当 `fn_pipeline` 为字符串时，使用对应的预设 pipeline，并使用 `**pipeline_kwargs` 作为参数

    Parameters
    ----------
    - `k` : `int`
        - 折数
    - `all_files` : `list[str] | Sequence[list[str]]`
        - 所有文件路径列表。
        - `list[str]`: 直接对整个文件列表进行 K 折交叉验证。
        - `Sequence[list[str]]`: 按组进行分层 K 折交叉验证。每个子列表被视为一个独立的组，
          K 折切分将独立应用于每个组，然后将每折的结果合并，以确保类别平衡。
    - `random_seed` : `int | None`, optional
        - 随机种子; 当 `random_seed` 为 `None` 时，不设置随机种子
    - `yield_split_files` : `bool`, optional
        - 是否返回训练集和验证集的文件路径列表; 当 `yield_split_files` 为 `True` 时，返回训练集和验证集的文件路径列表
    - `fn_pipeline` : `FuncPipeline | Literal['fileseg', 'fileloader', 'basic', 'simple']`, optional
        - 数据处理 pipeline; 默认为 `'fileseg'`，使用 `fileseg_pipeline`
        - `'fileseg'` / `'basic'`: 使用 `fileseg_pipeline`（信号切片流程）
        - `'fileloader'` / `'simple'`: 使用 `fileloader_pipeline`（灵活加载流程）
    - `**pipeline_kwargs` :
        - pipeline 函数的参数
        - 当使用 `'fileseg'` / `'basic'` 时，包括: `['fn_read', 'fn_segment', 'label', 'fn_before_segment', 'fn_after_segment', 'show_progress', 'fn_encode_label', 'random_shuffle']`
        - 当使用 `'fileloader'` / `'simple'` 时，包括: `['fn_fileloader', 'fn_encode_label', 'show_progress', 'random_shuffle']`

    Yields
    ----------
    - `X_train` : `np.ndarray`, shape = (N_train, C, T)
        - 训练数据
    - `X_val` : `np.ndarray`, shape = (N_val, C, T)
        - 验证数据
    - `y_train` : `np.ndarray`, shape = (N_train, n_classes)
        - 训练标签
    - `y_val` : `np.ndarray`, shape = (N_val, n_classes)
        - 验证标签
    - `train_files` : `list[str]`, optional
        - 训练集文件路径列表; 当 `yield_split_files` 为 `True` 时返回
    - `val_files` : `list[str]`, optional
        - 验证集文件路径列表; 当 `yield_split_files` 为 `True` 时返回
    """
    all_files = _normalize_file_groups(all_files)

    # 验证每个组的文件数量
    for group in all_files:
        if len(group) < k:
            raise ValueError(
                f'Number of files in a group ({len(group)}) is less than k ({k}). '
                'Cannot perform K-Fold splitting.'
            )

    # 为每个组预先生成所有的分割索引
    all_splits = []
    for group in all_files:
        kf = KFold(n_splits=k, shuffle=True, random_state=random_seed)
        splits = list(kf.split(group))  # 获取所有 k 折的分割索引
        all_splits.append(splits)

    pipeline_kwargs = pipeline_kwargs or {}
    if fn_pipeline in ('fileseg', 'basic'):
        fn_pipeline = lambda files: fileseg_pipeline(files, **pipeline_kwargs)
    elif fn_pipeline in ('fileloader', 'simple'):
        fn_pipeline = lambda files: fileloader_pipeline(files, **pipeline_kwargs)

    for fold_idx in range(k):
        print_info(f'Processing fold {fold_idx + 1}')
        train_files, val_files = [], []

        # 从每个组的预生成分割中获取当前折的索引
        for i, group in enumerate(all_files):
            train_idx, val_idx = all_splits[i][fold_idx]
            train_files.extend([group[j] for j in train_idx])
            val_files.extend([group[j] for j in val_idx])

        print_info('Processing train files...')
        X_train, y_train = fn_pipeline(train_files)
        print_info('Processing val files...')
        X_val, y_val = fn_pipeline(val_files)

        if yield_split_files:
            yield X_train, X_val, y_train, y_val, train_files, val_files
        else:
            yield X_train, X_val, y_train, y_val
