# -*- coding: utf-8 -*-
"""
Created on 2020/8/5
@project: SPAIC
@filename: Neuron
@author: Hong Chaofei
@contact: hongchf@gmail.com

@description:
定义神经集群和神经元模型。
神经元集群保存神经元数目、神经元编号、类型、模型、模型参数、神经元位置等信息，参与网络构建
"""
# print('__file__={0:<35} | __name__={1:<20} | __package__={2:<20}'.format(__file__,__name__,str(__package__)))
import numpy as np
from ..Network import Assembly
from abc import ABC, abstractmethod
from collections import OrderedDict
from ..Network.BaseModule import VariableAgent
from ..Network.Operator import Op
from ..IO.Initializer import BaseInitializer
from ..IO.Initializer import uniform
import re


# from brian2 import *

class NeuronGroup(Assembly):
    '''Class for a group of neurons.
    '''

    _class_label = '<neg>'
    _is_terminal = True
    def __init__(self, num=None,
                 model=None,
                 shape=None,
                 neuron_type=('excitatory', 'inhibitory', 'pyramidal', '...'),
                 neuron_position='x, y, z',
                 name=None,
                 parameter_variables: list = None,
                 varible_ranges: dict = None,
                 prefer_device=None,
                 **kwargs
                 ):
        super(NeuronGroup, self).__init__(name=name)
        self.set_num_shape(num=num, shape=shape)
        self.outlayer = kwargs.get("outlayer", False)

        # self.model = model
        if neuron_type == ('excitatory', 'inhibitory', 'pyramidal', '...'):
            self.type = ['nontype']
        elif isinstance(neuron_type, list):
            self.type = neuron_type
        else:
            self.type = [neuron_type]

        if neuron_position == 'x, y, z':
            self.position = []
        else:
            neuron_position = np.array(neuron_position)
            assert neuron_position.shape[0] == num, " Neuron_position not equal to neuron number"
            self.position = neuron_position

        self.parameters = kwargs
        if isinstance(model, str):
            self.model_class = NeuronModel.apply_model(model)
            self.model_name = model  # self.model -> self.model_name
            self.model = None
        elif isinstance(model, NeuronModel):
            self.model = model
            self.model_class = None
            self.model_name = 'custom_model'
        else:
            raise ValueError("only support set neuron model with string or NeuronModel class constructed by @custom_model()")
        # variables that can be trained
        if parameter_variables is None:
            self._parameter_var_names = []
        else:
            assert isinstance(parameter_variables, list)
            self._parameter_var_names = parameter_variables
        # variables that needs clamp
        if varible_ranges is None:
            self._variablel_ranges = dict()
        else:
            assert isinstance(varible_ranges, dict)
            self._variablel_ranges = varible_ranges
        self._var_names = list()
        self._var_dict = dict()
        self._operations = OrderedDict()
        self._init_operations = OrderedDict()
        self.prefer_device = prefer_device


    def set_num_shape(self, num, shape):
        self.num = num
        self.shape = shape
        if self.shape is not None:
            num = int(np.prod(self.shape))
            if self.num is None:
                self.num = num
            else:
                assert self.num == num, "the neuron number is not accord with neuron shape"
        elif self.num is not None:
            self.shape = [self.num]
        else:
            raise ValueError("neither neuron number nor neuron shape is defined")

    def set_parameter(self):
        pass

    def get_model(self):
        return self.model


    def add_neuron_label(self, key: str):
        if isinstance(key, str) or isinstance(key, VariableAgent):
            return self._add_label(key)
        elif isinstance(key, list) or isinstance(key, tuple):
            keys = []
            for k in key:
                keys.append(self._add_label(k))
            return keys

    def add_variable_kwargs(self, key):
        t_kwargs = dict()
        if key in self._parameter_var_names:
            t_kwargs['is_parameter'] = True
        if key in self._variablel_ranges:
            min, max = self._variablel_ranges[key]
            if min is not None:
                assert isinstance(min, float)
                t_kwargs['min'] = min
            if max is not None:
                assert isinstance(max, float)
                t_kwargs['max'] = max
        return t_kwargs

    def build(self, backend):
        '''
        Parameters
        ----------
        backend : Backend.Backend
        Returns
        -------
        '''
        self._backend = backend
        prefer_device=self.prefer_device if self.prefer_device != None else None
        # batch_size = self._backend.get_batch_size()
        # if(self.parameters is not None):
        #     self.model = self.model_class(**self.model_parameters)
        # else:
        #     self.model = self.model_class()
        if self.model_class is not None:
            if 'dt' not in self.parameters:
                self.parameters['dt'] = self._backend.dt
            self.parameters['neuron_num'] = self.num
            self.model = self.model_class(**self.parameters)


        dt = backend.dt
        for (key, tau_var) in self.model._tau_variables.items():
            t_kwargs = self.add_variable_kwargs(key)
            tau_var = self.sequence_to_tensor(tau_var)
            self.model._tau_variables[key] = tau_var
            key = self.add_neuron_label(key)
            tau_var = torch.exp(-dt / tau_var)
            # tau_var = np.exp(-dt / tau_var)
            # shape = ()
            if tau_var.numel() > 1:
                shape = list(tau_var.shape)
                assert tau_var.numel() == self.num, "The number of tau should equal the number of neurons."
            else:
                shape = ()
            # if tau_var.size > 1:
            #     shape = tau_var.shape
            #     assert tau_var.size == self.num, "The number of tau should equal the number of neurons."
            # else:
            #     shape = ()

            self.variable_to_backend(key, shape, value=tau_var, prefer_device=prefer_device,  **t_kwargs)

        for (key, membrane_tau_var) in self.model._membrane_variables.items():
            t_kwargs = self.add_variable_kwargs(key)
            membrane_tau_var = self.sequence_to_tensor(membrane_tau_var)
            self.model._membrane_variables[key] = membrane_tau_var
            key = self.add_neuron_label(key)
            # TODO: 计划把membrane tau反过来了变成membrane_tau/dt，需要把用的模型也改一下
            membrane_tau_var = dt/membrane_tau_var
            shape = (1, *self.shape)  # (1, neuron_num)
            self.variable_to_backend(key, shape, value=membrane_tau_var, prefer_device=prefer_device, **t_kwargs)


        for (key, var) in self.model._variables.items():
            # add the rule to extend new dimension before shape (for slif model)
            t_kwargs = self.add_variable_kwargs(key)
            var = self.sequence_to_tensor(var)
            self.model._variables[key] = var
            extend_tag = re.search("\[\d*\]", key)
            if extend_tag is not None:
                extend_tag = int(key[extend_tag.start() + 1:extend_tag.end() - 1])
            key = self.add_neuron_label(key)


            if extend_tag is not None:
                shape = (1, extend_tag, *self.shape)
            else:
                shape = (1, *self.shape)  # (batch_size, shape)
            if hasattr(var, 'shape') and len(var.shape) > 0:
                # adding more dimensions after the neuron shape, assuming neuron'shape in var.shape is 1
                if len(var.shape) > len(shape):
                    var_shape = list(var.shape)
                    for i, s in enumerate(shape):
                        if s != var_shape[i]:
                            assert var_shape[i] == 1
                            var_shape[i] = shape[i]
                    var = var.expand(var_shape)
                    shape = var.shape

                elif var.shape[0] > 1 and var.shape[0]==shape[1]:
                    # for batch dim
                    var = var.unsqueeze(dim=0)


            self.variable_to_backend(key, shape, value=var, prefer_device=prefer_device, **t_kwargs)

        for (key, var) in self.model._parameter_variables.items():
            t_kwargs = self.add_variable_kwargs(key)
            t_kwargs['is_parameter'] = True
            t_kwargs['min'] = 0.0
            t_kwargs['max'] = 1.0
            var = self.sequence_to_tensor(var)
            self.model._parameter_variables[key] = var
            key = self.add_neuron_label(key)

            # is pararmeter by default?
            if isinstance(var, np.ndarray):
                if var.size > 1:
                    shape = var.shape
                    # assert var.size == self.num, "The number of tau should equal the number of neurons."
                else:
                    shape = ()
            elif isinstance(var, torch.Tensor):
                shape = var.shape
            elif hasattr(var, '__iter__'):
                var = np.array(var)
                if var.size > 1:
                    shape = var.shape
                    assert var.size == self.num, "The number of tau should equal the number of neurons."
                else:
                    shape = ()
            else:
                shape = ()

            self.variable_to_backend(key, shape, value=var, prefer_device=prefer_device, **t_kwargs)

        for (key, var) in self.model._constant_variables.items():
            var = self.sequence_to_tensor(var)
            self.model._constant_variables[key] = var

            key = self.add_neuron_label(key)
            shape = None
            self.variable_to_backend(key, shape, value=var, is_constant=True, prefer_device=prefer_device)

        op_count = 0
        for op in self.model._operations:
            addcode_op = Op(owner=self)
            if isinstance(op[1], str):
                op_name = f'{op_count}:{op[1]}'
                addcode_op.func_name = op[1]
                addcode_op.output = self.add_neuron_label(op[0])
                if len(op) > 3:  # 为了解决历史的单一list格式的问题
                    addcode_op.input = self.add_neuron_label(op[2:])
                else:
                    addcode_op.input = self.add_neuron_label(op[2])
                addcode_op.place = prefer_device
                backend.add_operation(addcode_op)
            else:
                op_name = f'{op_count}:custom_function'
                addcode_op.func_name = op[1]
                addcode_op.output = self.add_neuron_label(op[0])
                if len(op) > 3:  # 为了解决历史的单一list格式的问题
                    addcode_op.input = self.add_neuron_label(op[2:])
                else:
                    addcode_op.input = self.add_neuron_label(op[2])
                addcode_op.place = prefer_device
                backend.register_standalone(addcode_op)
            op_count += 1

            self._operations[op_name] = addcode_op

        for op in self.model._init_operations:
            addcode_op = []
            for ind, value in enumerate(op):
                if ind != 1:
                    addcode_op.append(self.add_neuron_label(value))
                else:
                    addcode_op.append(value)
            if len(op) > 3:
                addcode_op[2] = addcode_op[2:]
                addcode_op = addcode_op[:3]
            op_count += 1
            self.init_op_to_backend(addcode_op[0], addcode_op[1], addcode_op[2], prefer_device)
            op_name = str(op_count) + ':' + 'custom_initial_function'
            self._init_operations[op_name] = addcode_op



        if self.model_name == "slif" or self.model_name == 'selif':
            self.model.build((1, *self.shape), backend)
            self.model.outlayer = self.outlayer
            update_code = self.model.update_op_code
            intital_code = self.model.initial_op_code
            self.init_op_to_backend(intital_code[0], intital_code[1], intital_code[2], prefer_device)
            backend.register_standalone(Op(self.add_neuron_label(update_code[0]), update_code[1],
                                          [self.add_neuron_label(update_code[2])], prefer_device, owner=self))
            backend.register_standalone(Op(self.add_neuron_label('V'), self.model.return_V, [], prefer_device, owner=self))
            backend.register_standalone(Op(self.add_neuron_label('S'), self.model.return_S, [], prefer_device, owner=self))


    @staticmethod
    def custom_model(input_vars, output_vars, new_vars_dict, equation_type=('iterative','euler_iterative','exp_euler_iterative','ode'), backend='torch', custom_function_name='custom', base_model=None, add_threshold=True):
        '''
        Examples:
            @NeuronGroup.custom_model(input_vars=['M', 'S', 'WgtSum'], output_vars=['V', 'M', 'S'],
            new_vars_dict={'V':0, 'M':0, 'S':0, 'WgtSum':0}, equation_type='exp_euler_iterative')
            def func(M, S, WgtSum):
                M = (WgtSum-M)/tau
                S = (WgtSum-S)/tau
                V = M - S
                return V, M, S

            NeuronGroup(...., model=func)
        '''
        assert backend == 'torch'
        if base_model is None:
            model = NeuronModel()
        elif isinstance(base_model, NeuronModel):
            model = base_model
        else:
            raise ValueError("base model is given wrong type")
        model.name = custom_function_name
        if equation_type == 'iterative' or equation_type == 'ode':
            for key, value in new_vars_dict.items():
                if '[constant]' in key:
                    model._constant_variables[key.replace('[constant]','')] = value
                else:
                    model._variables[key] = value
        elif equation_type == 'euler_iterative':
            for key, value in new_vars_dict.items():
                if '[constant]' in key:
                    model._constant_variables[key.replace('[constant]','')] = value
                elif 'tau' in key.lower():
                    model._membrane_variables[key] = value
                else:
                    model._variables[key] = value
        elif equation_type == 'exp_euler_iterative':
            for key, value in new_vars_dict.items():
                if '[constant]' in key:
                    model._constant_variables[key.replace('[constant]','')] = value
                elif 'tau' in key.lower():
                    model._tau_variables[key] = value
                else:
                    model._variables[key] = value
        new_vars_dict = dict()
        new_vars_dict.update(model._variables)
        new_vars_dict.update(model._tau_variables)
        new_vars_dict.update(model._membrane_variables)
        new_vars_dict.update(model._constant_variables)
        for var in input_vars:
            if isinstance(var, VariableAgent):
                continue
            elif var not in new_vars_dict:
                if '[updated]' in var:
                    if var.replace('[updated]', '') in new_vars_dict:
                        continue
                else:
                    raise ValueError("The variable %s is not in model variable dict and not a Variable of other modules"%var)
        for var in output_vars:
            if isinstance(var, VariableAgent):
                continue
            elif var not in new_vars_dict:
                if '[updated]' in var:
                    if var.replace('[updated]', '') in new_vars_dict:
                        continue
                else:
                    raise ValueError("The variable %s is not in model variable dict and not a Variable of other modules"%var)

        def model_function(func):
            op_code = [output_vars, func,input_vars ]
            model._operations.append(op_code)
            if add_threshold == True:
                model._operations.append(('O', 'threshold', 'V[updated]', 'Vth'))
            return model
        return model_function

    def sequence_to_tensor(self, var):
        # import collections.abc
        # if isinstance(var, collections.abc.Iterable):
        #     return torch.Tensor(var)
        # return var
        if isinstance(var, BaseInitializer):
            return var
        return torch.tensor(var, dtype=self._backend.data_type, device=self._backend.device0)


