import numpy as np
from scipy import signal


def _square_waveform(fs, iti=0.02, duration=0.01, repeat=1):
    n_samples = int(round(fs * iti))
    n_trig = int(round(fs * duration))

    trig = np.zeros(n_samples)
    trig[:n_trig] = 1
    trig = np.tile(trig, (repeat, 1))
    trig[1:] *= 0.5
    trig = trig.ravel()
    return trig


def _cos_waveform(fs, pulse_frequency=128, pulse_cycles=1, repeat=1):
    n_samp = int(round((1 / pulse_frequency) * fs)) * pulse_cycles
    t = np.arange(n_samp) / fs
    y = np.sin(2 * np.pi * pulse_frequency * t)
    y = np.tile(y, (repeat, 1))
    y[1:] *= 0.5
    return y.ravel()


def make_trigger(fs, n, shape='square', shape_settings=None):
    if shape_settings is None:
        shape_settings = {}
    fn = globals()[f'_{shape}_waveform']
    trig = fn(fs, **shape_settings)
    if len(trig) > n:
        n_extra = len(trig) - n
        raise ValueError(f'Trigger sequence is too long by {n_extra} samples')

    waveform = np.zeros(n)
    waveform[:len(trig)] = trig
    return waveform


def group_triggers(x, window=10):
    '''
    Group triggers if they fall within `window` of the previous trigger.

    Parameters
    ----------
    x : array
        Array of timestamps (e.g., in units of sample number or seconds).
    window : float
        Group triggers occuring within `window` of the previous trigger. The
        timestamp for the set of triggers is the timestamp of the first trigger
        in the train.

    Returns
    -------
    trigger_times : dict
        Grouped trigger times. Dictionary keys are the number of triggers in
        each group and values are a list of times where that trigger group was
        detected. The reported times are the time of the first trigger in the
        group.

    Note
    ----
    The units of `window` and `trigger_times` are the same units of `x` (e.g.,
    if `x` is specified as trigger timestamps in seconds, then `window` would
    be seconds; if `x` is specified as trigger timestamps in samples, then
    `window` would be number of samples).
    '''
    n_trig = 1

    # This handles the edge condition to make sure the very last set of
    # triggers gets included.
    x = np.pad(x, (0, 1), constant_values=x[-1] + window * 2)

    triggers = {}
    for i in range(1, len(x)):
        if (x[i] - x[i-1]) < window:
            n_trig += 1
        else:
            triggers.setdefault(n_trig, []).append(x[i-n_trig])
            n_trig = 1

    # Convert trigger lists to arrays
    return {k: np.array(v) for k, v in triggers.items()}


def extract_triggers(x, threshold=None, group_window=10, fs=None):
    '''
    Extract and group triggers generated by the trigger interface of subclasses
    of SoundDevice that support including triggers (e.g., Babyface).

    Parameters
    ----------
    x : array
        The waveform of the trigger channel (e.g., Erg2 in the Biosemi BDF
        file).
    threshold : {None, float}
        Threshold used to detect triggers. The trigger times will be the point
        at which the leading edge of the trigger crosses threshold. If None,
        will automatically be set to 20% of the maximum value in the signal.
    group_window : float
        Group triggers occuring within `window` of the previous trigger. The
        timestamp for the set of triggers is the timestamp of the first trigger
        in the train. If `fs` is not None, `group_window` should be in seconds.
        If `fs` is None, `group_window` should be in units of samples.
    fs : {None, float}
        Sampling rate of `x`.

    Returns
    -------
    trigger_times : dict
        Grouped trigger times. Dictionary keys are the number of triggers in
        each group and values are a list of times where that trigger group was
        detected. The reported times are the time of the first trigger in the
        group. Trigger times will be in units of seconds when `fs` is provided,
        samples otherwise.
    '''
    if threshold is None:
        threshold = x.max() * 0.2

    # Find all points in x that exceed the threshold value. The default return
    # type of a logical comparision is np.bool, but we need np.int so that we
    # can properly execute the next step.
    m = (x >= threshold).astype('int')

    # Compute the difference between adjacent samples. A value of 1 indicates
    # the rising edge of the trigger and -1 the falling edge (if m had been
    # kept as a boolean array, the value for both 1 and -1 would just be True).
    d = np.diff(m)

    # Pull out the indices where the difference is 1, indicating the rising
    # (leading) edge of the trigger.
    i = np.flatnonzero(d == 1)

    if fs is not None:
        i = i / fs
    return group_triggers(i, group_window)
