import functools

import pytest

from brian2 import *
from brian2.codegen.runtime.GSLcython_rt import IntegrationError
from brian2.core.preferences import PreferenceError
from brian2.stateupdaters.base import UnsupportedEquationsException
from brian2.tests.utils import exc_isinstance

max_difference = 0.1 * mV

pytestmark = pytest.mark.gsl


def skip_if_not_implemented(func):
    @functools.wraps(func)
    def wrapped():
        try:
            func()
        except (BrianObjectException, NotImplementedError) as exc:
            if not (
                isinstance(exc, NotImplementedError)
                or isinstance(exc.__cause__, NotImplementedError)
            ):
                raise
            pytest.skip("GSL support for numpy has not been implemented yet")

    return wrapped


@pytest.mark.standalone_compatible
@skip_if_not_implemented
def test_GSL_stateupdater_basic():
    # just the adaptive_threshold example: run for exponential_euler and GSL and see
    # if results are comparable (same amount of spikes and spikes > 0)
    eqs = """
    dv/dt = -v/(10*ms) : volt
    dvt/dt = (10*mV-vt)/(15*ms) : volt
    """
    reset = """
    v = 0*mV
    vt += 3*mV
    """
    neurons_conventional = NeuronGroup(
        1, model=eqs, reset=reset, threshold="v>vt", method="exponential_euler"
    )
    neurons_GSL = NeuronGroup(1, model=eqs, reset=reset, threshold="v>vt", method="gsl")
    neurons_conventional.vt = 10 * mV
    neurons_GSL.vt = 10 * mV
    # 50 'different' neurons so no neuron spikes more than once per dt
    P = SpikeGeneratorGroup(1, [0] * 50, np.arange(50) / 50.0 * 100 * ms)
    C_conventional = Synapses(P, neurons_conventional, on_pre="v += 3*mV")
    C_GSL = Synapses(P, neurons_GSL, on_pre="v += 3*mV")
    C_conventional.connect()
    C_GSL.connect()
    SM_conventional = SpikeMonitor(neurons_conventional, variables="v")
    SM_GSL = SpikeMonitor(neurons_GSL, variables="v")
    net = Network(
        neurons_conventional,
        neurons_GSL,
        P,
        C_conventional,
        C_GSL,
        SM_conventional,
        SM_GSL,
    )
    net.run(100 * ms)
    assert (
        SM_conventional.num_spikes > 0
    ), "simulation should produce spiking, but no spikes monitored"
    assert SM_conventional.num_spikes == SM_GSL.num_spikes, (
        "GSL_statupdater produced different number of spikes than integration with ",
        "exponential euler",
    )


@pytest.mark.standalone_compatible
@skip_if_not_implemented
def test_GSL_different_clocks():
    vt = 10 * mV
    eqs = "dv/dt = -v/(10*ms) : volt"
    neurons = NeuronGroup(1, model=eqs, threshold="v>vt", method="gsl", dt=0.2 * ms)
    # for this test just check if it compiles
    run(0 * ms)


@pytest.mark.standalone_compatible
@skip_if_not_implemented
def test_GSL_default_function():
    # phase_locking example
    tau = 20 * ms
    n = 100
    b = 1.2  # constant current mean, the modulation varies
    freq = 10 * Hz
    eqs = """
        dv/dt = (-v + a * sin(2 * pi * freq * t) + b) / tau : 1
        a : 1
    """
    vrand = rand(n)
    neurons_conventional = NeuronGroup(
        n, model=eqs, threshold="v > 1", reset="v = 0", method="exponential_euler"
    )
    neurons_GSL = NeuronGroup(
        n, model=eqs, threshold="v > 1", reset="v = 0", method="gsl"
    )
    neurons_conventional.v = vrand
    neurons_GSL.v = vrand
    neurons_conventional.a = "0.05 + 0.7*i/n"
    neurons_GSL.a = "0.05 + 0.7*i/n"

    trace_conventional = StateMonitor(neurons_conventional, "v", record=50)
    trace_GSL = StateMonitor(neurons_GSL, "v", record=50)
    net = Network(neurons_conventional, neurons_GSL, trace_conventional, trace_GSL)
    net.run(10 * ms)

    assert (
        max(trace_conventional.v[0] - trace_GSL.v[0]) < max_difference / mV
    ), "difference between conventional and GSL output is larger than max_difference"