class NeuronModel(ABC):
    '''
    op -> (return_name, operation_name, input_name1, input_name2...)
    '''

    #: A dictionary mapping neuron model names to `Model` objects
    neuron_models = dict()

    def __init__(self, **kwargs):
        super(NeuronModel, self).__init__()
        self.name = 'none'
        self._operations = []
        self._init_operations = []
        self._variables = dict()
        # self._tau_constant_variables = dict()
        self._tau_variables = dict()
        self._parameter_variables = dict()
        self._membrane_variables = dict()
        self._constant_variables = dict()
        # self.neuron_parameters = dict()

    @staticmethod
    def register(name, model):
        '''
        Register a neuron model. Registered neuron models can be referred to
        # via their name.
        Parameters
        ----------
        name : str
            A short name for the state updater (e.g. `'lif'`)
        model : `NeuronModel`
            The neuron model object, e.g. an `CLIFModel`, 'SLIFModel'.
        '''

        # only deal with lower case names -- we don't want to have 'LIF' and
        # 'lif', for example
        name = name.lower()
        if name in NeuronModel.neuron_models:
            raise ValueError(('A model with the name "%s" has already been registered') % name)

        if not issubclass(model, NeuronModel):
            raise ValueError(('Given model of type %s does not seem to be a valid NeuronModel.' % str(type(model))))

        NeuronModel.neuron_models[name] = model
        model.name = name

    @staticmethod
    def apply_model(model_name):
        '''
        Parameters
        ----------
        model_name : str
        Returns
        -------
        '''
        model_name = model_name.lower()
        if model_name not in NeuronModel.neuron_models:
            raise ValueError(('Given model name is not in the model list'))
        else:
            return NeuronModel.neuron_models[model_name]



    # @abstractmethod
    # def get_var(self):
    #     NotImplementedError()

    # @abstractmethod
    # def get_op(self):
    #     NotImplementedError()

    # @abstractmethod
    # def get_tau(self):
    #     NotImplementedError()


class CLIFModel(NeuronModel):
    """
    Current LIF 3-kernel model:
    V(t) = M(t) − S(t) − E(t)
    I^n[t] = V0 * Isyn^n[t-1] #sum(w * O^(n-1)[t])
    M^n[t] = betaM * M^n[t-1] + I^n[t-1]
    S^n[t] = betaS * S^n[t-1] + I^n[t-1]
    E^n[t] = betaM * E^n[t-1] + Vth * O^n[t-1]
    O^n[t] = spike_func(V^n[t-1])
    """

    def __init__(self, **kwargs):
        super(CLIFModel, self).__init__()
        # self.neuron_parameters['tau_p'] = kwargs.get('tau_p', 12.0)
        # self.neuron_parameters['tau_q'] = kwargs.get('tau_q', 8.0)
        # self.neuron_parameters['tau_m'] = kwargs.get('tau_m', 20.0)
        # self.neuron_parameters['v_th'] = kwargs.get('v_th', 1.0)

        self._variables['M'] = 0.0
        self._variables['S'] = 0.0
        self._variables['E'] = 0.0
        # self._variables['I_che'] = 0.0
        # self._variables['I_ele'] = 0.0
        self._variables['I'] = 0.0
        self._variables['O'] = 0.0
        self._variables['V'] = 0.0
        self._variables['Isyn'] = 0.0

        self._tau_variables['tauM'] = np.asarray(kwargs.get('tau_m', 20.0))  # self.neuron_parameters['tau_m']
        self._tau_variables['tauP'] = np.asarray(kwargs.get('tau_p', 12.0))  # self.neuron_parameters['tau_p']
        self._tau_variables['tauQ'] = np.asarray(kwargs.get('tau_q', 8.0))  # self.neuron_parameters['tau_q']

        beta = self._tau_variables['tauP'] / self._tau_variables['tauQ']
        V0 = (1 / (beta - 1)) * (beta ** (beta / (beta - 1)))
        self._parameter_variables['V0'] = V0

        self._constant_variables['Vth'] = kwargs.get('v_th', 1.0)

        # self._operations.append(('I_che', 'var_mult', 'V0', 'WgtSum[updated]'))
        # self._operations.append(('I', 'add', 'I_che[updated]', 'I_ele'))
        self._operations.append(('I', 'var_mult', 'V0', 'Isyn[updated]'))
        self._operations.append(('M', 'var_linear', 'tauP', 'M', 'I[updated]'))
        self._operations.append(('S', 'var_linear', 'tauQ', 'S', 'I[updated]'))
        self._operations.append(('PSP', 'minus', 'M[updated]', 'S[updated]'))
        self._operations.append(('V', 'minus', 'PSP', 'E'))
        self._operations.append(('O', 'threshold', 'V[updated]', 'Vth'))
        self._operations.append(('Resetting', 'var_mult', 'Vth', 'O[updated]'))
        self._operations.append(('E', 'var_linear', 'tauM', 'E', 'Resetting'))

NeuronModel.register("clif", CLIFModel)


class CLIF2kModel(NeuronModel):
    """
    Current LIF 2-kernel model:
    V(t) = M(t) − S(t) − E(t)
    I^n[t] = V0 * Isyn^n[t-1] #sum(w * O^(n-1)[t])
    M^n[t] = betaM * M^n[t-1] + I^n[t-1]
    S^n[t] = betaS * S^n[t-1] + I^n[t-1]
    E^n[t] = betaM * E^n[t-1] + Vth * O^n[t-1]
    O^n[t] = spike_func(V^n[t-1])
    """

    def __init__(self, **kwargs):
        super(CLIF2kModel, self).__init__()
        # self.neuron_parameters['tau_p'] = kwargs.get('tau_p', 12.0)
        # self.neuron_parameters['tau_q'] = kwargs.get('tau_q', 8.0)
        # self.neuron_parameters['tau_m'] = kwargs.get('tau_m', 20.0)
        # self.neuron_parameters['v_th'] = kwargs.get('v_th', 1.0)

        self._variables['M'] = 0.0
        self._variables['S'] = 0.0
        self._variables['E'] = 0.0
        # self._variables['I_che'] = 0.0
        # self._variables['I_ele'] = 0.0
        self._variables['I'] = 0.0
        self._variables['O'] = 0.0
        self._variables['V'] = 0.0
        self._variables['Isyn'] = 0.0

        self._tau_variables['tauM'] = np.asarray(kwargs.get('tau_m', 20.0))  # self.neuron_parameters['tau_m']
        self._tau_variables['tauS'] = np.asarray(kwargs.get('tau_s', 12.0))  # self.neuron_parameters['tau_s']

        beta = self._tau_variables['tauM'] / self._tau_variables['tauS']
        V0 = (1 / (beta - 1)) * (beta ** (beta / (beta - 1)))
        self._parameter_variables['V0'] = V0

        self._constant_variables['Vth'] = kwargs.get('v_th', 1.0)

        # self._operations.append(('I_che', 'var_mult', 'V0', 'WgtSum[updated]'))
        # self._operations.append(('I', 'add', 'I_che[updated]', 'I_ele'))
        self._operations.append(('I', 'var_mult', 'V0', 'Isyn[updated]'))
        self._operations.append(('M', 'var_linear', 'tauM', 'M', 'I[updated]'))
        self._operations.append(('S', 'var_linear', 'tauS', 'S', 'I[updated]'))
        self._operations.append(('PSP', 'minus', 'M[updated]', 'S[updated]'))
        self._operations.append(('V', 'minus', 'PSP', 'E'))
        self._operations.append(('O', 'threshold', 'V[updated]', 'Vth'))
        self._operations.append(('Resetting', 'var_mult', 'Vth', 'O[updated]'))
        self._operations.append(('E', 'var_linear', 'tauM', 'E', 'Resetting'))

NeuronModel.register("clif2k", CLIF2kModel)


class PLIF2kModel(NeuronModel):
    """
    Current LIF 2-kernel model:
    V(t) = M(t) − S(t) − E(t)
    I^n[t] = V0 * Isyn^n[t-1] #sum(w * O^(n-1)[t])
    M^n[t] = betaM * M^n[t-1] + I^n[t-1]
    S^n[t] = betaS * S^n[t-1] + I^n[t-1]
    E^n[t] = betaM * E^n[t-1] + Vth * O^n[t-1]
    O^n[t] = spike_func(V^n[t-1])
    """

    def __init__(self, **kwargs):
        super(PLIF2kModel, self).__init__()
        # self.neuron_parameters['tau_p'] = kwargs.get('tau_p', 12.0)
        # self.neuron_parameters['tau_q'] = kwargs.get('tau_q', 8.0)
        # self.neuron_parameters['tau_m'] = kwargs.get('tau_m', 20.0)
        # self.neuron_parameters['v_th'] = kwargs.get('v_th', 1.0)

        self._variables['M'] = 0.0
        self._variables['S'] = 0.0
        self._variables['E'] = 0.0
        # self._variables['I_che'] = 0.0
        # self._variables['I_ele'] = 0.0
        self._variables['I'] = 0.0
        self._variables['O'] = 0.0
        self._variables['V'] = 0.0
        self._variables['Isyn'] = 0.0

        self._parameter_variables['tauM'] = np.asarray(kwargs.get('tau_m', 0.9))  # self.neuron_parameters['tau_m']
        self._parameter_variables['tauS'] = np.asarray(kwargs.get('tau_s', 0.6))  # self.neuron_parameters['tau_s']

        # beta = self._tau_variables['tauM'] / self._tau_variables['tauS']
        # V0 = (1 / (beta - 1)) * (beta ** (beta / (beta - 1)))
        # self._parameter_variables['V0'] = V0

        self._constant_variables['Vth'] = kwargs.get('v_th', 1.0)

        # self._operations.append(('I_che', 'var_mult', 'V0', 'WgtSum[updated]'))
        # self._operations.append(('I', 'add', 'I_che[updated]', 'I_ele'))
        # self._operations.append(('I', 'var_mult', 'V0', 'Isyn[updated]'))
        self._operations.append(('Mtemp', 'add', 'M', 'Isyn[updated]'))
        self._operations.append(('M', 'var_mult', 'tauM', 'Mtemp'))
        self._operations.append(('Stemp', 'add', 'S', 'Isyn[updated]'))
        self._operations.append(('S', 'var_mult', 'tauS', 'Stemp'))
        self._operations.append(('PSP', 'minus', 'M[updated]', 'S[updated]'))
        self._operations.append(('Resetting', 'var_mult', 'Vth', 'O'))
        self._operations.append(('E', 'var_linear', 'tauM', 'E', 'Resetting'))
        self._operations.append(('V', 'minus', 'PSP', 'E[updated]'))
        self._operations.append(('O', 'threshold', 'V[updated]', 'Vth'))


NeuronModel.register("pclif", PLIF2kModel)


class AdaptiveCLIFModel(NeuronModel):
    """
    Current LIF 3-kernel model:
    V(t) = M(t) − S(t) − E(t)
    I^n[t] = V0 * Isyn^n[t-1] #sum(w * O^(n-1)[t])
    M^n[t] = betaM * M^n[t-1] + I^n[t-1]
    S^n[t] = betaS * S^n[t-1] + I^n[t-1]
    E^n[t] = betaM * E^n[t-1] + Vth * O^n[t-1]
    O^n[t] = spike_func(V^n[t-1])
    """

    def __init__(self, **kwargs):
        super(AdaptiveCLIFModel, self).__init__()
        # self.neuron_parameters['tau_p'] = kwargs.get('tau_p', 12.0)
        # self.neuron_parameters['tau_q'] = kwargs.get('tau_q', 8.0)
        # self.neuron_parameters['tau_m'] = kwargs.get('tau_m', 20.0)
        # self.neuron_parameters['v_th'] = kwargs.get('v_th', 1.0)

        self._variables['M'] = 0.0
        self._variables['S'] = 0.0
        self._variables['E'] = 0.0
        # self._variables['I_che'] = 0.0
        # self._variables['I_ele'] = 0.0
        self._variables['I'] = 0.0
        self._variables['O'] = 0.0
        self._variables['V'] = 0.0
        self._variables['Isyn'] = 0.0

        self._tau_variables['tauM'] = np.asarray(kwargs.get('tau_m', 20.0))  # self.neuron_parameters['tau_m']
        self._tau_variables['tauP'] = np.asarray(kwargs.get('tau_p', 12.0))  # self.neuron_parameters['tau_p']
        self._tau_variables['tauQ'] = np.asarray(kwargs.get('tau_q', 8.0))  # self.neuron_parameters['tau_q']
        self._tau_variables['tauA'] = np.asarray(kwargs.get('tau_a', 2000.0))  # self.neuron_parameters['tau_q']

        beta = self._tau_variables['tauP'] / self._tau_variables['tauQ']
        V0 = (1 / (beta - 1)) * (beta ** (beta / (beta - 1)))
        self._parameter_variables['V0'] = V0

        self._constant_variables['Vth'] = kwargs.get('v_th', 1.0)
        self._variables['Vth1[stay]'] = 0 #uniform(0, 1.0)
        self._variables['beta_adaptive'] = kwargs.get('beta_adaptive', 0.1)

        # self._operations.append(('I_che', 'var_mult', 'V0', 'WgtSum[updated]'))
        # self._operations.append(('I', 'add', 'I_che[updated]', 'I_ele'))

        self._operations.append(('Vth1[stay]', 'var_linear', 'tauA', 'Vth1[stay]', 'O'))
        # self._operations.append(('Vth', 'add', 'Vth1[stay][updated]', 'Vth0'))
        self._operations.append(('I', 'var_mult', 'V0', 'Isyn[updated]'))
        self._operations.append(('M', 'var_linear', 'tauP', 'M', 'I[updated]'))
        self._operations.append(('S', 'var_linear', 'tauQ', 'S', 'I[updated]'))
        self._operations.append(('PSP', 'minus', 'M[updated]', 'S[updated]'))
        self._operations.append(('EE', 'var_linear', 'beta_adaptive', 'Vth1[stay][updated]', 'E'))
        self._operations.append(('V', 'minus', 'PSP', 'EE'))
        self._operations.append(('O', 'threshold', 'V[updated]', 'Vth'))
        self._operations.append(('Resetting', 'var_mult', 'Vth', 'O[updated]'))
        self._operations.append(('E', 'var_linear', 'tauM', 'E', 'Resetting'))

