"""Preprocessors that work on Raw or Epochs objects."""

# Authors: Hubert Banville <hubert.jbanville@gmail.com>
#          Lukas Gemein <l.gemein@gmail.com>
#          Simon Brandt <simonbrandt@protonmail.com>
#          David Sabbagh <dav.sabbagh@gmail.com>
#          Bruno Aristimunha <b.aristimunha@gmail.com>
#
# License: BSD (3-clause)

from __future__ import annotations
from warnings import warn
from functools import partial
from collections.abc import Iterable
import sys
import platform

if sys.version_info < (3, 9):
    from typing import Callable
else:
    from collections.abc import Callable

import numpy as np
from numpy.typing import NDArray
import pandas as pd
from mne import create_info, BaseEpochs
from mne.io import BaseRaw
from joblib import Parallel, delayed

from braindecode.datasets.base import (
    BaseConcatDataset,
    BaseDataset,
    WindowsDataset,
    EEGWindowsDataset,
)
from braindecode.datautil.serialization import (
    load_concat_dataset,
    _check_save_dir_empty,
)


class Preprocessor(object):
    """Preprocessor for an MNE Raw or Epochs object.

    Applies the provided preprocessing function to the data of a Raw or Epochs
    object.
    If the function is provided as a string, the method with that name will be
    used (e.g., 'pick_channels', 'filter', etc.).
    If it is provided as a callable and `apply_on_array` is True, the
    `apply_function` method of Raw and Epochs object will be used to apply the
    function on the internal arrays of Raw and Epochs.
    If `apply_on_array` is False, the callable must directly modify the Raw or
    Epochs object (e.g., by calling its method(s) or modifying its attributes).

    Parameters
    ----------
    fn: str or callable
        If str, the Raw/Epochs object must have a method with that name.
        If callable, directly apply the callable to the object.
    apply_on_array : bool
        Ignored if `fn` is not a callable. If True, the `apply_function` of Raw
        and Epochs object will be used to run `fn` on the underlying arrays
        directly. If False, `fn` must directly modify the Raw or Epochs object.
    kwargs:
        Keyword arguments to be forwarded to the MNE function.
    """

    def __init__(self, fn: Callable | str, *, apply_on_array: bool = True, **kwargs):
        if hasattr(fn, "__name__") and fn.__name__ == "<lambda>":
            warn("Preprocessing choices with lambda functions cannot be saved.")
        if callable(fn) and apply_on_array:
            channel_wise = kwargs.pop("channel_wise", False)
            picks = kwargs.pop("picks", None)
            n_jobs = kwargs.pop("n_jobs", 1)
            kwargs = dict(
                fun=partial(fn, **kwargs),
                channel_wise=channel_wise,
                picks=picks,
                n_jobs=n_jobs,
            )
            fn = "apply_function"
        self.fn = fn
        self.kwargs = kwargs

    def apply(self, raw_or_epochs: BaseRaw | BaseEpochs):
        try:
            self._try_apply(raw_or_epochs)
        except RuntimeError:
            # Maybe the function needs the data to be loaded and the data was
            # not loaded yet. Not all MNE functions need data to be loaded,
            # most importantly the 'crop' function can be lazily applied
            # without preloading data which can make the overall preprocessing
            # pipeline substantially faster.
            raw_or_epochs.load_data()
            self._try_apply(raw_or_epochs)

    def _try_apply(self, raw_or_epochs):
        if callable(self.fn):
            self.fn(raw_or_epochs, **self.kwargs)
        else:
            if not hasattr(raw_or_epochs, self.fn):
                raise AttributeError(f"MNE object does not have a {self.fn} method.")
            getattr(raw_or_epochs, self.fn)(**self.kwargs)


