# -*- coding: utf-8 -*-
"""
Functions to generate spike trains from analog signals,
or to generate random spike trains.

Some functions are based on the NeuroTools stgen module, which was mostly
written by Eilif Muller, or from the NeuroTools signals.analogs module.

:copyright: Copyright 2015 by the Elephant team, see AUTHORS.txt.
:license: Modified BSD, see LICENSE.txt for details.
"""

from __future__ import division
import numpy as np
from quantities import ms, mV, Hz, Quantity, dimensionless
from neo import SpikeTrain
import random
from elephant.spike_train_surrogates import dither_spike_train


def threshold_detection(signal, threshold=0.0 * mV, sign='above'):
    """
    Returns the times when the analog signal crosses a threshold.
    Usually used for extracting spike times from a membrane potential.
    Adapted from version in NeuroTools.

    Parameters
    ----------
    signal : neo AnalogSignal object
        'signal' is an analog signal.
    threshold : A quantity, e.g. in mV
        'threshold' contains a value that must be reached
        for an event to be detected.
    sign : 'above' or 'below'
        'sign' determines whether to count thresholding crossings
        that cross above or below the threshold.
    format : None or 'raw'
        Whether to return as SpikeTrain (None)
        or as a plain array of times ('raw').

    Returns
    -------
    result_st : neo SpikeTrain object
        'result_st' contains the spike times of each of the events (spikes)
        extracted from the signal.
    """

    assert threshold is not None, "A threshold must be provided"

    if sign is 'above':
        cutout = np.where(signal > threshold)[0]
    elif sign in 'below':
        cutout = np.where(signal < threshold)[0]

    if len(cutout) <= 0:
        events = np.zeros(0)
    else:
        take = np.where(np.diff(cutout) > 1)[0] + 1
        take = np.append(0, take)

        time = signal.times
        events = time[cutout][take]

    events_base = events.base
    if events_base is None:
        # This occurs in some Python 3 builds due to some
        # bug in quantities.
        events_base = np.array([event.base for event in events])  # Workaround

    result_st = SpikeTrain(events_base, units=signal.times.units,
                           t_start=signal.t_start, t_stop=signal.t_stop)
    return result_st


def _homogeneous_process(interval_generator, args, mean_rate, t_start, t_stop,
                         as_array):
    """
    Returns a spike train whose spikes are a realization of a random process
    generated by the function `interval_generator` with the given rate,
    starting at time `t_start` and stopping `time t_stop`.
    """
    def rescale(x):
        return (x / mean_rate.units).rescale(t_stop.units)

    n = int(((t_stop - t_start) * mean_rate).simplified)
    number = np.ceil(n + 3 * np.sqrt(n))
    if number < 100:
        number = min(5 + np.ceil(2 * n), 100)
    assert number > 4  # if positive, number cannot be less than 5
    isi = rescale(interval_generator(*args, size=number))
    spikes = np.cumsum(isi)
    spikes += t_start

    i = spikes.searchsorted(t_stop)
    if i == len(spikes):
        # ISI buffer overrun
        extra_spikes = []
        t_last = spikes[-1] + rescale(interval_generator(*args, size=1))[0]
        while t_last < t_stop:
            extra_spikes.append(t_last)
            t_last = t_last + rescale(interval_generator(*args, size=1))[0]
        # np.concatenate does not conserve units
        spikes = Quantity(
            np.concatenate(
                (spikes, extra_spikes)).magnitude, units=spikes.units)
    else:
        spikes = spikes[:i]

    if as_array:
        spikes = spikes.magnitude
    else:
        spikes = SpikeTrain(
            spikes, t_start=t_start, t_stop=t_stop, units=spikes.units)

    return spikes


def homogeneous_poisson_process(rate, t_start=0.0 * ms, t_stop=1000.0 * ms,
                                as_array=False):
    """
    Returns a spike train whose spikes are a realization of a Poisson process
    with the given rate, starting at time `t_start` and stopping time `t_stop`.

    All numerical values should be given as Quantities, e.g. 100*Hz.

    Parameters
    ----------

    rate : Quantity scalar with dimension 1/time
           The rate of the discharge.
    t_start : Quantity scalar with dimension time
              The beginning of the spike train.
    t_stop : Quantity scalar with dimension time
             The end of the spike train.
    as_array : bool
               If True, a NumPy array of sorted spikes is returned,
               rather than a SpikeTrain object.

    Examples
    --------
        >>> from quantities import Hz, ms
        >>> spikes = homogeneous_poisson_process(50*Hz, 0*ms, 1000*ms)
        >>> spikes = homogeneous_poisson_process(
            20*Hz, 5000*ms, 10000*ms, as_array=True)

    """
    mean_interval = 1 / rate
    return _homogeneous_process(
        np.random.exponential, (mean_interval,), rate, t_start, t_stop,
        as_array)


