

"""
Tests for `SpikeGeneratorGroup`
"""
import os
import tempfile

import pytest
from numpy.testing import assert_equal, assert_array_equal

from brian2 import *
from brian2.core.network import schedule_propagation_offset
from brian2.devices.device import reinit_and_delete
from brian2.tests.utils import assert_allclose, exc_isinstance
from brian2.utils.logger import catch_logs


@pytest.mark.standalone_compatible
def test_spikegenerator_connected():
    """
    Test that `SpikeGeneratorGroup` connects properly.
    """
    G = NeuronGroup(10, 'v:1')
    mon = StateMonitor(G, 'v', record=True, when='end')
    indices = np.array([3, 2, 1, 1, 4, 5])
    times =   np.array([6, 5, 4, 3, 3, 1]) * ms
    SG = SpikeGeneratorGroup(10, indices, times)
    S = Synapses(SG, G, on_pre='v+=1')
    S.connect(j='i')
    run(7*ms)
    # The following neurons should not receive any spikes
    for idx in [0, 6, 7, 8, 9]:
        assert all(mon[idx].v == 0)
    offset = schedule_propagation_offset()
    # The following neurons should receive a single spike
    for idx, time in zip([2, 3, 4, 5], [5, 6, 3, 1]*ms):
        assert all(mon[idx].v[mon.t<time+offset] == 0)
        assert all(mon[idx].v[mon.t>=time+offset] == 1)
    # This neuron receives two spikes
    assert all(mon[1].v[mon.t<3*ms+offset] == 0)
    assert all(mon[1].v[(mon.t>=3*ms+offset) & (mon.t<4*ms+offset)] == 1)
    assert all(mon[1].v[(mon.t>=4*ms+offset)] == 2)

@pytest.mark.standalone_compatible
def test_spikegenerator_basic():
    """
    Basic test for `SpikeGeneratorGroup`.
    """
    indices = np.array([3, 2, 1, 1, 2, 3, 3, 2, 1])
    times   = np.array([1, 4, 4, 3, 2, 4, 2, 3, 2]) * ms
    SG = SpikeGeneratorGroup(5, indices, times)
    s_mon = SpikeMonitor(SG)
    run(5*ms)
    _compare_spikes(5, indices, times, s_mon)


@pytest.mark.standalone_compatible
def test_spikegenerator_basic_sorted():
    """
    Basic test for `SpikeGeneratorGroup` with already sorted spike events.
    """
    indices = np.array([3, 1, 2, 3, 1, 2, 1, 2, 3])
    times   = np.array([1, 2, 2, 2, 3, 3, 4, 4, 4]) * ms
    SG = SpikeGeneratorGroup(5, indices, times)
    s_mon = SpikeMonitor(SG)
    run(5*ms)
    _compare_spikes(5, indices, times, s_mon)


@pytest.mark.standalone_compatible
def test_spikegenerator_basic_sorted_with_sorted():
    """
    Basic test for `SpikeGeneratorGroup` with already sorted spike events.
    """
    indices = np.array([3, 1, 2, 3, 1, 2, 1, 2, 3])
    times   = np.array([1, 2, 2, 2, 3, 3, 4, 4, 4]) * ms
    SG = SpikeGeneratorGroup(5, indices, times, sorted=True)
    s_mon = SpikeMonitor(SG)
    run(5*ms)
    _compare_spikes(5, indices, times, s_mon)


@pytest.mark.standalone_compatible
def test_spikegenerator_period():
    """
    Basic test for `SpikeGeneratorGroup`.
    """
    indices = np.array([3, 2, 1, 1, 2, 3, 3, 2, 1])
    times   = np.array([1, 4, 4, 3, 2, 4, 2, 3, 2]) * ms
    SG = SpikeGeneratorGroup(5, indices, times, period=5*ms)

    s_mon = SpikeMonitor(SG)
    run(10*ms)
    for idx in range(5):
        generator_spikes = sorted([(idx, time) for time in times[indices==idx]] + [(idx, time+5*ms) for time in times[indices==idx]])
        recorded_spikes = sorted([(idx, time) for time in s_mon.t[s_mon.i==idx]])
        assert_allclose(generator_spikes, recorded_spikes)


@pytest.mark.codegen_independent
def test_spikegenerator_extreme_period():
    """
    Basic test for `SpikeGeneratorGroup`.
    """
    indices = np.array([0, 1, 2])
    times   = np.array([0, 1, 2]) * ms
    SG = SpikeGeneratorGroup(5, indices, times, period=1e6*second)
    s_mon = SpikeMonitor(SG)
    with catch_logs() as l:
        run(10*ms)

    assert_equal(s_mon.i, np.array([0, 1, 2]))
    assert_allclose(s_mon.t, [0, 1, 2]*ms)
    assert len(l) == 1 and l[0][1].endswith('spikegenerator_long_period')

