"""
检测脑电信号当中的坏通道和坏 epoch 等受污染严重的部分
"""

import typing

from typing import Iterable, List, Optional, Union, cast

import numpy as np
from sklearn.preprocessing import StandardScaler

from zyplib.utils import ensure_2dims

from .outlier_detect import _MethodArgs, _MethodName, detect_outliers


def find_bad_channels_simple(
    eeg: np.ndarray,
    fs: float,
    ch_names: Optional[typing.Sequence[str]] = None,
    method: _MethodName = 'lof',
    **method_args,
) -> tuple[List[str], np.ndarray]:
    """基于特征工程和异常检测算法识别坏道。

    该函数通过为每个通道提取一组统计特征，然后使用 `detect_outliers`
    函数来识别特征异常的通道，从而实现坏道检测。

    Parameters
    ----------
    - `eeg` : `np.ndarray`
        - EEG 数据，形状为 [C, T] (通道数, 时间点数)。
    - `fs` : `float`
        - 采样频率 (Hz)。
    - `ch_names` : `Optional[List[str]]`, optional
        - 通道名称列表。如果为 None，将自动生成 'eeg0', 'eeg1' 等名称。
    - `method` : `_MethodName`, optional
        - 使用的异常检测算法，默认为 'lof'。
    - `**method_args`
        - 传递给 `zyplib.preprocessing.outlier_detect.detect_outliers` 的参数。
        - 例如 `n_neighbors`, `contamination` 等。

    Returns
    -------
    - `tuple[List[str], np.ndarray]`
        - `bad_channels`: 检测到的坏通道名称列表。
        - `scores`: 每个坏通道对应的异常分数。
    """

    if fs <= 0:
        raise ValueError('采样频率必须为正数')

    eeg = cast(np.ndarray, ensure_2dims(eeg))
    n_channels, n_times = eeg.shape

    if n_times < fs:  # 至少需要1秒的数据
        raise ValueError('数据长度太短，至少需要1秒的数据进行坏通道检测')

    # 验证或生成通道名称
    if ch_names is None:
        ch_names = [f'eeg{i}' for i in range(n_channels)]
    elif len(ch_names) != n_channels:
        raise ValueError('ch_names 的长度必须与 eeg 数据中的通道数匹配')

    # 1. 特征提取
    features = []
    for i in range(n_channels):
        ch_data = eeg[i, :]

        # 基本统计特征
        std_dev = np.std(ch_data)
        ptp = np.ptp(ch_data)

        diff = np.diff(ch_data)
        diff_mean = np.mean(np.abs(diff))

        features.append([std_dev, ptp, diff_mean])

    features = np.array(features)

    is_bad_mask, scores = detect_outliers(
        features, method=method, method_args=cast(_MethodArgs, method_args)
    )

    bad_channels = [ch_names[i] for i, is_bad in enumerate(is_bad_mask) if is_bad]
    scores = scores[is_bad_mask]

    return bad_channels, scores


