import itertools
import os

import pytest
from numpy.testing import assert_equal

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

try:
    import scipy
except ImportError:
    scipy = None


numpy_needs_scipy = pytest.mark.skipif(
    # Using condition string, since we cannot yet know
    # prefs.codegen.target at module import time
    "prefs.codegen.target == 'numpy' and not scipy",
    reason="multi-compartmental models need scipy to run with numpy",
)


@pytest.mark.codegen_independent
@numpy_needs_scipy
def test_custom_events():
    # Set (could be moved in a setup)
    EL = -65 * mV
    gL = 0.0003 * siemens / cm**2
    ev = """
    Im = gL * (EL - v) : amp/meter**2
    event_time1 : second
    """
    # Create a three compartments morphology
    morpho = Soma(diameter=10 * um)
    morpho.dend1 = Cylinder(n=1, diameter=1 * um, length=10 * um)
    morpho.dend2 = Cylinder(n=1, diameter=1 * um, length=10 * um)
    G = SpatialNeuron(
        morphology=morpho, model=ev, events={"event1": "t>=i*ms and t<i*ms+dt"}
    )
    G.run_on_event("event1", "event_time1 = 0.1*ms")
    run(0.2 * ms)
    # Event has size three now because there are three compartments
    assert_allclose(G.event_time1[:], [0.1, 0, 0] * ms)


@pytest.mark.codegen_independent
def test_construction():
    BrianLogger.suppress_name("resolution_conflict")
    morpho = Soma(diameter=30 * um)
    morpho.L = Cylinder(length=10 * um, diameter=1 * um, n=10)
    morpho.LL = Cylinder(length=5 * um, diameter=2 * um, n=5)
    morpho.LR = Cylinder(length=5 * um, diameter=2 * um, n=10)
    morpho.right = Cylinder(length=3 * um, diameter=1 * um, n=7)
    morpho.right.nextone = Cylinder(length=2 * um, diameter=1 * um, n=3)
    gL = 1e-4 * siemens / cm**2
    EL = -70 * mV
    eqs = """
    Im=gL*(EL-v) : amp/meter**2
    I : meter (point current)
    """

    # Check units of currents
    with pytest.raises(DimensionMismatchError):
        SpatialNeuron(morphology=morpho, model=eqs)

    eqs = """
    Im=gL*(EL-v) : amp/meter**2
    """
    neuron = SpatialNeuron(
        morphology=morpho, model=eqs, Cm=1 * uF / cm**2, Ri=100 * ohm * cm
    )
    # Test initialization of values
    neuron.LL.v = EL
    assert_allclose(neuron.L.main.v, 0 * mV)
    assert_allclose(neuron.LL.v, EL)
    neuron.LL[1 * um : 3 * um].v = 0 * mV
    assert_allclose(neuron.LL.v, Quantity([EL, 0 * mV, 0 * mV, EL, EL]))
    assert_allclose(neuron.Cm, 1 * uF / cm**2)

    # Test morphological variables
    assert_allclose(neuron.L.main.distance, morpho.L.distance)
    assert_allclose(neuron.L.main.area, morpho.L.area)
    assert_allclose(neuron.L.main.length, morpho.L.length)
    assert_allclose(neuron.L.main.volume, morpho.L.volume)

    # Check basic consistency of the flattened representation
    assert all(
        neuron.diffusion_state_updater._ends[:].flat
        >= neuron.diffusion_state_updater._starts[:].flat
    )

    # Check that length and distances make sense
    assert_allclose(sum(morpho.L.length), 10 * um)
    assert_allclose(morpho.L.distance, (0.5 + np.arange(10)) * um)
    assert_allclose(sum(morpho.LL.length), 5 * um)
    assert_allclose(morpho.LL.distance, (10 + 0.5 + np.arange(5)) * um)
    assert_allclose(sum(morpho.LR.length), 5 * um)
    assert_allclose(morpho.LR.distance, (10 + 0.25 + np.arange(10) * 0.5) * um)
    assert_allclose(sum(morpho.right.length), 3 * um)
    assert_allclose(morpho.right.distance, (0.5 + np.arange(7)) * 3.0 / 7.0 * um)
    assert_allclose(sum(morpho.right.nextone.length), 2 * um)
    assert_allclose(
        morpho.right.nextone.distance, 3 * um + (0.5 + np.arange(3)) * 2.0 / 3.0 * um
    )