@pytest.mark.standalone_compatible
def test_spikegenerator_period_rounding():
    # See discussion in PR #1042
    # The last spike will be considered to be in the time step *after* 1s, due
    # to the way our rounding works. Although probably not what the user
    # expects, this should therefore raise an error. In previous versions of
    # Brian, this did not raise any error but silently discarded the spike.
    with pytest.raises(ValueError):
        SpikeGeneratorGroup(1, [0, 0, 0], [0*ms, .9*ms, .99999*ms],
                            period=1*ms, dt=0.1*ms)
    # This should also raise a ValueError, since the last two spikes fall into
    # the same bin
    s = SpikeGeneratorGroup(1, [0, 0, 0], [0*ms, .9*ms, .96*ms],
                            period=1*ms, dt=0.1*ms)
    net = Network(s)
    with pytest.raises(BrianObjectException) as exc:
        net.run(0*ms)
    assert exc_isinstance(exc, ValueError)


def test_spikegenerator_period_repeat():
    """
    Basic test for `SpikeGeneratorGroup`.
    """
    indices = np.zeros(10)
    times   = arange(0, 1, 0.1) * ms

    rec = np.rec.fromarrays([times, indices], names=['t', 'i'])
    rec.sort()
    sorted_times = np.ascontiguousarray(rec.t)*1000
    sorted_indices = np.ascontiguousarray(rec.i)
    SG = SpikeGeneratorGroup(1, indices, times, period=1*ms)
    s_mon = SpikeMonitor(SG)
    net   = Network(SG, s_mon)
    rate  = PopulationRateMonitor(SG)
    for idx in range(5):
        net.run(1*ms)
        assert (idx+1)*len(SG.spike_time) == s_mon.num_spikes


def _compare_spikes(N, indices, times, recorded, start_time=0*ms, end_time=1e100*second):
    for idx in range(N):
        generator_spikes = sorted([(idx, time) for time in times[indices==idx]])
        recorded_spikes = sorted([(idx, time) for time in recorded.t[recorded.i==idx] if time >= start_time and time < end_time])
        assert_allclose(generator_spikes, recorded_spikes)

@pytest.mark.standalone_compatible
@pytest.mark.multiple_runs
def test_spikegenerator_change_spikes():
    indices1 = np.array([3, 2, 1, 1, 2, 3, 3, 2, 1])
    times1   = np.array([1, 4, 4, 3, 2, 4, 2, 3, 2]) * ms
    SG = SpikeGeneratorGroup(5, indices1, times1)
    s_mon = SpikeMonitor(SG)
    net = Network(SG, s_mon)
    net.run(5*ms)

    indices2 = np.array([3, 2, 1, 1, 2, 3, 3, 2, 1, 3,   3,   3,   1,   2])
    times2   = np.array([1, 4, 4, 3, 2, 4, 2, 3, 2, 4.5, 4.7, 4.8, 4.5, 4.7])*ms + 5*ms

    SG.set_spikes(indices2, times2)
    net.run(5*ms)

    indices3 = np.array([4, 1, 0])
    times3   = np.array([1, 3, 4])*ms + 10*ms

    SG.set_spikes(indices3, times3)
    net.run(5*ms)
    device.build(direct_call=False, **device.build_options)
    _compare_spikes(5, indices1, times1, s_mon, 0*ms, 5*ms)
    _compare_spikes(5, indices2, times2, s_mon, 5*ms, 10*ms)
    _compare_spikes(5, indices3, times3, s_mon, 10*ms)

@pytest.mark.standalone_compatible
@pytest.mark.multiple_runs
def test_spikegenerator_change_period():
    """
    Basic test for `SpikeGeneratorGroup`.
    """
    indices1 = np.array([3, 2, 1, 1, 2, 3, 3, 2, 1])
    times1   = np.array([1, 4, 4, 3, 2, 4, 2, 3, 2]) * ms
    SG = SpikeGeneratorGroup(5, indices1, times1, period=5*ms)
    s_mon = SpikeMonitor(SG)
    net = Network(SG, s_mon)
    net.run(10*ms)

    indices2 = np.array([3, 2, 1, 1, 2, 3, 3, 2, 1, 3,   3,   3,   1,   2])
    times2   = np.array([1, 4, 4, 3, 2, 4, 2, 3, 2, 4.5, 4.7, 4.8, 4.5, 4.7])*ms + 10*ms

    SG.set_spikes(indices2, times2)
    net.run(10*ms)  # period should no longer be in effect
    device.build(direct_call=False, **device.build_options)

    _compare_spikes(5, np.hstack([indices1, indices1]),
                    np.hstack([times1, times1+5*ms]),
                    s_mon, 0*ms, 10*ms)
    _compare_spikes(5, indices2, times2, s_mon, 10*ms)

