"""
Check that the features of `Synapses` are available and correct.
"""

import numpy

from brian2 import *
from brian2.tests.features import FeatureTest, InaccuracyError


class SynapsesPre(FeatureTest):
    category = "Synapses"
    name = "Presynaptic code"
    tags = ["NeuronGroup", "run", "Synapses", "Presynaptic code"]

    def run(self):
        tau = 5 * ms
        eqs = """
        dV/dt = k/tau : 1
        k : 1
        """
        G = NeuronGroup(10, eqs, threshold="V>1", reset="V=0")
        G.k = linspace(1, 5, len(G))
        H = NeuronGroup(10, "V:1")
        S = Synapses(G, H, on_pre="V += 1")
        S.connect(j="i")
        self.H = H
        run(101 * ms)

    def results(self):
        return self.H.V[:]

    compare = FeatureTest.compare_arrays


class SynapsesPost(FeatureTest):
    category = "Synapses"
    name = "Postsynaptic code"
    tags = ["NeuronGroup", "run", "Synapses", "Postsynaptic code"]

    def run(self):
        tau = 5 * ms
        eqs = """
        dV/dt = k/tau : 1
        k : 1
        """
        G = NeuronGroup(10, eqs, threshold="V>1", reset="V=0")
        G.k = linspace(1, 5, len(G))
        H = NeuronGroup(10, "V:1")
        S = Synapses(H, G, on_post="V_pre += 1")
        S.connect(j="i")
        self.H = H
        run(101 * ms)

    def results(self):
        return self.H.V[:]

    compare = FeatureTest.compare_arrays


class SynapsesSTDP(FeatureTest):
    category = "Synapses"
    name = "STDP"
    tags = [
        "NeuronGroup",
        "Threshold",
        "Reset",
        "Refractory",
        "run",
        "Synapses",
        "Postsynaptic code",
        "Presynaptic code",
        "SpikeMonitor",
        "StateMonitor",
        "SpikeGeneratorGroup",
    ]

    def run(self):
        n_cells = 100
        n_recorded = 10
        numpy.random.seed(42)
        taum = 20 * ms
        taus = 5 * ms
        Vt = -50 * mV
        Vr = -60 * mV
        El = -49 * mV
        fac = 60 * 0.27 / 10
        gmax = 20 * fac
        dApre = 0.01
        taupre = 20 * ms
        taupost = taupre
        dApost = -dApre * taupre / taupost * 1.05
        dApost *= 0.1 * gmax
        dApre *= 0.1 * gmax

        connectivity = numpy.random.randn(n_cells, n_cells)
        sources = numpy.random.random_integers(0, n_cells - 1, 10 * n_cells)
        # Only use one spike per time step (to rule out that a single source neuron
        # has more than one spike in a time step)
        times = (
            numpy.random.choice(numpy.arange(10 * n_cells), 10 * n_cells, replace=False)
            * ms
        )
        v_init = Vr + numpy.random.rand(n_cells) * (Vt - Vr)

        eqs = Equations(
            """
            dv/dt = (g-(v-El))/taum : volt
            dg/dt = -g/taus         : volt
            """
        )

        P = NeuronGroup(
            n_cells, model=eqs, threshold="v>Vt", reset="v=Vr", refractory=5 * ms
        )
        Q = SpikeGeneratorGroup(n_cells, sources, times)
        P.v = v_init
        P.g = 0 * mV
        S = Synapses(
            P,
            P,
            model="""dApre/dt=-Apre/taupre    : 1 (event-driven)    
                     dApost/dt=-Apost/taupost : 1 (event-driven)
                     w                        : 1""",
            pre="""g     += w*mV
                   Apre  += dApre
                   w      = w + Apost""",
            post="""Apost += dApost
                    w      = w + Apre""",
        )
        S.connect()

        S.w = fac * connectivity.flatten()

        T = Synapses(Q, P, model="w : 1", on_pre="g += w*mV")
        T.connect(j="i")
        T.w = 10 * fac

        spike_mon = SpikeMonitor(P)
        state_mon = StateMonitor(S, "w", record=np.arange(n_recorded))
        v_mon = StateMonitor(P, "v", record=np.arange(n_recorded))

        self.state_mon = state_mon
        self.spike_mon = spike_mon
        self.v_mon = v_mon

        run(0.2 * second, report="text")

    def results(self):
        return self.state_mon.w[:], self.v_mon.v[:], self.spike_mon.num_spikes

    def compare(self, maxrelerr, res1, res2):
        w1, v1, n1 = res1
        w2, v2, n2 = res2
        FeatureTest.compare_arrays(self, maxrelerr, w1, w2)
        FeatureTest.compare_arrays(self, maxrelerr, v1, v2)
        FeatureTest.compare_arrays(
            self, maxrelerr, array([n1], dtype=float), array([n2], dtype=float)
        )


if __name__ == "__main__":
    for ftc in [SynapsesPre, SynapsesPost]:
        ft = ftc()
        ft.run()
        print(ft.results())