@pytest.mark.codegen_independent
def test_construction_coordinates():
    # Same as test_construction, but uses coordinates instead of lengths to
    # set up everything
    # Note that all coordinates here are relative to the origin of the
    # respective cylinder
    BrianLogger.suppress_name("resolution_conflict")
    morpho = Soma(diameter=30 * um)
    morpho.L = Cylinder(x=[0, 10] * um, diameter=1 * um, n=10)
    morpho.LL = Cylinder(y=[0, 5] * um, diameter=2 * um, n=5)
    morpho.LR = Cylinder(z=[0, 5] * um, diameter=2 * um, n=10)
    morpho.right = Cylinder(
        x=[0, sqrt(2) * 1.5] * um, y=[0, sqrt(2) * 1.5] * um, diameter=1 * um, n=7
    )
    morpho.right.nextone = Cylinder(
        y=[0, sqrt(2)] * um, z=[0, sqrt(2)] * um, diameter=1 * um, n=3
    )
    gL = 1e-4 * siemens / cm**2
    EL = -70 * mV
    eqs = """
    Im=gL*(EL-v) : amp/meter**2
    I : meter (point current)
    """

    # Check units of currents
    with pytest.raises(DimensionMismatchError):
        SpatialNeuron(morphology=morpho, model=eqs)

    eqs = """
    Im=gL*(EL-v) : amp/meter**2
    """
    neuron = SpatialNeuron(
        morphology=morpho, model=eqs, Cm=1 * uF / cm**2, Ri=100 * ohm * cm
    )

    # Test initialization of values
    neuron.LL.v = EL
    assert_allclose(neuron.L.main.v, 0 * mV)
    assert_allclose(neuron.LL.v, EL)
    neuron.LL[1 * um : 3 * um].v = 0 * mV
    assert_allclose(neuron.LL.v, Quantity([EL, 0 * mV, 0 * mV, EL, EL]))
    assert_allclose(neuron.Cm, 1 * uF / cm**2)

    # Test morphological variables
    assert_allclose(neuron.L.main.x, morpho.L.x)
    assert_allclose(neuron.LL.main.x, morpho.LL.x)
    assert_allclose(neuron.right.main.x, morpho.right.x)
    assert_allclose(neuron.L.main.distance, morpho.L.distance)
    # assert_allclose(neuron.L.main.diameter, morpho.L.diameter)
    assert_allclose(neuron.L.main.area, morpho.L.area)
    assert_allclose(neuron.L.main.length, morpho.L.length)

    # Check basic consistency of the flattened representation
    assert all(
        neuron.diffusion_state_updater._ends[:].flat
        >= neuron.diffusion_state_updater._starts[:].flat
    )

    # Check that length and distances make sense
    assert_allclose(sum(morpho.L.length), 10 * um)
    assert_allclose(morpho.L.distance, (0.5 + np.arange(10)) * um)
    assert_allclose(sum(morpho.LL.length), 5 * um)
    assert_allclose(morpho.LL.distance, (10 + 0.5 + np.arange(5)) * um)
    assert_allclose(sum(morpho.LR.length), 5 * um)
    assert_allclose(morpho.LR.distance, (10 + 0.25 + np.arange(10) * 0.5) * um)
    assert_allclose(sum(morpho.right.length), 3 * um)
    assert_allclose(morpho.right.distance, (0.5 + np.arange(7)) * 3.0 / 7.0 * um)
    assert_allclose(sum(morpho.right.nextone.length), 2 * um)
    assert_allclose(
        morpho.right.nextone.distance, 3 * um + (0.5 + np.arange(3)) * 2.0 / 3.0 * um
    )


