from collections import Counter

import pytest
from numpy.testing import assert_equal

from brian2.core.functions import timestep
from brian2.utils.logger import catch_logs
from brian2 import *
from brian2.equations.refractory import add_refractoriness
from brian2.devices.device import reinit_and_delete
from brian2.tests.utils import assert_allclose, exc_isinstance


@pytest.mark.codegen_independent
def test_add_refractoriness():
    eqs = Equations("""
    dv/dt = -x*v/second : volt (unless refractory)
    dw/dt = -w/second : amp
    x : 1
    """)
    # only make sure it does not throw an error
    eqs = add_refractoriness(eqs)
    # Check that the parameters were added
    assert 'not_refractory' in eqs
    assert 'lastspike' in eqs


@pytest.mark.codegen_independent
def test_missing_refractory_warning():
    # Forgotten refractory argument
    with catch_logs() as l:
        group = NeuronGroup(1, 'dv/dt = -v / (10*ms) : 1 (unless refractory)',
                            threshold='v > 1', reset='v = 0')
    assert len(l) == 1
    assert l[0][0] == 'WARNING' and l[0][1].endswith('no_refractory')


@pytest.mark.standalone_compatible
def test_refractoriness_basic():
    G = NeuronGroup(1, """
                       dv/dt = 99.999*Hz : 1 (unless refractory)
                       dw/dt = 99.999*Hz : 1
                       """,
                    threshold='v>1', reset='v=0;w=0',
                    refractory=5*ms)
    # It should take 10ms to reach the threshold, then v should stay at 0
    # for 5ms, while w continues to increase
    mon = StateMonitor(G, ['v', 'w'], record=True, when='end')
    run(20*ms)
    # No difference before the spike
    assert_allclose(mon[0].v[:timestep(10*ms, defaultclock.dt)],
                    mon[0].w[:timestep(10*ms, defaultclock.dt)])
    # v is not updated during refractoriness
    in_refractoriness = mon[0].v[timestep(10*ms, defaultclock.dt):timestep(15*ms, defaultclock.dt)]
    assert_equal(in_refractoriness, np.zeros_like(in_refractoriness))
    # w should evolve as before
    assert_allclose(mon[0].w[:timestep(5*ms, defaultclock.dt)],
                    mon[0].w[timestep(10*ms, defaultclock.dt)+1:timestep(15*ms, defaultclock.dt)+1])
    assert np.all(mon[0].w[timestep(10*ms, defaultclock.dt)+1:timestep(15*ms, defaultclock.dt)+1] > 0)
    # After refractoriness, v should increase again
    assert np.all(mon[0].v[timestep(15*ms, defaultclock.dt):timestep(20*ms, defaultclock.dt)] > 0)


@pytest.mark.standalone_compatible
@pytest.mark.parametrize('ref_time', ['5*ms',
                                      '(t-lastspike + 1e-3*dt) < 5*ms',
                                      'time_since_spike + 1e-3*dt < 5*ms',
                                      'ref_subexpression',
                                      '(t-lastspike + 1e-3*dt) < ref',
                                      'ref',
                                      'ref_no_unit*ms'])