NeuronModel.register("aclif", AdaptiveCLIFModel)
class IFModel(NeuronModel):
    """
    IF model:
    V(t) = V(t-1) * (1 - O(t-1)) + Isyn[t] - ConstantDecay

    O^n[t] = spike_func(V^n[t-1])
    """

    def __init__(self, **kwargs):
        super(IFModel, self).__init__()
        # self.neuron_parameters['ConstantDecay'] = kwargs.get('ConstantDecay', 0.0)
        # self.neuron_parameters['v_th'] = kwargs.get('v_th', 1.0)

        self._variables['O'] = 0.0
        self._variables['V'] = 0.0
        self._variables['Isyn'] = 0.0

        self._constant_variables['ConstantDecay'] = kwargs.get('ConstantDecay', 0.0)
        self._constant_variables['Vth'] = kwargs.get('v_th', 1.0)

        self._operations.append(('Vtemp', 'add', 'V', 'Isyn[updated]'))
        self._operations.append(('Vtemp1', 'minus', 'Vtemp', 'ConstantDecay'))
        self._operations.append(('O', 'threshold', 'Vtemp1', 'Vth'))
        # self._operations.append(('Resetting', 'var_mult', 'Vtemp1', 'O[updated]'))
        self._operations.append(('V', 'reset', 'Vtemp1', 'O[updated]'))

NeuronModel.register("if", IFModel)

class IFSoftModel(NeuronModel):
    """
    IF soft reset model:
    V(t) = V(t-1) * (1 - O(t-1)) + Isyn[t] - ConstantDecay
    O^n[t] = spike_func(V^n[t-1])
    """

    def __init__(self, **kwargs):
        super(IFSoftModel, self).__init__()

        self._variables['O'] = 0.0
        self._variables['V'] = 0.0
        self._variables['Isyn'] = 0.0
        self._constant_variables['ConstantDecay'] = kwargs.get('ConstantDecay', 0.0)
        self._constant_variables['Vth'] = kwargs.get('v_th', 1.0)
        self._operations.append(('Vtemp', 'add', 'V', 'Isyn[updated]'))
        self._operations.append(('Vtemp1', 'minus', 'Vtemp', 'ConstantDecay'))
        self._operations.append(('O', 'threshold', 'Vtemp1', 'Vth'))
        self._operations.append(('Resetting', 'var_mult', 'Vtemp1', 'O[updated]'))
        self._operations.append(('V', 'minus', 'Vtemp1', 'Resetting'))

NeuronModel.register("ifsoftreset", IFSoftModel)

class IFBModel(NeuronModel):
    """
    IF model:
    V(t) = V(t-1) * (1 - O(t-1)) + Isyn[t] - ConstantDecay

    O^n[t] = spike_func(V^n[t-1])
    """

    def __init__(self, **kwargs):
        super(IFBModel, self).__init__()
        # self.neuron_parameters['ConstantDecay'] = kwargs.get('ConstantDecay', 0.0)
        # self.neuron_parameters['v_th'] = kwargs.get('v_th', 1.0)

        self._variables['O'] = 0.0
        self._variables['V'] = 0.0
        self._variables['Isyn'] = 0.0

        self._constant_variables['Vth'] = kwargs.get('v_th', 1.0)
        self._constant_variables['bias'] = kwargs.get('bias', 0)

        # self._operations.append(('Vtemp', 'add', 'V', 'Isyn[updated]'))
        self._operations.append(('Itemp', 'add', 'Isyn[updated]', 'bias'))
        self._operations.append(('Vtemp', 'add', 'V', 'Itemp'))
        # self._operations.append(('Vtemp1', 'minus', 'Vtemp', 'ConstantDecay'))
        self._operations.append(('O', 'threshold', 'Vtemp', 'Vth'))
        self._operations.append(('Resetting', 'var_mult', 'Vtemp', 'O[updated]'))
        self._operations.append(('V', 'minus', 'Vtemp', 'Resetting'))


NeuronModel.register("ifb", IFBModel)

class NonSpikingIFBModel(NeuronModel):
    """
    IFB model:
    V(t) = V(t-1) * (1 - O(t-1)) + Isyn[t] - ConstantDecay

    O^n[t] = spike_func(V^n[t-1])
    """

    def __init__(self, **kwargs):
        super(NonSpikingIFBModel, self).__init__()
        # self.neuron_parameters['ConstantDecay'] = kwargs.get('ConstantDecay', 0.0)
        # self.neuron_parameters['v_th'] = kwargs.get('v_th', 1.0)

        self._variables['V'] = 0.0
        self._variables['Isyn'] = 0.0

        self._constant_variables['bias'] = kwargs.get('bias', 0)

        self._operations.append(('Itemp', 'add', 'Isyn[updated]', 'bias'))
        self._operations.append(('V', 'add', 'V', 'Itemp'))

NeuronModel.register("nonspikingifb", NonSpikingIFBModel)

class NullModel(NeuronModel):
    """
    return 'O'
    """

    def __init__(self, **kwargs):
        super(NullModel, self).__init__()
        # self.neuron_parameters['ConstantDecay'] = kwargs.get('ConstantDecay', 0.0)
        # self.neuron_parameters['v_th'] = kwargs.get('v_th', 1.0)
        self._variables['Isyn'] = 0.0
        self._variables['O'] = 0.0

        self._operations.append(('O', 'assign', 'Isyn[updated]'))

NeuronModel.register("null", NullModel)



import torch


class QLIFModel(NeuronModel):
    """
    LIF neuron model for Q-Backprop learning

    E[t] = beta_m*E[t-1] + Vth*O[t-1]
    V[t] = WgtSum(PSP[t]) - E[t]
    O[t] = spike_function(V[t])


    # Q for non-spiking
    Mp[t] = beta_m*Mp[t-1] + WpSum[t]
    Sp[t] = beta_s*Sp[t-1] + WpSum[t]
    P[t] = Mp[t] - Sp[t]

    # Q for spiking
    Mq[t] = beta_m*Mq[t-1] + WqSum[t]
    Sq[t] = beta_s*Sq[t-1] + WqSum[t]
    Q[t] = Mq[t] - Sq[t]

    R[t] = ~O[t]*beta_s*R[t-1] + O[t]
    F[t] = 1 - R[t]
    QP[t] = Q[t] - P[t]
    BQ[t] = R[t]*Q[t] + F[t]*P[t]

    TP[t] = R[t]*P[t] + F[t]*Reward - P[t-1]
    TQ[t] = F[t]*Q[t] + R[t]*Reward - Q[t-1]
    """

    # WgtSum(BQ_post[t - 1]) / WgtSum)
    def __init__(self, **kwargs):
        super(QLIFModel, self).__init__()
        self._tau_variables['Beta_m'] = kwargs.get('tau_m', 20.0)
        self._tau_variables['Beta_s'] = kwargs.get('tau_s', 8.0)
        self._parameter_variables['Vth'] = kwargs.get('v_th', 1.0)
        self._constant_variables['One'] = 1.0

        self._variables['E'] = 0.0
        self._variables['O'] = 0.0
        self._variables['V'] = 0.0
        self._variables['P'] = 0.0
        self._variables['Mp'] = 0.0
        self._variables['Sp'] = 0.0
        self._variables['Q'] = 0.0
        self._variables['Mq'] = 0.0
        self._variables['Sq'] = 0.0
        self._variables['R'] = 0.0
        self._variables['F'] = 0.0
        self._variables['QP'] = 0.0
        self._variables['BQ'] = 0.0
        self._variables['TP'] = 0.0
        self._variables['TQ'] = 0.0
        self._variables['PSP'] = 0.0
        self._variables['WpSum'] = 0.0
        self._variables['WqSum'] = 0.0
        self._variables['Reward'] = 0.0

        self._operations.append(('Resetting', 'var_mult', 'Vth', 'O'))
        self._operations.append(('E', 'var_linear', 'Beta_m', 'E', 'Resetting'))
        self._operations.append(('SumPSP', 'reduce_sum', 'PSP'))
        self._operations.append(('V', 'minus', 'SumPSP', 'E[updated]'))
        self._operations.append(('Mp', 'var_linear', 'Beta_m', 'Mp', 'WpSum'))
        self._operations.append(('Sp', 'var_linear', 'Beta_s', 'Sp', 'WpSum'))
        self._operations.append(('P', 'minus', 'Mp[updated]', 'Sp[updated]'))
        self._operations.append(('Mq', 'var_linear', 'Beta_m', 'Mq', 'WpSum'))
        self._operations.append(('Sq', 'var_linear', 'Beta_s', 'Sq', 'WpSum'))
        self._operations.append(('Q', 'minus', 'Mq[updated]', 'Sq[updated]'))
        self._operations.append(('RTmp', 'var_mult', 'Beta_s', 'R'))
        self._operations.append(('', 'var_linear', 'Beta_s', 'R'))



NeuronModel.register("qlif", QLIFModel)


class SELIFModel(NeuronModel):  #Exponential Model
    """
    SpikeProp LIF 3-kernel model:


    V[t] = WgtSum[t-1] - E[t-1]
    I[t] = spike_func(V[t])
    M[t] = betaM * M[t-1] + I[t]
    S[t] = betaS * S[t-1] + I[t]
    E[t] = betaM * E[t-1] + Vth * I[t]
    O[t] = M[t] − S[t]
    """

    def __init__(self,
                 tau_m=12.0,
                 tau_p=6.0,
                 tau_q=2.0,
                 tau_r=16.0,
                 v_th=1.0,
                 v_reset=2.0,
                 outlayer=False,
                 **kwargs
                 ):
        super(SELIFModel, self).__init__()
        # from spaic.Learning.TRUE_Learner import TRUE_SpikeProp
        # initial value for state variables
        self._variables['[2]O'] = 0.0
        self._variables['V'] = 0.0
        self._variables['dV'] = 0.0
        self._variables['S'] = 0.0
        self._variables['WgtSum'] = 0.0
        self.tau_m = tau_m
        self.tau_e = tau_p
        self.tau_s = tau_q
        self.tau_r = tau_r
        self.v_th = v_th
        self.v_reset = 1.0 * v_th
        self.outlayer = outlayer

        self.beta = tau_p / tau_q
        self.V0 = (1 / (self.beta - 1)) * (self.beta ** (self.beta / (self.beta - 1)))
        # self.delat_m = self.tau_m/(self.tau_m-self.tau_s)
        # self.delat_s = self.tau_s/(self.tau_m-self.tau_s)
        # self.delat_ms = self.tau_m*self.tau_s/(self.tau_m-self.tau_s)

        self.update_op_code = ('[2]O', self.update, 'WgtSum[updated]')  # 'WgtSum[updated]'

        self.return_op_code = (None, self.return_V, [])

        self.initial_op_code = (None, self.initial, [])

        '''
            V(t) = M(t) − S(t) − E(t)
            I^n[t] = V0 * WgtSum^n[t-1] #sum(w * O^(n-1)[t])
            M^n[t] = betaM * M^n[t-1] + I^n[t-1]
            S^n[t] = betaS * S^n[t-1] + I^n[t-1]
            E^n[t] = betaM * E^n[t-1] + Vth * O^n[t-1]

            O^n[t] = spike_func(V^n[t-1])
        '''

    def attach_learner(self, learner):
        self.learner = learner

    def build(self, shape, backend):
        self.dt = backend.dt
        self.M_initial = torch.zeros(shape, device=backend.device)
        self.S_initial = torch.zeros(shape, device=backend.device)
        self.R_initial = torch.zeros(shape, device=backend.device)
        self.V_initial = torch.zeros(shape, device=backend.device)
        self.beta_m = np.exp(-backend.dt / self.tau_m)
        self.beta_s = np.exp(-backend.dt / self.tau_s)
        self.beta_e = np.exp(-backend.dt / self.tau_e)
        self.beta_r = np.exp(-backend.dt / self.tau_r)
        self.running_var = None
        self.running_mean = None
        self.decay = 0.9999
        self.initial()
        self.rec_E = []

    def initial(self):
        self.M = self.M_initial
        self.S = self.S_initial
        self.R = self.R_initial
        self.V = self.V_initial
        self.O = None
        self.rec_E = []

    def norm_hook(self, grad):
        if self.running_var is None:
            self.running_var = torch.norm(grad, dim=0) * 0
            self.running_mean = torch.mean(grad, dim=0) * 0
        else:
            self.running_var = self.decay * self.running_var + (1 - self.decay) * torch.norm(grad, dim=0)
            self.running_mean = self.decay * self.running_mean + (1 - self.decay) * torch.mean(grad, dim=0)
        return (grad - self.running_mean) / (1.0e-10 + self.running_var)

    def update(self, WgtSum):
        # with torch.no_grad():
        #     self.dV = self.E / self.tau_m +self.S / self.tau_s - self.M / self.tau_m
        I = self.V0 * WgtSum
        # WgtSum.register_hook(self.norm_hook)

        if I.dim() == self.M.dim() + 1:
            Ii = I[:, 0, ...]
            I0 = I[:, 1, ...]
            self.M = self.beta_e * self.M + (Ii - I0 / self.tau_e)
            self.S = self.beta_s * self.S + (Ii - I0 / self.tau_s)
        else:
            self.M = self.beta_e * self.M + I
            self.S = self.beta_s * self.S + I

        if self.O is not None:
            Oi = self.O[:, 0, ...] #+ 0.9*self.O[:, 0, ...].detach()  # *self.O[:, 0, ...].gt(0.0)
            Ot = self.O[:, 1, ...] #+ 0.9*self.O[:, 1, ...].detach()
        else:
            Oi = 0.0
            Ot = 0.0


        # expv = torch.clamp_max(torch.exp(2.5 * self.V)-1, 12)
        expv = 2.0*torch.pow(torch.clamp(self.V,-10,self.v_th), 2.0)
        self.R = self.beta_r*self.R + (1-self.beta_r)*self.V + 10.0*(Oi-Ot/self.tau_r)
        self.dV = (expv - self.V + (self.M - self.S) + 0.5-0.5*self.R) / self.tau_m
        # with torch.no_grad():
        #     self.ddV = (self.dV * (0.2 * expv - 1) + self.S / self.tau_s - self.M / self.tau_e) / self.tau_m
        #     self.dV2 = self.dV + self.ddV * self.dt
        self.V = self.V + self.dV * self.dt - self.v_reset*(Oi-Ot/self.tau_m)# + self.ddV * self.dt ** 2 / 2.0
        # self.P = self.M + self.S


        self.O = self.learner.threshold(self.V, self.dV, self.v_th)
        #
        # self.rec_E.append(self.Vmax)
        # if (I is not None) and (I.requires_grad == True):
        #     I.retain_grad()
        #     self.rec_E.append(I)

        return self.O

    def return_V(self):
        return self.V

    def return_M(self):
        return self.M

    def return_S(self):
        return self.S

    def return_dV(self):
        return self.dV

    @property
    def E_values(self):
        return torch.stack(self.rec_E, dim=-1).cpu().detach().numpy()

    @property
    def E_grads(self):
        grads = []
        for v in self.rec_E:
            if v.grad is not None:
                grads.append(v.grad.cpu().numpy())
            else:
                grads.append(torch.zeros_like(v).cpu().numpy())
        grads = np.stack(grads[1:], axis=-1)
        return grads