def find_bad_epochs_simple(
    epochs: Union[np.ndarray, Iterable[np.ndarray]],
    fs: float,
    *,
    method: _MethodName = 'lof',
    method_args: Optional[_MethodArgs] = None,
    return_indices: bool = True,
) -> tuple[List[int | bool], np.ndarray]:
    """通过向量化操作高效检测坏的 epoch

    基于多种信号特征和 `zyplib.preprocessing.outlier_detect` 模块进行坏 epoch 检测。
    此版本经过向量化优化，可高效处理大量的 epoch。

    Parameters
    ----------
    - `epochs` : `Union[np.ndarray, Iterable[np.ndarray]]`
        - epoch 数据，可以是 [N, C, T] 的3D数组，或 `np.ndarray` 的可迭代对象。
    - `fs` : `float`
        - 采样频率 (Hz)，用于验证。
    - `method` : `_MethodName`, optional
        - 使用的异常检测算法，传递给 `detect_outliers` 函数, by default 'lof'.
    - `method_args` : `_MethodArgs`, optional
        - 传递给所选方法的参数字典, by default None.
    - `return_indices` : `bool`, optional
        - 如果为 True，返回坏 epoch 的索引；如果为 False，返回布尔掩码, by default True

    Returns
    -------
    - `List[int] | List[bool]`
        - 如果 return_indices=True: 坏 epoch 的索引列表
        - 如果 return_indices=False: 布尔掩码，True 表示坏 epoch
    - `np.ndarray`
        - 异常分数，值越大越可能是异常值

    Notes
    -----
    1.  所有 epoch 被统一转换为 [N, C, T] 的3D数组。
    2.  使用向量化操作提取每个 epoch 的多种统计特征，包括：
        - 差分特征 (信号变化率)
        - 功率特征
        - 峰峰值特征
        - 方差特征
        - 变异系数
    3.  对特征进行标准化处理。
    4.  调用 `detect_outliers` 函数进行异常检测。

    参考:
    - `zyplib.preprocessing.outlier_detect.detect_outliers`
    - `zyplib.utils.ensure.ensure_3dims`
    """
    from zyplib.utils import ensure_3dims

    from .outlier_detect import detect_outliers

    if fs <= 0:
        raise ValueError('采样频率必须为正数')

    # 1. 统一数据格式为 [N, C, T] 的3D数组
    if isinstance(epochs, np.ndarray) and epochs.ndim >= 3:
        epochs_array = ensure_3dims(epochs)
    else:  # 处理可迭代对象
        epochs_list = list(epochs)
        if not epochs_list:
            raise ValueError('epochs 不能为空')
        epochs_array = np.array([ensure_2dims(e) for e in epochs_list])

    if epochs_array.ndim != 3:
        raise ValueError(f'处理后的 epoch 数组应为3维，但得到 {epochs_array.ndim} 维')

    # 2. 向量化特征提取
    # 差分特征 (axis=2 是时间轴)
    diff = np.diff(epochs_array, axis=2)
    diff_mean_ch = np.mean(np.abs(diff), axis=2)  # [N, C]
    diff_std_ch = np.std(diff, axis=2)  # [N, C]
    f1 = np.hstack(
        [
            np.mean(diff_mean_ch, axis=1, keepdims=True),
            np.std(diff_mean_ch, axis=1, keepdims=True),
            np.mean(diff_std_ch, axis=1, keepdims=True),
            np.std(diff_std_ch, axis=1, keepdims=True),
        ]
    )  # [N, 4]

    # 功率特征
    power_ch = np.mean(epochs_array**2, axis=2)  # [N, C]
    f2 = np.hstack(
        [
            np.mean(power_ch, axis=1, keepdims=True),
            np.std(power_ch, axis=1, keepdims=True),
        ]
    )  # [N, 2]

    # 峰峰值特征
    ptp_ch = np.ptp(epochs_array, axis=2)  # [N, C]
    f3 = np.hstack(
        [np.mean(ptp_ch, axis=1, keepdims=True), np.std(ptp_ch, axis=1, keepdims=True)]
    )  # [N, 2]

    # 方差特征
    var_ch = np.var(epochs_array, axis=2)  # [N, C]
    f4 = np.hstack(
        [np.mean(var_ch, axis=1, keepdims=True), np.std(var_ch, axis=1, keepdims=True)]
    )  # [N, 2]

    # 变异系数
    mean_abs_ch = np.mean(np.abs(epochs_array), axis=2)  # [N, C]
    std_ch = np.std(epochs_array, axis=2)  # [N, C]
    cv_ch = std_ch / (mean_abs_ch + 1e-8)
    f5 = np.hstack(
        [np.mean(cv_ch, axis=1, keepdims=True), np.std(cv_ch, axis=1, keepdims=True)]
    )  # [N, 2]

    # 合并所有特征
    features = np.hstack([f1, f2, f3, f4, f5])  # [N, 12]

    # 注意：不再手动标准化，让 detect_outliers 内部处理
    # 这避免了二次标准化的问题

    # 4. 使用 outlier_detect 模块进行异常检测
    bad_epochs_mask, scores = detect_outliers(
        features, method=method, method_args=method_args
    )

    if return_indices:
        return np.where(bad_epochs_mask)[0].tolist(), scores
    else:
        return bad_epochs_mask.tolist(), scores