@pytest.mark.codegen_independent
def test_spikegenerator_incorrect_values():
    with pytest.raises(TypeError):
        SpikeGeneratorGroup(0, [], []*second)
    # Floating point value for N
    with pytest.raises(TypeError):
        SpikeGeneratorGroup(1.5, [], []*second)
    # Negative index
    with pytest.raises(ValueError):
        SpikeGeneratorGroup(5, [0, 3, -1], [0, 1, 2]*ms)
    # Too high index
    with pytest.raises(ValueError):
        SpikeGeneratorGroup(5, [0, 5, 1], [0, 1, 2]*ms)
    # Negative time
    with pytest.raises(ValueError):
        SpikeGeneratorGroup(5, [0, 1, 2], [0, -1, 2]*ms)

@pytest.mark.codegen_independent
def test_spikegenerator_incorrect_period():
    """
    Test that you cannot provide incorrect period arguments or combine
    inconsistent period and dt arguments.
    """
    # Period is negative
    with pytest.raises(ValueError):
        SpikeGeneratorGroup(1, [], []*second, period=-1*ms)

    # Period is smaller than the highest spike time
    with pytest.raises(ValueError):
        SpikeGeneratorGroup(1, [0], [2]*ms, period=1*ms)
    # Period is not an integer multiple of dt
    SG = SpikeGeneratorGroup(1, [], []*second, period=1.25*ms, dt=0.1*ms)
    net = Network(SG)
    with pytest.raises(BrianObjectException) as exc:
        net.run(0*ms)
    assert exc_isinstance(exc, NotImplementedError)

    SG = SpikeGeneratorGroup(1, [], [] * second, period=0.101 * ms, dt=0.1 * ms)
    net = Network(SG)
    with pytest.raises(BrianObjectException) as exc:
        net.run(0*ms)
    assert exc_isinstance(exc, NotImplementedError)

    SG = SpikeGeneratorGroup(1, [], [] * second, period=3.333 * ms, dt=0.1 * ms)
    net = Network(SG)
    with pytest.raises(BrianObjectException) as exc:
        net.run(0*ms)
    assert exc_isinstance(exc, NotImplementedError)

    # This should not raise an error (see #1041)
    SG = SpikeGeneratorGroup(1, [], []*ms, period=150*ms, dt=0.1*ms)
    net = Network(SG)
    net.run(0*ms)

    # Period is smaller than dt
    SG = SpikeGeneratorGroup(1, [], []*second, period=1*ms, dt=2*ms)
    net = Network(SG)
    with pytest.raises(BrianObjectException) as exc:
        net.run(0*ms)
    assert exc_isinstance(exc, ValueError)


def test_spikegenerator_rounding():
    # all spikes should fall into the first time bin
    indices = np.arange(100)
    times = np.linspace(0, 0.1, 100, endpoint=False)*ms
    SG = SpikeGeneratorGroup(100, indices, times, dt=0.1*ms)
    mon = SpikeMonitor(SG)
    net = Network(SG, mon)
    net.run(0.1*ms)
    assert_equal(mon.count, np.ones(100))

    # all spikes should fall in separate bins
    dt = 0.1*ms
    indices = np.zeros(10000)
    times = np.arange(10000)*dt
    SG = SpikeGeneratorGroup(1, indices, times, dt=dt)
    target = NeuronGroup(1, 'count : 1',
                         threshold='True', reset='count=0')  # set count to zero at every time step
    syn = Synapses(SG, target, on_pre='count+=1')
    syn.connect()
    mon = StateMonitor(target, 'count', record=0, when='end')
    net = Network(SG, target, syn, mon)
    # change the schedule so that resets are processed before synapses
    net.schedule = ['start', 'groups', 'thresholds', 'resets', 'synapses', 'end']
    net.run(10000*dt)
    assert_equal(mon[0].count, np.ones(10000))