NeuronModel.register("selif", SELIFModel)

class SLIFModel(NeuronModel):
    """
    SpikeProp LIF 3-kernel model:


    V[t] = WgtSum[t-1] - E[t-1]
    I[t] = spike_func(V[t])
    M[t] = betaM * M[t-1] + I[t]
    S[t] = betaS * S[t-1] + I[t]
    E[t] = betaM * E[t-1] + Vth * I[t]
    O[t] = M[t] − S[t]
    """

    def __init__(self,
                 tau_m=20.0,
                 tau_p=20.0,
                 tau_q=8.0,
                 v_th=1.0,
                 v_reset=2.0,
                 outlayer=False,
                 ):
        super(SLIFModel, self).__init__()
        # from spaic.Learning.TRUE_Learner import TRUE_SpikeProp
        # initial value for state variables
        self._variables['[2]O'] = 0.0
        self._variables['V'] = 0.0
        self._variables['dV'] = 0.0
        self._variables['S'] = 0.0
        self._variables['Isyn'] = 0.0
        self.tau_m = tau_m
        self.tau_e = tau_p
        self.tau_s = tau_q
        self.v_th = v_th
        self.v_reset = 5.0 * v_th
        self.outlayer = outlayer

        self.beta = tau_m / tau_q
        self.V0 = (1 / (self.beta - 1)) * (self.beta ** (self.beta / (self.beta - 1)))
        # self.delat_m = self.tau_m/(self.tau_m-self.tau_s)
        # self.delat_s = self.tau_s/(self.tau_m-self.tau_s)
        # self.delat_ms = self.tau_m*self.tau_s/(self.tau_m-self.tau_s)

        self.update_op_code = ('[2]O', self.update, 'Isyn[updated]')  # 'WgtSum[updated]'

        self.return_op_code = (None, self.return_V, [])

        self.initial_op_code = (None, self.initial, [])

        '''
            V(t) = M(t) − S(t) − E(t)
            I^n[t] = V0 * WgtSum^n[t-1] #sum(w * O^(n-1)[t])
            M^n[t] = betaM * M^n[t-1] + I^n[t-1]
            S^n[t] = betaS * S^n[t-1] + I^n[t-1]
            E^n[t] = betaM * E^n[t-1] + Vth * O^n[t-1]

            O^n[t] = spike_func(V^n[t-1])
        '''

    def attach_learner(self, learner):
        self.learner = learner

    def build(self, shape, backend):
        self.dt = backend.dt
        self.M_initial = torch.zeros(shape, device=backend.device)
        self.S_initial = torch.zeros(shape, device=backend.device)
        self.E_initial = torch.zeros(shape, device=backend.device)
        self.V_initial = torch.zeros(shape, device=backend.device)
        self.O_initial = torch.zeros((1,2,1), device=backend.device)
        self.beta_m = np.exp(-backend.dt / self.tau_m)
        self.beta_s = np.exp(-backend.dt / self.tau_s)
        self.beta_e = np.exp(-backend.dt / self.tau_e)
        self.deta_m = (1 - self.dt/(2*self.tau_m))/self.tau_m
        self.deta_s = (1 - self.dt/(2*self.tau_s))/self.tau_s
        self.running_var = None
        self.running_mean = None
        self.decay = 0.9999
        self.initial()
        self.rec_E = []

    def initial(self):
        self.M = self.M_initial
        self.S = self.S_initial
        self.E = self.E_initial
        self.V = self.V_initial
        self.O = None
        self.rec_E = []

    def norm_hook(self, grad):
        if self.running_var is None:
            self.running_var = torch.norm(grad, dim=0) * 0
            self.running_mean = torch.mean(grad, dim=0) * 0
        else:
            self.running_var = self.decay * self.running_var + (1 - self.decay) * torch.norm(grad, dim=0)
            self.running_mean = self.decay * self.running_mean + (1 - self.decay) * torch.mean(grad, dim=0)
        return (grad - self.running_mean) / (1.0e-10 + self.running_var)

    def update(self, WgtSum):
        # with torch.no_grad():
        #     self.dV = self.E / self.tau_m +self.S / self.tau_s - self.M / self.tau_m
        I = self.V0 * WgtSum
        # WgtSum.register_hook(self.norm_hook)
        # Oi = self.O[:, 0, ...]
        # Ot = self.O[:, 1, ...]
        if self.O is not None:
            Oi = 0.1*self.O[:, 0, ...] + 0.9*self.O[:, 0, ...].detach()
            Ot = 0.1*self.O[:, 1, ...] + 0.9*self.O[:, 1, ...].detach()
        else:
            Oi = 0
            Ot = 0

        if I.dim() == self.M.dim() + 1:
            Ii = I[:, 0, ...]
            I0 = I[:, 1, ...]
            self.M = self.beta_m * self.M + Ii - I0 / self.tau_m - self.v_reset*(Oi-Ot/self.tau_m)
            self.S = self.beta_s * self.S + Ii - I0 / self.tau_s
        else:
            self.M = self.beta_m * self.M + I - self.v_reset*(Oi-Ot/self.tau_m)
            self.S = self.beta_s * self.S + I




        self.V = self.M - self.S
        # self.P = self.M + self.S
        self.dV = self.S*self.deta_s - self.M*self.deta_m

        #
        # if self.O is not None:
        #     self.E = self.E*Oi.lt(1.0).float()# + (0.5*self.E.detach() - 0.5*self.E)*Oi.ge(1.0).float()
        # MSbase = torch.clamp_min(self.M*self.tau_s/(self.S*self.tau_m+1.0e-20), 0)
        # self.Vmax = self.M*MSbase**self.delat_s - self.S*MSbase**self.delat_m

        self.O = self.learner.threshold(self.V, self.dV, self.v_th)
        #
        # self.rec_E.append(self.Vmax)
        # if (I is not None) and (I.requires_grad == True):
        #     I.retain_grad()
        #     self.rec_E.append(I)

        return self.O

    def return_V(self):
        return self.V

    def return_M(self):
        return self.M

    def return_S(self):
        return self.S

    def return_dV(self):
        return self.dV

    @property
    def E_values(self):
        return torch.stack(self.rec_E, dim=-1).cpu().detach().numpy()

    @property
    def E_grads(self):
        grads = []
        for v in self.rec_E:
            if v.grad is not None:
                grads.append(v.grad.cpu().numpy())
            else:
                grads.append(torch.zeros_like(v).cpu().numpy())
        grads = np.stack(grads[1:], axis=-1)
        return grads
# for ii in range(2):
#     out.append(test.update(I))

NeuronModel.register("slif", SLIFModel)


class SELIFDebugModel(NeuronModel):

    def __init__(self,
                 tau_m=20.0, tau_p=20.0, tau_q=8.0, v_th=1.0,
                 outlayer=False
                 ):
        super(SELIFDebugModel, self).__init__()
        # from spaic.Learning.TRUE_Learner import TRUE_SpikeProp
        # initial value for state variables
        self._variables['[2]O'] = 0.0
        self._variables['V'] = 0.0
        self._variables['dV'] = 0.0
        self._variables['S'] = 0.0
        self._variables['WgtSum'] = 0.0
        self._variables['cumV'] = 0.0
        self.tau_m = tau_m
        self.tau_e = tau_p
        self.tau_s = tau_q
        self.v_th = v_th
        self.v_reset = 5.0 * v_th
        self.outlayer = outlayer

        self.beta = tau_m / tau_q
        self.V0 = (1 / (self.beta - 1)) * (self.beta ** (self.beta / (self.beta - 1)))


        self.update_op_code = ('[2]O', self.update, 'WgtSum[updated]')  # 'WgtSum[updated]'
        self.return_op_code = (None, self.return_V, [])
        self.initial_op_code = (None, self.initial, [])

    def attach_learner(self, learner):
        self.learner = learner

    def build(self, shape, backend):
        self.dt = backend.dt
        self.M_initial = torch.zeros(shape, device=backend.device)
        self.S_initial = torch.zeros(shape, device=backend.device)
        self.V_initial = torch.zeros(shape, device=backend.device)
        self.beta_m = np.exp(-backend.dt / self.tau_m)
        self.beta_s = np.exp(-backend.dt / self.tau_s)
        self.beta_e = np.exp(-backend.dt / self.tau_e)
        self.initial()
        self.rec_E = []

    def initial(self):
        self.M = self.M_initial
        self.S = self.S_initial
        self.V = self.V_initial
        self.cumV = self.V_initial
        self.O = None
        self.rec_E = []

    def update(self, WgtSum):

        I = self.V0 * WgtSum

        if I.dim() == self.M.dim() + 1:
            Ii = I[:, 0, ...]
            I0 = I[:, 1, ...]
            self.M = self.beta_e * self.M + (Ii - I0 / self.tau_e)
            self.S = self.beta_s * self.S + (Ii - I0 / self.tau_s)
        else:
            self.M = self.beta_e * self.M + I
            self.S = self.beta_s * self.S + I

        if self.O is not None:
            Oi = self.O[:, 0, ...] #+ 0.9*self.O[:, 0, ...].detach()  # *self.O[:, 0, ...].gt(0.0)
            Ot = self.O[:, 1, ...] #+ 0.9*self.O[:, 1, ...].detach()
        else:
            Oi = 0.0
            Ot = 0.0


        # expv = torch.clamp_max(torch.exp(2.5 * self.V)-1, 12)
        expv = 2.0*torch.pow(torch.clamp(self.V,-10,self.v_th), 2.0)
        self.dV = (expv - self.V + (self.M - self.S)) / self.tau_m
        self.V = self.V + self.dV * self.dt - self.v_reset*(Oi-Ot/self.tau_m)



        self.O = self.learner.threshold(self.V, self.dV, self.v_th)

        return self.O

    def return_V(self):
        return self.V

class LIFModel(NeuronModel):
    """
    LIF model:
    # V(t) = tuaM * V^n[t-1] + Isyn[t]   # tauM: constant membrane time (tauM=RmCm)
    O^n[t] = spike_func(V^n[t-1])
    """

    def __init__(self, **kwargs):
        super(LIFModel, self).__init__()
        # initial value for state variables
        self._variables['V'] = 0.0
        self._variables['O'] = 0.0
        self._variables['Isyn'] = 0.0


        self._constant_variables['Vth'] = kwargs.get('v_th', 1.0)
        self._constant_variables['Vreset'] = kwargs.get('v_reset', 0.0)

        self._tau_variables['tauM'] = kwargs.get('tau_m', 8.0)

        self._operations.append(('Vtemp', 'var_linear', 'tauM', 'V', 'Isyn[updated]'))
        self._operations.append(('O', 'threshold', 'Vtemp', 'Vth'))
        self._operations.append(('V', 'reset', 'Vtemp',  'O[updated]'))

NeuronModel.register("lif", LIFModel)


class PLIFModel(NeuronModel):
    """
    LIF model:
    # V(t) = tuaM * V^n[t-1] + Isyn[t]   # tauM: constant membrane time (tauM=RmCm)
    O^n[t] = spike_func(V^n[t-1])
    """

    def __init__(self, **kwargs):
        super(PLIFModel, self).__init__()
        # initial value for state variables

        # self.neuron_parameters['tau_m'] = kwargs.get('tau_m', 20.0)  # 20
        # self.neuron_parameters['v_th'] = kwargs.get('v_th', 1)
        # self.neuron_parameters['v_reset'] = kwargs.get('v_reset', 0.0)

        self._variables['V'] = 0.0
        self._variables['O'] = 0.0
        self._variables['Isyn'] = 0.0
        # self._variables['b'] = 0.0
        # self._variables['I_che'] = 0.0
        # self._variables['I_ele'] = 0.0
        # self._variables['I'] = 0.0

        self._constant_variables['Vth'] = kwargs.get('v_th', 1)
        self._constant_variables['Vreset'] = kwargs.get('v_reset', 0.0)


        import math
        tau = kwargs.get('tau_m', 2.0)  # dt/taum
        tau = -math.log(tau - 1.)
        self._parameter_variables['tau'] = tau  # dt/taum

        # self.mem = self.mem + ((inputs - self.mem) * self.w.sigmoid()) * self.dt
        self._operations.append(('tau_temp', 'sigmoid', 'tau'))
        self._operations.append(('Vtemp', 'minus', 'Isyn[updated]', 'V'))
        self._operations.append(('Vtemp1', 'var_linear', 'tau_temp', 'Vtemp', 'V'))
        # self._operations.append(('Vtemp', 'var_linear', 'tauM', 'V', 'Isyn[updated]'))
        self._operations.append(('O', 'threshold', 'Vtemp1', 'Vth'))
        self._operations.append(('V', 'reset', 'Vtemp1',  'O'))
        # self._operations.append(('tau_temp', 'sigmoid', 'tau'))
        # self._operations.append(('Vtemp', 'var_linear', 'tau_temp', 'V', 'Isyn[updated]'))
        # self._operations.append(('O', 'threshold', 'Vtemp', 'Vth'))
        # self._operations.append(('V', 'reset', 'Vtemp', 'O'))

NeuronModel.register("plif", PLIFModel)


# class linearDecayLIFModel(NeuronModel):
#     """
#     LIF model:
#     # V(t) = tuaM * V^n[t-1] + Isyn[t]   # tauM: constant membrane time (tauM=RmCm)
#     O^n[t] = spike_func(V^n[t-1])
#     """
#
#     def __init__(self, **kwargs):
#         super(ConstantDecayLIFModel, self).__init__()
#         # initial value for state variables
#         self._variables['V'] = 0.0
#         self._variables['O'] = 0.0
#         self._variables['Isyn'] = 0.0
#
#         self._parameter_variables['Vth'] = kwargs.get('v_th', 1)
#         self._constant_variables['Vreset'] = kwargs.get('v_reset', 0.0)
#         self._constant_variables['Vdecay'] = kwargs.get('v_decay', 1.0)
#
#         self._operations.append(('Vtemp', 'var_linear', 'tauM', 'V', 'Isyn[updated]'))
#         self._operations.append(('O', 'threshold', 'Vtemp', 'Vth'))
#         self._operations.append(('Resetting', 'var_mult', 'Vtemp', 'O[updated]'))
#         self._operations.append(('V', 'minus', 'Vtemp', 'Resetting'))
#
# NeuronModel.register("constantdecaylif", ConstantDecayLIFModel)