def homogeneous_gamma_process(a, b, t_start=0.0 * ms, t_stop=1000.0 * ms,
                              as_array=False):
    """
    Returns a spike train whose spikes are a realization of a gamma process
    with the given parameters, starting at time `t_start` and stopping time
    `t_stop` (average rate will be b/a).

    All numerical values should be given as Quantities, e.g. 100*Hz.

    Parameters
    ----------

    a : int or float
        The shape parameter of the gamma distribution.
    b : Quantity scalar with dimension 1/time
        The rate parameter of the gamma distribution.
    t_start : Quantity scalar with dimension time
              The beginning of the spike train.
    t_stop : Quantity scalar with dimension time
             The end of the spike train.
    as_array : bool
               If True, a NumPy array of sorted spikes is returned,
               rather than a SpikeTrain object.

    Examples
    --------
        >>> from quantities import Hz, ms
        >>> spikes = homogeneous_gamma_process(2.0, 50*Hz, 0*ms, 1000*ms)
        >>> spikes = homogeneous_gamma_process(
                5.0, 20*Hz, 5000*ms, 10000*ms, as_array=True)

    """
    rate = b / a
    k, theta = a, (1 / b)
    return _homogeneous_process(np.random.gamma, (k, theta), rate, t_start, t_stop, as_array)


def _n_poisson(rate, t_stop, t_start=0.0 * ms, n=1):
    """
    Generates one or more independent Poisson spike trains.

    Parameters
    ----------
    rate : Quantity or Quantity array
        Expected firing rate (frequency) of each output SpikeTrain.
        Can be one of:
        *  a single Quantity value: expected firing rate of each output
           SpikeTrain
        *  a Quantity array: rate[i] is the expected firing rate of the i-th
           output SpikeTrain
    t_stop : Quantity
        Single common stop time of each output SpikeTrain. Must be > t_start.
    t_start : Quantity (optional)
        Single common start time of each output SpikeTrain. Must be < t_stop.
        Default: 0 s.
    n: int (optional)
        If rate is a single Quantity value, n specifies the number of
        SpikeTrains to be generated. If rate is an array, n is ignored and the
        number of SpikeTrains is equal to len(rate).
        Default: 1


    Returns
    -------
    list of neo.SpikeTrain
        Each SpikeTrain contains one of the independent Poisson spike trains,
        either n SpikeTrains of the same rate, or len(rate) SpikeTrains with
        varying rates according to the rate parameter. The time unit of the
        SpikeTrains is given by t_stop.
    """
    # Check that the provided input is Hertz of return error
    try:
        for r in rate.reshape(-1, 1):
            r.rescale('Hz')
    except AttributeError:
        raise ValueError('rate argument must have rate unit (1/time)')

    # Check t_start < t_stop and create their strip dimensions
    if not t_start < t_stop:
        raise ValueError(
            't_start (=%s) must be < t_stop (=%s)' % (t_start, t_stop))

    # Set number n of output spike trains (specified or set to len(rate))
    if not (type(n) == int and n > 0):
        raise ValueError('n (=%s) must be a positive integer' % str(n))
    rate_dl = rate.simplified.magnitude.flatten()

    # Check rate input parameter
    if len(rate_dl) == 1:
        if rate_dl < 0:
            raise ValueError('rate (=%s) must be non-negative.' % rate)
        rates = np.array([rate_dl] * n)
    else:
        rates = rate_dl.flatten()
        if any(rates < 0):
            raise ValueError('rate must have non-negative elements.')
    sts = []
    for r in rates:
        sts.append(homogeneous_poisson_process(r * Hz, t_start, t_stop))
    return sts