def test_refractoriness_variables(ref_time):
    # Try a string evaluating to a quantity, and an explicit boolean
    # condition -- all should do the same thing
    G = NeuronGroup(1, """
                    dv/dt = 99.999*Hz : 1 (unless refractory)
                    dw/dt = 99.999*Hz : 1
                    ref : second
                    ref_no_unit : 1
                    time_since_spike = (t - lastspike) +1e-3*dt : second
                    ref_subexpression = (t - lastspike + 1e-3*dt) < ref : boolean
                    """,
                    threshold='v>1', reset='v=0;w=0',
                    refractory=ref_time,
                    dtype={'ref': defaultclock.variables['t'].dtype,
                           'ref_no_unit': defaultclock.variables['t'].dtype,
                           'lastspike': defaultclock.variables['t'].dtype,
                           'time_since_spike': defaultclock.variables['t'].dtype})
    G.ref = 5*ms
    G.ref_no_unit = 5
    # It should take 10ms to reach the threshold, then v should stay at 0
    # for 5ms, while w continues to increase
    mon = StateMonitor(G, ['v', 'w'], record=True, when='end')
    run(20*ms)
    try:
        # No difference before the spike
        assert_allclose(mon[0].v[:timestep(10*ms, defaultclock.dt)],
                        mon[0].w[:timestep(10*ms, defaultclock.dt)])
        # v is not updated during refractoriness
        in_refractoriness = mon[0].v[timestep(10*ms, defaultclock.dt):timestep(15*ms, defaultclock.dt)]
        assert_allclose(in_refractoriness, np.zeros_like(in_refractoriness))
        # w should evolve as before
        assert_allclose(mon[0].w[:timestep(5*ms, defaultclock.dt)],
                     mon[0].w[timestep(10*ms, defaultclock.dt)+1:timestep(15*ms, defaultclock.dt)+1])
        assert np.all(mon[0].w[timestep(10*ms, defaultclock.dt)+1:timestep(15*ms, defaultclock.dt)+1] > 0)
        # After refractoriness, v should increase again
        assert np.all(mon[0].v[timestep(15*ms, defaultclock.dt):timestep(20*ms, defaultclock.dt)] > 0)
    except AssertionError as ex:
        raise
        raise AssertionError(f'Assertion failed when using {ref_time!r} as refractory argument:\n{ex}')


@pytest.mark.standalone_compatible
def test_refractoriness_threshold_basic():
    G = NeuronGroup(1, """
    dv/dt = 199.99*Hz : 1
    """, threshold='v > 1', reset='v=0', refractory=10*ms)
    # The neuron should spike after 5ms but then not spike for the next
    # 10ms. The state variable should continue to integrate so there should
    # be a spike after 15ms
    spike_mon = SpikeMonitor(G)
    run(16*ms)
    assert_allclose(spike_mon.t, [5, 15] * ms)


@pytest.mark.standalone_compatible
def test_refractoriness_repeated():
    # Create a group that spikes whenever it can
    group = NeuronGroup(1, '', threshold='True', refractory=10*defaultclock.dt)
    spike_mon = SpikeMonitor(group)
    run(10000*defaultclock.dt)
    assert spike_mon.t[0] == 0*ms
    assert_allclose(np.diff(spike_mon.t), 10*defaultclock.dt)


@pytest.mark.standalone_compatible
def test_refractoriness_repeated_legacy():
    if prefs.core.default_float_dtype == np.float32:
        pytest.skip('Not testing legacy refractory mechanism with single '
                       'precision floats.')
    # Switch on behaviour from versions <= 2.1.2
    prefs.legacy.refractory_timing = True
    # Create a group that spikes whenever it can
    group = NeuronGroup(1, '', threshold='True', refractory=10*defaultclock.dt)
    spike_mon = SpikeMonitor(group)
    run(10000*defaultclock.dt)
    assert spike_mon.t[0] == 0*ms

    # Empirical values from running with earlier Brian versions
    assert_allclose(np.diff(spike_mon.t)[:10],
                    [1.1, 1, 1.1, 1, 1.1, 1.1, 1.1, 1.1, 1, 1.1]*ms)
    steps = Counter(np.diff(np.int_(np.round(spike_mon.t / defaultclock.dt))))
    assert len(steps) == 2 and steps[10] == 899 and steps[11] == 91
    prefs.legacy.refractory_timing = False


@pytest.mark.standalone_compatible
@pytest.mark.parametrize('ref_time', [10*ms, '10*ms',
                                      'timestep(t-lastspike, dt) < timestep(10*ms, dt)',
                                      'timestep(t-lastspike, dt) < timestep(ref, dt)',
                                      'ref',
                                      'ref_no_unit*ms'])
def test_refractoriness_threshold(ref_time):
    # Try a quantity, a string evaluating to a quantity, and an explicit boolean
    # condition -- all should do the same thing
    G = NeuronGroup(1, """
                    dv/dt = 199.999*Hz : 1
                    ref : second
                    ref_no_unit : 1
                    """, threshold='v > 1',
                    reset='v=0', refractory=ref_time,
                    dtype={'ref': defaultclock.variables['t'].dtype,
                           'ref_no_unit': defaultclock.variables['t'].dtype})
    G.ref = 10*ms
    G.ref_no_unit = 10
    # The neuron should spike after 5ms but then not spike for the next
    # 10ms. The state variable should continue to integrate so there should
    # be a spike after 15ms
    spike_mon = SpikeMonitor(G)
    run(16*ms)
    assert_allclose(spike_mon.t, [5, 15] * ms)