class ConstantCurrentLIFModel(NeuronModel):
    """
    ConstantCurrentLIF model:
    V(t) = V^n[t-1] + (dt/taum) * (Ureset-V^n[t-1]+I)  # tauM: constant membrane time (tauM=RmCm)  Isyn = I*Weight
    O^n[t] = spike_func(V^n[t-1])
    """

    def __init__(self, **kwargs):
        super(ConstantCurrentLIFModel, self).__init__()
        # initial value for state variables
        # self.neuron_parameters['tau_m'] = kwargs.get('tau_m', 20.0)
        # self.neuron_parameters['v_th'] = kwargs.get('v_th', 1.0)
        # self.neuron_parameters['v_reset'] = kwargs.get('v_reset', 0.0)

        self._variables['V'] = 0.0
        self._variables['O'] = 0.0
        self._variables['I'] = 0.0
        self._variables['Isyn'] = 0.0
        # beta = self.neuron_parameters['tau_p'] / self.neuron_parameters['tau_q']
        # V0 = 1.0  # (1 / (beta - 1)) * (beta ** (beta / (beta - 1)))
        # self._constant_variables['V0'] = V0

        self._constant_variables['Vth'] = kwargs.get('v_th', 1.0)
        self._constant_variables['Vreset'] = kwargs.get('v_reset', 0.0)

        self._membrane_variables['tauM'] = kwargs.get('tau_m', 20.0)

        # self._operations.append(('I', 'var_mult', 'V0', 'I_synapse'))
        self._operations.append(('decayV', 'minus', 'Isyn', 'V'))
        self._operations.append(('Vtemp', 'var_linear', 'tauM', 'decayV', 'V'))
        self._operations.append(('O', 'threshold', 'Vtemp', 'Vth'))
        self._operations.append(('Resetting', 'var_mult', 'Vtemp', 'O'))
        self._operations.append(('V', 'minus', 'Vtemp', 'Resetting'))


NeuronModel.register("constantcurrentlif", ConstantCurrentLIFModel)


class NonSpikingLIFModel(NeuronModel):
    """
    NonSpikingLIF model:
    # V(t) = -tuaM * V^n[t-1] + M^n[t] - S^n[t]   # tauM: constant membrane time (tauM=RmCm)
    V(t) = V^n[t-1] + (dt/taum) * (PSP-V^n[t-1])  # tauM: constant membrane time (tauM=RmCm)
    I^n[t] = V0 * Isyn^n[t-1]                # sum(w * O^(n-1)[t])
    M^n[t] = tauP * M^n[t-1] + I^n[t-1]        # tauP: decaying time constants of membrane integration
    S^n[t] = tauQ * S^n[t-1] + I^n[t-1]        # tauQ: decaying time constants of synaptic currents
    PSP = M - S
    """

    def __init__(self, **kwargs):
        super(NonSpikingLIFModel, self).__init__()
        # initial value for state variables
        # self.neuron_parameters['tau_p'] = kwargs.get('tau_p', 4.0)
        # self.neuron_parameters['tau_q'] = kwargs.get('tau_q', 1.0)
        # self.neuron_parameters['tau_m'] = kwargs.get('tau_m', 1.0)

        self._variables['I'] = 0.0
        self._variables['M'] = 0.0
        self._variables['S'] = 0.0
        self._variables['V'] = 0.0
        self._variables['O'] = 0.0
        self._variables['Isyn'] = 0.0

        self._tau_variables['tauP'] = np.asarray(kwargs.get('tau_p', 4.0))
        self._tau_variables['tauQ'] = np.asarray(kwargs.get('tau_q', 1.0))
        beta = self._tau_variables['tauP'] / self._tau_variables['tauQ']
        V0 = (1 / (beta - 1)) * (beta ** (beta / (beta - 1)))
        self._constant_variables['V0'] = V0

        self._membrane_variables['tauM'] = kwargs.get('tau_m', 1.0)
        # self._tau_constant_variables['tauM'] = kwargs.get('tau_m', 1.0)



        self._operations.append(('I', 'var_mult', 'V0', 'Isyn'))
        self._operations.append(('M', 'var_linear', 'tauP', 'M', 'I'))
        self._operations.append(('S', 'var_linear', 'tauQ', 'S', 'I'))
        self._operations.append(('PSP', 'minus', 'M', 'S'))
        self._operations.append(('decayV', 'minus', 'PSP', 'V'))
        self._operations.append(('V', 'var_linear', 'tauM', 'decayV', 'V'))


NeuronModel.register("nonspikinglif", NonSpikingLIFModel)


class LIFMModel(NeuronModel):
    """
    LIF model:
    # I_che = tauP*I + Isyn^n[t-1] + b^n                         # sum(w * O^(n-1)[t])
    # I = I_che + I_ele
    # F = tauM * exp(-O^n[t-1] / tauM)
    # V(t) = V^n[t-1] * F + I
    # O^(n)[t] = spike_func(V^n(t))
    """

    def __init__(self, **kwargs):
        super(LIFMModel, self).__init__()
        # initial value for state variables
        # self.neuron_parameters['tau_p'] = kwargs.get('tau_p', 1.0)
        # self.neuron_parameters['tau_m'] = kwargs.get('tau_m', 10.0)
        # self.neuron_parameters['v_th']  = kwargs.get('v_th', 1.0)

        self._variables['V'] = 0.0
        self._variables['O'] = 0.0
        self._variables['Isyn'] = 0.0
        # self._variables['b'] = 0.0
        # self._variables['I_che'] = 0.0
        # self._variables['I_ele'] = 0.0
        self._variables['I'] = 0.0

        self._constant_variables['Vth'] = kwargs.get('v_th', 1.0)
        # self._constant_variables['Vreset'] = v_reset

        self._tau_variables['tauM'] = kwargs.get('tau_m', 10.0)
        self._tau_variables['tauP'] = kwargs.get('tau_p', 1.0)

        self._operations.append(('I', 'var_linear', 'tauP', 'I', 'Isyn[updated]'))

        # self._operations.append(('I_che', 'var_linear', 'tauP', 'I_che', 'PSP'))
        # self._operations.append(('I', 'add', 'I_che[updated]', 'I_ele'))
        self._operations.append(('Vtemp', 'var_linear', 'V', 'tauM', 'I[updated]'))
        self._operations.append(('O', 'threshold', 'Vtemp', 'Vth'))
        self._operations.append(('Vreset', 'var_mult', 'Vtemp', 'O[updated]'))
        self._operations.append(('V', 'minus', 'Vtemp', 'Vreset'))


NeuronModel.register("lifm", LIFMModel)


class IZHModel(NeuronModel):
    """
    IZH model:

    .. math::
        V = V + dt / tauM * (C1 * V * V + C2 * V + C3 - U + I)  # tauM=1 此处加tauM是为了添加op时和LIF模型保存一致 \\
        V = V + dt / tauM * (V* (C1 * V + C2) + C3 - U + I)     # 由上式拆分而来 \\
        U = U + a. * (b. * V - U) \\

        O^n[t] = spike\_func(V^n[t-1])

        if V > Vth, \\
        then V = C, U = U + d

    References:
        Izhikevich, E. M. (2003). Simple model of spiking neurons. IEEE Transactions on neural networks, 14(6), 1569-1572.
    """

    def __init__(self, **kwargs):
        super(IZHModel, self).__init__()

        self._variables['I'] = 0.0
        self._variables['O'] = 0.0
        self._variables['V'] = -65.0
        self._variables['U'] = 1.0
        self._variables['Isyn'] = 0.0  # 1.8

        self._constant_variables['a'] = kwargs.get('a', 0.02)
        self._constant_variables['b'] = kwargs.get('b', 0.2)
        self._constant_variables['Vth'] = np.asarray(kwargs.get('v_th', 30))# 30.0
        self._constant_variables['Vreset'] = np.asarray(kwargs.get('Vreset', -65.0))
        self._constant_variables['d'] = kwargs.get('d', 8.0)  # 8.0
        self._constant_variables['C1'] = 0.04
        self._constant_variables['C2'] = 5
        self._constant_variables['C3'] = 140

        self._membrane_variables['tauM'] = 1.0

        # V = V + dt / tauM * (C1 * V * V + C2 * V + C3 - U + I)
        # V = V + dt / tauM * (V* (C1 * V + C2) + C3 - U + I)
        # U = U + dt /tauM * a. * (b. * V - U)

        self._operations.append(('temp_V1', 'var_linear', 'C1', 'V', 'C2'))
        self._operations.append(('temp_V2', 'var_linear', 'temp_V1', 'V', 'C3'))
        self._operations.append(('temp_V3', 'minus', 'temp_V2', 'U'))
        self._operations.append(('temp_V4', 'add', 'temp_V3', 'Isyn[updated]'))
        self._operations.append(('temp_V', 'var_linear', 'tauM', 'temp_V4', 'V'))

        self._operations.append(('temp_U1', 'var_mult', 'b', 'V'))
        self._operations.append(('temp_U2', 'minus', 'temp_U1', 'U'))
        self._operations.append(('temp_U', 'var_linear', 'a', 'temp_U2', 'U'))

        self._operations.append(('O', 'threshold', 'temp_V', 'Vth'))

        # if V > Vth,
        # then V <- C, U <- U + d
        self._operations.append(('VResetting', 'minus', 'Vreset', 'temp_V'))
        self._operations.append(('V', 'var_linear', 'O[updated]', 'VResetting', 'temp_V'))
        self._operations.append(('U', 'var_linear', 'd', 'O[updated]', 'temp_U'))


NeuronModel.register("izh", IZHModel)


class aEIFModel(NeuronModel):
    """
    aEIF model:

    .. math::
        V = V + dt / tauM * (EL - V + EXP - U + I^n[t]) \\
        U = U + dt / tauW * (a * (V - EL) - U) \\
        EXP = delta\_t * delta\_t2 * exp(dv\_th/delta\_t2) \\
        dv = V - EL \\
        dv\_th = V - Vth

        O^n[t] = spike\_func(V^n[t-1]) \\

        If V > 20: \\
        then V = EL, U = U + b

    References:
        Brette, R., & Gerstner, W. (2005). Adaptive exponential integrate-and-fire model as an
        effective description of neuronal activity. Journal of neurophysiology, 94(5), 3637-3642.
    """

    def __init__(self, **kwargs):
        super(aEIFModel, self).__init__()

        self._variables['I'] = 0.0

        self._variables['O'] = 0.0
        self._variables['V'] = -70.6
        # self._variables['Vt'] = -70.6
        self._variables['U'] = 0.0
        self._variables['Isyn'] = 0.0
        self._variables['EXP'] = 0.0

        self._constant_variables['EL'] = kwargs.get('EL', -70.6)
        self._constant_variables['a'] = kwargs.get('a', 4) #0.05
        self._constant_variables['b'] = kwargs.get('b', 0.0805)
        self._constant_variables['Vth'] = kwargs.get('v_th', -50.4)
        self._constant_variables['Vspk'] = kwargs.get('Vspk', 20)

        self._constant_variables['delta_t'] = kwargs.get('delta_t', 2.0)

        self._membrane_variables['C'] = kwargs.get('C', 281)
        self._membrane_variables['tauW'] = kwargs.get('tau_w', 144) #144

        self._constant_variables['gL'] = kwargs.get('gL', 30)

        # V = V + dt / tauM * (gL * (EL - V + EXP) - U + I ^ n[t])
        # U = U + dt / tauW * (a * (V - EL) - U)
        # EXP = delta_t * exp(du_th / delta_t2)
        # du = V - EL
        # du_th = V - Vth
        # I ^ n[t] = V0 * WgtSum ^ n[t - 1]
        # O ^ n[t] = spike_func(V ^ n[t - 1])

        self._operations.append(('dv', 'minus', 'V', 'EL'))
        self._operations.append(('dv_th', 'minus', 'V', 'Vth'))

        self._operations.append(('EXP_T1', 'div', 'dv_th', 'delta_t'))
        self._operations.append(('EXP_T2', 'exp', 'EXP_T1'))
        self._operations.append(('EXP', 'var_mult', 'delta_t', 'EXP_T2'))

        self._operations.append(('temp_V0', 'minus', 'EXP[updated]', 'dv'))
        self._operations.append(('temp_V1', 'var_mult', 'gL', 'temp_V0'))
        self._operations.append(('temp_V2', 'minus', 'temp_V1', 'U'))
        self._operations.append(('temp_V3', 'add', 'temp_V2', 'Isyn[updated]'))
        self._operations.append(('Vt', 'var_linear', 'C', 'temp_V3', 'V'))

        self._operations.append(('temp_U1', 'var_mult', 'a', 'dv'))
        self._operations.append(('temp_U2', 'minus', 'temp_U1', 'U'))
        self._operations.append(('Ut', 'var_linear', 'tauW', 'temp_U2', 'U'))

        self._operations.append(('O', 'threshold', 'Vt', 'Vspk'))

        self._operations.append(('Vtemp2', 'var_mult', 'Vt', 'O[updated]'))
        self._operations.append(('Vtemp3', 'minus', 'Vt', 'Vtemp2'))
        self._operations.append(('V', 'var_linear', 'EL', 'O[updated]', 'Vtemp3'))
        self._operations.append(('U', 'var_linear', 'b', 'O[updated]', 'Ut'))


NeuronModel.register("aeif", aEIFModel)
NeuronModel.register("adex", aEIFModel)


