import numpy as np
import random

from collections import namedtuple
import numpy.random as rd
import tensorflow.compat.v1 as tf
from tensorflow.python.ops.variables import Variable

import os
import warnings
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' 
warnings.filterwarnings("ignore",category=FutureWarning)
from tensorflow.python.util import deprecation
deprecation._PRINT_DEPRECATION_WARNINGS = False
tf.disable_v2_behavior()

os.environ['CUDA_VISIBLE_DEVICES'] = '-1' #use CPU for this experiment
#######################################################################################
# adaptive LIF neuron model 
######################################################################################
@tf.custom_gradient
def SpikeFunction(v, v_scaled, dampening_factor):
    z_ = tf.greater(v, 0.)
    z_ = tf.cast(z_, dtype=tf.float32)

    def grad(dy):
        dE_dz = dy
        dz_dv_scaled = tf.maximum(1 - tf.abs(v_scaled), 0)
        dz_dv_scaled *= dampening_factor

        dE_dv_scaled = dE_dz * dz_dv_scaled

        return [dE_dv_scaled,
                tf.zeros_like(dampening_factor)]

    return tf.identity(z_, name="SpikeFunction"), grad

#########################################################
ALIFStateTuple = namedtuple('ALIFState', ('u','s','w','ref', 'I1', 'I2', 'out'))
Cell = tf.nn.rnn_cell.BasicRNNCell
class ALIF(Cell):
    def __init__(self, network_param, neuron_param, synapse_param, dtype=tf.float32):

        # Parameters
        self.n_in = network_param[0]     
        self.n_rec = network_param[1]    
        self.n_out = network_param[2]       
        
        self.Vrest = neuron_param[0]
        self.Vth = neuron_param[1]
        self.Vahp = neuron_param[2]
        self.Vmax = neuron_param[3]
        
        self.t_ref = neuron_param[4]
        self.R = neuron_param[5]
        self.tau_m = neuron_param[6]
        self.R_adp = neuron_param[7]
        self.tau_w = neuron_param[8]
        self.a = neuron_param[9] 
        self.delT = neuron_param[10]
        
        self.p = neuron_param[11]
        self.q = neuron_param[12]
        self.r = neuron_param[13]
        self.s = neuron_param[14]
        
        self.tau_recsyn = synapse_param[0]
        self.tau_postsyn = synapse_param[1]
        self.inW = synapse_param[2]
        self.inBias = synapse_param[3]
        self.recW = synapse_param[4]
        self.outW = synapse_param[5]
        
        self.dampening_factor = 0.3

        # Trainable parameters
        with tf.variable_scope('InputWeights'): #initial input weights
            self.W_in = tf.Variable(self.inW, dtype=dtype, name='InputWeight', trainable=False)
            self.B_in = tf.Variable(self.inBias, dtype=dtype, name='InputBias', trainable=False)

        with tf.variable_scope('RecWeights'): #initial recurrent weights
            self.W_rec = tf.Variable(self.recW, dtype=dtype, name='RecurrentWeight', trainable=False)
        
        with tf.variable_scope('outputWeight'): #initial output weights
            self.W_out = tf.Variable(self.outW, dtype=dtype, name='RecurrentWeight', trainable=False)
            
    @property
    def output_size(self):
        return [self.n_rec, self.n_rec, self.n_rec, self.n_rec, self.n_rec, self.n_out]

    @property
    def state_size(self):
        return ALIFStateTuple(u=self.n_rec, s=self.n_rec, w=self.n_rec, ref=self.n_rec, I1=self.n_rec, I2=self.n_rec, out=self.n_out)

    def zero_state(self, batch_size, dtype):
        n_rec = self.n_rec
        u0 = tf.ones(shape=(batch_size, n_rec), dtype=dtype)*self.Vrest #membrane potential
        s0 = tf.zeros(shape=(batch_size, n_rec), dtype=dtype) #spike pattern
        w0 = tf.zeros(shape=(batch_size, n_rec), dtype=dtype) #adaptive property
        ref0 = tf.zeros(shape=(batch_size, n_rec), dtype=dtype) #refractory period
        I01 = tf.zeros(shape=(batch_size, n_rec), dtype=dtype) #previous current
        I02 = tf.zeros(shape=(batch_size, n_rec), dtype=dtype) #previous current
        out0 = tf.zeros(shape=(batch_size, self.n_out), dtype=dtype)
        return ALIFStateTuple(u=u0, s=s0, w=w0, ref=ref0, I1=I01, I2=I02, out=out0)

    def __call__(self, inputs, state, scope=None, dtype=tf.float32):
        with tf.name_scope('ALIFcall'):
            
            # current 
            I1 = tf.matmul(inputs, self.W_in) + self.B_in #current generated by input stimuli
            I2 = state.I2*tf.exp(-1.0/self.tau_recsyn) + tf.matmul(state.s, self.W_rec)*(1-tf.exp(-1.0/self.tau_recsyn)) #synapse currrent
            w = state.w*(1 - 1.0/self.tau_w) + self.a*state.s
            
            # membrane potential 
            u = state.u + (-1*(state.u-tf.ones_like(state.u)*self.Vrest) \
                           + self.delT*tf.exp((state.u-tf.ones_like(state.u)*self.Vth)/self.delT) \
                           + self.R*(I1+I2) - self.R_adp*state.w) / self.tau_m
            u = tf.where(tf.greater(state.ref, 0.0), tf.ones_like(u)*self.Vahp, u) #voltage at Vrh during refractory period
            u = tf.where(tf.greater(u, self.Vmax), tf.ones_like(u)*self.Vmax+1, u) #set voltage upper limit
            u = tf.where(tf.greater(self.Vrest, u), tf.ones_like(u)*self.Vrest, u) #set voltage lower 
            
            # generate spike 
            u_scale = (u-self.Vth)/abs(self.Vth) #normalized membrane different, used for backpropogation
            s = SpikeFunction( u-self.Vmax, u_scale, self.dampening_factor)
            
            # refractory period                                                                              
            ref = tf.where(tf.greater(state.ref, 0.0), tf.ones_like(u)*(state.ref-1), state.ref) #refractory timing reduce 1 step
            ref = tf.where(tf.greater(u, self.Vmax), tf.ones_like(u)*self.t_ref, ref) #enter refractory period
            
            # output filter 
            out = state.out*tf.exp(-1.0/self.tau_postsyn) + tf.matmul(s, self.W_out)*(1-tf.exp(-1.0/self.tau_postsyn))          
            
            new_state = ALIFStateTuple(u, s, w, ref, I1, I2, out)
            
        return [u, s, w, I1, I2, out], new_state

