
import pytest

from brian2 import *
from brian2.devices.device import reinit_and_delete
from brian2.tests.utils import assert_allclose
from brian2.utils.caching import _hashable


@pytest.mark.codegen_independent
def test_timedarray_direct_use():
    ta = TimedArray(np.linspace(0, 10, 11), 1*ms)
    assert ta(-1*ms) == 0
    assert ta(5*ms) == 5
    assert ta(10*ms) == 10
    assert ta(15*ms) == 10
    ta = TimedArray(np.linspace(0, 10, 11)*amp, 1*ms)
    assert ta(-1*ms) == 0*amp
    assert ta(5*ms) == 5*amp
    assert ta(10*ms) == 10*amp
    assert ta(15*ms) == 10*amp
    ta2d = TimedArray((np.linspace(0, 11, 12)*amp).reshape(4, 3), 1*ms)
    assert ta2d(-1*ms, 0) == 0*amp
    assert ta2d(0*ms, 0) == 0*amp
    assert ta2d(0*ms, 1) == 1*amp
    assert ta2d(1*ms, 1) == 4*amp
    assert_allclose(ta2d(1*ms, [0, 1, 2]), [3, 4, 5]*amp)
    assert_allclose(ta2d(15*ms, [0, 1, 2]), [9, 10, 11]*amp)


@pytest.mark.standalone_compatible
def test_timedarray_semantics():
    # Make sure that timed arrays are interpreted as specifying the values
    # between t and t+dt (not between t-dt/2 and t+dt/2 as in Brian1)
    ta = TimedArray(array([0, 1]), dt=0.4*ms)
    G = NeuronGroup(1, 'value = ta(t) : 1', dt=0.1*ms)
    mon = StateMonitor(G, 'value', record=0)
    run(0.8*ms)
    assert_allclose(mon[0].value, [0, 0, 0, 0, 1, 1, 1, 1])
    assert_allclose(mon[0].value, ta(mon.t))

@pytest.mark.standalone_compatible
def test_timedarray_no_units():
    ta = TimedArray(np.arange(10), dt=0.1*ms)
    G = NeuronGroup(1, 'value = ta(t) + 1: 1', dt=0.1*ms)
    mon = StateMonitor(G, 'value', record=True, dt=0.1*ms)
    run(1.1*ms)
    assert_allclose(mon[0].value_, np.clip(np.arange(len(mon[0].t)), 0, 9) + 1)

@pytest.mark.standalone_compatible
def test_timedarray_with_units():
    ta = TimedArray(np.arange(10)*amp, dt=0.1*ms)
    G = NeuronGroup(1, 'value = ta(t) + 2*nA: amp', dt=0.1*ms)
    mon = StateMonitor(G, 'value', record=True, dt=0.1*ms)
    run(1.1*ms)
    assert_allclose(mon[0].value, np.clip(np.arange(len(mon[0].t)), 0, 9)*amp + 2*nA)

@pytest.mark.standalone_compatible
def test_timedarray_2d():
    # 4 time steps, 3 neurons
    ta2d = TimedArray(np.arange(12).reshape(4, 3), dt=0.1*ms)
    G = NeuronGroup(3, 'value = ta2d(t, i) + 1: 1', dt=0.1*ms)
    mon = StateMonitor(G, 'value', record=True, dt=0.1*ms)
    run(0.5*ms)
    assert_allclose(mon[0].value_, np.array([0, 3, 6, 9, 9]) + 1)
    assert_allclose(mon[1].value_, np.array([1, 4, 7, 10, 10]) + 1)
    assert_allclose(mon[2].value_, np.array([2, 5, 8, 11, 11]) + 1)


@pytest.mark.codegen_independent
def test_timedarray_incorrect_use():
    ta = TimedArray(np.linspace(0, 10, 11), 1*ms)
    ta2d = TimedArray((np.linspace(0, 11, 12)*amp).reshape(4, 3), 1*ms)
    G = NeuronGroup(3, 'I : amp')
    with pytest.raises(ValueError):
        setattr(G, 'I', 'ta2d(t)*amp')
    with pytest.raises(ValueError):
        setattr(G, 'I', 'ta(t, i)*amp')
    with pytest.raises(ValueError):
        setattr(G, 'I', 'ta()*amp')
    with pytest.raises(ValueError):
        setattr(G, 'I', 'ta*amp')


@pytest.mark.standalone_compatible
def test_timedarray_no_upsampling():
    # Test a TimedArray where no upsampling is necessary because the monitor's
    # dt is bigger than the TimedArray's
    ta = TimedArray(np.arange(10), dt=0.01*ms)
    G = NeuronGroup(1, 'value = ta(t): 1', dt=0.1*ms)
    mon = StateMonitor(G, 'value', record=True, dt=1*ms)
    run(2.1*ms)
    assert_allclose(mon[0].value, [0, 9, 9])


#@pytest.mark.standalone_compatible  # see FIXME comment below
def test_long_timedarray():
    """
    Use a very long timedarray (with a big dt), where the upsampling can lead
    to integer overflow.
    """
    ta = TimedArray(np.arange(16385), dt=1*second)
    G = NeuronGroup(1, 'value = ta(t) : 1')
    mon = StateMonitor(G, 'value', record=True)
    net = Network(G, mon)
    # We'll start the simulation close to the critical boundary
    # FIXME: setting the time like this does not work for standalone
    net.t_ = float(16384*second - 5*ms)
    net.run(10*ms)
    assert_allclose(mon[0].value[mon.t < 16384*second], 16383)
    assert_allclose(mon[0].value[mon.t >= 16384*second], 16384)


def test_timedarray_repeated_use():
    # Check that recreating a TimedArray with different values does work
    # correctly (no issues with caching)
    values = np.array([[1, 2, 3], [2, 4, 6]])
    for run_idx in range(2):
        ta = TimedArray(values[run_idx], dt=defaultclock.dt, name='ta')
        G = NeuronGroup(1, 'dx/dt = ta(t)/dt : 1', name='G')
        run(3*defaultclock.dt)
        assert G.x[0] == 6 * (run_idx + 1)


if __name__ == '__main__':
    test_timedarray_direct_use()
    test_timedarray_semantics()
    test_timedarray_no_units()
    test_timedarray_with_units()
    test_timedarray_2d()
    test_timedarray_incorrect_use()
    test_timedarray_no_upsampling()
    test_long_timedarray()
    test_timedarray_repeated_use()