class GLIFModel(NeuronModel):
    """
    Current GLIF5 model:

        V = V + dt / C * (I + I1 + I2 - (V - E_L) / R)
        Theta_s = Theta_s - dt * b_s * Theta_s
        I_j = I_j - dt * k_j * I_j (j = 1,2)
        Theta_v = Theta_v + dt * (a_v * (V - E_L) - b_v * Theta_v)

        v_th = Theta_v + Theta_s + Theta_inf
        O = spike_func(V)

        Reset function:

        V = E_L + f_v * (V - E_L) - delta_v
        Theta_s = Theta_s + delta_Theta_s
        I_j = f_j * I_j + delta_I_j (j = 1, 2)
        Theta_v = Theta_v

    References:
        Teeter, C., Iyer, R., Menon, V., Gouwens, N., Feng, D., Berg, J., ... & Mihalas, S. (2018). Generalized leaky
        integrate-and-fire models classify multiple neuron types. Nature communications, 9(1), 1-15.
    """

    def __init__(self, **kwargs):
        super(GLIFModel, self).__init__()

        self._variables['V']       = 0.0
        self._variables['Theta_s'] = 0.1
        self._variables['Theta_v'] = 0.2
        self._variables['I1']      = 0.08
        self._variables['I2']      = 0.06
        self._variables['I']       = 0.0
        self._variables['S']       = 0.0
        self._variables['O']       = 0.0
        self._variables['Vth']    = 1.0
        self._variables['Isyn'] = 0.0

        # self._constant_variables['deltaV'] = kwargs.get('E_L', 0.0) - kwargs.get('delta_v', 0.05)
        self._membrane_variables['C'] = kwargs.get('C', 60.0)
        self._membrane_variables['b_s'] = -1 / kwargs.get('b_s', 0.02)
        self._membrane_variables['k1'] = -1 / kwargs.get('k_1', 0.8)
        self._membrane_variables['k2'] = -1 / kwargs.get('k_2', 0.6)

        self._membrane_variables['a_v'] = 1 / kwargs.get('a_v', 0.12)
        self._membrane_variables['b_v'] = -1 / kwargs.get('b_v', 0.3)

        self._constant_variables['R'] = -1 / kwargs.get('R', 2)
        self._constant_variables['E_L'] = kwargs.get('E_L', 0.0)

        self._constant_variables['Theta_inf'] = kwargs.get('Theta_inf', 1.0)

        self._constant_variables['f_v'] = kwargs.get('f_v', 0.1)
        self._constant_variables['delta_v'] = kwargs.get('delta_v', 0.05)

        self._constant_variables['delta_Theta_s'] = kwargs.get('delta_Theta_s', 0.2)

        self._constant_variables['f1'] = kwargs.get('f1', 1.0)
        self._constant_variables['f2'] = kwargs.get('f2', 1.0)
        self._constant_variables['delta_I1'] = kwargs.get('delta_I1', 0.1)
        self._constant_variables['delta_I2'] = kwargs.get('delta_I2', 0.1)

        # I_j = I_j - dt * k_j * I_j
        self._operations.append(('I1temp', 'var_linear', 'k1', 'I1', 'I1'))
        self._operations.append(('I2temp', 'var_linear', 'k2', 'I2', 'I2'))

        # I_sum = I + I1 + I2
        self._operations.append(('I_sum1', 'add', 'I1temp', 'I2temp'))
        self._operations.append(('I_sum', 'add', 'I_sum1', 'Isyn[updated]'))

        # dv = V - E_L
        self._operations.append(('dv', 'minus', 'V', 'E_L'))

        # V_up = I + I1 + I2 - (V - E_L) / R
        self._operations.append(('V_up', 'var_linear', 'dv', 'R', 'I_sum'))

        # V = V + dt / C * (I + I1 + I2 - (V - E_L) / R)
        self._operations.append(('Vtemp', 'var_linear', 'C', 'V_up', 'V'))

        # Theta_s = Theta_s - dt * b_s * Theta_s
        self._operations.append(('Theta_s_temp', 'var_linear', 'b_s', 'Theta_s', 'Theta_s'))

        # Theta_v = Theta_v + dt * (a_v * (V - E_L) - b_v * Theta_v)
        # Theta_temp =Theta_v + dt * a_v * (V - E_L)
        # Theta_v = Theta_temp - dt * b_v * Theta_v

        self._operations.append(('Theta_temp', 'var_linear', 'a_v', 'dv', 'Theta_v'))
        self._operations.append(('Theta_v', 'var_linear', 'b_v', 'Theta_v', 'Theta_temp'))

        # v_th = Theta_v + Theta_s + Theta_inf
        self._operations.append(('Theta', 'add', 'Theta_v[updated]', 'Theta_s_temp'))
        self._operations.append(('Vth', 'add', 'Theta', 'Theta_inf'))

        self._operations.append(('O', 'threshold', 'Vtemp', 'Vth[updated]'))

        # V = E_L + f_v * (V - E_L) - delta_v
        self._operations.append(('dv_reset', 'minus', 'Vtemp', 'E_L'))
        self._operations.append(('Vreset1', 'var_linear', 'f_v', 'dv_reset', 'E_L'))
        self._operations.append(('Vreset2', 'minus', 'Vreset1', 'delta_v'))
        self._operations.append(('Vreset', 'minus', 'Vreset2', 'Vtemp'))

        self._operations.append(('V', 'var_linear', 'Vreset', 'O[updated]', 'Vtemp'))

        # Theta_s = Theta_s + delta_Theta_s
        self._operations.append(('Theta_s', 'var_linear', 'delta_Theta_s', 'O[updated]', 'Theta_s_temp'))

        # I_j = f_j * I_j + delta_I_j (j = 1, 2)
        self._operations.append(('I1_temp1', 'var_linear', 'f1', 'I1temp', 'delta_I1'))
        self._operations.append(('I1_temp2', 'minus', 'I1_temp1', 'I1temp'))
        self._operations.append(('I1', 'var_linear', 'I1_temp1', 'O[updated]', 'I1temp'))
        self._operations.append(('I2_temp1', 'var_linear', 'f2', 'I2temp', 'delta_I2'))
        self._operations.append(('I2_temp2', 'minus', 'I2_temp1', 'I2temp'))
        self._operations.append(('I2', 'var_linear', 'I2_temp1', 'O[updated]', 'I2temp'))


NeuronModel.register("glif", GLIFModel)


class HodgkinHuxleyModel(NeuronModel):
    """
    Hodgkin-Huxley model:

    V = V + dt/tau_v * (I - Ik)
    Ik = NA + K + L
    NA = g_NA * m^3 * h * (V - V_NA)
    K = g_K * n^4 * (V - V_K)
    L = g_L * (V - V_L)

    Na activation:
    m = m + dt/tau_m * (alpha_m * (1-m) - beta_m * m)

    K activation:
    n = n + dt/tau_n * (alpha_n * (1-n) - beta_n * n)

    Na inactivation:
    h = h + dt/tau_h * (alpha_h * (1-h) - beta_h * h)

    original function1:
    alpha_m = 0.1 * (-V + 25) / (exp((-V+25)/10) - 1)
    beta_m = 4 * exp(-V/18)
    alpha_n = 0.01 * (-V + 10) / (exp((-V+10)/10) - 1)
    beta_n = 0.125 * exp(-V/80)
    alpha_h = 0.07 * exp(-V/20)
    beta_h = 1/(exp((-V+30)/10) + 1)

    O^n[t] = spike_func(V^n[t-1])
    """

    def __init__(self, **kwargs):
        super(HodgkinHuxleyModel, self).__init__()

        self._variables['I'] = 0.0
        self._variables['O'] = 0.0
        self._variables['V'] = kwargs.get('V', 0.0)

        self._variables['m'] = kwargs.get('m', 0.5)
        self._variables['n'] = kwargs.get('n', 0.5)
        self._variables['h'] = kwargs.get('h', 0.06)

        self._variables['Isyn'] = 0.0

        self._constant_variables['Vth'] = kwargs.get('v_th', 1.0)
        # self._constant_variables['Vreset'] = 0.0

        self._membrane_variables['tauM'] = 1.0
        self._membrane_variables['tauN'] = 1.0
        self._membrane_variables['tauH'] = 1.0
        self._membrane_variables['tauV'] = 1.0

        self._constant_variables['1'] = 1.0
        # self._constant_variables['65'] = kwargs.get('V65', 0.0)
        self._constant_variables['Vreset'] = kwargs.get('Vreset', 0.0)
        self._constant_variables['dt'] = kwargs.get('dt', 0.1)

        self._constant_variables['g_NA'] = kwargs.get('g_NA', 120.0)
        self._constant_variables['g_K'] = kwargs.get('g_K', 36.0)
        self._constant_variables['g_L'] = kwargs.get('g_L', 0.3)

        self._constant_variables['E_NA'] = kwargs.get('E_NA', 120.0)
        self._constant_variables['E_K'] = kwargs.get('E_K', -12.0)
        self._constant_variables['E_L'] = kwargs.get('E_L', 10.6)

        self._constant_variables['a_m1'] = kwargs.get('alpha_m1', 0.1)
        self._constant_variables['a_m2'] = kwargs.get('alpha_m2', 25.0)
        self._constant_variables['a_m3'] = kwargs.get('alpha_m3', 10.0)
        self._constant_variables['b_m1'] = kwargs.get('beta_m1', 4.0)
        self._constant_variables['b_m2'] = kwargs.get('beta_m2', 0.0)
        self._constant_variables['b_m3'] = kwargs.get('beta_m3', 18.0)

        self._constant_variables['a_n1'] = kwargs.get('alpha_n1', 0.01)
        self._constant_variables['a_n2'] = kwargs.get('alpha_n2', 10.0)
        self._constant_variables['a_n3'] = kwargs.get('alpha_n3', 10.0)
        self._constant_variables['b_n1'] = kwargs.get('beta_n1', 0.125)
        self._constant_variables['b_n2'] = kwargs.get('beta_n2', 0.0)
        self._constant_variables['b_n3'] = kwargs.get('beta_n3', 80.0)

        self._constant_variables['a_h1'] = kwargs.get('alpha_h1', 0.07)
        self._constant_variables['a_h2'] = kwargs.get('alpha_h2', 0.0)
        self._constant_variables['a_h3'] = kwargs.get('alpha_h3', 20.0)
        self._constant_variables['b_h1'] = kwargs.get('beta_h1', 1.0)
        self._constant_variables['b_h2'] = kwargs.get('beta_h2', 30.0)
        self._constant_variables['b_h3'] = kwargs.get('beta_h3', 10.0)

        self._operations.append(('V65', 'add', 'V', 'Vreset'))

        # a_m = a_m1 * (-V + a_m2) / (exp((-V + a_m2) / a_m3) - 1)
        # b_m = b_m1 * exp((-V + b_m2) / b_m3)
        # alpha_m
        self._operations.append(('Vam', 'minus', 'a_m2', 'V65'))
        self._operations.append(('Vamd', 'div', 'Vam', 'a_m3'))
        self._operations.append(('expVamd1', 'exp', 'Vamd'))
        self._operations.append(('expVamd', 'minus', 'expVamd1', '1'))
        self._operations.append(('amtemp', 'div', 'Vam', 'expVamd'))
        self._operations.append(('a_m', 'var_mult', 'a_m1', 'amtemp'))

        # beta_m
        self._operations.append(('Vbm', 'minus', 'b_m2', 'V65'))
        self._operations.append(('Vbmd', 'div', 'Vbm', 'b_m3'))
        self._operations.append(('expVbmd', 'exp', 'Vbmd'))
        self._operations.append(('b_m', 'var_mult', 'b_m1', 'expVbmd'))

        # a_n = a_n1 * (-V + a_n2) / (exp((-V + a_n2) / a_n3) - 1)
        # b_n = b_n1 * exp((-V + b_n2) / b_n3)
        # alpha_n
        self._operations.append(('Van', 'minus', 'a_n2', 'V65'))
        self._operations.append(('Vand', 'div', 'Van', 'a_n3'))
        self._operations.append(('expVand1', 'exp', 'Vand'))
        self._operations.append(('expVand', 'minus', 'expVand1', '1'))
        self._operations.append(('antemp', 'div', 'Van', 'expVand'))
        self._operations.append(('a_n', 'var_mult', 'a_n1', 'antemp'))

        # beta_n
        self._operations.append(('Vbn', 'minus', 'b_n2', 'V65'))
        self._operations.append(('Vbnd', 'div', 'Vbn', 'b_n3'))
        self._operations.append(('expVbnd', 'exp', 'Vbnd'))
        self._operations.append(('b_n', 'var_mult', 'b_n1', 'expVbnd'))

        # a_h = a_h1 * exp((-V + a_h2) / a_h3)
        # b_h = b_h1 / (exp((-V + b_h2) / b_h3 + 1)
        # alpha_h
        self._operations.append(('Vah', 'minus', 'a_h2', 'V65'))
        self._operations.append(('Vahd', 'div', 'Vah', 'a_h3'))
        self._operations.append(('expVahd', 'exp', 'Vahd'))
        self._operations.append(('a_h', 'var_mult', 'a_h1', 'expVahd'))

        # beta_h
        self._operations.append(('Vbh', 'minus', 'b_h2', 'V65'))
        self._operations.append(('Vbhd', 'div', 'Vbh', 'b_h3'))
        self._operations.append(('expVbhd1', 'exp', 'Vbhd'))
        self._operations.append(('expVbhd', 'add', 'expVbhd1', '1'))
        self._operations.append(('b_h', 'div', 'b_h1', 'expVbhd'))

        # m = m + alpha_m * (1 - m) - beta_m * m
        # n = n + alpha_n * (1 - n) - beta_n * n
        # h = h + alpha_h * (1 - h) - beta_h * h

        # m
        self._operations.append(('mtemp1', 'minus', '1', 'm'))
        self._operations.append(('mtemp2', 'var_mult', 'a_m', 'mtemp1'))
        self._operations.append(('betam', 'var_mult', 'b_m', 'm'))
        self._operations.append(('mtemp3', 'minus', 'mtemp2', 'betam'))
        self._operations.append(('mtemp4', 'var_mult', 'tauM', 'mtemp3'))
        self._operations.append(('m', 'add', 'mtemp4', 'm'))
        # n
        self._operations.append(('ntemp1', 'minus', '1', 'n'))
        self._operations.append(('ntemp2', 'var_mult', 'a_n', 'ntemp1'))
        self._operations.append(('betan', 'var_mult', 'b_n', 'n'))
        self._operations.append(('ntemp3', 'minus', 'ntemp2', 'betan'))
        self._operations.append(('ntemp4', 'var_mult', 'tauN', 'ntemp3'))
        self._operations.append(('n', 'add', 'ntemp4', 'n'))
        # h
        self._operations.append(('htemp1', 'minus', '1', 'h'))
        self._operations.append(('htemp2', 'var_mult', 'a_h', 'htemp1'))
        self._operations.append(('betah', 'var_mult', 'b_h', 'h'))
        self._operations.append(('htemp3', 'minus', 'htemp2', 'betah'))
        self._operations.append(('htemp4', 'var_mult', 'tauH', 'htemp3'))
        self._operations.append(('h', 'add', 'htemp4', 'h'))

        # g_NAm3h
        self._operations.append(('m2', 'var_mult', 'm[updated]', 'm[updated]'))
        self._operations.append(('m3', 'var_mult', 'm2', 'm[updated]'))
        self._operations.append(('m3h', 'var_mult', 'm3', 'h[updated]'))
        self._operations.append(('g_NAm3h', 'var_mult', 'g_NA', 'm3h'))

        # g_Kn4
        self._operations.append(('n2', 'var_mult', 'n[updated]', 'n[updated]'))
        self._operations.append(('n4', 'var_mult', 'n2', 'n2'))
        self._operations.append(('g_Kn4', 'var_mult', 'g_K', 'n4'))

        self._operations.append(('d_NA', 'minus', 'V', 'E_NA'))
        self._operations.append(('d_K', 'minus', 'V', 'E_K'))
        self._operations.append(('d_L', 'minus', 'V', 'E_L'))

        # Ik, NA, K, L
        self._operations.append(('NA', 'var_mult', 'g_NAm3h', 'd_NA'))
        self._operations.append(('K', 'var_mult', 'g_Kn4', 'd_K'))
        self._operations.append(('L', 'var_mult', 'g_L', 'd_L'))

        self._operations.append(('Ik1', 'add', 'NA', 'K'))
        self._operations.append(('Ik2', 'add', 'Ik1', 'L'))
        self._operations.append(('Ik', 'var_mult', 'tauV', 'Ik2'))

        # I
        self._operations.append(('Vtemp', 'minus', 'Isyn[updated]', 'Ik'))
        self._operations.append(('V', 'add', 'V', 'Vtemp'))

        # O
        self._operations.append(('O', 'threshold', 'V[updated]', 'Vth'))