@pytest.mark.standalone_compatible
@skip_if_not_implemented
def test_GSL_user_defined_function():
    # phase_locking example with user_defined sin
    eqs = """
    dv/dt = (-v + a * sin(2 * pi * freq * t) + b) / tau : 1
    a : 1
    """

    @implementation(
        "cpp",
        """
        double user_sin(double phase)
        {
            return sin(phase);
        }
        """,
    )
    @implementation(
        "cython",
        """
        cdef double user_sin(double phase):
            return sin(phase)
        """,
    )
    @check_units(phase=1, result=1)
    def user_sin(phase):
        raise Exception

    tau = 20 * ms
    n = 100
    b = 1.2  # constant current mean, the modulation varies
    freq = 10 * Hz
    eqs = """
    dv/dt = (-v + a * user_sin(2 * pi * freq * t) + b) / tau : 1
    a : 1
    """
    vrand = rand(n)
    neurons_conventional = NeuronGroup(
        n, model=eqs, threshold="v > 1", reset="v = 0", method="exponential_euler"
    )
    neurons_GSL = NeuronGroup(
        n, model=eqs, threshold="v > 1", reset="v = 0", method="gsl"
    )
    neurons_conventional.v = vrand
    neurons_GSL.v = vrand
    neurons_conventional.a = "0.05 + 0.7*i/n"
    neurons_GSL.a = "0.05 + 0.7*i/n"

    trace_conventional = StateMonitor(neurons_conventional, "v", record=50)
    trace_GSL = StateMonitor(neurons_GSL, "v", record=50)
    net = Network(neurons_conventional, neurons_GSL, trace_conventional, trace_GSL)
    net.run(10 * ms)

    assert (
        max(trace_conventional.v[0] - trace_GSL.v[0]) < max_difference / mV
    ), "difference between conventional and GSL output is larger than max_difference"
    # assert not all(trace_conventional.v[0]==trace_GSL.v[0]), \
    #         ('output of GSL stateupdater is exactly the same as Brians stateupdater (unlikely to be right)')


@pytest.mark.standalone_compatible
@skip_if_not_implemented
def test_GSL_x_variable():
    neurons = NeuronGroup(
        2, "dx/dt = 300*Hz : 1", threshold="x>1", reset="x=0", method="gsl"
    )
    # just testing compilation
    run(0 * ms)


@pytest.mark.codegen_independent
def test_GSL_failing_directory():
    def set_dir(arg):
        prefs.GSL.directory = arg

    with pytest.raises(PreferenceError):
        set_dir(1)
    with pytest.raises(PreferenceError):
        set_dir("/usr/")
    with pytest.raises(PreferenceError):
        set_dir("/usr/blablabla/")


@pytest.mark.codegen_independent
@skip_if_not_implemented
def test_GSL_stochastic():
    tau = 20 * ms
    sigma = 0.015
    eqs = """
    dx/dt = (1.1 - x) / tau + sigma * (2 / tau)**.5 * xi : 1
    """
    neuron = NeuronGroup(1, eqs, method="gsl")
    net = Network(neuron)
    with pytest.raises(BrianObjectException) as exc:
        net.run(0 * ms, namespace={"tau": tau, "sigma": sigma})
    assert exc_isinstance(
        exc, UnsupportedEquationsException, raise_not_implemented=True
    )


@pytest.mark.standalone_compatible
@skip_if_not_implemented
def test_GSL_error_dimension_mismatch_unit():
    eqs = """
    dv/dt = (v0 - v)/(10*ms) : volt
    v0 : volt
    """
    options = {"absolute_error_per_variable": {"v": 1 * nS}}
    neuron = NeuronGroup(
        1,
        eqs,
        threshold="v > 10*mV",
        reset="v = 0*mV",
        method="gsl",
        method_options=options,
    )
    net = Network(neuron)
    with pytest.raises(BrianObjectException) as exc:
        net.run(0 * ms, namespace={})
    assert exc_isinstance(exc, DimensionMismatchError, raise_not_implemented=True)


