import numpy as np
from scipy import signal


def _square_waveform(fs, duration=0.01):
    return np.ones(int(round(fs * duration)))


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, repeat=1, iti=0.02, shape='square',
                 shape_settings=None):

    n_trig = int(round(fs * iti))
    trig = np.zeros(n_trig)
    if shape_settings is None:
        shape_settings = {}
    fn = globals()[f'_{shape}_waveform']
    w = fn(fs, **shape_settings)
    if len(w) > n_trig:
        raise ValueError('ITI too short')
    trig[:len(w)] = w

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

    if len(trig) > n:
        raise ValueError('Trigger sequence is too long')

    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.

    Note
    ----
    The units of window 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
    return triggers


def extract_triggers(x, group_window=10):
    '''
    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).
    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.
    '''
    peaks, _ = signal.find_peaks(x, prominence=0.5)
    return group_triggers(peaks, group_window)