def preprocess(
    concat_ds: BaseConcatDataset,
    preprocessors: list[Preprocessor],
    save_dir: str | None = None,
    overwrite: bool = False,
    n_jobs: int | None = None,
    offset: int = 0,
):
    """Apply preprocessors to a concat dataset.

    Parameters
    ----------
    concat_ds: BaseConcatDataset
        A concat of BaseDataset or WindowsDataset datasets to be preprocessed.
    preprocessors: list(Preprocessor)
        List of Preprocessor objects to apply to the dataset.
    save_dir : str | None
        If a string, the preprocessed data will be saved under the specified
        directory and the datasets in ``concat_ds`` will be reloaded with
        `preload=False`.
    overwrite : bool
        When `save_dir` is provided, controls whether to delete the old
        subdirectories that will be written to under `save_dir`. If False and
        the corresponding subdirectories already exist, a ``FileExistsError``
        will be raised.
    n_jobs : int | None
        Number of jobs for parallel execution. See `joblib.Parallel` for
        a more detailed explanation.
    offset : int
        If provided, the integer is added to the id of the dataset in the
        concat. This is useful in the setting of very large datasets, where
        one dataset has to be processed and saved at a time to account for
        its original position.

    Returns
    -------
    BaseConcatDataset:
        Preprocessed dataset.
    """
    # In case of serialization, make sure directory is available before
    # preprocessing
    if save_dir is not None and not overwrite:
        _check_save_dir_empty(save_dir)

    if not isinstance(preprocessors, Iterable):
        raise ValueError("preprocessors must be a list of Preprocessor objects.")
    for elem in preprocessors:
        assert hasattr(elem, "apply"), "Preprocessor object needs an `apply` method."

    parallel_processing = (n_jobs is not None) and (n_jobs != 1)

    job_prefer = "threads" if platform.system() == "Windows" else None
    list_of_ds = Parallel(n_jobs=n_jobs, prefer=job_prefer)(
        delayed(_preprocess)(
            ds,
            i + offset,
            preprocessors,
            save_dir,
            overwrite,
            copy_data=(parallel_processing and (save_dir is None)),
        )
        for i, ds in enumerate(concat_ds.datasets)
    )

    if save_dir is not None:  # Reload datasets and replace in concat_ds
        ids_to_load = [i + offset for i in range(len(concat_ds.datasets))]
        concat_ds_reloaded = load_concat_dataset(
            save_dir,
            preload=False,
            target_name=None,
            ids_to_load=ids_to_load,
        )
        _replace_inplace(concat_ds, concat_ds_reloaded)
    else:
        if parallel_processing:  # joblib made copies
            _replace_inplace(concat_ds, BaseConcatDataset(list_of_ds))
        else:  # joblib did not make copies, the
            # preprocessing happened in-place
            # Recompute cumulative sizes as transforms might have changed them
            concat_ds.cumulative_sizes = concat_ds.cumsum(concat_ds.datasets)

    return concat_ds


def _replace_inplace(concat_ds, new_concat_ds):
    """Replace subdatasets and preproc_kwargs of a BaseConcatDataset inplace.

    Parameters
    ----------
    concat_ds : BaseConcatDataset
        Dataset to modify inplace.
    new_concat_ds : BaseConcatDataset
        Dataset to use to modify ``concat_ds``.
    """
    if len(concat_ds.datasets) != len(new_concat_ds.datasets):
        raise ValueError("Both inputs must have the same length.")
    for i in range(len(new_concat_ds.datasets)):
        concat_ds.datasets[i] = new_concat_ds.datasets[i]

    concat_kind = "raw" if hasattr(concat_ds.datasets[0], "raw") else "window"
    preproc_kwargs_attr = concat_kind + "_preproc_kwargs"
    if hasattr(new_concat_ds, preproc_kwargs_attr):
        setattr(
            concat_ds, preproc_kwargs_attr, getattr(new_concat_ds, preproc_kwargs_attr)
        )


def _preprocess(
    ds, ds_index, preprocessors, save_dir=None, overwrite=False, copy_data=False
):
    """Apply preprocessor(s) to Raw or Epochs object.

    Parameters
    ----------
    ds: BaseDataset | WindowsDataset
        Dataset object to preprocess.
    ds_index : int
        Index of the BaseDataset in its BaseConcatDataset. Ignored if save_dir
        is None.
    preprocessors: list(Preprocessor)
        List of preprocessors to apply to the dataset.
    save_dir : str | None
        If provided, save the preprocessed BaseDataset in the
        specified directory.
    overwrite : bool
        If True, overwrite existing file with the same name.
    copy_data : bool
        First copy the data in case it is preloaded. Necessary for parallel processing to work.
    """

    def _preprocess_raw_or_epochs(raw_or_epochs, preprocessors):
        # Copying the data necessary in some scenarios for parallel processing
        # to work when data is in memory (else error about _data not being writeable)
        if raw_or_epochs.preload and copy_data:
            raw_or_epochs._data = raw_or_epochs._data.copy()
        for preproc in preprocessors:
            preproc.apply(raw_or_epochs)

    if hasattr(ds, "raw"):
        if isinstance(ds, EEGWindowsDataset):
            warn(
                f"Applying preprocessors {preprocessors} to the mne.io.Raw of an EEGWindowsDataset."
            )
        _preprocess_raw_or_epochs(ds.raw, preprocessors)
    elif hasattr(ds, "windows"):
        _preprocess_raw_or_epochs(ds.windows, preprocessors)
    else:
        raise ValueError(
            "Can only preprocess concatenation of BaseDataset or "
            "WindowsDataset, with either a `raw` or `windows` attribute."
        )

    # Store preprocessing keyword arguments in the dataset
    _set_preproc_kwargs(ds, preprocessors)

    if save_dir is not None:
        concat_ds = BaseConcatDataset([ds])
        concat_ds.save(save_dir, overwrite=overwrite, offset=ds_index)
    else:
        return ds