def single_interaction_process(
        rate, rate_c, t_stop, n=2, jitter=0 * ms, coincidences='deterministic',
        t_start=0 * ms, min_delay=0 * ms, return_coinc=False):
    """
    Generates a multidimensional Poisson SIP (single interaction process)
    plus independent Poisson processes

    A Poisson SIP consists of Poisson time series which are independent
    except for simultaneous events in all of them. This routine generates
    a SIP plus additional parallel independent Poisson processes.

    See [1].

    Parameters
    -----------
    t_stop: quantities.Quantity
        Total time of the simulated processes. The events are drawn between
        0 and `t_stop`.
    rate: quantities.Quantity
        Overall mean rate of the time series to be generated (coincidence
        rate `rate_c` is subtracted to determine the background rate). Can be:
        * a float, representing the overall mean rate of each process. If
          so, it must be higher than `rate_c`.
        * an iterable of floats (one float per process), each float
          representing the overall mean rate of a process. If so, all the
          entries must be larger than `rate_c`.
    rate_c: quantities.Quantity
        Coincidence rate (rate of coincidences for the n-dimensional SIP).
        The SIP spike trains will have coincident events with rate `rate_c`
        plus independent 'background' events with rate `rate-rate_c`.
    n: int, optional
        If `rate` is a single Quantity value, `n` specifies the number of
        SpikeTrains to be generated. If rate is an array, `n` is ignored and
        the number of SpikeTrains is equal to `len(rate)`.
        Default: 1
    jitter: quantities.Quantity, optional
        Jitter for the coincident events. If `jitter == 0`, the events of all
        n correlated processes are exactly coincident. Otherwise, they are
        jittered around a common time randomly, up to +/- `jitter`.
    coincidences: string, optional
        Whether the total number of injected coincidences must be determin-
        istic (i.e. rate_c is the actual rate with which coincidences are
        generated) or stochastic (i.e. rate_c is the mean rate of coincid-
        ences):
        * 'deterministic': deterministic rate
        * 'stochastic': stochastic rate
        Default: 'deterministic'
    t_start: quantities.Quantity, optional
        Starting time of the series. If specified, it must be lower than
        t_stop
        Default: 0 * ms
    min_delay: quantities.Quantity, optional
        Minimum delay between consecutive coincidence times.
        Default: 0 * ms
    return_coinc: bool, optional
        Whether to return the coincidence times for the SIP process
        Default: False


    Returns
    --------
    output: list
        Realization of a SIP consisting of n Poisson processes characterized
        by synchronous events (with the given jitter)
        If `return_coinc` is `True`, the coincidence times are returned as a
        second output argument. They also have an associated time unit (same
        as `t_stop`).

    References
    ----------
    [1] Kuhn, Aertsen, Rotter (2003) Neural Comput 15(1):67-101

    EXAMPLE:

    >>> import quantities as qt
    >>> import jelephant.core.stocmod as sm
    >>> sip, coinc = sm.sip_poisson(n=10, n=0, t_stop=1*qt.sec, \
            rate=20*qt.Hz,  rate_c=4, return_coinc = True)

    *************************************************************************
    """

    # Check if n is a positive integer
    if not (isinstance(n, int) and n > 0):
        raise ValueError('n (=%s) must be a positive integer' % str(n))

    # Assign time unit to jitter, or check that its existing unit is a time
    # unit
    jitter = abs(jitter)

    # Define the array of rates from input argument rate. Check that its length
    # matches with n
    if rate.ndim == 0:
        if rate < 0 * Hz:
            raise ValueError(
                'rate (=%s) must be non-negative.' % str(rate))
        rates_b = np.array(
            [rate.magnitude for _ in range(n)]) * rate.units
    else:
        rates_b = np.array(rate).flatten() * rate.units
        if not all(rates_b >= 0. * Hz):
            raise ValueError('*rate* must have non-negative elements')

    # Check: rate>=rate_c
    if np.any(rates_b < rate_c):
        raise ValueError('all elements of *rate* must be >= *rate_c*')

    # Check min_delay < 1./rate_c
    if not (rate_c == 0 * Hz or min_delay < 1. / rate_c):
        raise ValueError(
            "'*min_delay* (%s) must be lower than 1/*rate_c* (%s)." %
            (str(min_delay), str((1. / rate_c).rescale(min_delay.units))))

    # Generate the n Poisson processes there are the basis for the SIP
    # (coincidences still lacking)
    embedded_poisson_trains = _n_poisson(
        rate=rates_b - rate_c, t_stop=t_stop, t_start=t_start)
    # Convert the trains from neo SpikeTrain objects to simpler Quantity
    # objects
    embedded_poisson_trains = [
        emb.view(Quantity) for emb in embedded_poisson_trains]

    # Generate the array of times for coincident events in SIP, not closer than
    # min_delay. The array is generated as a quantity from the Quantity class
    # in the quantities module
    if coincidences == 'deterministic':
        Nr_coinc = int(((t_stop - t_start) * rate_c).rescale(dimensionless))
        while True:
            coinc_times = t_start + \
                np.sort(np.random.random(Nr_coinc)) * (t_stop - t_start)
            if len(coinc_times) < 2 or min(np.diff(coinc_times)) >= min_delay:
                break
    elif coincidences == 'stochastic':
        while True:
            coinc_times = homogeneous_poisson_process(
                rate=rate_c, t_stop=t_stop, t_start=t_start)
            if len(coinc_times) < 2 or min(np.diff(coinc_times)) >= min_delay:
                break
        # Convert coinc_times from a neo SpikeTrain object to a Quantity object
        # pq.Quantity(coinc_times.base)*coinc_times.units
        coinc_times = coinc_times.view(Quantity)
        # Set the coincidence times to T-jitter if larger. This ensures that
        # the last jittered spike time is <T
        for i in range(len(coinc_times)):
            if coinc_times[i] > t_stop - jitter:
                coinc_times[i] = t_stop - jitter

    # Replicate coinc_times n times, and jitter each event in each array by
    # +/- jitter (within (t_start, t_stop))
    embedded_coinc = coinc_times + \
        np.random.random(
            (len(rates_b), len(coinc_times))) * 2 * jitter - jitter
    embedded_coinc = embedded_coinc + \
        (t_start - embedded_coinc) * (embedded_coinc < t_start) - \
        (t_stop - embedded_coinc) * (embedded_coinc > t_stop)

    # Inject coincident events into the n SIP processes generated above, and
    # merge with the n independent processes
    sip_process = [
        np.sort(np.concatenate((
            embedded_poisson_trains[m].rescale(t_stop.units),
            embedded_coinc[m].rescale(t_stop.units))) * t_stop.units)
        for m in range(len(rates_b))]

    # Convert back sip_process and coinc_times from Quantity objects to
    # neo.SpikeTrain objects
    sip_process = [
        SpikeTrain(t, t_start=t_start, t_stop=t_stop).rescale(t_stop.units)
        for t in sip_process]
    coinc_times = [
        SpikeTrain(t, t_start=t_start, t_stop=t_stop).rescale(t_stop.units)
        for t in embedded_coinc]

    # Return the processes in the specified output_format
    if not return_coinc:
        output = sip_process
    else:
        output = sip_process, coinc_times

    return output