@pytest.mark.standalone_compatible
@skip_if_not_implemented
def test_GSL_error_dimension_mismatch_dimensionless1():
    eqs = """
    dv/dt = (v0 - v)/(10*ms) : 1
    v0 : 1
    """
    options = {"absolute_error_per_variable": {"v": 1 * mV}}
    neuron = NeuronGroup(
        1, eqs, threshold="v > 10", reset="v = 0", method="gsl", method_options=options
    )
    net = Network(neuron)
    with pytest.raises(BrianObjectException) as exc:
        net.run(0 * ms, namespace={})
    assert exc_isinstance(exc, DimensionMismatchError, raise_not_implemented=True)


@pytest.mark.standalone_compatible
@skip_if_not_implemented
def test_GSL_error_dimension_mismatch_dimensionless2():
    eqs = """
    dv/dt = (v0 - v)/(10*ms) : volt
    v0 : volt
    """
    options = {"absolute_error_per_variable": {"v": 1e-3}}
    neuron = NeuronGroup(
        1,
        eqs,
        threshold="v > 10*mV",
        reset="v = 0*mV",
        method="gsl",
        method_options=options,
    )
    net = Network(neuron)
    with pytest.raises(BrianObjectException) as exc:
        net.run(0 * ms, namespace={})
    assert exc_isinstance(exc, DimensionMismatchError, raise_not_implemented=True)


@pytest.mark.standalone_compatible
@skip_if_not_implemented
def test_GSL_error_nonexisting_variable():
    eqs = """
    dv/dt = (v0 - v)/(10*ms) : volt
    v0 : volt
    """
    options = {"absolute_error_per_variable": {"dummy": 1e-3 * mV}}
    neuron = NeuronGroup(
        1,
        eqs,
        threshold="v > 10*mV",
        reset="v = 0*mV",
        method="gsl",
        method_options=options,
    )
    net = Network(neuron)
    with pytest.raises(BrianObjectException) as exc:
        net.run(0 * ms, namespace={})
    assert exc_isinstance(exc, KeyError, raise_not_implemented=True)


@pytest.mark.standalone_compatible
@skip_if_not_implemented
def test_GSL_error_incorrect_error_format():
    eqs = """
    dv/dt = (v0 - v)/(10*ms) : volt
    v0 : volt
    """
    options = {"absolute_error_per_variable": object()}
    neuron = NeuronGroup(
        1,
        eqs,
        threshold="v > 10*mV",
        reset="v = 0*mV",
        method="gsl",
        method_options=options,
    )
    net = Network(neuron)
    options2 = {"absolute_error": "not a float"}
    neuron2 = NeuronGroup(
        1,
        eqs,
        threshold="v > 10*mV",
        reset="v = 0*mV",
        method="gsl",
        method_options=options2,
    )
    net2 = Network(neuron2)
    with pytest.raises(BrianObjectException) as exc:
        net.run(0 * ms, namespace={})
    assert exc_isinstance(exc, TypeError, raise_not_implemented=True)
    with pytest.raises(BrianObjectException) as exc:
        net2.run(0 * ms, namespace={})
    assert exc_isinstance(exc, TypeError, raise_not_implemented=True)


@pytest.mark.standalone_compatible
@skip_if_not_implemented
def test_GSL_error_nonODE_variable():
    eqs = """
    dv/dt = (v0 - v)/(10*ms) : volt
    v0 : volt
    """
    options = {"absolute_error_per_variable": {"v0": 1e-3 * mV}}
    neuron = NeuronGroup(
        1,
        eqs,
        threshold="v > 10*mV",
        reset="v = 0*mV",
        method="gsl",
        method_options=options,
    )
    net = Network(neuron)
    with pytest.raises(BrianObjectException) as exc:
        net.run(0 * ms, namespace={})
    assert exc_isinstance(exc, KeyError, raise_not_implemented=True)