NeuronModel.register("hh", HodgkinHuxleyModel)


class LIFSTDPEXModel(NeuronModel):
    """
    LIF model:
    V(t) = decay_v * (v - v_rest) + v_rest + I^n[t]
    I^n[t] = V0 * Isyn^n[t]  #V0 = 1
    theta(t) = decay_th * theta[t-1]
    if v >= (vth + theta) then s_out = 1; else s_out = 0;
    Reset:
    V(t) = s_out * v_reset + (1 - s_out) * v; theta = theta + s_out * th_inc
    O^n[t] = spike_func(V^n[t-1])
    """
    def __init__(self, **kwargs):
        super(LIFSTDPEXModel, self).__init__()
        # initial value for state variables
        # self.neuron_parameters['decay_v'] = kwargs.get('decay_v', np.exp(-1/100))
        # self.neuron_parameters['decay_th'] = kwargs.get('decay_th', np.exp(-1/1e7))
        # self.neuron_parameters['th_inc'] = kwargs.get('th_inc', 0.05)
        # self.neuron_parameters['v_th'] = kwargs.get('v_th', -52.0)
        # self.neuron_parameters['v_rest'] = kwargs.get('v_rest', -65.0)
        # self.neuron_parameters['v_reset'] = kwargs.get('v_reset', -60.0)

        self._variables['I'] = 0.0
        self._variables['V'] = -65.0
        self._variables['O'] = 0.0
        self._variables['Isyn'] = 0.0
        self._variables['theta[stay]'] = 0.0
        self._variables['Vth_theta'] = 0.0

        # self._constant_variables['V0'] = 1

        self._constant_variables['Vth'] = kwargs.get('v_th', -52.0)
        self._constant_variables['Vreset'] = kwargs.get('v_reset', -60.0)
        self._constant_variables['Vrest'] = kwargs.get('v_rest', -65.0)
        self._constant_variables['th_inc'] = kwargs.get('th_inc', 0.05)
        self._constant_variables['decay_th'] = kwargs.get('decay_th', np.exp(-1/1e7))
        self._constant_variables['decay_v'] = kwargs.get('decay_v', np.exp(-1/100))


        # self._operations.append(('I', 'var_mult', 'V0', 'I_synapse[updated]'))
        self._operations.append(('PSP1', 'minus', 'V', 'Vrest'))
        self._operations.append(('PSP2', 'var_linear', 'decay_v', 'PSP1', 'Vrest'))
        self._operations.append(('Vtemp', 'add', 'PSP2', 'Isyn[updated]'))
        self._operations.append(('theta_temp', 'var_mult', 'decay_th', 'theta[stay]'))
        self._operations.append(('Vth_theta', 'add', 'Vth', 'theta_temp'))
        self._operations.append(('O', 'threshold', 'Vtemp', 'Vth_theta'))
        self._operations.append(('Resetting1', 'var_mult', 'Vreset', 'O[updated]'))
        self._operations.append(('Resetting2', 'var_mult', 'Vtemp', 'O[updated]'))
        self._operations.append(('Resetting3', 'minus', 'Vtemp', 'Resetting2'))
        self._operations.append(('V', 'add', 'Resetting1', 'Resetting3'))
        self._operations.append(('Resetting_theta', 'var_mult', 'O[updated]', 'th_inc'))
        self._operations.append(('theta[stay]', 'add', 'theta_temp', 'Resetting_theta'))


NeuronModel.register("lifstdp_ex", LIFSTDPEXModel)

class ALIFSTDPEXModel(NeuronModel):
    """
    LIF model:
    V(t) = decay_v * (v - v_rest) + v_rest + I^n[t]
    I^n[t] = V0 * Isyn^n[t]  #V0 = 1
    theta(t) = decay_th * theta[t-1]
    if v >= (vth + theta) then s_out = 1; else s_out = 0;
    Reset:
    V(t) = s_out * v_reset + (1 - s_out) * v; theta = theta + s_out * th_inc

    O^n[t] = spike_func(V^n[t-1])
    """
    def __init__(self, **kwargs):
        super(ALIFSTDPEXModel, self).__init__()
        # initial value for state variables
        # self.neuron_parameters['decay_v'] = kwargs.get('decay_v', np.exp(-1/100))
        # self.neuron_parameters['decay_th'] = kwargs.get('decay_th', np.exp(-1/1e7))
        # self.neuron_parameters['th_inc'] = kwargs.get('th_inc', 0.05)
        # self.neuron_parameters['v_th'] = kwargs.get('v_th', -52.0)
        # self.neuron_parameters['v_rest'] = kwargs.get('v_rest', -65.0)
        # self.neuron_parameters['v_reset'] = kwargs.get('v_reset', -60.0)

        self._variables['I'] = 0.0
        self._variables['V'] = 0.0
        self._variables['O'] = 0.0
        self._variables['Isyn'] = 0.0
        self._variables['theta[stay]'] = 0.0
        self._variables['Vth_theta'] = 2.0

        # self._constant_variables['V0'] = 1

        self._constant_variables['Vth'] = kwargs.get('v_th', 2.0)
        self._constant_variables['Vreset'] = kwargs.get('v_reset', 2.0)
        self._constant_variables['Vrest'] = kwargs.get('v_rest', 0.0)
        self._constant_variables['th_inc'] = kwargs.get('th_inc', 0.001)
        self._constant_variables['decay_th'] = kwargs.get('decay_th', np.exp(-1/100000.0))
        self._constant_variables['decay_v'] = kwargs.get('decay_v', np.exp(-1/100.0))


        # self._operations.append(('I', 'var_mult', 'V0', 'I_synapse[updated]'))
        self._operations.append(('PSP1', 'minus', 'V', 'Vrest'))
        self._operations.append(('PSP2', 'var_linear', 'decay_v', 'PSP1', 'Vrest'))
        self._operations.append(('Vtemp', 'add', 'PSP2', 'Isyn[updated]'))
        self._operations.append(('theta_temp', 'var_mult', 'decay_th', 'theta[stay]'))
        self._operations.append(('Vth_theta', 'add', 'Vth', 'theta_temp'))
        self._operations.append(('O', 'threshold', 'Vtemp', 'Vth_theta'))
        self._operations.append(('Resetting1', 'var_mult', 'Vreset', 'O[updated]'))
        self._operations.append(('V', 'minus', 'Vtemp', 'Resetting1'))
        self._operations.append(('Resetting_theta1', 'var_mult', 'O[updated]', 'th_inc'))
        self._operations.append(('Resetting_theta', 'var_mult', 'Resetting_theta1', 'Vth_theta'))
        self._operations.append(('theta[stay]', 'add', 'theta_temp', 'Resetting_theta'))


NeuronModel.register("alifstdp_ex", ALIFSTDPEXModel)

class LIFSTDPIHModel(NeuronModel):
    """
    LIF model:
    V(t) = decay_v * (v - v_rest) + v_rest + I^n[t]
    I^n[t] = V0 * Isyn^n[t]  #V0 = 1

    Reset:
    V(t) = s_out * v_reset + (1 - s_out) * v;

    O^n[t] = spike_func(V^n[t-1])
    """
    def __init__(self, **kwargs):
        super(LIFSTDPIHModel, self).__init__()
        # initial value for state variables
        # self.neuron_parameters['decay_v'] = kwargs.get('decay_v', np.exp(-1/10))
        # self.neuron_parameters['v_th'] = kwargs.get('v_th', -40.0)
        # self.neuron_parameters['v_rest'] = kwargs.get('v_rest', -60.0)
        # self.neuron_parameters['v_reset'] = kwargs.get('v_reset', -45.0)

        self._variables['I'] = 0.0
        self._variables['V'] = -60.0
        self._variables['O'] = 0.0
        self._variables['Isyn'] = 0.0

        # self._constant_variables['V0'] = 1

        self._constant_variables['Vth'] = kwargs.get('v_th', -40.0)
        self._constant_variables['Vreset'] = kwargs.get('v_reset', -45.0)
        self._constant_variables['Vrest'] = kwargs.get('v_rest', -60.0)
        self._constant_variables['decay_v'] = kwargs.get('decay_v', np.exp(-1/10))


        # self._operations.append(('I', 'var_mult', 'V0', 'I_synapse[updated]'))
        self._operations.append(('PSP1', 'minus', 'V', 'Vrest'))
        self._operations.append(('PSP2', 'var_linear', 'decay_v', 'PSP1', 'Vrest'))
        self._operations.append(('Vtemp', 'add', 'PSP2', 'Isyn[updated]'))
        self._operations.append(('O', 'threshold', 'Vtemp', 'Vth'))
        self._operations.append(('Resetting1', 'var_mult', 'Vreset', 'O[updated]'))
        self._operations.append(('Resetting2', 'var_mult', 'Vtemp', 'O[updated]'))
        self._operations.append(('Resetting3', 'minus', 'Vtemp', 'Resetting2'))
        self._operations.append(('V', 'add', 'Resetting1', 'Resetting3'))


NeuronModel.register("lifstdp_ih", LIFSTDPIHModel)


class DiehlAndCookModelInt(NeuronModel):
    """
    Layer of leaky integrate-and-fire (LIF) neurons with adaptive thresholds (modified for Diehl & Cook 2015
    replication), based on integer arithmetic:
    if thresh > max_threshold: thresh = thresh/2
    # Decay voltages.
    V(t) = ((self.decay * (v - self.rest)) >> shift_voltage) + self.rest  # decay = 1013;  v_rest = 0;  shift_voltage = 10
    # Integrate inputs.
    V(t) = V(t) + (self.refrac_count <= 0).int() * ((Isyn^n[t] * self.weight_sum_coef) >> self.shift_weight_sum_coef)
    refrac_count=0;  weight_sum_coef = 154;  shift_weight_sum_coef = 10
    # Decrement refractory counters.
    self.refrac_count -= 1
    # Check for spiking neurons.
    O^n[t] = spike_func(V^n[t-1])
    # Refractoriness, voltage reset, and adaptive thresholds.
    self.refrac_count.masked_fill_(self.s, self.reset_refrac)   self.reset_refrac = 5
    self.v.masked_fill_(self.s, self.reset)
    self.thresh = self.thresh + self.thresh_plus * self.s.int().sum(0)
    # Voltage clipping to lower bound.
    if self.lbound is not None:
        self.v.masked_fill_(self.v < self.lbound, self.lbound)
    """
    def __init__(self, **kwargs):
        super(DiehlAndCookModelInt, self).__init__()
        # initial value for state variables
        self._variables['V'] = 0
        self._variables['O'] = 0
        self._variables['Isyn'] = 0
        self._variables['thresh'] = kwargs.get('thresh', 213)
        self._variables['refrac_count'] = kwargs.get('refrac_count', 0)

        self._constant_variables['max_threshold'] = kwargs.get('max_threshold', 4915)
        self._constant_variables['decay'] = kwargs.get('decay', 1013)
        self._constant_variables['shift_voltage'] = kwargs.get('shift_voltage', 10)
        self._constant_variables['weight_sum_coef'] = kwargs.get('weight_sum_coef', 154)
        self._constant_variables['shift_weight_sum_coef'] = kwargs.get('shift_weight_sum_coef', 10)
        self._constant_variables['weight_sum_min'] = kwargs.get('weight_sum_min', -32767)
        self._constant_variables['weight_sum_max'] = kwargs.get('weight_sum_max', 32767)
        self._constant_variables['Vreset'] = kwargs.get('v_reset', 0)
        self._constant_variables['Vrest'] = kwargs.get('v_rest', 0)
        self._constant_variables['init_refrac'] = kwargs.get('init_refrac', 0)
        self._constant_variables['reset_refrac'] = kwargs.get('reset_refrac', 5)
        self._constant_variables['delta'] = kwargs.get('delta', 1)
        self._constant_variables['thresh_plus'] = kwargs.get('thresh_plus', 66)
        self._constant_variables['lbound'] = kwargs.get('lbound', -32767)

        # self._operations.append(('halve', 'gt', 'thresh', 'max_threshold'))
        # self._operations.append(('shiftthresh', 'right_shift', 'thresh', 'halve'))
        # self._operations.append(('decayV', 'var_mult', 'decay', 'V'))
        # self._operations.append(('shiftDecayV', 'right_shift', 'decayV', 'shift_voltage'))
        # self._operations.append(('refrac_mask', 'le', 'refrac_count', 'init_refrac'))
        # self._operations.append(('decayIsyn', 'var_mult', 'weight_sum_coef', 'Isyn[updated]'))
        # self._operations.append(('shiftDecayIsyn', 'right_shift', 'decayIsyn', 'shift_weight_sum_coef'))
        # self._operations.append(('refrac_Isyn', 'var_mult', 'refrac_mask', 'shiftDecayIsyn'))
        # self._operations.append(('Vtemp', 'add', 'shiftDecayV', 'refrac_Isyn'))
        # self._operations.append(('refrac_count_temp', 'minus', 'refrac_count', 'delta'))
        # self._operations.append(('O', 'threshold', 'Vtemp', 'shiftthresh'))
        # self._operations.append(('refrac_count', 'reset', 'refrac_count_temp', 'O[updated]', 'reset_refrac'))
        # self._operations.append(('resetV', 'reset', 'Vtemp', 'O[updated]', 'Vreset'))
        # self._operations.append(('sumO', 'sum', 'O[updated]'))
        # self._operations.append(('thresh', 'var_linear', 'thresh_plus', 'sumO', 'shiftthresh'))
        # self._operations.append(('v_mask', 'lt', 'resetV', 'lbound'))
        # self._operations.append(('V', 'reset', 'resetV', 'v_mask', 'lbound'))

        self._operations.append((['V', 'O', 'thresh', 'refrac_count'], self.update,
                           ['Isyn[updated]', 'V', 'thresh', 'refrac_count', 'decay', 'Vrest', 'shift_voltage', 'max_threshold',
                            'weight_sum_coef', 'shift_weight_sum_coef', 'reset_refrac', 'Vreset', 'thresh_plus',
                            'lbound', 'weight_sum_min', 'weight_sum_max']))

    def update(self, x, v, thresh, refrac_count, decay, rest, shift_voltage, max_threshold, weight_sum_coef,
               shift_weight_sum_coef, refrac, reset, thresh_plus, lbound, weight_sum_min, weight_sum_max):

        x_ = x.int().clamp(min=weight_sum_min, max=weight_sum_max)
        # Halve the firing threshold (including the initial firing threshold and
        # the adaptive firing threshold) if the adaptive firing threshold exceeds the
        # maximum adaptive firing threshold.
        # At every time step, neurons halve their firing thresholds before further
        # updating the adaptive firing threshold, in response to excessive adaptive
        # firing thresholds computed at the last time step. This ensures that weight
        # updates at the last time step can access the excessive adaptive firing
        # thresholds before they are halved.

        halve = (thresh > max_threshold)
        thresh >>= halve  # Note: thresh >>= halve 不能改写为 thresh = thresh >> halve，会导致量化版本的STDP训练没效果
        del halve
        # Decay voltages.
        v = ((decay * (v - rest)) >> shift_voltage) + rest
        # Integrate inputs.
        v += (refrac_count <= 0).int() * ((x_ * weight_sum_coef) >> shift_weight_sum_coef)
        # Decrement refractory counters.
        refrac_count -= 1
        # Check for spiking neurons.
        s = v > thresh

        # Refractoriness, voltage reset, and adaptive thresholds.
        refrac_count.masked_fill_(s, refrac)
        v.masked_fill_(s, reset)
        thresh += thresh_plus * s.int().sum(0)

        # Voltage clipping to lower bound.
        v.masked_fill_(v < lbound, lbound)
        return v, s.to(dtype=torch.int32), thresh, refrac_count


