import typing
from typing import Protocol

import numpy as np
import pandas as pd

from zyplib.utils.ensure import ensure_3dims, ensure_npy, ensure_scalar
from zyplib.utils.progress_bar import rich_progress


# @runtime_checkable
class EEGFeatureFunction(Protocol):
    """输入脑电信号 [N, C, T], 输出:

    1. 标量
    2. 向量 [N,]
    3. 矩阵 [N, C] 或 [N, K]
    4. 三维张量 [N, C, K]
    """

    def __call__(self, signal: np.ndarray) -> np.ndarray | float: ...


def _feature_name(func: object) -> str:
    return getattr(func, '__name__', func.__class__.__name__)


def _transform_feature_shape(raw_feat: np.ndarray) -> np.ndarray:
    """
    Preprocesses the raw feature array by removing trailing singleton dimensions
    to match expected 1D, 2D, or 3D structures (after batch dimension).
    """
    processed_feat = raw_feat

    # Step 1: Handle cases like (N, C, K, 1, ..., 1) -> (N, C, K)
    # Or (N, C, K, M, 1, ..., 1) -> (N, C, K, M) if M != 1
    if processed_feat.ndim > 3:
        temp_shape_list = list(processed_feat.shape)
        # Remove trailing 1s as long as the dimension is still > 3
        while len(temp_shape_list) > 3 and temp_shape_list[-1] == 1:
            temp_shape_list.pop()
        if tuple(temp_shape_list) != processed_feat.shape:
            processed_feat = processed_feat.reshape(tuple(temp_shape_list))

    # Step 2: Handle cases like (N, K_or_C, 1) -> (N, K_or_C)
    # This applies if processed_feat is 3D and its last dimension is 1
    if processed_feat.ndim == 3 and processed_feat.shape[-1] == 1:
        processed_feat = processed_feat.reshape(
            processed_feat.shape[:-1]
        )  # remove last dimension

    # Step 3: Handle cases like (N, 1) -> (N,)
    # This applies if processed_feat is 2D and its last dimension is 1
    if processed_feat.ndim == 2 and processed_feat.shape[-1] == 1:
        processed_feat = processed_feat.reshape(
            processed_feat.shape[:-1]
        )  # remove last dimension

    return processed_feat


