import typing
from typing import Literal, Optional

import numpy as np
from mne.decoding import CSP

from zyplib.utils.ensure import ensure_3dims


def csp_feature(
    signal: np.ndarray,
    n_components: int = 4,
    transform_into: Literal['average_power', 'csp_space'] = 'average_power',
    y: Optional[np.ndarray] = None,
    csp_kwargs: Optional[typing.Mapping] = None,
) -> np.ndarray:
    """Extract CSP features from EEG signals

    本函数使用 mne 库计算 CSP 特征, 需要安装 mne 库

    Parameters
    ----------
    - `signal` : `np.ndarray`
        - 输入信号，形状为 [Batch, Channel, Times]
    - `n_components` : `int`, optional
        - CSP组件数量，默认为4。实际特征数量将是该值的2倍
    - `y` : `np.ndarray`, optional
        - 标签，形状为 [Batch]，必须提供
    - `csp_kwargs` : `dict`, optional
        - CSP参数，默认为None

    Returns
    -------
    - `np.ndarray`
        - CSP特征
    """
    signal = typing.cast(np.ndarray, ensure_3dims(signal, 'batch'))
    # Apply CSP
    csp = CSP(
        n_components=n_components, transform_into=transform_into, **(csp_kwargs or {})
    )
    features = csp.fit_transform(signal, y)

    return features
