from brian import *
from nose.tools import *
from operator import itemgetter
from brian.utils.approximatecomparisons import is_approx_equal

def test():
    """
    :class:`Threshold`
    ~~~~~~~~~~~~~~~~~~
    
    Initialised as ``Threshold(threshold[,state=0])``
    
    Causes a spike whenever the given state variable is above
    the threshold value.
    
    :class:`NoThreshold`
    ~~~~~~~~~~~~~~~~~~~~
    
    Does nothing, initialised as ``NoThreshold()``
    
    Functional thresholds
    ~~~~~~~~~~~~~~~~~~~~~
    
    Initialised as::
    
        FunThreshold(thresholdfun)
        SimpleFunThreshold(thresholdfun[,state=0])
    
    Threshold functions return a boolean array the same size as the
    number of neurons in the group, where if the returned array is
    True at index i then neuron i fires.
    
    The arguments passed to the :class:`FunThreshold` function are the
    full array of state variables for the group in order.
    
    The argument passed to the :class:`SimpleFunThreshold` function is the
    array of length N corresponding to the given state variable.
    
    :class:`VariableThreshold`
    ~~~~~~~~~~~~~~~~~~~~~~~~~~
    
    Initialised as ``VariableThreshold(threshold_state[,state=0])``
    
    Causes a spike whenever the state variable defined by state
    is above the state variable defined by threshold_state.
    
    :class:`EmpiricalThreshold`
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~
    
    Initialised as::
    
        EmpiricalThreshold(threshold[,refractory=1*msecond[,state=0[,clock]]])
    
    Causes a spike when the given state variable exceeds the threshold value,
    but only if there has been no spike within the refractory period. Will
    use the given clock if specified, otherwise the standard guessing procedure
    is used.
    
    Poisson thresholds
    ~~~~~~~~~~~~~~~~~~
    
    Initialised as::
    
        PoissonThreshold([state=0])
        HomogeneousPoissonThreshold([state=0])
    
    The Poisson process gets the rates from the specified state variable, the
    homogeneous version uses the rates from the specified variable of the first
    neuron in the group.      
    """
    reinit_default_clock()

    # test that Threshold works as expected with default state
    G = NeuronGroup(3, model=LazyStateUpdater(), reset=Reset(0.), threshold=Threshold(1.), init=(0.,))
    M = SpikeMonitor(G, True)
    net = Network(G, M)
    net.run(1 * msecond)
    assert (len(M.spikes) == 0)
    G.state(0)[:] = array([0.5, 1.5, 2.5])
    net.run(1 * msecond)
    i, t = zip(*sorted(M.spikes, key=itemgetter(0)))
    assert (i == (1, 2))
    for s in t: assert (is_approx_equal(s, 1 * msecond))

    # test that Threshold works as expected with specified state
    G = NeuronGroup(3, model=LazyStateUpdater(numstatevariables=2), reset=Reset(0., state=1), threshold=Threshold(1., state=1), init=(0., 0.))
    M = SpikeMonitor(G, True)
    net = Network(G, M)
    net.run(1 * msecond)
    assert (len(M.spikes) == 0)
    net.reinit()
    G.state(0)[:] = array([0.5, 1.5, 2.5])
    net.run(1 * msecond)
    assert (len(M.spikes) == 0)
    net.reinit()
    G.state(1)[:] = array([0.5, 1.5, 2.5])
    net.run(1 * msecond)
    i, t = zip(*sorted(M.spikes, key=itemgetter(0)))
    assert (i == (1, 2))
    for s in t: assert (is_approx_equal(s, 0 * msecond))

    # test that VariableThreshold works as expected
    G = NeuronGroup(3, model=LazyStateUpdater(numstatevariables=3), reset=Reset(0., state=1), threshold=VariableThreshold(2, state=1), init=(0., 0., 0.))
    M = SpikeMonitor(G, True)
    net = Network(G, M)
    get_default_clock().reinit()
    G.state(2)[:] = array([1., 2., 3.]) # the thresholds
    G.state(1)[:] = array([4., 1., 2.]) # the values
    net.run(1 * msecond)
    i, t = zip(*sorted(M.spikes, key=itemgetter(0)))
    assert (i == (0,))
    assert (is_approx_equal(t[0], 0 * second))

    # test that FunThreshold works as expected
    def f(S0, S1):
        return S0 > S1 * S1
    G = NeuronGroup(3, model=LazyStateUpdater(numstatevariables=2), reset=Reset(0.), threshold=FunThreshold(f), init=(0., 0.))
    G.state(0)[:] = array([2., 3., 10.])
    G.state(1)[:] = array([1., 2., 3.]) # the square root of the threshold values
    M = SpikeMonitor(G, True)
    net = Network(G, M)
    get_default_clock().reinit()
    net.run(1 * msecond)
    i, t = zip(*sorted(M.spikes, key=itemgetter(0)))
    assert (i == (0, 2))
    for s in t: assert (is_approx_equal(s, 0 * msecond))

    # test that SimpleFunThreshold works as expected
    def f(S):
        return S > 1.
    G = NeuronGroup(3, model=LazyStateUpdater(), reset=Reset(0.), threshold=SimpleFunThreshold(f), init=(0.,))
    G.state(0)[:] = array([0.5, 1.5, 2.5])
    M = SpikeMonitor(G, True)
    net = Network(G, M)
    get_default_clock().reinit()
    net.run(1 * msecond)
    i, t = zip(*sorted(M.spikes, key=itemgetter(0)))
    assert (i == (1, 2))
    for s in t: assert (is_approx_equal(s, 0 * msecond))

    # test that EmpiricalThreshold works as expected
    G = NeuronGroup(1, model=LazyStateUpdater(numstatevariables=2), reset=NoReset(), threshold=EmpiricalThreshold(1., refractory=0.5 * msecond, state=1), init=(0., 2.))
    M = SpikeMonitor(G, True)
    net = Network(G, M)
    get_default_clock().reinit()
    net.run(1.6 * msecond)
    i, t = zip(*sorted(M.spikes, key=itemgetter(1)))
    assert (i == (0, 0, 0, 0))
    for i, s in enumerate(t): assert (is_approx_equal(s, i * 0.5 * msecond))

    # test that PoissonThreshold works
    init = float(1. / get_default_clock().dt) # should cause spiking at every time interval
    G = NeuronGroup(3, model=LazyStateUpdater(), reset=NoReset(), threshold=PoissonThreshold())
    G.state(0)[:] = array([0., init, 0.])
    M = SpikeMonitor(G, True)
    net = Network(G, M)
    net.run(1 * msecond)
    assert (len(M.spikes))
    i, t = zip(*sorted(M.spikes, key=itemgetter(1)))
    assert (all(j == 1 for j in i))

    # test that HomogeneousPoissonThreshold works
    init = float(1. / get_default_clock().dt) # should cause spiking at every time interval
    G = NeuronGroup(3, model=LazyStateUpdater(), reset=NoReset(), threshold=HomogeneousPoissonThreshold())
    M = SpikeMonitor(G, True)
    net = Network(G, M)
    G.state(0)[:] = array([0., init, 0.]) # should do nothing, because only first neuron is looked at 
    net.run(1 * msecond)
    assert (len(M.spikes) == 0)
    G.state(0)[:] = array([init, 0., 0.]) # should do nothing, because only first neuron is looked at 
    net.run(1 * msecond)
    # we actually cannot make any assertion about the behaviour of this system, other than
    # that it should run correctly    

if __name__ == '__main__':
    test()