@pytest.mark.standalone_compatible
@skip_if_not_implemented
def test_GSL_error_bounds():
    runtime = 50 * ms
    error1 = 1e-2 * volt
    error2 = 1e-4 * volt
    error3 = 1e-6 * volt  # default error
    eqs = """
    dv/dt = (stimulus(t) + -v)/(.1*ms) : volt
    """
    stimulus = TimedArray(rand(int(runtime / (10 * ms))) * 3 * volt, dt=5 * ms)
    neuron1 = NeuronGroup(
        1,
        model=eqs,
        reset="v=0*mV",
        threshold="v>10*volt",
        method="gsl",
        method_options={"absolute_error_per_variable": {"v": error1}},
        dt=1 * ms,
    )
    neuron2 = NeuronGroup(
        1,
        model=eqs,
        reset="v=0*mV",
        threshold="v>10*volt",
        method="gsl",
        method_options={"absolute_error_per_variable": {"v": error2}},
        dt=1 * ms,
    )
    neuron3 = NeuronGroup(
        1,
        model=eqs,
        reset="v=0*mV",
        threshold="v>10*volt",
        method="gsl",
        method_options={"absolute_error_per_variable": {}},
        dt=1 * ms,
    )  # Uses default error
    neuron_control = NeuronGroup(1, model=eqs, method="linear", dt=1 * ms)
    mon1 = StateMonitor(neuron1, "v", record=True)
    mon2 = StateMonitor(neuron2, "v", record=True)
    mon3 = StateMonitor(neuron3, "v", record=True)
    mon_control = StateMonitor(neuron_control, "v", record=True)
    run(runtime)
    err1 = abs(mon1.v[0] - mon_control.v[0])
    err2 = abs(mon2.v[0] - mon_control.v[0])
    err3 = abs(mon3.v[0] - mon_control.v[0])
    assert (
        max(err1) < error1
    ), f"Error bound exceeded, error bound: {error1:e}, obtained error: {max(err1):e}"
    assert max(err2) < error2, "Error bound exceeded"
    assert max(err3) < error3, "Error bound exceeded"
    assert max(err1) > max(
        err2
    ), "The simulation with smaller error bound produced a bigger maximum error"
    assert max(err2) > max(
        err3
    ), "The simulation with smaller error bound produced a bigger maximum error"


@pytest.mark.standalone_compatible
@skip_if_not_implemented
def test_GSL_non_autonomous():
    eqs = """dv/dt = sin(2*pi*freq*t)/ms : 1
             freq : Hz"""
    neuron = NeuronGroup(10, eqs, method="gsl")
    neuron.freq = "i*10*Hz + 10*Hz"
    neuron2 = NeuronGroup(10, eqs, method="euler")
    neuron2.freq = "i*10*Hz + 10*Hz"
    mon = StateMonitor(neuron, "v", record=True)
    mon2 = StateMonitor(neuron2, "v", record=True)
    run(20 * ms)
    abs_err = np.abs(mon.v.T - mon2.v.T)
    max_allowed = 1000 * np.finfo(prefs.core.default_float_dtype).eps
    assert np.max(abs_err) < max_allowed


@pytest.mark.standalone_compatible
@skip_if_not_implemented
def test_GSL_non_autonomous():
    eqs = """dv/dt = sin(2*pi*freq*t)/ms : 1
             freq : Hz"""
    neuron = NeuronGroup(10, eqs, method="gsl")
    neuron.freq = "i*10*Hz + 10*Hz"
    neuron2 = NeuronGroup(10, eqs, method="euler")
    neuron2.freq = "i*10*Hz + 10*Hz"
    mon = StateMonitor(neuron, "v", record=True)
    mon2 = StateMonitor(neuron2, "v", record=True)
    run(20 * ms)
    abs_err = np.abs(mon.v.T - mon2.v.T)
    max_allowed = 1000 * np.finfo(prefs.core.default_float_dtype).eps
    assert np.max(abs_err) < max_allowed


@pytest.mark.standalone_compatible
@skip_if_not_implemented
def test_GSL_refractory():
    eqs = """dv/dt = 99.99*Hz : 1 (unless refractory)"""
    neuron = NeuronGroup(
        1, eqs, method="gsl", threshold="v>1", reset="v=0", refractory=3 * ms
    )
    neuron2 = NeuronGroup(
        1, eqs, method="euler", threshold="v>1", reset="v=0", refractory=3 * ms
    )
    mon = SpikeMonitor(neuron, "v")
    mon2 = SpikeMonitor(neuron2, "v")
    run(20 * ms)
    assert mon.count[0] == mon2.count[0]


