from typing import Iterable, List, Literal, Union

import numpy as np
from typing import cast

from zyplib.utils import ensure_2dims


def apply_eog_regression(
    eeg: np.ndarray,
    eog: np.ndarray,
    fs: float,
    return_type: Literal['raw', 'eeg'] = 'raw',
):
    """基于 MNE 进行 EOG 回归去除眼电伪迹

    1. [Reduce EOG artifacts through regression](https://mne.tools/stable/auto_examples/preprocessing/eog_regression.html)
    2. [mne.preprocessing.EOGRegression API Documentation](https://mne.tools/stable/generated/mne.preprocessing.EOGRegression.html)

    Parameters
    ----------
    - `eeg` : `np.ndarray`
        - _description_
    - `eog` : `np.ndarray`
        - _description_
    - `fs` : `float`
        - _description_
    - `return_type` : `Literal[&#39;raw&#39;, &#39;eeg&#39;]`, optional
        - _description_, by default 'raw'

    Returns
    ----------
    - `_type_`
        - _description_

    Raises
    ----------
    - `ValueError`
        - _description_
    """
    from mne.preprocessing import EOGRegression

    from zyplib.io._base import to_mne_raw

    eeg = cast(np.ndarray, ensure_2dims(eeg))
    eog = cast(np.ndarray, ensure_2dims(eog))
    eeg_ch, eeg_time = eeg.shape
    eog_ch, eog_time = eog.shape
    assert eeg_time == eog_time, 'EEG 和 EOG 通道必须具有相同的时间长度'

    ch_type = ['eeg'] * eeg_ch + ['eog'] * eog_ch
    eeg_chnames = [f'eeg{i}' for i in range(eeg_ch)]
    eog_chnames = [f'eog{i}' for i in range(eog_ch)]
    channels = eeg_chnames + eog_chnames
    data = np.concatenate((eeg, eog), axis=0)
    raw = to_mne_raw(data, fs, channels, ch_types=ch_type)

    model = EOGRegression(picks='eeg', picks_artifact='eog')
    weight = model.fit(raw)
    clean_raw = weight.apply(raw, copy=True)

    if return_type == 'raw':
        return clean_raw
    elif return_type == 'eeg':
        return clean_raw.get_data(picks=eeg_chnames)
    else:
        raise ValueError(f'不支持的返回类型: {return_type}')