def _pool_two_spiketrains(a, b, extremes='inner'):
    """
    Pool the spikes of two spike trains a and b into a unique spike train.

    Parameters
    ----------
    a, b : neo.SpikeTrains
        Spike trains to be pooled

    extremes: str, optional
        Only spikes of a and b in the specified extremes are considered.
        * 'inner': pool all spikes from max(a.tstart_ b.t_start) to
           min(a.t_stop, b.t_stop)
        * 'outer': pool all spikes from min(a.tstart_ b.t_start) to
           max(a.t_stop, b.t_stop)
        Default: 'inner'

    Output
    ------
    neo.SpikeTrain containing all spikes in a and b falling in the
    specified extremes
    """

    unit = a.units
    times_a_dimless = list(a.view(Quantity).magnitude)
    times_b_dimless = list(b.rescale(unit).view(Quantity).magnitude)
    times = (times_a_dimless + times_b_dimless) * unit

    if extremes == 'outer':
        t_start = min(a.t_start, b.t_start)
        t_stop = max(a.t_stop, b.t_stop)
    elif extremes == 'inner':
        t_start = max(a.t_start, b.t_start)
        t_stop = min(a.t_stop, b.t_stop)
        times = times[times > t_start]
        times = times[times < t_stop]

    else:
        raise ValueError(
            'extremes (%s) can only be "inner" or "outer"' % extremes)
    pooled_train = SpikeTrain(