@skip_if_not_implemented
def test_GSL_save_step_count():
    eqs = """
    dv/dt = -v/(.1*ms) : volt
    """
    neuron = NeuronGroup(
        1, model=eqs, method="gsl", method_options={"save_step_count": True}, dt=1 * ms
    )
    run(0 * ms)
    mon = StateMonitor(neuron, "_step_count", record=True, when="end")
    run(10 * ms)
    assert mon._step_count[0][0] > 0, "Monitor did not save GSL step count"


HH_namespace = {
    "Cm": 1 * ufarad * cm**-2,
    "gl": 5e-5 * siemens * cm**-2,
    "El": -65 * mV,
    "EK": -90 * mV,
    "ENa": 50 * mV,
    "g_na": 100 * msiemens * cm**-2,
    "g_kd": 30 * msiemens * cm**-2,
    "VT": -63 * mV,
}

HH_eqs = Equations(
    """
    dv/dt = (gl*(El-v) - g_na*(m*m*m)*h*(v-ENa) - g_kd*(n*n*n*n)*(v-EK) + I)/Cm : volt
    dm/dt = 0.32*(mV**-1)*(13.*mV-v+VT)/
        (exp((13.*mV-v+VT)/(4.*mV))-1.)/ms*(1-m)-0.28*(mV**-1)*(v-VT-40.*mV)/
        (exp((v-VT-40.*mV)/(5.*mV))-1.)/ms*m : 1
    dn/dt = 0.032*(mV**-1)*(15.*mV-v+VT)/
        (exp((15.*mV-v+VT)/(5.*mV))-1.)/ms*(1.-n)-.5*exp((10.*mV-v+VT)/(40.*mV))/ms*n : 1
    dh/dt = 0.128*exp((17.*mV-v+VT)/(18.*mV))/ms*(1.-h)-4./(1+exp((40.*mV-v+VT)/(5.*mV)))/ms*h : 1
    I : amp/metre**2
    """
)


@pytest.mark.standalone_compatible
@skip_if_not_implemented
def test_GSL_fixed_timestep_big_dt_small_error():
    # should raise integration error
    neuron = NeuronGroup(
        1,
        model=HH_eqs,
        threshold="v > -40*mV",
        refractory="v > -40*mV",
        method="gsl",
        method_options={"adaptable_timestep": False, "absolute_error": 1e-12},
        dt=0.001 * ms,
        namespace=HH_namespace,
    )
    neuron.I = 0.7 * nA / (20000 * umetre**2)
    neuron.v = HH_namespace["El"]
    net = Network(neuron)
    with pytest.raises((BrianObjectException, RuntimeError)):
        net.run(10 * ms)


@pytest.mark.codegen_independent
@skip_if_not_implemented
def test_GSL_internal_variable():
    with pytest.raises(SyntaxError):
        Equations("d_p/dt = 300*Hz : 1")


@pytest.mark.standalone_compatible
@skip_if_not_implemented
def test_GSL_method_options_neurongroup():
    neuron1 = NeuronGroup(
        1,
        model="dp/dt = 300*Hz : 1",
        method="gsl",
        method_options={"adaptable_timestep": True},
    )
    neuron2 = NeuronGroup(
        1,
        model="dp/dt = 300*Hz : 1",
        method="gsl",
        method_options={"adaptable_timestep": False},
    )
    run(0 * ms)
    assert "if (gsl_odeiv2_driver_apply_fixed_step" not in str(
        neuron1.state_updater.codeobj.code
    ), "This neuron should not call gsl_odeiv2_driver_apply_fixed_step()"
    assert "if (gsl_odeiv2_driver_apply_fixed_step" in str(
        neuron2.state_updater.codeobj.code
    ), "This neuron should call gsl_odeiv2_driver_apply_fixed_step()"