def _get_preproc_kwargs(preprocessors):
    preproc_kwargs = []
    for p in preprocessors:
        # in case of a mne function, fn is a str, kwargs is a dict
        func_name = p.fn
        func_kwargs = p.kwargs
        # in case of another function
        # if apply_on_array=False
        if callable(p.fn):
            func_name = p.fn.__name__
        # if apply_on_array=True
        else:
            if "fun" in p.fn:
                func_name = p.kwargs["fun"].func.__name__
                func_kwargs = p.kwargs["fun"].keywords
        preproc_kwargs.append((func_name, func_kwargs))
    return preproc_kwargs


def _set_preproc_kwargs(ds, preprocessors):
    """Record preprocessing keyword arguments in BaseDataset or WindowsDataset.

    Parameters
    ----------
    ds : BaseDataset | WindowsDataset
        Dataset in which to record preprocessing keyword arguments.
    preprocessors : list
        List of preprocessors.
    """
    preproc_kwargs = _get_preproc_kwargs(preprocessors)
    if isinstance(ds, WindowsDataset):
        kind = "window"
    if isinstance(ds, EEGWindowsDataset):
        kind = "raw"
    elif isinstance(ds, BaseDataset):
        kind = "raw"
    else:
        raise TypeError(f"ds must be a BaseDataset or a WindowsDataset, got {type(ds)}")
    setattr(ds, kind + "_preproc_kwargs", preproc_kwargs)


def exponential_moving_standardize(
    data: NDArray,
    factor_new: float = 0.001,
    init_block_size: int | None = None,
    eps: float = 1e-4,
):
    r"""Perform exponential moving standardization.

    Compute the exponental moving mean :math:`m_t` at time `t` as
    :math:`m_t=\mathrm{factornew} \cdot mean(x_t) + (1 - \mathrm{factornew}) \cdot m_{t-1}`.

    Then, compute exponential moving variance :math:`v_t` at time `t` as
    :math:`v_t=\mathrm{factornew} \cdot (m_t - x_t)^2 + (1 - \mathrm{factornew}) \cdot v_{t-1}`.

    Finally, standardize the data point :math:`x_t` at time `t` as:
    :math:`x'_t=(x_t - m_t) / max(\sqrt{->v_t}, eps)`.


    Parameters
    ----------
    data: np.ndarray (n_channels, n_times)
    factor_new: float
    init_block_size: int
        Standardize data before to this index with regular standardization.
    eps: float
        Stabilizer for division by zero variance.

    Returns
    -------
    standardized: np.ndarray (n_channels, n_times)
        Standardized data.
    """
    data = data.T
    df = pd.DataFrame(data)
    meaned = df.ewm(alpha=factor_new).mean()
    demeaned = df - meaned
    squared = demeaned * demeaned
    square_ewmed = squared.ewm(alpha=factor_new).mean()
    standardized = demeaned / np.maximum(eps, np.sqrt(np.array(square_ewmed)))
    standardized = np.array(standardized)
    if init_block_size is not None:
        i_time_axis = 0
        init_mean = np.mean(data[0:init_block_size], axis=i_time_axis, keepdims=True)
        init_std = np.std(data[0:init_block_size], axis=i_time_axis, keepdims=True)
        init_block_standardized = (data[0:init_block_size] - init_mean) / np.maximum(
            eps, init_std
        )
        standardized[0:init_block_size] = init_block_standardized
    return standardized.T