@pytest.mark.codegen_independent
def test_refractoriness_types():
    # make sure that using a wrong type of refractoriness does not work
    group = NeuronGroup(1, '', refractory='3*Hz')
    with pytest.raises(BrianObjectException) as exc:
        Network(group).run(0*ms)
    assert exc_isinstance(exc, TypeError)
    group = NeuronGroup(1, 'ref: Hz', refractory='ref')
    with pytest.raises(BrianObjectException) as exc:
        Network(group).run(0*ms)
    assert exc_isinstance(exc, TypeError)
    group = NeuronGroup(1, '', refractory='3')
    with pytest.raises(BrianObjectException) as exc:
        Network(group).run(0*ms)
    assert exc_isinstance(exc, TypeError)
    group = NeuronGroup(1, 'ref: 1', refractory='ref')
    with pytest.raises(BrianObjectException) as exc:
        Network(group).run(0*ms)
    assert exc_isinstance(exc, TypeError)

@pytest.mark.codegen_independent
def test_conditional_write_set():
    """
    Test that the conditional_write attribute is set correctly
    """
    G = NeuronGroup(1, """dv/dt = 10*Hz : 1 (unless refractory)
                          dw/dt = 10*Hz : 1""", refractory=2*ms)
    assert G.variables['v'].conditional_write is G.variables['not_refractory']
    assert G.variables['w'].conditional_write is None

@pytest.mark.standalone_compatible
def test_conditional_write_behaviour():
    H = NeuronGroup(1, 'v:1', threshold='v>-1')

    tau = 1*ms
    eqs = """
    dv/dt = (2-v)/tau : 1 (unless refractory)
    dx/dt = 0/tau : 1 (unless refractory)
    dy/dt = 0/tau : 1
    """
    reset = """
    v = 0
    x -= 0.05
    y -= 0.05
    """
    G = NeuronGroup(1, eqs, threshold='v>1', reset=reset, refractory=1*ms)

    Sx = Synapses(H, G, on_pre='x += dt*100*Hz')
    Sx.connect(True)

    Sy = Synapses(H, G, on_pre='y += dt*100*Hz')
    Sy.connect(True)

    M = StateMonitor(G, variables=True, record=True)

    run(10*ms)

    assert G.x[0] < 0.2
    assert G.y[0] > 0.2
    assert G.v[0] < 1.1


@pytest.mark.standalone_compatible
def test_conditional_write_automatic_and_manual():
    source = NeuronGroup(1, '', threshold='True')  # spiking all the time
    target = NeuronGroup(2, """dv/dt = 0/ms : 1 (unless refractory)
                               dw/dt = 0/ms : 1""",
                         threshold='t == 0*ms',
                         refractory='False')  # only refractory for the first time step
    # Cell is spiking/refractory only in the first time step
    syn = Synapses(source, target, on_pre="""v += 1
                                             w += 1 * int(not_refractory_post)""")
    syn.connect()
    mon = StateMonitor(target, ['v', 'w'], record=True, when='end')
    run(2*defaultclock.dt)

    # Synapse should not have been effective in the first time step
    assert_allclose(mon.v[:, 0], 0)
    assert_allclose(mon.v[:, 1], 1)
    assert_allclose(mon.w[:, 0], 0)
    assert_allclose(mon.w[:, 1], 1)


if __name__ == '__main__':
    test_add_refractoriness()
    test_missing_refractory_warning()
    test_refractoriness_basic()
    test_refractoriness_variables()
    test_refractoriness_threshold()
    test_refractoriness_threshold_basic()
    test_refractoriness_repeated()
    test_refractoriness_repeated_legacy()
    test_refractoriness_types()
    test_conditional_write_set()
    test_conditional_write_behaviour()
    test_conditional_write_automatic_and_manual()