@pytest.mark.standalone_compatible
@skip_if_not_implemented
def test_GSL_method_options_spatialneuron():
    morpho = Soma(30 * um)
    eqs = """
    Im = g * v : amp/meter**2
    dg/dt = siemens/metre**2/second : siemens/metre**2
    """
    neuron1 = SpatialNeuron(
        morphology=morpho,
        model=eqs,
        Cm=1 * uF / cm**2,
        Ri=100 * ohm * cm,
        method="gsl_rkf45",
        method_options={"adaptable_timestep": True},
    )
    neuron2 = SpatialNeuron(
        morphology=morpho,
        model=eqs,
        Cm=1 * uF / cm**2,
        Ri=100 * ohm * cm,
        method="gsl_rkf45",
        method_options={"adaptable_timestep": False},
    )
    run(0 * ms)
    assert "if (gsl_odeiv2_driver_apply_fixed_step" not in str(
        neuron1.state_updater.codeobj.code
    ), "This neuron should not call gsl_odeiv2_driver_apply_fixed_step()"
    assert "if (gsl_odeiv2_driver_apply_fixed_step" in str(
        neuron2.state_updater.codeobj.code
    ), "This neuron should call gsl_odeiv2_driver_apply_fixed_step()"


@skip_if_not_implemented
def test_GSL_method_options_synapses():
    N = 1000
    taum = 10 * ms
    taupre = 20 * ms
    taupost = taupre
    Ee = 0 * mV
    vt = -54 * mV
    vr = -60 * mV
    El = -74 * mV
    taue = 5 * ms
    F = 15 * Hz
    gmax = 0.01
    dApre = 0.01
    dApost = -dApre * taupre / taupost * 1.05
    dApost *= gmax
    dApre *= gmax
    eqs_neurons = """
    dv/dt = (ge * (Ee-vr) + El - v) / taum : volt
    dge/dt = -ge / taue : 1
    """
    poisson_input = PoissonGroup(N, rates=F)
    neurons = NeuronGroup(
        1, eqs_neurons, threshold="v>vt", reset="v = vr", method="gsl_rkf45"
    )
    S1 = Synapses(
        poisson_input,
        neurons,
        """
        w : 1
        dApre/dt = -Apre / taupre : 1 (clock-driven)
        dApost/dt = -Apost / taupost : 1 (clock-driven)
        """,
        method="gsl_rkf45",
        method_options={"adaptable_timestep": True},
    )
    S2 = Synapses(
        poisson_input,
        neurons,
        """
        w : 1
        dApre/dt = -Apre / taupre : 1 (clock-driven)
        dApost/dt = -Apost / taupost : 1 (clock-driven)
        """,
        method="gsl_rkf45",
        method_options={"adaptable_timestep": False},
    )
    run(0 * ms)
    assert "if (gsl_odeiv2_driver_apply_fixed_step" not in str(
        S1.state_updater.codeobj.code
    ), "This state_updater should not call gsl_odeiv2_driver_apply_fixed_step()"
    assert "if (gsl_odeiv2_driver_apply_fixed_step" in str(
        S2.state_updater.codeobj.code
    ), "This state_updater should call gsl_odeiv2_driver_apply_fixed_step()"


if __name__ == "__main__":
    from _pytest.outcomes import Skipped

    for t in [
        test_GSL_stateupdater_basic,
        test_GSL_different_clocks,
        test_GSL_default_function,
        test_GSL_user_defined_function,
        test_GSL_x_variable,
        test_GSL_failing_directory,
        test_GSL_stochastic,
        test_GSL_error_dimension_mismatch_unit,
        test_GSL_error_dimension_mismatch_dimensionless1,
        test_GSL_error_dimension_mismatch_dimensionless2,
        test_GSL_error_nonexisting_variable,
        test_GSL_error_incorrect_error_format,
        test_GSL_error_nonODE_variable,
        test_GSL_error_bounds,
        test_GSL_non_autonomous,
        test_GSL_refractory,
        test_GSL_save_step_count,
        test_GSL_fixed_timestep_big_dt_small_error,
        test_GSL_method_options_neurongroup,
        test_GSL_method_options_spatialneuron,
        test_GSL_method_options_synapses,
    ]:
        try:
            t()
        except Skipped as ex:
            print(f"Skipped: {t.__name__} ({str(ex)})")