def exponential_moving_demean(
    data: NDArray, factor_new: float = 0.001, init_block_size: int | None = None
):
    r"""Perform exponential moving demeanining.

    Compute the exponental moving mean :math:`m_t` at time `t` as
    :math:`m_t=\mathrm{factornew} \cdot mean(x_t) + (1 - \mathrm{factornew}) \cdot m_{t-1}`.

    Deman the data point :math:`x_t` at time `t` as:
    :math:`x'_t=(x_t - m_t)`.

    Parameters
    ----------
    data: np.ndarray (n_channels, n_times)
    factor_new: float
    init_block_size: int
        Demean data before to this index with regular demeaning.

    Returns
    -------
    demeaned: np.ndarray (n_channels, n_times)
        Demeaned data.
    """
    data = data.T
    df = pd.DataFrame(data)
    meaned = df.ewm(alpha=factor_new).mean()
    demeaned = df - meaned
    demeaned = np.array(demeaned)
    if init_block_size is not None:
        i_time_axis = 0
        init_mean = np.mean(data[0:init_block_size], axis=i_time_axis, keepdims=True)
        demeaned[0:init_block_size] = data[0:init_block_size] - init_mean
    return demeaned.T


def filterbank(
    raw: BaseRaw,
    frequency_bands: list[tuple[float, float]],
    drop_original_signals: bool = True,
    order_by_frequency_band: bool = False,
    **mne_filter_kwargs,
):
    """Applies multiple bandpass filters to the signals in raw. The raw will be
    modified in-place and number of channels in raw will be updated to
    len(frequency_bands) * len(raw.ch_names) (-len(raw.ch_names) if
    drop_original_signals).

    Parameters
    ----------
    raw: mne.io.Raw
        The raw signals to be filtered.
    frequency_bands: list(tuple)
        The frequency bands to be filtered for (e.g. [(4, 8), (8, 13)]).
    drop_original_signals: bool
        Whether to drop the original unfiltered signals
    order_by_frequency_band: bool
        If True will return channels ordered by frequency bands, so if there
        are channels Cz, O1 and filterbank ranges [(4,8), (8,13)], returned
        channels will be [Cz_4-8, O1_4-8, Cz_8-13, O1_8-13]. If False, order
        will be [Cz_4-8, Cz_8-13, O1_4-8, O1_8-13].
    mne_filter_kwargs: dict
        Keyword arguments for filtering supported by mne.io.Raw.filter().
        Please refer to mne for a detailed explanation.
    """
    if not frequency_bands:
        raise ValueError(f"Expected at least one frequency band, got {frequency_bands}")
    if not all([len(ch_name) < 8 for ch_name in raw.ch_names]):
        warn(
            "Try to use shorter channel names, since frequency band "
            "annotation requires an estimated 4-8 chars depending on the "
            "frequency ranges. Will truncate to 15 chars (mne max)."
        )
    original_ch_names = raw.ch_names
    all_filtered = []
    for l_freq, h_freq in frequency_bands:
        filtered = raw.copy()
        filtered.filter(l_freq=l_freq, h_freq=h_freq, **mne_filter_kwargs)
        # mne automatically changes the highpass/lowpass info values
        # when applying filters and channels can't be added if they have
        # different such parameters. Not needed when making picks as
        # high pass is not modified by filter if pick is specified

        ch_names = filtered.info.ch_names
        ch_types = filtered.info.get_channel_types()
        sampling_freq = filtered.info["sfreq"]

        info = create_info(ch_names=ch_names, ch_types=ch_types, sfreq=sampling_freq)

        filtered.info = info

        # add frequency band annotation to channel names
        # truncate to a max of 15 characters, since mne does not allow for more
        filtered.rename_channels(
            {
                old_name: (old_name + f"_{l_freq}-{h_freq}")[-15:]
                for old_name in filtered.ch_names
            }
        )
        all_filtered.append(filtered)
    raw.add_channels(all_filtered)
    if not order_by_frequency_band:
        # order channels by name and not by frequency band:
        # index the list with a stepsize of the number of channels for each of
        # the original channels
        chs_by_freq_band = []
        for i in range(len(original_ch_names)):
            chs_by_freq_band.extend(raw.ch_names[i :: len(original_ch_names)])
        raw.reorder_channels(chs_by_freq_band)
    if drop_original_signals:
        raw.drop_channels(original_ch_names)