from typing import List, Optional

import matplotlib as mpl

mpl.rcParams["axes.formatter.useoffset"] = False
mpl.rcParams["axes.grid"] = True
mpl.rcParams["grid.color"] = "k"
mpl.rcParams["grid.linestyle"] = ":"
mpl.rcParams["grid.linewidth"] = 0.5
mpl.rcParams["figure.dpi"] = 120
mpl.rcParams["figure.figsize"] = [8., 3.]

import matplotlib.pyplot as plt
import nest
import numpy as np
import os
import random
import re

from pynestml.codegeneration.nest_code_generator_utils import NESTCodeGeneratorUtils
from pynestml.codegeneration.nest_tools import NESTTools

nestml_stdp_dopa_model = """
model neuromodulated_stdp_synapse:
    state:
        w real = 1. [[w >= 0]]
        n real = 0.   # Neuromodulator concentration
        c real = 0.   # Eligibility trace
        pre_tr real = 0.
        post_tr real = 0.

    parameters:
        d ms = 1 ms
        tau_tr_pre ms = 20 ms    # STDP time constant for weight changes caused by pre-before-post spike pairings.
        tau_tr_post ms = 20 ms    # STDP time constant for weight changes caused by post-before-pre spike pairings.
        tau_c ms = 1000 ms    # Time constant of eligibility trace
        tau_n ms = 200 ms   # Time constant of dopaminergic trace
        b real = 0.   # Dopaminergic baseline concentration
        Wmax real = 200. [[Wmax >= 0]]    # Maximal synaptic weight
        Wmin real = 0. [[Wmin >= 0]]    # Minimal synaptic weight
        A_plus real = 1.    # Multiplier applied to weight changes caused by pre-before-post spike pairings. If b (dopamine baseline concentration) is zero, then A_plus is simply the multiplier for facilitation (as in the stdp_synapse model). If b is not zero, then A_plus will be the multiplier for facilitation only if n - b is positive, where n is the instantenous dopamine concentration in the volume transmitter. If n - b is negative, A_plus will be the multiplier for depression.
        A_minus real = 1.5    # Multiplier applied to weight changes caused by post-before-pre spike pairings. If b (dopamine baseline concentration) is zero, then A_minus is simply the multiplier for depression (as in the stdp_synapse model). If b is not zero, then A_minus will be the multiplier for depression only if n - b is positive, where n is the instantenous dopamine concentration in the volume transmitter. If n - b is negative, A_minus will be the multiplier for facilitation.

    equations:
        pre_tr' = -pre_tr / tau_tr_pre
        post_tr' = -post_tr / tau_tr_post

    internals:
        tau_s 1/ms = (tau_c + tau_n) / (tau_c * tau_n)

    input:
        pre_spikes <- spike
        post_spikes <- spike
        mod_spikes <- spike

    output:
        spike

    onReceive(mod_spikes):
        n += 1. / tau_n

    onReceive(post_spikes):
        post_tr += 1.

        # facilitation
        c += A_plus * pre_tr

    onReceive(pre_spikes):
        pre_tr += 1.

        # depression
        c -= A_minus * post_tr

        # deliver spike to postsynaptic partner
        emit_spike()

    # update from time t to t + timestep()
    update:
        # timestep() returns the timestep to be made (in units of time)
        # the sequence here matters: the update step for w requires the "old" values of c and n
        w -= c * (n / tau_s * expm1(-tau_s * timestep()) \
                 - b * tau_c * expm1(-timestep() / tau_c))
        c = c * exp(-timestep() / tau_c)
        n = n * exp(-timestep() / tau_n)

        integrate_odes(pre_tr, post_tr)
"""

module_name, neuron_model_name, synapse_model_name = \
    NESTCodeGeneratorUtils.generate_code_for("./iaf_psc_delta_neuron.nestml",
                                             nestml_stdp_dopa_model,
                                             post_ports=["post_spikes"],
                                             mod_ports=["mod_spikes"],
                                             logging_level="WARNING",
                                             codegen_opts={
                                             "weight_variable": {"neuromodulated_stdp_synapse": "w"},
                                             "delay_variable": {"neuromodulated_stdp_synapse": "d"}})