@pytest.mark.long
@numpy_needs_scipy
def test_infinitecable():
    """
    Test simulation of an infinite cable vs. theory for current pulse (Green function)
    """
    BrianLogger.suppress_name("resolution_conflict")
    defaultclock.dt = 0.001 * ms

    # Morphology
    diameter = 1 * um
    Cm = 1 * uF / cm**2
    Ri = 100 * ohm * cm
    N = 500
    morpho = Cylinder(diameter=diameter, length=3 * mm, n=N)

    # Passive channels
    gL = 1e-4 * siemens / cm**2
    eqs = """
    Im=-gL*v : amp/meter**2
    I : amp (point current)
    """

    neuron = SpatialNeuron(morphology=morpho, model=eqs, Cm=Cm, Ri=Ri)

    # Monitors
    mon = StateMonitor(neuron, "v", record=N / 2 - 20)

    neuron.I[len(neuron) // 2] = 1 * nA  # injecting in the middle
    run(0.02 * ms)
    neuron.I = 0 * amp
    run(3 * ms)
    t = mon.t
    v = mon[N // 2 - 20].v
    # Theory (incorrect near cable ends)
    x = 20 * morpho.length[0]
    la = neuron.space_constant[0]
    taum = Cm / gL  # membrane time constant
    theory = (
        1.0
        / (la * Cm * pi * diameter)
        * sqrt(taum / (4 * pi * (t + defaultclock.dt)))
        * exp(
            -(t + defaultclock.dt) / taum
            - taum / (4 * (t + defaultclock.dt)) * (x / la) ** 2
        )
    )
    theory = theory * 1 * nA * 0.02 * ms
    assert_allclose(
        v[t > 0.5 * ms], theory[t > 0.5 * ms], atol=float(6.32 * uvolt)
    )  # high error tolerance (not exact because not infinite cable)


@pytest.mark.standalone_compatible
@numpy_needs_scipy
def test_finitecable():
    """
    Test simulation of short cylinder vs. theory for constant current.
    """
    BrianLogger.suppress_name("resolution_conflict")

    defaultclock.dt = 0.01 * ms

    # Morphology
    diameter = 1 * um
    length = 300 * um
    Cm = 1 * uF / cm**2
    Ri = 150 * ohm * cm
    N = 200
    morpho = Cylinder(diameter=diameter, length=length, n=N)

    # Passive channels
    gL = 1e-4 * siemens / cm**2
    EL = -70 * mV
    eqs = """
    Im=gL*(EL-v) : amp/meter**2
    I : amp (point current)
    """

    neuron = SpatialNeuron(morphology=morpho, model=eqs, Cm=Cm, Ri=Ri)
    neuron.v = EL

    neuron.I[0] = 0.02 * nA  # injecting at the left end

    run(100 * ms)

    # Theory
    x = neuron.distance
    v = neuron.v
    la = neuron.space_constant[0]
    ra = la * 4 * Ri / (pi * diameter**2)
    theory = EL + ra * neuron.I[0] * cosh((length - x) / la) / sinh(length / la)
    assert_allclose(v - EL, theory - EL, atol=1e-6)


@pytest.mark.standalone_compatible
@numpy_needs_scipy
def test_rallpack1():
    """
    Rallpack 1
    """
    if prefs.core.default_float_dtype is np.float32:
        pytest.skip("Need double precision for this test")
    defaultclock.dt = 0.05 * ms

    # Morphology
    diameter = 1 * um
    length = 1 * mm
    Cm = 1 * uF / cm**2
    Ri = 100 * ohm * cm
    N = 1000
    morpho = Cylinder(diameter=diameter, length=length, n=N)

    # Passive channels
    gL = 1.0 / (40000 * ohm * cm**2)
    EL = -65 * mV
    eqs = """
    Im = gL*(EL - v) : amp/meter**2
    I : amp (point current, constant)
    """
    neuron = SpatialNeuron(morphology=morpho, model=eqs, Cm=Cm, Ri=Ri)
    neuron.v = EL

    neuron.I[0] = 0.1 * nA  # injecting at the left end

    # Record at the two ends
    mon = StateMonitor(neuron, "v", record=[0, 999], when="start", dt=0.05 * ms)

    run(250 * ms + defaultclock.dt)

    # Load the theoretical results
    basedir = os.path.join(os.path.abspath(os.path.dirname(__file__)), "rallpack_data")
    data_0 = np.loadtxt(os.path.join(basedir, "ref_cable.0"))
    data_x = np.loadtxt(os.path.join(basedir, "ref_cable.x"))

    scale_0 = max(data_0[:, 1] * volt) - min(data_0[:, 1] * volt)
    scale_x = max(data_x[:, 1] * volt) - min(data_x[:, 1] * volt)
    squared_diff_0 = (data_0[:, 1] * volt - mon[0].v) ** 2
    squared_diff_x = (data_x[:, 1] * volt - mon[999].v) ** 2
    rel_RMS_0 = sqrt(mean(squared_diff_0)) / scale_0
    rel_RMS_x = sqrt(mean(squared_diff_x)) / scale_x
    max_rel_0 = sqrt(max(squared_diff_0)) / scale_0
    max_rel_x = sqrt(max(squared_diff_x)) / scale_x

    # sanity check: times are the same
    assert_allclose(mon.t / second, data_0[:, 0])
    assert_allclose(mon.t / second, data_x[:, 0])

    # RMS error should be < 0.1%, maximum error along the curve should be < 0.5%
    assert 100 * rel_RMS_0 < 0.1
    assert 100 * rel_RMS_x < 0.1
    assert 100 * max_rel_0 < 0.5
    assert 100 * max_rel_x < 0.5


@pytest.mark.standalone_compatible
@numpy_needs_scipy
def test_rallpack2():
    """
    Rallpack 2
    """
    if prefs.core.default_float_dtype is np.float32:
        pytest.skip("Need double precision for this test")
    defaultclock.dt = 0.1 * ms

    # Morphology
    diameter = 32 * um
    length = 16 * um
    Cm = 1 * uF / cm**2
    Ri = 100 * ohm * cm

    # Construct binary tree according to Rall's formula
    morpho = Cylinder(n=1, diameter=diameter, y=[0, float(length)] * meter)
    endpoints = {morpho}
    for depth in range(1, 10):
        diameter /= 2.0 ** (1.0 / 3.0)
        length /= 2.0 ** (2.0 / 3.0)
        new_endpoints = set()
        for endpoint in endpoints:
            new_L = Cylinder(n=1, diameter=diameter, length=length)
            new_R = Cylinder(n=1, diameter=diameter, length=length)
            new_endpoints.add(new_L)
            new_endpoints.add(new_R)
            endpoint.L = new_L
            endpoint.R = new_R
        endpoints = new_endpoints

    # Passive channels
    gL = 1.0 / (40000 * ohm * cm**2)
    EL = -65 * mV
    eqs = """
    Im = gL*(EL - v) : amp/meter**2
    I : amp (point current, constant)
    """
    neuron = SpatialNeuron(morphology=morpho, model=eqs, Cm=Cm, Ri=Ri, method="rk4")
    neuron.v = EL

    neuron.I[0] = 0.1 * nA  # injecting at the origin

    endpoint_indices = [endpoint.indices[0] for endpoint in endpoints]
    mon = StateMonitor(
        neuron, "v", record=[0] + endpoint_indices, when="start", dt=0.1 * ms
    )

    run(250 * ms + defaultclock.dt)

    # Load the theoretical results
    basedir = os.path.join(os.path.abspath(os.path.dirname(__file__)), "rallpack_data")
    # Only use very second time step, since we run with 0.1ms instead of 0.05ms
    data_0 = np.loadtxt(os.path.join(basedir, "ref_branch.0"))[::2]
    data_x = np.loadtxt(os.path.join(basedir, "ref_branch.x"))[::2]

    # sanity check: times are the same
    assert_allclose(mon.t / second, data_0[:, 0])
    assert_allclose(mon.t / second, data_x[:, 0])

    # Check that all endpoints are the same:
    for endpoint in endpoints:
        assert_allclose(mon[endpoint].v, mon[endpoint[0]].v)

    scale_0 = max(data_0[:, 1] * volt) - min(data_0[:, 1] * volt)
    scale_x = max(data_x[:, 1] * volt) - min(data_x[:, 1] * volt)
    squared_diff_0 = (data_0[:, 1] * volt - mon[0].v) ** 2

    # One endpoint
    squared_diff_x = (data_x[:, 1] * volt - mon[endpoint_indices[0]].v) ** 2
    rel_RMS_0 = sqrt(mean(squared_diff_0)) / scale_0
    rel_RMS_x = sqrt(mean(squared_diff_x)) / scale_x
    max_rel_0 = sqrt(max(squared_diff_0)) / scale_0
    max_rel_x = sqrt(max(squared_diff_x)) / scale_x

    # RMS error should be < 0.25%, maximum error along the curve should be < 0.5%
    assert 100 * rel_RMS_0 < 0.25
    assert 100 * rel_RMS_x < 0.25
    assert 100 * max_rel_0 < 0.5
    assert 100 * max_rel_x < 0.5


@pytest.mark.standalone_compatible
@pytest.mark.long
@numpy_needs_scipy
def test_rallpack3():
    """
    Rallpack 3
    """
    if prefs.core.default_float_dtype is np.float32:
        pytest.skip("Need double precision for this test")
    defaultclock.dt = 1 * usecond

    # Morphology
    diameter = 1 * um
    length = 1 * mm
    N = 1000
    morpho = Cylinder(diameter=diameter, length=length, n=N)
    # Passive properties
    gl = 1.0 / (40000 * ohm * cm**2)
    El = -65 * mV
    Cm = 1 * uF / cm**2
    Ri = 100 * ohm * cm
    # Active properties
    ENa = 50 * mV
    EK = -77 * mV
    gNa = 120 * msiemens / cm**2
    gK = 36 * msiemens / cm**2
    eqs = """
    Im = gl * (El-v) + gNa * m**3 * h * (ENa-v) + gK * n**4 * (EK-v) : amp/meter**2
    dm/dt = alpham * (1-m) - betam * m : 1
    dn/dt = alphan * (1-n) - betan * n : 1
    dh/dt = alphah * (1-h) - betah * h : 1
    v_shifted = v - El : volt
    alpham = (0.1/mV) * (-v_shifted+25*mV) / (exp((-v_shifted+25*mV) / (10*mV)) - 1)/ms : Hz
    betam = 4 * exp(-v_shifted/(18*mV))/ms : Hz
    alphah = 0.07 * exp(-v_shifted/(20*mV))/ms : Hz
    betah = 1/(exp((-v_shifted+30*mV) / (10*mV)) + 1)/ms : Hz
    alphan = (0.01/mV) * (-v_shifted+10*mV) / (exp((-v_shifted+10*mV) / (10*mV)) - 1)/ms : Hz
    betan = 0.125*exp(-v_shifted/(80*mV))/ms : Hz
    I : amp (point current, constant)
    """
    axon = SpatialNeuron(
        morphology=morpho, model=eqs, Cm=Cm, Ri=Ri, method="exponential_euler"
    )
    axon.v = El
    # Pre-calculated equilibrium values at v = El
    axon.m = 0.0529324852572
    axon.n = 0.317676914061
    axon.h = 0.596120753508
    axon.I[0] = 0.1 * nA  # injecting at the left end

    # Record at the two ends
    mon = StateMonitor(axon, "v", record=[0, 999], when="start", dt=0.05 * ms)

    run(250 * ms + defaultclock.dt)

    # Load the theoretical results
    basedir = os.path.join(os.path.abspath(os.path.dirname(__file__)), "rallpack_data")
    data_0 = np.loadtxt(os.path.join(basedir, "ref_axon.0.neuron"))
    data_x = np.loadtxt(os.path.join(basedir, "ref_axon.x.neuron"))

    # sanity check: times are the same
    assert_allclose(mon.t / second, data_0[:, 0])
    assert_allclose(mon.t / second, data_x[:, 0])

    scale_0 = max(data_0[:, 1] * volt) - min(data_0[:, 1] * volt)
    scale_x = max(data_x[:, 1] * volt) - min(data_x[:, 1] * volt)
    squared_diff_0 = (data_0[:, 1] * volt - mon[0].v) ** 2
    squared_diff_x = (data_x[:, 1] * volt - mon[999].v) ** 2

    rel_RMS_0 = sqrt(mean(squared_diff_0)) / scale_0
    rel_RMS_x = sqrt(mean(squared_diff_x)) / scale_x
    max_rel_0 = sqrt(max(squared_diff_0)) / scale_0
    max_rel_x = sqrt(max(squared_diff_x)) / scale_x

    # RMS error should be < 0.1%, maximum error along the curve should be < 0.5%
    # Note that this is much stricter than the original Rallpack evaluation, but
    # with the 1us time step, the voltage traces are extremely similar
    assert 100 * rel_RMS_0 < 0.1
    assert 100 * rel_RMS_x < 0.1
    assert 100 * max_rel_0 < 0.5
    assert 100 * max_rel_x < 0.5


@pytest.mark.standalone_compatible
@numpy_needs_scipy
def test_rall():
    """
    Test simulation of a cylinder plus two branches, with diameters according to Rall's formula
    """
    BrianLogger.suppress_name("resolution_conflict")

    defaultclock.dt = 0.01 * ms

    # Passive channels
    gL = 1e-4 * siemens / cm**2
    EL = -70 * mV

    # Morphology
    diameter = 1 * um
    length = 300 * um
    Cm = 1 * uF / cm**2
    Ri = 150 * ohm * cm
    N = 500
    rm = 1 / (gL * pi * diameter)  # membrane resistance per unit length
    ra = (4 * Ri) / (pi * diameter**2)  # axial resistance per unit length
    la = sqrt(rm / ra)  # space length
    morpho = Cylinder(diameter=diameter, length=length, n=N)
    d1 = 0.5 * um
    L1 = 200 * um
    rm = 1 / (gL * pi * d1)  # membrane resistance per unit length
    ra = (4 * Ri) / (pi * d1**2)  # axial resistance per unit length
    l1 = sqrt(rm / ra)  # space length
    morpho.L = Cylinder(diameter=d1, length=L1, n=N)
    d2 = (diameter**1.5 - d1**1.5) ** (1.0 / 1.5)
    rm = 1 / (gL * pi * d2)  # membrane resistance per unit length
    ra = (4 * Ri) / (pi * d2**2)  # axial resistance per unit length
    l2 = sqrt(rm / ra)  # space length
    L2 = (L1 / l1) * l2
    morpho.R = Cylinder(diameter=d2, length=L2, n=N)

    eqs = """
    Im=gL*(EL-v) : amp/meter**2
    I : amp (point current)
    """

    neuron = SpatialNeuron(morphology=morpho, model=eqs, Cm=Cm, Ri=Ri)
    neuron.v = EL

    neuron.I[0] = 0.02 * nA  # injecting at the left end
    run(100 * ms)

    # Check space constant calculation
    assert_allclose(la, neuron.space_constant[0])
    assert_allclose(l1, neuron.L.space_constant[0])
    assert_allclose(l2, neuron.R.space_constant[0])

    # Theory
    x = neuron.main.distance
    ra = la * 4 * Ri / (pi * diameter**2)
    l = length / la + L1 / l1
    theory = EL + ra * neuron.I[0] * cosh(l - x / la) / sinh(l)
    v = neuron.main.v
    assert_allclose(v - EL, theory - EL, atol=2e-6)
    x = neuron.L.distance
    theory = EL + ra * neuron.I[0] * cosh(
        l - neuron.main.distance[-1] / la - (x - neuron.main.distance[-1]) / l1
    ) / sinh(l)
    v = neuron.L.v
    assert_allclose(v - EL, theory - EL, atol=2e-6)
    x = neuron.R.distance
    theory = EL + ra * neuron.I[0] * cosh(
        l - neuron.main.distance[-1] / la - (x - neuron.main.distance[-1]) / l2
    ) / sinh(l)
    v = neuron.R.v
    assert_allclose(v - EL, theory - EL, atol=2e-6)


@pytest.mark.standalone_compatible
@numpy_needs_scipy
def test_basic_diffusion():
    # A very basic test that shows that propagation is working in a very basic
    # sense, testing all morphological classes

    defaultclock.dt = 0.01 * ms

    EL = -70 * mV
    gL = 1e-4 * siemens / cm**2
    target = -10 * mV
    eqs = """
    Im = gL*(EL-v) + gClamp*(target-v): amp/meter**2
    gClamp : siemens/meter**2
    """

    morph = Soma(diameter=30 * um)
    morph.axon = Cylinder(n=10, diameter=10 * um, length=100 * um)
    morph.dend = Section(
        n=10,
        diameter=[10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0.1] * um,
        length=np.ones(10) * 10 * um,
    )

    neuron = SpatialNeuron(morph, eqs)
    neuron.v = EL
    neuron.axon.gClamp[0] = 100 * siemens / cm**2

    mon = StateMonitor(neuron, "v", record=True)

    run(0.25 * ms)
    assert all(abs(mon.v[:, -1] / mV + 10) < 0.25), mon.v[:, -1] / mV


@pytest.mark.codegen_independent
def test_allowed_integration():
    morph = Soma(diameter=30 * um)
    EL = -70 * mV
    gL = 1e-4 * siemens / cm**2
    ENa = 115 * mV
    gNa = 120 * msiemens / cm**2
    VT = -50.4 * mV
    DeltaT = 2 * mV
    ENMDA = 0.0 * mV

    @check_units(voltage=volt, result=volt)
    def user_fun(voltage):
        return voltage  # could be an arbitrary function and is therefore unsafe

    allowed_eqs = [
        "Im = gL*(EL-v) : amp/meter**2",
        """
        Im = gl * (El-v) + gNa * m**3 * h * (ENa-v) : amp/meter**2
        dm/dt = alpham * (1-m) - betam * m : 1
        dh/dt = alphah * (1-h) - betah * h : 1
        alpham = (0.1/mV) * (-v+25*mV) / (exp((-v+25*mV) / (10*mV)) - 1)/ms : Hz
        betam = 4 * exp(-v/(18*mV))/ms : Hz
        alphah = 0.07 * exp(-v/(20*mV))/ms : Hz
        betah = 1/(exp((-v+30*mV) / (10*mV)) + 1)/ms : Hz
        """,
        """
        Im = gl * (El-v) : amp/meter**2
        I_ext = 1*nA + sin(2*pi*100*Hz*t)*nA : amp (point current)
        """,
        """
        Im = I_leak + I_spike : amp/meter**2
        I_leak = gL*(EL - v) : amp/meter**2
        I_spike = gL*DeltaT*exp((v - VT)/DeltaT): amp/meter**2 (constant over dt)
        """,
        """
        Im = gL*(EL-v) : amp/meter**2
        I_NMDA = gNMDA*(ENMDA-v)*Mgblock : amp (point current)
        gNMDA : siemens
        Mgblock = 1./(1. +  exp(-0.062*v/mV)/3.57) : 1 (constant over dt)
        """,
        "Im = gL*(EL - v) + gL*DeltaT*exp((v - VT)/DeltaT) : amp/meter**2",
        """
        Im = I_leak + I_spike : amp/meter**2
        I_leak = gL*(EL - v) : amp/meter**2
        I_spike = gL*DeltaT*exp((v - VT)/DeltaT): amp/meter**2
        """,
        """
        Im = gL*(EL-v) : amp/meter**2
        I_NMDA = gNMDA*(ENMDA-v)*Mgblock : amp (point current)
        gNMDA : siemens
        Mgblock = 1./(1. +  exp(-0.062*v/mV)/3.57) : 1
        """,
    ]
    forbidden_eqs = [
        """Im = gl * (El-v + user_fun(v)) : amp/meter**2""",
        """Im = gl * clip(El-v, -100*mV, 100*mV) : amp/meter**2""",
    ]
    for eqs in allowed_eqs:
        # Should not raise an error
        neuron = SpatialNeuron(morph, eqs)

    for eqs in forbidden_eqs:
        # Should raise an error
        with pytest.raises(TypeError):
            SpatialNeuron(morph, eqs)


@pytest.mark.codegen_independent
def test_spatialneuron_indexing():
    sec = Cylinder(length=50 * um, diameter=10 * um, n=1)
    sec.sec1 = Cylinder(length=50 * um, diameter=10 * um, n=2)
    sec.sec1.sec11 = Cylinder(length=50 * um, diameter=10 * um, n=4)
    sec.sec1.sec12 = Cylinder(length=50 * um, diameter=10 * um, n=8)
    sec.sec2 = Cylinder(length=50 * um, diameter=10 * um, n=16)
    sec.sec2.sec21 = Cylinder(length=50 * um, diameter=10 * um, n=32)
    neuron = SpatialNeuron(sec, "Im = 0*amp/meter**2 : amp/meter**2")
    neuron.v = "i*volt"
    # Accessing indices/variables of a subtree refers to the full subtree
    assert len(neuron.indices[:]) == 1 + 2 + 4 + 8 + 16 + 32
    assert len(neuron.sec1.indices[:]) == 2 + 4 + 8
    assert len(neuron.sec1.sec11.indices[:]) == 4
    assert len(neuron.sec1.sec12.indices[:]) == 8
    assert len(neuron.sec2.indices[:]) == 16 + 32
    assert len(neuron.sec2.sec21.indices[:]) == 32
    assert len(neuron.v[:]) == 1 + 2 + 4 + 8 + 16 + 32
    assert len(neuron.sec1.v[:]) == 2 + 4 + 8
    assert len(neuron.sec1.sec11.v[:]) == 4
    assert len(neuron.sec1.sec12.v[:]) == 8
    assert len(neuron.sec2.v[:]) == 16 + 32
    assert len(neuron.sec2.sec21.v[:]) == 32
    # Accessing indices/variables with ".main" only refers to the section
    assert len(neuron.main.indices[:]) == 1
    assert len(neuron.sec1.main.indices[:]) == 2
    assert len(neuron.sec1.sec11.main.indices[:]) == 4
    assert len(neuron.sec1.sec12.main.indices[:]) == 8
    assert len(neuron.sec2.main.indices[:]) == 16
    assert len(neuron.sec2.sec21.main.indices[:]) == 32
    assert len(neuron.main.v[:]) == 1
    assert len(neuron.sec1.main.v[:]) == 2
    assert len(neuron.sec1.sec11.main.v[:]) == 4
    assert len(neuron.sec1.sec12.main.v[:]) == 8
    assert len(neuron.sec2.main.v[:]) == 16
    assert len(neuron.sec2.sec21.main.v[:]) == 32
    # Accessing subgroups
    assert len(neuron[0].indices[:]) == 1
    assert len(neuron[0 * um : 50 * um].indices[:]) == 1
    assert len(neuron[0:1].indices[:]) == 1
    assert len(neuron[sec.sec2.indices[:]]) == 16
    assert len(neuron[sec.sec2]) == 16
    assert_equal(neuron.sec1.sec11.v, [3, 4, 5, 6] * volt)
    assert_equal(neuron.sec1.sec11[1].v, neuron.sec1.sec11.v[1])
    assert_equal(neuron.sec1.sec11[1:3].v, neuron.sec1.sec11.v[1:3])
    assert_equal(neuron.sec1.sec11[1:3].v, [4, 5] * volt)


@pytest.mark.codegen_independent
def test_tree_index_consistency():
    # Test all possible trees with depth 3 and a maximum of 3 branches subtree
    # (a total of 84 trees)
    # This tests whether the indices (i.e. where the compartments are placed in
    # the overall flattened 1D structure) make sense: for the `SpatialSubgroup`
    # mechanism to work correctly, each subtree has to have contiguous indices.
    # Separate subtrees should of course have non-overlapping indices.
    for tree_description in itertools.product(
        [1, 2, 3],  # children of root
        [0, 1, 2, 3],  # children of first branch
        [0, 1, 2, 3],  # children of second branch
        [0, 1, 2, 3],  # children of third branch
    ):
        sec = Cylinder(length=50 * um, diameter=10 * um, n=1)
        root_children = tree_description[0]
        if not all([tree_description[x] == 0 for x in range(root_children + 1, 4)]):
            # skip redundant descriptions (differing number of branches in a
            # subtree that does not exist)
            continue

        # Create a tree according to the description
        for idx in range(root_children):
            setattr(
                sec,
                f"sec{int(idx + 1)}",
                Cylinder(length=50 * um, diameter=10 * um, n=2 * (idx + 1)),
            )
        for child in range(root_children):
            subsec = getattr(sec, f"sec{int(child + 1)}")
            subsec_children = tree_description[child + 1]
            for idx in range(subsec_children):
                setattr(
                    subsec,
                    f"sec{int(child + 1)}{int(idx + 1)}",
                    Cylinder(length=50 * um, diameter=10 * um, n=1 + (child + 1) * idx),
                )

        neuron = SpatialNeuron(sec, "Im = 0*amp/meter**2 : amp/meter**2")
        # Check the indicies for the full neuron:
        assert_equal(neuron.indices[:], np.arange(sec.total_compartments))

        all_subsec_indices = []
        for child in range(root_children):
            subsec = getattr(neuron, f"sec{int(child + 1)}")
            sub_indices = set(subsec.main.indices[:])
            subsec_children = tree_description[child + 1]
            for idx in range(subsec_children):
                subsubsec = getattr(subsec, f"sec{int(child + 1)}{int(idx + 1)}")
                sub_indices |= set(subsubsec.main.indices[:])
            # The indices for a full subtree should be the union of the indices
            # for all subsections within that subtree
            assert sub_indices == set(subsec.indices[:])
            all_subsec_indices.extend(subsec.indices[:])
        # Separate subtrees should not overlap
        assert len(all_subsec_indices) == len(set(all_subsec_indices))


@pytest.mark.codegen_independent
def test_spatialneuron_subtree_assignment():
    sec = Cylinder(length=50 * um, diameter=10 * um, n=2)
    sec.sec1 = Cylinder(length=50 * um, diameter=10 * um, n=2)
    sec.sec1.sec11 = Cylinder(length=50 * um, diameter=10 * um, n=2)
    sec.sec1.sec12 = Cylinder(length=50 * um, diameter=10 * um, n=2)
    sec.sec2 = Cylinder(length=50 * um, diameter=10 * um, n=2)
    sec.sec2.sec21 = Cylinder(length=50 * um, diameter=10 * um, n=2)
    neuron = SpatialNeuron(sec, "Im = 0*amp/meter**2 : amp/meter**2")

    neuron.v = 1 * volt
    assert_allclose(neuron.v[:], np.ones(12) * volt)
    neuron.sec1.v += 1 * volt
    assert_allclose(neuron.main.v[:], np.ones(2) * volt)
    assert_allclose(neuron.sec1.v[:], np.ones(6) * 2 * volt)
    assert_allclose(neuron.sec1.main.v[:], np.ones(2) * 2 * volt)
    assert_allclose(neuron.sec1.sec11.v[:], np.ones(2) * 2 * volt)
    assert_allclose(neuron.sec1.sec12.v[:], np.ones(2) * 2 * volt)
    assert_allclose(neuron.sec2.v[:], np.ones(4) * volt)
    neuron.sec2.v = 5 * volt
    assert_allclose(neuron.sec2.v[:], np.ones(4) * 5 * volt)
    assert_allclose(neuron.sec2.main.v[:], np.ones(2) * 5 * volt)
    assert_allclose(neuron.sec2.sec21.v[:], np.ones(2) * 5 * volt)


@pytest.mark.codegen_independent
def test_spatialneuron_morphology_assignment():
    sec = Cylinder(length=50 * um, diameter=10 * um, n=2)
    sec.sec1 = Cylinder(length=50 * um, diameter=10 * um, n=2)
    sec.sec1.sec11 = Cylinder(length=50 * um, diameter=10 * um, n=2)
    sec.sec1.sec12 = Cylinder(length=50 * um, diameter=10 * um, n=2)
    sec.sec2 = Cylinder(length=50 * um, diameter=10 * um, n=2)
    sec.sec2.sec21 = Cylinder(length=50 * um, diameter=10 * um, n=2)
    neuron = SpatialNeuron(sec, "Im = 0*amp/meter**2 : amp/meter**2")

    neuron.v[sec.sec1.sec11] = 1 * volt
    assert_allclose(neuron.sec1.sec11.v[:], np.ones(2) * volt)
    assert_allclose(neuron.sec1.sec12.v[:], np.zeros(2) * volt)
    assert_allclose(neuron.sec1.main.v[:], np.zeros(2) * volt)
    assert_allclose(neuron.main.v[:], np.zeros(2) * volt)
    assert_allclose(neuron.sec2.v[:], np.zeros(4) * volt)

    neuron.v[sec.sec2[25 * um :]] = 2 * volt
    neuron.v[sec.sec1[: 25 * um]] = 3 * volt
    assert_allclose(neuron.main.v[:], np.zeros(2) * volt)
    assert_allclose(neuron.sec2.main.v[:], [0, 2] * volt)
    assert_allclose(neuron.sec2.sec21.v[:], np.zeros(2) * volt)
    assert_allclose(neuron.sec1.main.v[:], [3, 0] * volt)
    assert_allclose(neuron.sec1.sec11.v[:], np.ones(2) * volt)
    assert_allclose(neuron.sec1.sec12.v[:], np.zeros(2) * volt)


@pytest.mark.standalone_compatible
@pytest.mark.multiple_runs
@numpy_needs_scipy
def test_spatialneuron_capacitive_currents():
    if prefs.core.default_float_dtype is np.float32:
        pytest.skip("Need double precision for this test")
    defaultclock.dt = 0.1 * ms
    morpho = Cylinder(x=[0, 10] * cm, diameter=2 * 238 * um, n=200, type="axon")

    El = 10.613 * mV
    ENa = 115 * mV
    EK = -12 * mV
    gl = 0.3 * msiemens / cm**2
    gNa0 = 120 * msiemens / cm**2
    gK = 36 * msiemens / cm**2

    # Typical equations
    eqs = """
    # The same equations for the whole neuron, but possibly different parameter values
    # distributed transmembrane current
    Im = gl * (El-v) + gNa * m**3 * h * (ENa-v) + gK * n**4 * (EK-v) : amp/meter**2
    I : amp (point current) # applied current
    dm/dt = alpham * (1-m) - betam * m : 1
    dn/dt = alphan * (1-n) - betan * n : 1
    dh/dt = alphah * (1-h) - betah * h : 1
    alpham = (0.1/mV) * (-v+25*mV) / (exp((-v+25*mV) / (10*mV)) - 1)/ms : Hz
    betam = 4 * exp(-v/(18*mV))/ms : Hz
    alphah = 0.07 * exp(-v/(20*mV))/ms : Hz
    betah = 1/(exp((-v+30*mV) / (10*mV)) + 1)/ms : Hz
    alphan = (0.01/mV) * (-v+10*mV) / (exp((-v+10*mV) / (10*mV)) - 1)/ms : Hz
    betan = 0.125*exp(-v/(80*mV))/ms : Hz
    gNa : siemens/meter**2
    """

    neuron = SpatialNeuron(
        morphology=morpho,
        model=eqs,
        Cm=1 * uF / cm**2,
        Ri=35.4 * ohm * cm,
        method="exponential_euler",
    )
    mon = StateMonitor(neuron, ["Im", "Ic"], record=True, when="end")
    run(10 * ms)
    neuron.I[0] = 1 * uA  # current injection at one end
    run(3 * ms)
    neuron.I = 0 * amp
    run(10 * ms)
    device.build(direct_call=False, **device.build_options)
    assert_allclose(
        (mon.Im - mon.Ic).sum(axis=0) / (mA / cm**2), np.zeros(230), atol=1e6
    )


@pytest.mark.codegen_independent
def test_point_current():
    soma = Soma(10 * um)
    eqs = """Im = 0*nA/cm**2 : amp/meter**2
             I1 = 1*nA : amp (point current)
             I2 = 1*nA : amp (point current, constant over dt)"""
    neuron = SpatialNeuron(soma, eqs)
    assert "I1/area" in neuron.equations["Im"].expr.code
    assert "I2/area" in neuron.equations["Im"].expr.code  # see issue #1160


@pytest.mark.standalone_compatible
@pytest.mark.multiple_runs
@numpy_needs_scipy
def test_spatialneuron_threshold_location():
    morpho = Soma(10 * um)
    morpho.axon = Cylinder(1 * um, n=2, length=20 * um)
    model = """
    Im = 0*nA/cm**2 : amp/meter**2
    should_spike : boolean (constant)
    """
    neuron = SpatialNeuron(
        morpho, model, threshold_location=morpho.axon[15 * um], threshold="should_spike"
    )
    # Different variants that should do the same thing
    neuron2 = SpatialNeuron(
        morpho,
        model,
        threshold_location=morpho.axon.indices[15 * um],
        threshold="should_spike",
    )
    neuron3 = SpatialNeuron(
        morpho, model, threshold_location=2, threshold="should_spike"
    )
    # Cannot use multiple compartments
    with pytest.raises(AttributeError):
        SpatialNeuron(
            morpho, model, threshold_location=[2, 3], threshold="should_spike"
        )
    with pytest.raises(AttributeError):
        SpatialNeuron(
            morpho,
            model,
            threshold_location=morpho.axon[5 * um : 15 * um],
            threshold="should_spike",
        )
    neurons = [neuron, neuron2, neuron3]
    monitors = [SpikeMonitor(n) for n in neurons]

    net = Network(neurons, monitors)
    for n in neurons:
        n.should_spike = True  # all compartments want to spike
    net.run(defaultclock.dt)
    for n in neurons:
        n.should_spike = False  # no compartment wants to spike
    net.run(defaultclock.dt)
    for n in neurons:
        n.should_spike = [False, False, True]
    net.run(defaultclock.dt)
    for n in neurons:
        n.should_spike = [True, True, False]
    net.run(defaultclock.dt)
    device.build(direct_call=False, **device.build_options)
    for mon in monitors:
        assert len(mon.i) == 2
        assert all(mon.i == 2)
        assert_allclose(mon.t, [0 * ms, 2 * defaultclock.dt])


@pytest.mark.standalone_compatible
@numpy_needs_scipy
def test_spatialneuron_timedarray():
    # See GitHub issue 1427
    ta = TimedArray([0, 1] * nA, dt=1 * ms)
    morpho = Soma(diameter=10 * um)
    neuron = SpatialNeuron(morpho, "Im = ta(t)/area : amp/meter**2", method="euler")
    mon = StateMonitor(neuron, "v", record=0, when="after_groups")
    run(2 * ms)
    assert_allclose(
        np.diff(mon.v_[0]),
        np.r_[
            np.zeros(9),
            np.array(
                np.ones(10) * 1 * nA / neuron.area[0] / neuron.Cm * defaultclock.dt
            ),
        ],
    )


if __name__ == "__main__":
    test_custom_events()
    test_construction()
    test_construction_coordinates()
    test_infinitecable()
    test_finitecable()
    test_rallpack1()
    test_rallpack2()
    test_rallpack3()
    test_rall()
    test_basic_diffusion()
    test_allowed_integration()
    test_spatialneuron_indexing()
    test_tree_index_consistency()
    test_spatialneuron_subtree_assignment()
    test_spatialneuron_morphology_assignment()
    test_spatialneuron_capacitive_currents()
    test_spatialneuron_timedarray()