@pytest.mark.standalone_compatible
@pytest.mark.long
def test_spikegenerator_rounding_long():
    # all spikes should fall in separate bins
    dt = 0.1*ms
    N = 1000000
    indices = np.zeros(N)
    times = np.arange(N)*dt
    SG = SpikeGeneratorGroup(1, indices, times, dt=dt)
    target = NeuronGroup(1, 'count : 1')
    syn = Synapses(SG, target, on_pre='count+=1')
    syn.connect()
    spikes = SpikeMonitor(SG)
    mon = StateMonitor(target, 'count', record=0, when='end')
    run(N*dt, report='text')
    assert spikes.count[0] == N, f'expected {int(N)} spikes, got {int(spikes.count[0])}'
    assert all(np.diff(mon[0].count[:]) == 1)


@pytest.mark.standalone_compatible
@pytest.mark.long
def test_spikegenerator_rounding_period():
    # all spikes should fall in separate bins
    dt = 0.1*ms
    N = 100
    repeats = 10000
    indices = np.zeros(N)
    times = np.arange(N)*dt
    SG = SpikeGeneratorGroup(1, indices, times, dt=dt, period=N*dt)
    target = NeuronGroup(1, 'count : 1')
    syn = Synapses(SG, target, on_pre='count+=1')
    syn.connect()
    spikes = SpikeMonitor(SG)
    mon = StateMonitor(target, 'count', record=0, when='end')
    run(N*repeats*dt, report='text')
    #print np.int_(np.round(spikes.t/dt))
    assert_equal(spikes.count[0], N*repeats)
    assert all(np.diff(mon[0].count[:]) == 1)


@pytest.mark.codegen_independent
def test_spikegenerator_multiple_spikes_per_bin():
    # Multiple spikes per bin are of course fine if they don't belong to the
    # same neuron
    SG = SpikeGeneratorGroup(2, [0, 1], [0, 0.05]*ms, dt=0.1*ms)
    net = Network(SG)
    net.run(0*ms)

    # This should raise an error
    SG = SpikeGeneratorGroup(2, [0, 0], [0, 0.05]*ms, dt=0.1*ms)
    net = Network(SG)
    with pytest.raises(BrianObjectException) as exc:
        net.run(0*ms)
    print(exc.value.__cause__)
    assert exc_isinstance(exc, ValueError)

    # More complicated scenario where dt changes between runs
    defaultclock.dt = 0.1*ms
    SG = SpikeGeneratorGroup(2, [0, 0], [0.05, 0.15]*ms)
    net = Network(SG)
    net.run(0*ms)  # all is fine
    defaultclock.dt = 0.2*ms  # Now the two spikes fall into the same bin
    with pytest.raises(BrianObjectException) as exc:
        net.run(0*ms)
    assert exc_isinstance(exc, ValueError)


@pytest.mark.standalone_compatible
@pytest.mark.multiple_runs
def test_spikegenerator_multiple_runs():
    indices = np.zeros(5)
    times = np.arange(5)*ms
    spike_gen = SpikeGeneratorGroup(1, indices, times)  # all good
    spike_mon = SpikeMonitor(spike_gen)
    run(5*ms)
    # Setting the same spike times again should not do anything, since they are
    # before the start of the current simulation
    spike_gen.set_spikes(indices, times)
    # however, a warning should be raised
    with catch_logs() as l:
        run(5*ms)
        device.build(direct_call=False, **device.build_options)
    assert len(l) == 1 and l[0][1].endswith('ignored_spikes')
    assert spike_mon.num_spikes == 5


def test_spikegenerator_restore():
    # Check whether SpikeGeneratorGroup works with store/restore
    # See github issue #1084
    gen = SpikeGeneratorGroup(1, [0, 0, 0], [0, 1, 2]*ms)
    mon = SpikeMonitor(gen)
    store()
    run(3*ms)
    assert_array_equal(mon.i, [0, 0, 0])
    assert_allclose(mon.t, [0, 1, 2] * ms)
    restore()
    run(3 * ms)
    assert_array_equal(mon.i, [0, 0, 0])
    assert_allclose(mon.t, [0, 1, 2]*ms)


if __name__ == '__main__':
    import time
    start = time.time()

    test_spikegenerator_connected()
    test_spikegenerator_basic()
    test_spikegenerator_basic_sorted()
    test_spikegenerator_basic_sorted_with_sorted()
    test_spikegenerator_period()
    test_spikegenerator_period_rounding()
    test_spikegenerator_extreme_period()
    test_spikegenerator_period_repeat()
    test_spikegenerator_change_spikes()
    test_spikegenerator_change_period()
    test_spikegenerator_incorrect_values()
    test_spikegenerator_incorrect_period()
    test_spikegenerator_rounding()
    test_spikegenerator_rounding_long()
    test_spikegenerator_rounding_period()
    test_spikegenerator_multiple_spikes_per_bin()
    test_spikegenerator_multiple_runs()
    test_spikegenerator_restore()
    print('Tests took', time.time()-start)