NeuronModel.register("diehlAndCook", DiehlAndCookModelInt)


class LIFInt(NeuronModel):
    """
    Layer of leaky integrate-and-fire (LIF) neurons based on integer arithmetic:
    """
    def __init__(self, **kwargs):
        super(LIFInt, self).__init__()
        # initial value for state variables
        self._variables['V'] = 0
        self._variables['O'] = 0
        self._variables['Isyn'] = 0
        self._variables['refrac_count'] = kwargs.get('refrac_count', 0)

        self._constant_variables['thresh'] = kwargs.get('thresh', 1)
        self._constant_variables['decay'] = kwargs.get('decay', 1)
        self._constant_variables['Vreset'] = kwargs.get('v_reset', 0)
        self._constant_variables['Vrest'] = kwargs.get('v_rest', 0)
        self._constant_variables['reset_refrac'] = kwargs.get('reset_refrac', 0)
        self._constant_variables['delta'] = kwargs.get('delta', 1)
        self._constant_variables['reset_isyn'] = kwargs.get('reset_isyn', 0.0)

        # self._operations.append(('Vdifference', 'minus', 'V', 'Vrest'))
        # self._operations.append(('decayV', 'var_linear', 'decay', 'Vdifference', 'Vrest'))
        # self._operations.append(('refrac_mask', 'gt', 'refrac_count', 'reset_refrac'))
        # self._operations.append(('Isynreset', 'reset', 'Isyn[updated]', 'refrac_mask', 'reset_isyn'))
        # self._operations.append(('refrac_count_temp', 'minus', 'refrac_count', 'delta'))
        # self._operations.append(('Vtemp', 'add', 'decayV', 'Isynreset'))
        # self._operations.append(('O', 'threshold', 'Vtemp', 'thresh'))
        # self._operations.append(('refrac_count', 'reset', 'refrac_count_temp', 'O[updated]', 'reset_refrac'))
        # self._operations.append(('V', 'reset', 'Vtemp', 'O[updated]', 'Vreset'))

        self._operations.append((['V', 'O', 'refrac_count'], self.update,
                           ['Isyn[updated]', 'V', 'thresh', 'refrac_count', 'decay', 'Vrest', 'reset_refrac', 'Vreset']))

    def update(self, x, v, thresh, refrac_count, decay, rest, refrac, reset):
         # Decay voltages.
        v = decay * (v - rest) + rest

        # Integrate inputs.
        x.masked_fill_(refrac_count > 0, 0.0)

        # Decrement refractory counters.
        refrac_count -= 1

        v += x  # interlaced

        # Check for spiking neurons.
        s = v >= thresh

        # Refractoriness and voltage reset.
        refrac_count.masked_fill_(s, refrac)
        v.masked_fill_(s, reset)

        return v, s.to(dtype=torch.int32), refrac_count
NeuronModel.register("lifint", LIFInt)


class CANN_MeanFieldModel(NeuronModel):
    """
    Mean Field Model used in "Fung CC, Wong KY, Wu S. A moving bump in a continuous manifold: a comprehensive study of
    the tracking dynamics of continuous attractor neural networks. Neural Comput. 2010 Mar;22(3):752-92. doi: 10.1162/neco.2009.07-08-824. "
    "Wu S, Wong KY, Fung CC, Mi Y, Zhang W. Continuous Attractor Neural Networks: Candidate of a Canonical Model for
    Neural Information Representation.F1000Res. 2016 Feb 10;5:F1000 Faculty Rev-156. doi: 10.12688/f1000research.7387.1. "

    U = U + dt/tau * (Iext + rho*WgtSum - U)
    O = U^2/(1 + k*rho*sum(U^2))
    (WgtSum = weight*O_pre)
    """
    def __init__(self, **kwargs):
        super(CANN_MeanFieldModel, self).__init__()

        self._constant_variables['rho'] = kwargs.get('rho', 0.02)
        self._constant_variables['k_rho'] = kwargs.get('k', 0.01)
        self._constant_variables['1'] = 1
        self._constant_variables['2'] = 2

        self._membrane_variables['tau'] = kwargs.get('tau', 1.0)

        self._variables['Iext'] = 0.0
        self._variables['Isyn'] = 0.0
        self._variables['U'] = 0.0
        self._variables['O'] = 0.0

        self._operations.append(('Isum', 'var_linear', 'rho', 'Isyn', 'Iext'))
        self._operations.append(('dU', 'minus', 'Isum', 'U'))
        self._operations.append(('U', 'var_linear', 'tau', 'dU', 'U'))
        self._operations.append(('ReU', 'relu', 'U[updated]'))
        self._operations.append(('U2', 'var_mult', 'ReU', 'ReU'))
        self._operations.append(('SumU2', 'reduce_sum', 'U2', '1'))
        self._operations.append(('RBase','var_linear', 'k_rho', 'SumU2', '1'))
        self._operations.append(('O', 'div', 'U2', 'RBase'))


NeuronModel.register("cann_field", CANN_MeanFieldModel)


class MeanFieldModel(NeuronModel):
    """
    Mean Field Model of LIF neuron "

    U = U + dt/tau * (rho*(Iext + Isyn) - U)
    O = relu(U)
    (WgtSum = weight*O_pre)
    """
    def __init__(self, **kwargs):
        super(MeanFieldModel, self).__init__()

        self._constant_variables['rho'] = kwargs.get('rho', 0.1)
        self._constant_variables['1'] = 1
        self._constant_variables['2'] = 2
        self._membrane_variables['tau'] = kwargs.get('tau', 1.0)

        self._variables['Iext'] = 0.0
        self._variables['Isyn'] = 0.0
        self._variables['U'] = 0.0
        self._variables['O'] = 0.0

        self._operations.append(('Isum', 'var_linear', 'rho', 'Isyn', 'Iext'))
        self._operations.append(('dU', 'minus', 'Isum', 'U'))
        self._operations.append(('U', 'var_linear', 'tau', 'dU', 'U'))
        self._operations.append(('O', 'relu', 'U[updated]'))


NeuronModel.register("meanfield", MeanFieldModel)


class SimpleRateModel(NeuronModel):
    """
    Rate model  "

    U = U + dt/tau * (sigmoid(Iext + WgtSum) - U)
    (WgtSum = weight*O_pre)
    """
    def __init__(self, **kwargs):
        super(MeanFieldModel, self).__init__()
        self._membrane_variables['tau'] = kwargs.get('tau', 1.0)

        self._variables['Iext'] = 0.0
        self._variables['WgtSum'] = 0.0
        self._variables['U'] = 0.0

        self._operations.append(('Isum', 'add', 'WgtSum', 'Iext'))
        self._operations.append(('F', 'sigmoid', 'Isum'))
        self._operations.append(('dU', 'minus', 'F', 'U'))
        self._operations.append(('U', 'var_linear', 'tau', 'dU', 'U'))


class Darwin_CLIF(NeuronModel):
    """
    I = I*P4 + Wgt_sum
    V = V*P0 + I + P1
    """

    def __init__(self, **kwargs):
        super(Darwin_CLIF, self).__init__()
        self.dt = kwargs.get('dt', 0.1)
        self.tau_m = kwargs.get('tau_m', 12.0)
        self.tau_s = kwargs.get('tau_s', 8.0)
        self.bias = kwargs.get('bias', 0.0)
        self.v_th = kwargs.get('v_th', 16384)

        self._constant_variables['P0'] = np.round(np.exp(-self.dt / self.tau_m) * 2.0 ** 8) / 2.0 ** 8 #non-negative
        self._constant_variables['P4'] = np.round(np.exp(-self.dt / self.tau_s) * 2.0 ** 8) / 2.0 ** 8 #non-negative
        self._constant_variables['P1'] = np.round(self.bias*2.0**15)/2.0*15
        self._constant_variables['Vth'] = self.v_th

        self._variables['V'] = uniform(10000, 15000)
        self._variables['I'] = 0.0
        self._variables['O'] = 0.0
        self._variables['Isyn'] = 0.0 # named WgtSum in Darwin

        self._operations.append((['V', 'I', 'O'], self.update, ['V', 'I', 'Isyn', 'P0', 'P4', 'P1', 'Vth']))

    def update(self, V, I, Isyn, P0, P4, P1, Vth):
        I = self.quantize_16(I*P4 + Isyn)
        V = self.quantize_16(V*P0 + I + P1)
        O = (V > Vth).to(torch.float)
        V = V - V*O
        return V, I, O

    def quantize_16(self, x):
        x = torch.round(torch.clamp(x+2**15, 0, 2**16)) - 2**15
        return x

NeuronModel.register("darwin_clif", Darwin_CLIF)


class Darwin_ALIF(NeuronModel):
    """
    Adaptive threshold LIF
    I = I*P4 + Wgt_sum
    V = V*P0 + I + P1
    Vth = Vth*P2 + C1
    if spike:
    Vth = Vth + C2
    """

    def __init__(self, **kwargs):
        super(Darwin_ALIF, self).__init__()
        self.dt = kwargs.get('dt', 0.1)
        self.tau_m = kwargs.get('tau_m', 12.0)
        self.tau_s = kwargs.get('tau_s', 8.0)
        self.bias = kwargs.get('bias', 0.0)
        self.v_th = kwargs.get('v_th', 16384)

        self._constant_variables['P0'] = np.round(np.exp(-self.dt / self.tau_m) * 2.0 ** 8) / 2.0 ** 8  # non-negative
        self._constant_variables['P4'] = np.round(np.exp(-self.dt / self.tau_s) * 2.0 ** 8) / 2.0 ** 8  # non-negative
        self._constant_variables['P1'] = np.round(self.bias * 2.0 ** 15) / 2.0 * 15
        self._constant_variables['P2'] = 0.999
        self._constant_variables['C1'] = np.round(self.v_th*0.001)
        self._constant_variables['C2'] = np.round(self.v_th*0.2)

        self._variables['Vth'] = self.v_th
        self._variables['V'] = uniform(-30000, 15000)
        self._variables['I'] = 0.0
        self._variables['O'] = 0.0
        self._variables['Isyn'] = 0.0  # named WgtSum in Darwin

        self._operations.append((['V', 'I', 'O', 'Vth'], self.update, ['V', 'I', 'Isyn', 'P0', 'P4', 'P1', 'Vth', 'P2', 'C1', 'C2']))

    def update(self, V, I, Isyn, P0, P4, P1, Vth, P2, C1, C2):
        I = self.quantize_16(I * P4 + Isyn)
        V = self.quantize_16(V * P0 + I + P1)
        O = (V > Vth).to(torch.float)
        V = V - V * O
        Vth = P2*Vth + C1 + O*C2
        return V, I, O, Vth

    def quantize_16(self, x):
        x = torch.round(torch.clamp(x + 2 ** 15, 0, 2 ** 16)) - 2 ** 15
        return x


NeuronModel.register("darwin_alif", Darwin_ALIF)
class Darwin_Random(NeuronModel):

    def __init__(self, **kwargs):
        super(Darwin_Random, self).__init__()
        self.dt = kwargs.get('dt', 0.1)
        self.tau_m = kwargs.get("tau_m", 10.0)
        self.std = kwargs.get('std', 0.0)
        self.bias = kwargs.get('bias', 0.0)
        self.v_th = kwargs.get('v_th', 15000)

        self._variables['V'] = uniform(10000, 15000)
        self._variables['O'] = 0.0

        self._constant_variables['P0'] = np.round(np.exp(-self.dt / self.tau_m) * 2.0 ** 16) / 2.0 ** 16  # non-negative
        self._constant_variables['bias'] = np.round(self.bias)
        self._constant_variables['std'] = np.round(self.std)
        self._constant_variables['Vth'] = self.v_th

        self._operations.append((['V', 'O'], self.update, ['V', 'P0', 'bias', 'std', 'Vth']))

    def update(self, V, P0, bias, std, Vth):
        V = P0*V + torch.round(std*torch.rand_like(V)-std/2.0) + bias
        O = (V > Vth).to(torch.float)
        V = V - V * O
        return V, O


NeuronModel.register("darwin_random", Darwin_Random)