def make_tabular(
    signals: np.ndarray,
    fn_features: typing.Sequence[EEGFeatureFunction]
    | typing.Mapping[str, EEGFeatureFunction],
    return_feat_names: bool = False,
    squeeze_high_dim_feat: bool = False,
    show_progress: bool = True,
    progress_desc: str = 'Making tabular...',
) -> np.ndarray | tuple[np.ndarray, list[str]]:
    """将脑电信号转换为表格数据

    此函数接收脑电信号数据和一系列特征提取函数，将每个特征函数应用于输入信号，
    并将结果组合成一个表格形式的特征矩阵。

    Parameters
    ----------
    - `signals` : `np.ndarray`
        - 脑电信号, 形状为 (n_samples, n_channels, n_times)
        - 如果输入是2维数组, 将被视为 `(1, n_channels, n_times)`
    - `fn_features` : `list[EEGFeatureFunction] | dict[str, EEGFeatureFunction]`， 特征函数, 可以是以下两种形式之一
        1. 函数列表：每个函数应接受形状为 [n_samples, n_channels, n_times] 的 numpy 数组，
           并返回标量、一维数组或与输入兼容的多维数组
        2. 字典：键为特征名称，值为特征函数
    - `return_feat_names` : `bool`, optional
        - 是否返回特征名称, 默认为 False
    - `squeeze_high_dim_feat` : `bool`, optional
        - 是否压缩高维特征（维度>3）, 默认为 False
        - 如果为 True, 将高维特征压缩为:
          * [N, C, X, Y, ...] -> [N, C, K] (K = X * Y * ...)
          * [N, X, Y, ...] -> [N, K] (K = X * Y * ...)
    - `show_progress` : `bool`, optional
        - 是否显示特征提取进度条, 默认为 True
    - `progress_desc` : `str`, optional
        - 进度条的描述文本, 默认为 "Extracting features"

    Returns
    ----------
    - `np.ndarray | tuple[np.ndarray, list[str]]`
        - 如果 `return_feat_names` 为 False：
            返回表格数据, 形状为 (n_samples, n_features)
        - 如果 `return_feat_names` 为 True：
            返回一个元组, 包含表格数据和特征名称列表

    Notes
    ----------
    1. 特征函数的返回值处理：
       - 标量：被广播到所有样本
       - 一维数组 `[N,]`：必须与样本数量匹配，每个元素对应一个样本的特征
       - 二维数组 `[N, K]` 或 `[N, C]`：
         * 如果第二维等于通道数，假定为每个通道的特征
         * 否则，假定为每个样本的 K 个特征
       - 三维数组 `[N, C, K]`：被解释为每个样本、每个通道的 K 个特征

    2. 特征名称生成规则：
       - 标量或一维数组：使用函数名或字典键
       - 二维数组：
         * 通道特征：`'name_ch{j}'`，j 为通道索引
         * 多特征：`'name{i}'`，i 为特征索引
       - 三维数组：`'name{i}_ch{j}'`，i 为特征索引，j 为通道索引

    Example
    ----------
    ```py
    import numpy as np
    from zyplib.feature.tabular import make_tabular

    # 创建示例数据
    signals = np.random.rand(10, 4, 1000)  # 10 个样本，4 个通道，1000 个时间点

    # 定义一些特征函数
    def mean_feature(x):
        return np.mean(x, axis=-1)

    def std_feature(x):
        return np.std(x, axis=-1)

    # 使用函数列表
    features, names = make_tabular(signals, [mean_feature, std_feature], return_feat_names=True)
    print(features.shape)  # 输出: (10, 8)
    print(names)  # 输出: ['mean_feature_ch0', 'mean_feature_ch1', ..., 'std_feature_ch3']

    # 使用字典
    feature_dict = {'Mean': mean_feature, 'Std': std_feature}
    features, names = make_tabular(signals, feature_dict, return_feat_names=True)
    print(features.shape)  # 输出: (10, 8)
    print(names)  # 输出: ['Mean_ch0', 'Mean_ch1', ..., 'Std_ch3']
    ```
    """
    signals = ensure_npy(signals)
    signals = ensure_3dims(signals, newaxis='batch')

    n_batch, n_channels, n_times = signals.shape

    # 特征名称, 如果为字典，特征名称就为 key name, 否则为 function 的名称
    if isinstance(fn_features, dict):
        feature_list: list[tuple[str, EEGFeatureFunction]] = list(fn_features.items())
    else:
        sequence_features = typing.cast(typing.Sequence[EEGFeatureFunction], fn_features)
        feature_list: list[tuple[str, EEGFeatureFunction]] = [
            (_feature_name(func), func) for func in sequence_features
        ]

    # 使用 fn_features 计算特征
    features: list[np.ndarray] = []
    feature_names: list[str] = []

    iterable: typing.Iterable[tuple[str, EEGFeatureFunction]] = feature_list
    if show_progress:
        iterable = rich_progress(
            feature_list, description=progress_desc, fn_get_status=lambda x: x[0]
        )

    for name, func in iterable:
        feat_raw = func(signals)

        if isinstance(feat_raw, tuple):
            raise ValueError(f'{func} 返回多个值, 不符合要求!')

        scalar = ensure_scalar(feat_raw)
        if scalar is not None:
            # fn 返回标量，则直接加入特征向量中
            features.append(np.full((n_batch, 1), scalar))
            feature_names.append(name)
            continue

        if not isinstance(feat_raw, np.ndarray):
            raise TypeError(f"特征函数 '{name}' 必须返回 numpy.ndarray 或标量")

        original_feat_shape = feat_raw.shape
        _feat = _transform_feature_shape(feat_raw)

        if _feat.ndim == 1:
            # fn 返回向量 [N,]，则将向量展平后加入特征向量中
            if not (_feat.shape[0] == n_batch):
                raise ValueError(
                    f"特征函数 '{name}' (原始形状 {original_feat_shape}, 处理后形状 {_feat.shape}) "
                    f'返回了一个一维数组，其维度 ({_feat.shape[0]}) '
                    f'与批次大小 n_batch ({n_batch}) 不匹配。'
                )
            features.append(_feat.reshape(-1, 1))
            feature_names.append(name)
        elif _feat.ndim == 2:
            # fn 返回 [N, K] 或者 [N, C] 的二维数组，则将特征拼接起来
            # feature name: name{i...k}_ch{j...c}
            if not (_feat.shape[0] == n_batch):
                raise ValueError(
                    f"特征函数 '{name}' (原始形状 {original_feat_shape}, 处理后形状 {_feat.shape}) "
                    f'返回了一个二维数组，其第一维度 ({_feat.shape[0]}) '
                    f'与批次大小 n_batch ({n_batch}) 不匹配。'
                )

            if _feat.shape[1] == n_channels:
                features.append(_feat)
                feature_names.extend([f'{name}_ch{j}' for j in range(n_channels)])
            else:
                features.append(_feat)
                feature_names.extend([f'{name}{i}' for i in range(_feat.shape[1])])
        elif _feat.ndim == 3:
            # fn 返回 [N, C, K] 的三维数组，则将特征拼接起来
            # feature name: name{i...k}_ch{j...c}
            if not (_feat.shape[0] == n_batch and _feat.shape[1] == n_channels):
                raise ValueError(
                    f"特征函数 '{name}' (原始形状 {original_feat_shape}, 处理后形状 {_feat.shape}) "
                    f'返回了一个三维数组。期望形状为 ({n_batch}, {n_channels}, K_features)，'
                    f'但得到了 ({_feat.shape[0]}, {_feat.shape[1]}, {_feat.shape[2]})。'
                )
            k_features = _feat.shape[2]
            feat_reshaped = _feat.reshape(n_batch, -1)  # Flatten to (N, C * K_features)
            features.append(feat_reshaped)
            feature_names.extend(
                [
                    f'{name}{k_idx}_ch{ch_idx}'  # Feature k_idx of channel ch_idx
                    for ch_idx in range(n_channels)
                    for k_idx in range(k_features)
                ]
            )
        else:
            # 处理高维特征 (ndim > 3)
            if squeeze_high_dim_feat:
                if _feat.shape[0] == n_batch:
                    if _feat.shape[1] == n_channels:
                        # 情况1: [N, C, X, Y, ...] -> [N, C, K]
                        # 将后面的维度压缩为一个维度
                        k_features = int(np.prod(_feat.shape[2:]))
                        feat_reshaped = _feat.reshape(n_batch, n_channels, k_features)
                        # 然后再压缩为 [N, C*K]
                        feat_reshaped = feat_reshaped.reshape(n_batch, -1)
                        features.append(feat_reshaped)
                        feature_names.extend(
                            [
                                f'{name}{k_idx}_ch{ch_idx}'
                                for ch_idx in range(n_channels)
                                for k_idx in range(k_features)
                            ]
                        )
                    else:
                        # 情况2: [N, X, Y, ...] -> [N, K]
                        # 将后面的维度压缩为一个维度
                        k_features = int(np.prod(_feat.shape[1:]))
                        feat_reshaped = _feat.reshape(n_batch, k_features)
                        features.append(feat_reshaped)
                        feature_names.extend([f'{name}{i}' for i in range(k_features)])
                else:
                    raise ValueError(
                        f"特征函数 '{name}' (原始形状 {original_feat_shape}, 处理后形状 {_feat.shape}) "
                        f'返回了一个高维数组，但第一维度 ({_feat.shape[0]}) '
                        f'与批次大小 n_batch ({n_batch}) 不匹配。'
                    )
            else:
                raise ValueError(
                    f"不支持的特征形状，特征函数 '{name}': "
                    f'原始形状 {original_feat_shape}，处理后形状 {_feat.shape}。'
                    f'处理后期望得到一维、二维或三维特征，或设置 squeeze_high_dim_feat=True 压缩高维特征。'
                )

    # 最后全部堆叠起来
    if not features:  # 如果 features 列表为空
        # n_batch 在函数前面已经从 signals.shape[0] 获取
        final_features_array = np.empty(
            (n_batch, 0), dtype=float
        )  # 或者选择一个更合适的默认dtype
    else:
        final_features_array = np.hstack(features)

    # 输出 (n_batch, n_features)
    if return_feat_names:
        return final_features_array, feature_names
    else:
        return final_features_array
