# Authors: The MNE-Python contributors.
# License: BSD-3-Clause
# Copyright the MNE-Python contributors.

import numpy as np
from scipy.interpolate import interp1d
from scipy.signal.windows import hann

from .._fiff.pick import _picks_to_idx
from ..epochs import BaseEpochs
from ..event import find_events
from ..evoked import Evoked
from ..io import BaseRaw
from ..utils import _check_option, _check_preload, fill_doc


def _get_window(start, end):
    """Return window which has length as much as parameter start - end."""
    window = 1 - np.r_[hann(4)[:2], np.ones(np.abs(end - start) - 4), hann(4)[-2:]].T
    return window


def _fix_artifact(data, window, picks, first_samp, last_samp, mode):
    """Modify original data by using parameter data."""
    if mode == "linear":
        x = np.array([first_samp, last_samp])
        f = interp1d(x, data[:, (first_samp, last_samp)][picks])
        xnew = np.arange(first_samp, last_samp)
        interp_data = f(xnew)
        data[picks, first_samp:last_samp] = interp_data
    if mode == "window":
        data[picks, first_samp:last_samp] = (
            data[picks, first_samp:last_samp] * window[np.newaxis, :]
        )


@fill_doc
def fix_stim_artifact(
    inst,
    events=None,
    event_id=None,
    tmin=0.0,
    tmax=0.01,
    mode="linear",
    stim_channel=None,
    picks=None,
):
    """Eliminate stimulation's artifacts from instance.

    .. note:: This function operates in-place, consider passing
              ``inst.copy()`` if this is not desired.

    Parameters
    ----------
    inst : instance of Raw or Epochs or Evoked
        The data.
    events : array, shape (n_events, 3)
        The list of events. Required only when inst is Raw.
    event_id : int
        The id of the events generating the stimulation artifacts.
        If None, read all events. Required only when inst is Raw.
    tmin : float
        Start time of the interpolation window in seconds.
    tmax : float
        End time of the interpolation window in seconds.
    mode : 'linear' | 'window'
        Way to fill the artifacted time interval.
        'linear' does linear interpolation
        'window' applies a (1 - hanning) window.
    stim_channel : str | None
        Stim channel to use.
    %(picks_all_data)s

    Returns
    -------
    inst : instance of Raw or Evoked or Epochs
        Instance with modified data.
    """
    _check_option("mode", mode, ["linear", "window"])
    s_start = int(np.ceil(inst.info["sfreq"] * tmin))
    s_end = int(np.ceil(inst.info["sfreq"] * tmax))
    if (mode == "window") and (s_end - s_start) < 4:
        raise ValueError(
            'Time range is too short. Use a larger interval or set mode to "linear".'
        )
    window = None
    if mode == "window":
        window = _get_window(s_start, s_end)

    picks = _picks_to_idx(inst.info, picks, "data", exclude=())

    _check_preload(inst, "fix_stim_artifact")
    if isinstance(inst, BaseRaw):
        if events is None:
            events = find_events(inst, stim_channel=stim_channel)
        if len(events) == 0:
            raise ValueError("No events are found")
        if event_id is None:
            events_sel = np.arange(len(events))
        else:
            events_sel = events[:, 2] == event_id
        event_start = events[events_sel, 0]
        data = inst._data
        for event_idx in event_start:
            first_samp = int(event_idx) - inst.first_samp + s_start
            last_samp = int(event_idx) - inst.first_samp + s_end
            _fix_artifact(data, window, picks, first_samp, last_samp, mode)
    elif isinstance(inst, BaseEpochs):
        if inst.reject is not None:
            raise RuntimeError(
                "Reject is already applied. Use reject=None in the constructor."
            )
        e_start = int(np.ceil(inst.info["sfreq"] * inst.tmin))
        first_samp = s_start - e_start
        last_samp = s_end - e_start
        data = inst._data
        for epoch in data:
            _fix_artifact(epoch, window, picks, first_samp, last_samp, mode)

    elif isinstance(inst, Evoked):
        first_samp = s_start - inst.first
        last_samp = s_end - inst.first
        data = inst.data
        _fix_artifact(data, window, picks, first_samp, last_samp, mode)

    else:
        raise TypeError(f"Not a Raw or Epochs or Evoked (got {type(inst)}).")

    return inst