
import torch, time, os

import torch.nn as nn
import torch.nn.functional as F

#from graphviz import Digraph
import scipy.io as sio
import numpy as np

os.environ['CUDA_VISIBLE_DEVICES'] = "0"
lens = 0.25

thresh = 0.3
decay = 0.3

num_classes = 10
batch_size = 1
num_epochs = 20
learning_rate = 1e-4#########default 1e-4
time_window = 100
names = 'animales_stbp'

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print (device)


# load datasets need (num,pol, size,size, step)

train_dataset=sio.loadmat('C:\\Users\\Ajay\\Documents\\IBMdataset\\DvsGesture\\ibmframes.mat')
tr=train_dataset['myframes']

tr_images = torch.from_numpy(tr)

print(tr_images.shape) #64,64,100,1219





train_labels=sio.loadmat('C:\\Users\\Ajay\\Documents\\IBMdataset\\DvsGesture\\ibmlabels.mat')
trl=train_labels['labels']

tr_labels = torch.from_numpy(trl)


class ActFun(torch.autograd.Function):

    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input)
        #print('\nActFun forward called',input.size())
        return input.gt(thresh).float()

    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors
        grad_input = grad_output.clone()
        temp = abs(input - thresh) < lens
        #print('start \n',temp.size(),'\n', grad_input.size(), torch.mean(grad_input), 'fin \n')
        #print('actfun backward called')
        return grad_input * temp.float() / (2 * lens)        
        #return grad_input * 1 / (2 * lens)
        




def lr_scheduler(optimizer, epoch, init_lr=0.1, lr_decay_epoch=100):
    """Decay learning rate by a factor of 0.1 every lr_decay_epoch epochs."""

    if epoch % lr_decay_epoch == 0 and epoch > 1:
        for param_group in optimizer.param_groups:
            param_group['lr'] = param_group['lr'] * 0.1

    return optimizer


best_acc = 0
acc_record = list([])


class SNN_Model(nn.Module):

    def __init__(self, num_classes=10):
        super(SNN_Model, self).__init__()
        self.conv0 = nn.Conv2d(1,10, kernel_size=5, stride=1, padding=2)
        self.conv2 = nn.Conv2d(1,10, kernel_size=5, stride=1, padding=2)

        self.conv1 = nn.Conv2d(20, 25, kernel_size=5, stride=1, padding=2)
        self.conv3 = nn.Conv2d(20, 25, kernel_size=5, stride=1, padding=2)

        #conv0,conv1 -> upper with delay in conn, 
        #conv2,3 -> lower no delay
        
        #in_planes, out_planes, stride = cfg_cnn[2]
        #self.conv2 = nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=1, padding=1)

        self.fc1 = nn.Linear(16*16 * 50, 100)
        self.fc3 = nn.Linear(16*16 * 50, 100)

        self.fc2 = nn.Linear(200, 10)

    def forward(self, input, win, bsize):
        #input to be bsize x 28 x 28
        c0_mem = c0_spike = torch.zeros( 10, 64, 64, device=device)
        p0_mem = p0_spike = torch.zeros( 10, 32, 32, device=device)
                
        c2_mem = c2_spike = torch.zeros( 10, 64, 64, device=device)
        p2_mem = p2_spike = torch.zeros( 10, 32, 32, device=device)
        
        
        
        
        
        

        c1_mem = c1_spike = torch.zeros( 25, 32, 32, device=device)
        p1_mem = p1_spike = torch.zeros( 25, 16, 16, device=device)
        
        c3_mem = c3_spike = torch.zeros( 25, 32, 32, device=device)
        p3_mem = p3_spike = torch.zeros( 25, 16, 16, device=device)
        
        
        
        
        
        

        #c2_mem = c2_spike = torch.zeros(bsize, cfg_cnn[2][1], 7, 7, device=device)
        #p2_mem = p2_spike = torch.zeros(bsize, cfg_cnn[2][1], 7, 7, device=device)

        h1_mem = h1_spike = h1_sumspike = torch.zeros( 100, device=device)
        h3_mem = h3_spike = h3_sumspike = torch.zeros( 100, device=device)

        h2_mem = h2_spike = h2_sumspike = torch.zeros( 10, device=device)

        x=torch.zeros(input.size()).float()

        #firstspikesidx=torch.zeros(batch_size,5)
        #spkcnt=torch.zeros(batch_size).byte()
        #input = input.float()
        for step in range(win):
            #x = input[:, :, :, :, step]
            x = input[:, :, step]
            x=x[:,:,None,None]
            x=x.permute(2,3,0,1)

            c0_mem, c0_spike = mem_update(self.conv0, x, c0_mem, c0_spike,decay,step)
            p0_mem, p0_spike = mem_update_pool(F.avg_pool2d, c0_spike, p0_mem, p0_spike,decay,step)
            
            c2_mem, c2_spike = mem_update(self.conv2, x, c2_mem, c2_spike,decay,step)
            p2_mem, p2_spike = mem_update_pool(F.avg_pool2d, c2_spike, p2_mem, p2_spike,decay,step)
            
            if step==0:
                p0spk_back5=torch.zeros_like(p0_spike)
                p0spk_back4=torch.zeros_like(p0_spike)
                p0spk_back3=torch.zeros_like(p0_spike)
                p0spk_back2=torch.zeros_like(p0_spike)
                p0spk_back1=torch.zeros_like(p0_spike)
            if step>5:
                p0spk_back5=p0spk_back4
                p0spk_back4=p0spk_back3
                p0spk_back3=p0spk_back2
                p0spk_back2=p0spk_back1
                p0spk_back1=p0_spike
                
                
            p0spk_output=torch.cat([p0spk_back5,p2_spike],dim=1)
            
            
            

            c1_mem, c1_spike = mem_update(self.conv1, p0spk_output, c1_mem, c1_spike,decay,step)
            p1_mem, p1_spike = mem_update_pool(F.avg_pool2d, c1_spike, p1_mem, p1_spike,decay,step)
            
            c3_mem, c3_spike = mem_update(self.conv3, p0spk_output, c3_mem, c3_spike,decay,step)
            p3_mem, p3_spike = mem_update_pool(F.avg_pool2d, c3_spike, p3_mem, p3_spike,decay,step)
            
            
            
            if step==0:
                p1spk_back5=torch.zeros_like(p1_spike)
                p1spk_back4=torch.zeros_like(p1_spike)
                p1spk_back3=torch.zeros_like(p1_spike)
                p1spk_back2=torch.zeros_like(p1_spike)
                p1spk_back1=torch.zeros_like(p1_spike)
            if step>5:
                p1spk_back5=p1spk_back4
                p1spk_back4=p1spk_back3
                p1spk_back3=p1spk_back2
                p1spk_back2=p1spk_back1
                p1spk_back1=p1_spike
                
                
            p1spk_output=torch.cat([p1spk_back5,p3_spike],dim=1)
            
            
            
            
            
            
            #x = p1_spike.view( -1)
            x = p1spk_output.view( -1)

            h1_mem, h1_spike = mem_update(self.fc1, x, h1_mem, h1_spike,decay,step)
            h3_mem, h3_spike = mem_update(self.fc3, x, h3_mem, h3_spike,decay,step)
            #h1_sumspike += h1_spike

            #h2_mem = mem_update_last(self.fc2, h1_spike, h2_mem,decay,step)
            h2_mem, h2_spike = mem_update(self.fc2, torch.cat([h1_spike,h3_spike],dim=0), h2_mem, h2_spike,decay,step)
            h2_sumspike += h2_spike
  
        outputs = h2_sumspike / time_window
        
        return outputs




def mem_update(conv, x, mem, spike,curr_decay,curr_step):
    #mem = mem * curr_decay *float(time_window-curr_step)/float(time_window)* (1 - spike) + conv(x)
    mem = mem * (1 - spike) *curr_decay + conv(x)
    spike = act_fun(mem)
    return mem, spike

def mem_update_last(conv, x, mem,curr_decay,curr_step):
    mem = mem *curr_decay + conv(x)
    return mem




def mem_update_pool(opts, x, mem, spike,curr_decay,curr_step):
    mem = mem * float(time_window-curr_step)/float(time_window)*curr_decay*(1 - spike) + opts(x, 2)
    #mem = mem * (1 - spike) *curr_decay + opts(x, 2)
    spike = act_fun(mem)
    return mem, spike


snn = SNN_Model()
print(torch.sum(snn.conv0.weight.data))
print(torch.max(snn.conv0.weight.data))

snn.to(device)
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(snn.parameters(), lr=learning_rate)

act_fun = ActFun.apply

tr_images=tr_images.float()
ts_images = tr_images.float()
tr_labels=tr_labels.long()

for kfold in range(4):
    trnum=914
    tsnum=305
    
    if kfold==0:
        tr_images_currset=tr_images[:,:,:,0:914]
        ts_images_currset=tr_images[:,:,:,914:1219]
        tr_labels_currset=tr_labels[0:914]
        ts_labels_currset=tr_labels[914:1219]
        
    if kfold==1:
        tr_images_currset=torch.cat((tr_images[:,:,:,0:305],tr_images[:,:,:,610:1219]),3)
        ts_images_currset=tr_images[:,:,:,305:610]
        tr_labels_currset=torch.cat((tr_labels[0:305],tr_labels[610:1219]),0)
        ts_labels_currset=tr_labels[305:610]
        
    if kfold==2:
        tr_images_currset=torch.cat((tr_images[:,:,:,0:610],tr_images[:,:,:,915:1219]),3)
        ts_images_currset=tr_images[:,:,:,610:915]
        tr_labels_currset=torch.cat((tr_labels[0:610],tr_labels[915:1219]),0)
        ts_labels_currset=tr_labels[610:915]
        
    if kfold==3:
        tr_images_currset=tr_images[:,:,:,305:1219]
        ts_images_currset=tr_images[:,:,:,0:305]
        tr_labels_currset=tr_labels[305:1219]
        ts_labels_currset=tr_labels[0:305]
        
    
    preds=np.zeros_like(ts_labels_currset)
    bestresult=0.0
    for epoch in range(num_epochs):

        running_loss = 0
        start_time = time.time()

        #numiters = 1219#60000//batch_size
        #shortiter=numiters

        for i in range(trnum):
            #images = tr_images[i,:,:,:,:,:]
            images = tr_images_currset[:,:,:,i]
            #images = images[:,:,:,None,None]
            #images.permute(2,3,0,1)




            #print('\n','iteration num',i,'\n')
            snn.zero_grad()
            optimizer.zero_grad()
            #print(images[0,0,:,:])
            images = images.float().to(device)
            outputs = snn(images, time_window,1)
            ops = outputs.cpu()

            #print(outputs)
            lbls=torch.FloatTensor(10)
            lbls.zero_()
            lbls[tr_labels_currset[i]]=1.
            loss = criterion(outputs.cpu(), lbls)

            running_loss += loss.item()
            loss.backward()
            optimizer.step()

            if (i + 1) % 50 == 0:
                #print (torch.sum(outputs,1))
                print('Epoch [%d/%d], Step [%d/%d], Loss: %.5f'
                      % (epoch + 1, num_epochs, i + 1, trnum, running_loss))
                running_loss = 0
                print('Time elasped:', time.time() - start_time)
        correct = 0
        total = 0

        optimizer = lr_scheduler(optimizer, epoch, learning_rate, 40)


        for i in range(tsnum):
            images = ts_images_currset[:,:,:,i]
            
            
            lbls=torch.FloatTensor(10)
            lbls.zero_()
            lbls[ts_labels_currset[i]]=1.

            images = images.float().to(device)
            outputs = snn(images, time_window,1) #/ time_window
            
            _,preds[i]=torch.max(outputs.data,0)

            #_, predicted = torch.max(outputs.data, 0)
            #_, lbls = torch.max(lbls.data, 0)
            #total += lbls.size(0)
            #correct += (predicted.cpu() == lbls).sum()
            #correct += (predicted.cpu() == tr_labels).sum()
            #print(correct,total)


            #i +=1
        ct = (preds[0:tsnum]==ts_labels_currset[0:tsnum])
        #print(preds[0:shortiter])
        print('\n\n\n')
        #print(tr_labels[0:shortiter])
        #print('\n\n\n')
        #print(preds.shape,tr_labels.shape)
        #print(ct.shape)
        print('Iters:', epoch, '\n')
        correct=ct.sum()
        print(correct)
        print('Test Accuracy of the model: %.3f  ' % (100 * correct.float() / tsnum))
        
        if correct.float()>=bestresult:
            bestresult=correct.float()
            if kfold==0:
                sio.savemat('bestresults_preds_splitnet_k0.mat',{'preds':preds})
                print('saved')
            if kfold==1:
                sio.savemat('bestresults_preds_splitnet_k1.mat',{'preds':preds})
                print('saved')
            if kfold==2:
                sio.savemat('bestresults_preds_splitnet_k2.mat',{'preds':preds})
                print('saved')
            if kfold==3:
                sio.savemat('bestresults_preds_splitnet_k3.mat',{'preds':preds})
                print('saved')
            
        
        print(torch.sum(snn.conv0.weight.data))
        print(torch.max(snn.conv0.weight.data))








91c41294e4464ce349b33df92df84a1306f1476e
<mask token>


class ActFun(torch.autograd.Function):

    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input)
        return input.gt(thresh).float()

    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors
        grad_input = grad_output.clone()
        temp = abs(input - thresh) < lens
        return grad_input * temp.float() / (2 * lens)


<mask token>


class SNN_Model(nn.Module):

    def __init__(self, num_classes=10):
        super(SNN_Model, self).__init__()
        self.conv0 = nn.Conv2d(1, 10, kernel_size=5, stride=1, padding=2)
        self.conv2 = nn.Conv2d(1, 10, kernel_size=5, stride=1, padding=2)
        self.conv1 = nn.Conv2d(20, 25, kernel_size=5, stride=1, padding=2)
        self.conv3 = nn.Conv2d(20, 25, kernel_size=5, stride=1, padding=2)
        self.fc1 = nn.Linear(16 * 16 * 50, 100)
        self.fc3 = nn.Linear(16 * 16 * 50, 100)
        self.fc2 = nn.Linear(200, 10)

    def forward(self, input, win, bsize):
        c0_mem = c0_spike = torch.zeros(10, 64, 64, device=device)
        p0_mem = p0_spike = torch.zeros(10, 32, 32, device=device)
        c2_mem = c2_spike = torch.zeros(10, 64, 64, device=device)
        p2_mem = p2_spike = torch.zeros(10, 32, 32, device=device)
        c1_mem = c1_spike = torch.zeros(25, 32, 32, device=device)
        p1_mem = p1_spike = torch.zeros(25, 16, 16, device=device)
        c3_mem = c3_spike = torch.zeros(25, 32, 32, device=device)
        p3_mem = p3_spike = torch.zeros(25, 16, 16, device=device)
        h1_mem = h1_spike = h1_sumspike = torch.zeros(100, device=device)
        h3_mem = h3_spike = h3_sumspike = torch.zeros(100, device=device)
        h2_mem = h2_spike = h2_sumspike = torch.zeros(10, device=device)
        x = torch.zeros(input.size()).float()
        for step in range(win):
            x = input[:, :, step]
            x = x[:, :, None, None]
            x = x.permute(2, 3, 0, 1)
            c0_mem, c0_spike = mem_update(self.conv0, x, c0_mem, c0_spike,
                decay, step)
            p0_mem, p0_spike = mem_update_pool(F.avg_pool2d, c0_spike,
                p0_mem, p0_spike, decay, step)
            c2_mem, c2_spike = mem_update(self.conv2, x, c2_mem, c2_spike,
                decay, step)
            p2_mem, p2_spike = mem_update_pool(F.avg_pool2d, c2_spike,
                p2_mem, p2_spike, decay, step)
            if step == 0:
                p0spk_back5 = torch.zeros_like(p0_spike)
                p0spk_back4 = torch.zeros_like(p0_spike)
                p0spk_back3 = torch.zeros_like(p0_spike)
                p0spk_back2 = torch.zeros_like(p0_spike)
                p0spk_back1 = torch.zeros_like(p0_spike)
            if step > 5:
                p0spk_back5 = p0spk_back4
                p0spk_back4 = p0spk_back3
                p0spk_back3 = p0spk_back2
                p0spk_back2 = p0spk_back1
                p0spk_back1 = p0_spike
            p0spk_output = torch.cat([p0spk_back5, p2_spike], dim=1)
            c1_mem, c1_spike = mem_update(self.conv1, p0spk_output, c1_mem,
                c1_spike, decay, step)
            p1_mem, p1_spike = mem_update_pool(F.avg_pool2d, c1_spike,
                p1_mem, p1_spike, decay, step)
            c3_mem, c3_spike = mem_update(self.conv3, p0spk_output, c3_mem,
                c3_spike, decay, step)
            p3_mem, p3_spike = mem_update_pool(F.avg_pool2d, c3_spike,
                p3_mem, p3_spike, decay, step)
            if step == 0:
                p1spk_back5 = torch.zeros_like(p1_spike)
                p1spk_back4 = torch.zeros_like(p1_spike)
                p1spk_back3 = torch.zeros_like(p1_spike)
                p1spk_back2 = torch.zeros_like(p1_spike)
                p1spk_back1 = torch.zeros_like(p1_spike)
            if step > 5:
                p1spk_back5 = p1spk_back4
                p1spk_back4 = p1spk_back3
                p1spk_back3 = p1spk_back2
                p1spk_back2 = p1spk_back1
                p1spk_back1 = p1_spike
            p1spk_output = torch.cat([p1spk_back5, p3_spike], dim=1)
            x = p1spk_output.view(-1)
            h1_mem, h1_spike = mem_update(self.fc1, x, h1_mem, h1_spike,
                decay, step)
            h3_mem, h3_spike = mem_update(self.fc3, x, h3_mem, h3_spike,
                decay, step)
            h2_mem, h2_spike = mem_update(self.fc2, torch.cat([h1_spike,
                h3_spike], dim=0), h2_mem, h2_spike, decay, step)
            h2_sumspike += h2_spike
        outputs = h2_sumspike / time_window
        return outputs


<mask token>

<mask token>


class ActFun(torch.autograd.Function):

    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input)
        return input.gt(thresh).float()

    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors
        grad_input = grad_output.clone()
        temp = abs(input - thresh) < lens
        return grad_input * temp.float() / (2 * lens)


def lr_scheduler(optimizer, epoch, init_lr=0.1, lr_decay_epoch=100):
    """Decay learning rate by a factor of 0.1 every lr_decay_epoch epochs."""
    if epoch % lr_decay_epoch == 0 and epoch > 1:
        for param_group in optimizer.param_groups:
            param_group['lr'] = param_group['lr'] * 0.1
    return optimizer


<mask token>


class SNN_Model(nn.Module):

    def __init__(self, num_classes=10):
        super(SNN_Model, self).__init__()
        self.conv0 = nn.Conv2d(1, 10, kernel_size=5, stride=1, padding=2)
        self.conv2 = nn.Conv2d(1, 10, kernel_size=5, stride=1, padding=2)
        self.conv1 = nn.Conv2d(20, 25, kernel_size=5, stride=1, padding=2)
        self.conv3 = nn.Conv2d(20, 25, kernel_size=5, stride=1, padding=2)
        self.fc1 = nn.Linear(16 * 16 * 50, 100)
        self.fc3 = nn.Linear(16 * 16 * 50, 100)
        self.fc2 = nn.Linear(200, 10)

    def forward(self, input, win, bsize):
        c0_mem = c0_spike = torch.zeros(10, 64, 64, device=device)
        p0_mem = p0_spike = torch.zeros(10, 32, 32, device=device)
        c2_mem = c2_spike = torch.zeros(10, 64, 64, device=device)
        p2_mem = p2_spike = torch.zeros(10, 32, 32, device=device)
        c1_mem = c1_spike = torch.zeros(25, 32, 32, device=device)
        p1_mem = p1_spike = torch.zeros(25, 16, 16, device=device)
        c3_mem = c3_spike = torch.zeros(25, 32, 32, device=device)
        p3_mem = p3_spike = torch.zeros(25, 16, 16, device=device)
        h1_mem = h1_spike = h1_sumspike = torch.zeros(100, device=device)
        h3_mem = h3_spike = h3_sumspike = torch.zeros(100, device=device)
        h2_mem = h2_spike = h2_sumspike = torch.zeros(10, device=device)
        x = torch.zeros(input.size()).float()
        for step in range(win):
            x = input[:, :, step]
            x = x[:, :, None, None]
            x = x.permute(2, 3, 0, 1)
            c0_mem, c0_spike = mem_update(self.conv0, x, c0_mem, c0_spike,
                decay, step)
            p0_mem, p0_spike = mem_update_pool(F.avg_pool2d, c0_spike,
                p0_mem, p0_spike, decay, step)
            c2_mem, c2_spike = mem_update(self.conv2, x, c2_mem, c2_spike,
                decay, step)
            p2_mem, p2_spike = mem_update_pool(F.avg_pool2d, c2_spike,
                p2_mem, p2_spike, decay, step)
            if step == 0:
                p0spk_back5 = torch.zeros_like(p0_spike)
                p0spk_back4 = torch.zeros_like(p0_spike)
                p0spk_back3 = torch.zeros_like(p0_spike)
                p0spk_back2 = torch.zeros_like(p0_spike)
                p0spk_back1 = torch.zeros_like(p0_spike)
            if step > 5:
                p0spk_back5 = p0spk_back4
                p0spk_back4 = p0spk_back3
                p0spk_back3 = p0spk_back2
                p0spk_back2 = p0spk_back1
                p0spk_back1 = p0_spike
            p0spk_output = torch.cat([p0spk_back5, p2_spike], dim=1)
            c1_mem, c1_spike = mem_update(self.conv1, p0spk_output, c1_mem,
                c1_spike, decay, step)
            p1_mem, p1_spike = mem_update_pool(F.avg_pool2d, c1_spike,
                p1_mem, p1_spike, decay, step)
            c3_mem, c3_spike = mem_update(self.conv3, p0spk_output, c3_mem,
                c3_spike, decay, step)
            p3_mem, p3_spike = mem_update_pool(F.avg_pool2d, c3_spike,
                p3_mem, p3_spike, decay, step)
            if step == 0:
                p1spk_back5 = torch.zeros_like(p1_spike)
                p1spk_back4 = torch.zeros_like(p1_spike)
                p1spk_back3 = torch.zeros_like(p1_spike)
                p1spk_back2 = torch.zeros_like(p1_spike)
                p1spk_back1 = torch.zeros_like(p1_spike)
            if step > 5:
                p1spk_back5 = p1spk_back4
                p1spk_back4 = p1spk_back3
                p1spk_back3 = p1spk_back2
                p1spk_back2 = p1spk_back1
                p1spk_back1 = p1_spike
            p1spk_output = torch.cat([p1spk_back5, p3_spike], dim=1)
            x = p1spk_output.view(-1)
            h1_mem, h1_spike = mem_update(self.fc1, x, h1_mem, h1_spike,
                decay, step)
            h3_mem, h3_spike = mem_update(self.fc3, x, h3_mem, h3_spike,
                decay, step)
            h2_mem, h2_spike = mem_update(self.fc2, torch.cat([h1_spike,
                h3_spike], dim=0), h2_mem, h2_spike, decay, step)
            h2_sumspike += h2_spike
        outputs = h2_sumspike / time_window
        return outputs


def mem_update(conv, x, mem, spike, curr_decay, curr_step):
    mem = mem * (1 - spike) * curr_decay + conv(x)
    spike = act_fun(mem)
    return mem, spike


<mask token>


def mem_update_pool(opts, x, mem, spike, curr_decay, curr_step):
    mem = mem * float(time_window - curr_step) / float(time_window
        ) * curr_decay * (1 - spike) + opts(x, 2)
    spike = act_fun(mem)
    return mem, spike


<mask token>

<mask token>
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
lens = 0.25
thresh = 0.3
decay = 0.3
num_classes = 10
batch_size = 1
num_epochs = 20
learning_rate = 0.0001
time_window = 100
names = 'animales_stbp'
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(device)
train_dataset = sio.loadmat(
    'C:\\Users\\Ajay\\Documents\\IBMdataset\\DvsGesture\\ibmframes.mat')
tr = train_dataset['myframes']
tr_images = torch.from_numpy(tr)
print(tr_images.shape)
train_labels = sio.loadmat(
    'C:\\Users\\Ajay\\Documents\\IBMdataset\\DvsGesture\\ibmlabels.mat')
trl = train_labels['labels']
tr_labels = torch.from_numpy(trl)


class ActFun(torch.autograd.Function):

    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input)
        return input.gt(thresh).float()

    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors
        grad_input = grad_output.clone()
        temp = abs(input - thresh) < lens
        return grad_input * temp.float() / (2 * lens)


def lr_scheduler(optimizer, epoch, init_lr=0.1, lr_decay_epoch=100):
    """Decay learning rate by a factor of 0.1 every lr_decay_epoch epochs."""
    if epoch % lr_decay_epoch == 0 and epoch > 1:
        for param_group in optimizer.param_groups:
            param_group['lr'] = param_group['lr'] * 0.1
    return optimizer


best_acc = 0
acc_record = list([])


class SNN_Model(nn.Module):

    def __init__(self, num_classes=10):
        super(SNN_Model, self).__init__()
        self.conv0 = nn.Conv2d(1, 10, kernel_size=5, stride=1, padding=2)
        self.conv2 = nn.Conv2d(1, 10, kernel_size=5, stride=1, padding=2)
        self.conv1 = nn.Conv2d(20, 25, kernel_size=5, stride=1, padding=2)
        self.conv3 = nn.Conv2d(20, 25, kernel_size=5, stride=1, padding=2)
        self.fc1 = nn.Linear(16 * 16 * 50, 100)
        self.fc3 = nn.Linear(16 * 16 * 50, 100)
        self.fc2 = nn.Linear(200, 10)

    def forward(self, input, win, bsize):
        c0_mem = c0_spike = torch.zeros(10, 64, 64, device=device)
        p0_mem = p0_spike = torch.zeros(10, 32, 32, device=device)
        c2_mem = c2_spike = torch.zeros(10, 64, 64, device=device)
        p2_mem = p2_spike = torch.zeros(10, 32, 32, device=device)
        c1_mem = c1_spike = torch.zeros(25, 32, 32, device=device)
        p1_mem = p1_spike = torch.zeros(25, 16, 16, device=device)
        c3_mem = c3_spike = torch.zeros(25, 32, 32, device=device)
        p3_mem = p3_spike = torch.zeros(25, 16, 16, device=device)
        h1_mem = h1_spike = h1_sumspike = torch.zeros(100, device=device)
        h3_mem = h3_spike = h3_sumspike = torch.zeros(100, device=device)
        h2_mem = h2_spike = h2_sumspike = torch.zeros(10, device=device)
        x = torch.zeros(input.size()).float()
        for step in range(win):
            x = input[:, :, step]
            x = x[:, :, None, None]
            x = x.permute(2, 3, 0, 1)
            c0_mem, c0_spike = mem_update(self.conv0, x, c0_mem, c0_spike,
                decay, step)
            p0_mem, p0_spike = mem_update_pool(F.avg_pool2d, c0_spike,
                p0_mem, p0_spike, decay, step)
            c2_mem, c2_spike = mem_update(self.conv2, x, c2_mem, c2_spike,
                decay, step)
            p2_mem, p2_spike = mem_update_pool(F.avg_pool2d, c2_spike,
                p2_mem, p2_spike, decay, step)
            if step == 0:
                p0spk_back5 = torch.zeros_like(p0_spike)
                p0spk_back4 = torch.zeros_like(p0_spike)
                p0spk_back3 = torch.zeros_like(p0_spike)
                p0spk_back2 = torch.zeros_like(p0_spike)
                p0spk_back1 = torch.zeros_like(p0_spike)
            if step > 5:
                p0spk_back5 = p0spk_back4
                p0spk_back4 = p0spk_back3
                p0spk_back3 = p0spk_back2
                p0spk_back2 = p0spk_back1
                p0spk_back1 = p0_spike
            p0spk_output = torch.cat([p0spk_back5, p2_spike], dim=1)
            c1_mem, c1_spike = mem_update(self.conv1, p0spk_output, c1_mem,
                c1_spike, decay, step)
            p1_mem, p1_spike = mem_update_pool(F.avg_pool2d, c1_spike,
                p1_mem, p1_spike, decay, step)
            c3_mem, c3_spike = mem_update(self.conv3, p0spk_output, c3_mem,
                c3_spike, decay, step)
            p3_mem, p3_spike = mem_update_pool(F.avg_pool2d, c3_spike,
                p3_mem, p3_spike, decay, step)
            if step == 0:
                p1spk_back5 = torch.zeros_like(p1_spike)
                p1spk_back4 = torch.zeros_like(p1_spike)
                p1spk_back3 = torch.zeros_like(p1_spike)
                p1spk_back2 = torch.zeros_like(p1_spike)
                p1spk_back1 = torch.zeros_like(p1_spike)
            if step > 5:
                p1spk_back5 = p1spk_back4
                p1spk_back4 = p1spk_back3
                p1spk_back3 = p1spk_back2
                p1spk_back2 = p1spk_back1
                p1spk_back1 = p1_spike
            p1spk_output = torch.cat([p1spk_back5, p3_spike], dim=1)
            x = p1spk_output.view(-1)
            h1_mem, h1_spike = mem_update(self.fc1, x, h1_mem, h1_spike,
                decay, step)
            h3_mem, h3_spike = mem_update(self.fc3, x, h3_mem, h3_spike,
                decay, step)
            h2_mem, h2_spike = mem_update(self.fc2, torch.cat([h1_spike,
                h3_spike], dim=0), h2_mem, h2_spike, decay, step)
            h2_sumspike += h2_spike
        outputs = h2_sumspike / time_window
        return outputs


def mem_update(conv, x, mem, spike, curr_decay, curr_step):
    mem = mem * (1 - spike) * curr_decay + conv(x)
    spike = act_fun(mem)
    return mem, spike


def mem_update_last(conv, x, mem, curr_decay, curr_step):
    mem = mem * curr_decay + conv(x)
    return mem


def mem_update_pool(opts, x, mem, spike, curr_decay, curr_step):
    mem = mem * float(time_window - curr_step) / float(time_window
        ) * curr_decay * (1 - spike) + opts(x, 2)
    spike = act_fun(mem)
    return mem, spike


snn = SNN_Model()
print(torch.sum(snn.conv0.weight.data))
print(torch.max(snn.conv0.weight.data))
snn.to(device)
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(snn.parameters(), lr=learning_rate)
act_fun = ActFun.apply
tr_images = tr_images.float()
ts_images = tr_images.float()
tr_labels = tr_labels.long()
for kfold in range(4):
    trnum = 914
    tsnum = 305
    if kfold == 0:
        tr_images_currset = tr_images[:, :, :, 0:914]
        ts_images_currset = tr_images[:, :, :, 914:1219]
        tr_labels_currset = tr_labels[0:914]
        ts_labels_currset = tr_labels[914:1219]
    if kfold == 1:
        tr_images_currset = torch.cat((tr_images[:, :, :, 0:305], tr_images
            [:, :, :, 610:1219]), 3)
        ts_images_currset = tr_images[:, :, :, 305:610]
        tr_labels_currset = torch.cat((tr_labels[0:305], tr_labels[610:1219
            ]), 0)
        ts_labels_currset = tr_labels[305:610]
    if kfold == 2:
        tr_images_currset = torch.cat((tr_images[:, :, :, 0:610], tr_images
            [:, :, :, 915:1219]), 3)
        ts_images_currset = tr_images[:, :, :, 610:915]
        tr_labels_currset = torch.cat((tr_labels[0:610], tr_labels[915:1219
            ]), 0)
        ts_labels_currset = tr_labels[610:915]
    if kfold == 3:
        tr_images_currset = tr_images[:, :, :, 305:1219]
        ts_images_currset = tr_images[:, :, :, 0:305]
        tr_labels_currset = tr_labels[305:1219]
        ts_labels_currset = tr_labels[0:305]
    preds = np.zeros_like(ts_labels_currset)
    bestresult = 0.0
    for epoch in range(num_epochs):
        running_loss = 0
        start_time = time.time()
        for i in range(trnum):
            images = tr_images_currset[:, :, :, i]
            snn.zero_grad()
            optimizer.zero_grad()
            images = images.float().to(device)
            outputs = snn(images, time_window, 1)
            ops = outputs.cpu()
            lbls = torch.FloatTensor(10)
            lbls.zero_()
            lbls[tr_labels_currset[i]] = 1.0
            loss = criterion(outputs.cpu(), lbls)
            running_loss += loss.item()
            loss.backward()
            optimizer.step()
            if (i + 1) % 50 == 0:
                print('Epoch [%d/%d], Step [%d/%d], Loss: %.5f' % (epoch + 
                    1, num_epochs, i + 1, trnum, running_loss))
                running_loss = 0
                print('Time elasped:', time.time() - start_time)
        correct = 0
        total = 0
        optimizer = lr_scheduler(optimizer, epoch, learning_rate, 40)
        for i in range(tsnum):
            images = ts_images_currset[:, :, :, i]
            lbls = torch.FloatTensor(10)
            lbls.zero_()
            lbls[ts_labels_currset[i]] = 1.0
            images = images.float().to(device)
            outputs = snn(images, time_window, 1)
            _, preds[i] = torch.max(outputs.data, 0)
        ct = preds[0:tsnum] == ts_labels_currset[0:tsnum]
        print('\n\n\n')
        print('Iters:', epoch, '\n')
        correct = ct.sum()
        print(correct)
        print('Test Accuracy of the model: %.3f  ' % (100 * correct.float() /
            tsnum))
        if correct.float() >= bestresult:
            bestresult = correct.float()
            if kfold == 0:
                sio.savemat('bestresults_preds_splitnet_k0.mat', {'preds':
                    preds})
                print('saved')
            if kfold == 1:
                sio.savemat('bestresults_preds_splitnet_k1.mat', {'preds':
                    preds})
                print('saved')
            if kfold == 2:
                sio.savemat('bestresults_preds_splitnet_k2.mat', {'preds':
                    preds})
                print('saved')
            if kfold == 3:
                sio.savemat('bestresults_preds_splitnet_k3.mat', {'preds':
                    preds})
                print('saved')
        print(torch.sum(snn.conv0.weight.data))
        print(torch.max(snn.conv0.weight.data))

import torch, time, os
import torch.nn as nn
import torch.nn.functional as F
import scipy.io as sio
import numpy as np
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
lens = 0.25
thresh = 0.3
decay = 0.3
num_classes = 10
batch_size = 1
num_epochs = 20
learning_rate = 0.0001
time_window = 100
names = 'animales_stbp'
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(device)
train_dataset = sio.loadmat(
    'C:\\Users\\Ajay\\Documents\\IBMdataset\\DvsGesture\\ibmframes.mat')
tr = train_dataset['myframes']
tr_images = torch.from_numpy(tr)
print(tr_images.shape)
train_labels = sio.loadmat(
    'C:\\Users\\Ajay\\Documents\\IBMdataset\\DvsGesture\\ibmlabels.mat')
trl = train_labels['labels']
tr_labels = torch.from_numpy(trl)


class ActFun(torch.autograd.Function):

    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input)
        return input.gt(thresh).float()

    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors
        grad_input = grad_output.clone()
        temp = abs(input - thresh) < lens
        return grad_input * temp.float() / (2 * lens)


def lr_scheduler(optimizer, epoch, init_lr=0.1, lr_decay_epoch=100):
    """Decay learning rate by a factor of 0.1 every lr_decay_epoch epochs."""
    if epoch % lr_decay_epoch == 0 and epoch > 1:
        for param_group in optimizer.param_groups:
            param_group['lr'] = param_group['lr'] * 0.1
    return optimizer


best_acc = 0
acc_record = list([])


class SNN_Model(nn.Module):

    def __init__(self, num_classes=10):
        super(SNN_Model, self).__init__()
        self.conv0 = nn.Conv2d(1, 10, kernel_size=5, stride=1, padding=2)
        self.conv2 = nn.Conv2d(1, 10, kernel_size=5, stride=1, padding=2)
        self.conv1 = nn.Conv2d(20, 25, kernel_size=5, stride=1, padding=2)
        self.conv3 = nn.Conv2d(20, 25, kernel_size=5, stride=1, padding=2)
        self.fc1 = nn.Linear(16 * 16 * 50, 100)
        self.fc3 = nn.Linear(16 * 16 * 50, 100)
        self.fc2 = nn.Linear(200, 10)

    def forward(self, input, win, bsize):
        c0_mem = c0_spike = torch.zeros(10, 64, 64, device=device)
        p0_mem = p0_spike = torch.zeros(10, 32, 32, device=device)
        c2_mem = c2_spike = torch.zeros(10, 64, 64, device=device)
        p2_mem = p2_spike = torch.zeros(10, 32, 32, device=device)
        c1_mem = c1_spike = torch.zeros(25, 32, 32, device=device)
        p1_mem = p1_spike = torch.zeros(25, 16, 16, device=device)
        c3_mem = c3_spike = torch.zeros(25, 32, 32, device=device)
        p3_mem = p3_spike = torch.zeros(25, 16, 16, device=device)
        h1_mem = h1_spike = h1_sumspike = torch.zeros(100, device=device)
        h3_mem = h3_spike = h3_sumspike = torch.zeros(100, device=device)
        h2_mem = h2_spike = h2_sumspike = torch.zeros(10, device=device)
        x = torch.zeros(input.size()).float()
        for step in range(win):
            x = input[:, :, step]
            x = x[:, :, None, None]
            x = x.permute(2, 3, 0, 1)
            c0_mem, c0_spike = mem_update(self.conv0, x, c0_mem, c0_spike,
                decay, step)
            p0_mem, p0_spike = mem_update_pool(F.avg_pool2d, c0_spike,
                p0_mem, p0_spike, decay, step)
            c2_mem, c2_spike = mem_update(self.conv2, x, c2_mem, c2_spike,
                decay, step)
            p2_mem, p2_spike = mem_update_pool(F.avg_pool2d, c2_spike,
                p2_mem, p2_spike, decay, step)
            if step == 0:
                p0spk_back5 = torch.zeros_like(p0_spike)
                p0spk_back4 = torch.zeros_like(p0_spike)
                p0spk_back3 = torch.zeros_like(p0_spike)
                p0spk_back2 = torch.zeros_like(p0_spike)
                p0spk_back1 = torch.zeros_like(p0_spike)
            if step > 5:
                p0spk_back5 = p0spk_back4
                p0spk_back4 = p0spk_back3
                p0spk_back3 = p0spk_back2
                p0spk_back2 = p0spk_back1
                p0spk_back1 = p0_spike
            p0spk_output = torch.cat([p0spk_back5, p2_spike], dim=1)
            c1_mem, c1_spike = mem_update(self.conv1, p0spk_output, c1_mem,
                c1_spike, decay, step)
            p1_mem, p1_spike = mem_update_pool(F.avg_pool2d, c1_spike,
                p1_mem, p1_spike, decay, step)
            c3_mem, c3_spike = mem_update(self.conv3, p0spk_output, c3_mem,
                c3_spike, decay, step)
            p3_mem, p3_spike = mem_update_pool(F.avg_pool2d, c3_spike,
                p3_mem, p3_spike, decay, step)
            if step == 0:
                p1spk_back5 = torch.zeros_like(p1_spike)
                p1spk_back4 = torch.zeros_like(p1_spike)
                p1spk_back3 = torch.zeros_like(p1_spike)
                p1spk_back2 = torch.zeros_like(p1_spike)
                p1spk_back1 = torch.zeros_like(p1_spike)
            if step > 5:
                p1spk_back5 = p1spk_back4
                p1spk_back4 = p1spk_back3
                p1spk_back3 = p1spk_back2
                p1spk_back2 = p1spk_back1
                p1spk_back1 = p1_spike
            p1spk_output = torch.cat([p1spk_back5, p3_spike], dim=1)
            x = p1spk_output.view(-1)
            h1_mem, h1_spike = mem_update(self.fc1, x, h1_mem, h1_spike,
                decay, step)
            h3_mem, h3_spike = mem_update(self.fc3, x, h3_mem, h3_spike,
                decay, step)
            h2_mem, h2_spike = mem_update(self.fc2, torch.cat([h1_spike,
                h3_spike], dim=0), h2_mem, h2_spike, decay, step)
            h2_sumspike += h2_spike
        outputs = h2_sumspike / time_window
        return outputs


def mem_update(conv, x, mem, spike, curr_decay, curr_step):
    mem = mem * (1 - spike) * curr_decay + conv(x)
    spike = act_fun(mem)
    return mem, spike


def mem_update_last(conv, x, mem, curr_decay, curr_step):
    mem = mem * curr_decay + conv(x)
    return mem


def mem_update_pool(opts, x, mem, spike, curr_decay, curr_step):
    mem = mem * float(time_window - curr_step) / float(time_window
        ) * curr_decay * (1 - spike) + opts(x, 2)
    spike = act_fun(mem)
    return mem, spike


snn = SNN_Model()
print(torch.sum(snn.conv0.weight.data))
print(torch.max(snn.conv0.weight.data))
snn.to(device)
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(snn.parameters(), lr=learning_rate)
act_fun = ActFun.apply
tr_images = tr_images.float()
ts_images = tr_images.float()
tr_labels = tr_labels.long()
for kfold in range(4):
    trnum = 914
    tsnum = 305
    if kfold == 0:
        tr_images_currset = tr_images[:, :, :, 0:914]
        ts_images_currset = tr_images[:, :, :, 914:1219]
        tr_labels_currset = tr_labels[0:914]
        ts_labels_currset = tr_labels[914:1219]
    if kfold == 1:
        tr_images_currset = torch.cat((tr_images[:, :, :, 0:305], tr_images
            [:, :, :, 610:1219]), 3)
        ts_images_currset = tr_images[:, :, :, 305:610]
        tr_labels_currset = torch.cat((tr_labels[0:305], tr_labels[610:1219
            ]), 0)
        ts_labels_currset = tr_labels[305:610]
    if kfold == 2:
        tr_images_currset = torch.cat((tr_images[:, :, :, 0:610], tr_images
            [:, :, :, 915:1219]), 3)
        ts_images_currset = tr_images[:, :, :, 610:915]
        tr_labels_currset = torch.cat((tr_labels[0:610], tr_labels[915:1219
            ]), 0)
        ts_labels_currset = tr_labels[610:915]
    if kfold == 3:
        tr_images_currset = tr_images[:, :, :, 305:1219]
        ts_images_currset = tr_images[:, :, :, 0:305]
        tr_labels_currset = tr_labels[305:1219]
        ts_labels_currset = tr_labels[0:305]
    preds = np.zeros_like(ts_labels_currset)
    bestresult = 0.0
    for epoch in range(num_epochs):
        running_loss = 0
        start_time = time.time()
        for i in range(trnum):
            images = tr_images_currset[:, :, :, i]
            snn.zero_grad()
            optimizer.zero_grad()
            images = images.float().to(device)
            outputs = snn(images, time_window, 1)
            ops = outputs.cpu()
            lbls = torch.FloatTensor(10)
            lbls.zero_()
            lbls[tr_labels_currset[i]] = 1.0
            loss = criterion(outputs.cpu(), lbls)
            running_loss += loss.item()
            loss.backward()
            optimizer.step()
            if (i + 1) % 50 == 0:
                print('Epoch [%d/%d], Step [%d/%d], Loss: %.5f' % (epoch + 
                    1, num_epochs, i + 1, trnum, running_loss))
                running_loss = 0
                print('Time elasped:', time.time() - start_time)
        correct = 0
        total = 0
        optimizer = lr_scheduler(optimizer, epoch, learning_rate, 40)
        for i in range(tsnum):
            images = ts_images_currset[:, :, :, i]
            lbls = torch.FloatTensor(10)
            lbls.zero_()
            lbls[ts_labels_currset[i]] = 1.0
            images = images.float().to(device)
            outputs = snn(images, time_window, 1)
            _, preds[i] = torch.max(outputs.data, 0)
        ct = preds[0:tsnum] == ts_labels_currset[0:tsnum]
        print('\n\n\n')
        print('Iters:', epoch, '\n')
        correct = ct.sum()
        print(correct)
        print('Test Accuracy of the model: %.3f  ' % (100 * correct.float() /
            tsnum))
        if correct.float() >= bestresult:
            bestresult = correct.float()
            if kfold == 0:
                sio.savemat('bestresults_preds_splitnet_k0.mat', {'preds':
                    preds})
                print('saved')
            if kfold == 1:
                sio.savemat('bestresults_preds_splitnet_k1.mat', {'preds':
                    preds})
                print('saved')
            if kfold == 2:
                sio.savemat('bestresults_preds_splitnet_k2.mat', {'preds':
                    preds})
                print('saved')
            if kfold == 3:
                sio.savemat('bestresults_preds_splitnet_k3.mat', {'preds':
                    preds})
                print('saved')
        print(torch.sum(snn.conv0.weight.data))
        print(torch.max(snn.conv0.weight.data))


import torch, time, os

import torch.nn as nn
import torch.nn.functional as F

#from graphviz import Digraph
import scipy.io as sio
import numpy as np

os.environ['CUDA_VISIBLE_DEVICES'] = "0"
lens = 0.25

thresh = 0.3
decay = 0.3

num_classes = 10
batch_size = 1
num_epochs = 20
learning_rate = 1e-4#########default 1e-4
time_window = 100
names = 'animales_stbp'

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print (device)


# load datasets need (num,pol, size,size, step)

train_dataset=sio.loadmat('C:\\Users\\Ajay\\Documents\\IBMdataset\\DvsGesture\\ibmframes.mat')
tr=train_dataset['myframes']

tr_images = torch.from_numpy(tr)

print(tr_images.shape) #64,64,100,1219





train_labels=sio.loadmat('C:\\Users\\Ajay\\Documents\\IBMdataset\\DvsGesture\\ibmlabels.mat')
trl=train_labels['labels']

tr_labels = torch.from_numpy(trl)


class ActFun(torch.autograd.Function):

    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input)
        #print('\nActFun forward called',input.size())
        return input.gt(thresh).float()

    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors
        grad_input = grad_output.clone()
        temp = abs(input - thresh) < lens
        #print('start \n',temp.size(),'\n', grad_input.size(), torch.mean(grad_input), 'fin \n')
        #print('actfun backward called')
        return grad_input * temp.float() / (2 * lens)        
        #return grad_input * 1 / (2 * lens)
        




def lr_scheduler(optimizer, epoch, init_lr=0.1, lr_decay_epoch=100):
    """Decay learning rate by a factor of 0.1 every lr_decay_epoch epochs."""

    if epoch % lr_decay_epoch == 0 and epoch > 1:
        for param_group in optimizer.param_groups:
            param_group['lr'] = param_group['lr'] * 0.1

    return optimizer


best_acc = 0
acc_record = list([])


class SNN_Model(nn.Module):

    def __init__(self, num_classes=10):
        super(SNN_Model, self).__init__()
        self.conv0 = nn.Conv2d(1,10, kernel_size=5, stride=1, padding=2)
        self.conv2 = nn.Conv2d(1,10, kernel_size=5, stride=1, padding=2)

        self.conv1 = nn.Conv2d(20, 25, kernel_size=5, stride=1, padding=2)
        self.conv3 = nn.Conv2d(20, 25, kernel_size=5, stride=1, padding=2)

        #conv0,conv1 -> upper with delay in conn, 
        #conv2,3 -> lower no delay
        
        #in_planes, out_planes, stride = cfg_cnn[2]
        #self.conv2 = nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=1, padding=1)

        self.fc1 = nn.Linear(16*16 * 50, 100)
        self.fc3 = nn.Linear(16*16 * 50, 100)

        self.fc2 = nn.Linear(200, 10)

    def forward(self, input, win, bsize):
        #input to be bsize x 28 x 28
        c0_mem = c0_spike = torch.zeros( 10, 64, 64, device=device)
        p0_mem = p0_spike = torch.zeros( 10, 32, 32, device=device)
                
        c2_mem = c2_spike = torch.zeros( 10, 64, 64, device=device)
        p2_mem = p2_spike = torch.zeros( 10, 32, 32, device=device)
        
        
        
        
        
        

        c1_mem = c1_spike = torch.zeros( 25, 32, 32, device=device)
        p1_mem = p1_spike = torch.zeros( 25, 16, 16, device=device)
        
        c3_mem = c3_spike = torch.zeros( 25, 32, 32, device=device)
        p3_mem = p3_spike = torch.zeros( 25, 16, 16, device=device)
        
        
        
        
        
        

        #c2_mem = c2_spike = torch.zeros(bsize, cfg_cnn[2][1], 7, 7, device=device)
        #p2_mem = p2_spike = torch.zeros(bsize, cfg_cnn[2][1], 7, 7, device=device)

        h1_mem = h1_spike = h1_sumspike = torch.zeros( 100, device=device)
        h3_mem = h3_spike = h3_sumspike = torch.zeros( 100, device=device)

        h2_mem = h2_spike = h2_sumspike = torch.zeros( 10, device=device)

        x=torch.zeros(input.size()).float()

        #firstspikesidx=torch.zeros(batch_size,5)
        #spkcnt=torch.zeros(batch_size).byte()
        #input = input.float()
        for step in range(win):
            #x = input[:, :, :, :, step]
            x = input[:, :, step]
            x=x[:,:,None,None]
            x=x.permute(2,3,0,1)

            c0_mem, c0_spike = mem_update(self.conv0, x, c0_mem, c0_spike,decay,step)
            p0_mem, p0_spike = mem_update_pool(F.avg_pool2d, c0_spike, p0_mem, p0_spike,decay,step)
            
            c2_mem, c2_spike = mem_update(self.conv2, x, c2_mem, c2_spike,decay,step)
            p2_mem, p2_spike = mem_update_pool(F.avg_pool2d, c2_spike, p2_mem, p2_spike,decay,step)
            
            if step==0:
                p0spk_back5=torch.zeros_like(p0_spike)
                p0spk_back4=torch.zeros_like(p0_spike)
                p0spk_back3=torch.zeros_like(p0_spike)
                p0spk_back2=torch.zeros_like(p0_spike)
                p0spk_back1=torch.zeros_like(p0_spike)
            if step>5:
                p0spk_back5=p0spk_back4
                p0spk_back4=p0spk_back3
                p0spk_back3=p0spk_back2
                p0spk_back2=p0spk_back1
                p0spk_back1=p0_spike
                
                
            p0spk_output=torch.cat([p0spk_back5,p2_spike],dim=1)
            
            
            

            c1_mem, c1_spike = mem_update(self.conv1, p0spk_output, c1_mem, c1_spike,decay,step)
            p1_mem, p1_spike = mem_update_pool(F.avg_pool2d, c1_spike, p1_mem, p1_spike,decay,step)
            
            c3_mem, c3_spike = mem_update(self.conv3, p0spk_output, c3_mem, c3_spike,decay,step)
            p3_mem, p3_spike = mem_update_pool(F.avg_pool2d, c3_spike, p3_mem, p3_spike,decay,step)
            
            
            
            if step==0:
                p1spk_back5=torch.zeros_like(p1_spike)
                p1spk_back4=torch.zeros_like(p1_spike)
                p1spk_back3=torch.zeros_like(p1_spike)
                p1spk_back2=torch.zeros_like(p1_spike)
                p1spk_back1=torch.zeros_like(p1_spike)
            if step>5:
                p1spk_back5=p1spk_back4
                p1spk_back4=p1spk_back3
                p1spk_back3=p1spk_back2
                p1spk_back2=p1spk_back1
                p1spk_back1=p1_spike
                
                
            p1spk_output=torch.cat([p1spk_back5,p3_spike],dim=1)
            
            
            
            
            
            
            #x = p1_spike.view( -1)
            x = p1spk_output.view( -1)

            h1_mem, h1_spike = mem_update(self.fc1, x, h1_mem, h1_spike,decay,step)
            h3_mem, h3_spike = mem_update(self.fc3, x, h3_mem, h3_spike,decay,step)
            #h1_sumspike += h1_spike

            #h2_mem = mem_update_last(self.fc2, h1_spike, h2_mem,decay,step)
            h2_mem, h2_spike = mem_update(self.fc2, torch.cat([h1_spike,h3_spike],dim=0), h2_mem, h2_spike,decay,step)
            h2_sumspike += h2_spike
  
        outputs = h2_sumspike / time_window
        
        return outputs




def mem_update(conv, x, mem, spike,curr_decay,curr_step):
    #mem = mem * curr_decay *float(time_window-curr_step)/float(time_window)* (1 - spike) + conv(x)
    mem = mem * (1 - spike) *curr_decay + conv(x)
    spike = act_fun(mem)
    return mem, spike

def mem_update_last(conv, x, mem,curr_decay,curr_step):
    mem = mem *curr_decay + conv(x)
    return mem




def mem_update_pool(opts, x, mem, spike,curr_decay,curr_step):
    mem = mem * float(time_window-curr_step)/float(time_window)*curr_decay*(1 - spike) + opts(x, 2)
    #mem = mem * (1 - spike) *curr_decay + opts(x, 2)
    spike = act_fun(mem)
    return mem, spike


snn = SNN_Model()
print(torch.sum(snn.conv0.weight.data))
print(torch.max(snn.conv0.weight.data))

snn.to(device)
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(snn.parameters(), lr=learning_rate)

act_fun = ActFun.apply

tr_images=tr_images.float()
ts_images = tr_images.float()
tr_labels=tr_labels.long()

for kfold in range(4):
    trnum=914
    tsnum=305
    
    if kfold==0:
        tr_images_currset=tr_images[:,:,:,0:914]
        ts_images_currset=tr_images[:,:,:,914:1219]
        tr_labels_currset=tr_labels[0:914]
        ts_labels_currset=tr_labels[914:1219]
        
    if kfold==1:
        tr_images_currset=torch.cat((tr_images[:,:,:,0:305],tr_images[:,:,:,610:1219]),3)
        ts_images_currset=tr_images[:,:,:,305:610]
        tr_labels_currset=torch.cat((tr_labels[0:305],tr_labels[610:1219]),0)
        ts_labels_currset=tr_labels[305:610]
        
    if kfold==2:
        tr_images_currset=torch.cat((tr_images[:,:,:,0:610],tr_images[:,:,:,915:1219]),3)
        ts_images_currset=tr_images[:,:,:,610:915]
        tr_labels_currset=torch.cat((tr_labels[0:610],tr_labels[915:1219]),0)
        ts_labels_currset=tr_labels[610:915]
        
    if kfold==3:
        tr_images_currset=tr_images[:,:,:,305:1219]
        ts_images_currset=tr_images[:,:,:,0:305]
        tr_labels_currset=tr_labels[305:1219]
        ts_labels_currset=tr_labels[0:305]
        
    
    preds=np.zeros_like(ts_labels_currset)
    bestresult=0.0
    for epoch in range(num_epochs):

        running_loss = 0
        start_time = time.time()

        #numiters = 1219#60000//batch_size
        #shortiter=numiters

        for i in range(trnum):
            #images = tr_images[i,:,:,:,:,:]
            images = tr_images_currset[:,:,:,i]
            #images = images[:,:,:,None,None]
            #images.permute(2,3,0,1)




            #print('\n','iteration num',i,'\n')
            snn.zero_grad()
            optimizer.zero_grad()
            #print(images[0,0,:,:])
            images = images.float().to(device)
            outputs = snn(images, time_window,1)
            ops = outputs.cpu()

            #print(outputs)
            lbls=torch.FloatTensor(10)
            lbls.zero_()
            lbls[tr_labels_currset[i]]=1.
            loss = criterion(outputs.cpu(), lbls)

            running_loss += loss.item()
            loss.backward()
            optimizer.step()

            if (i + 1) % 50 == 0:
                #print (torch.sum(outputs,1))
                print('Epoch [%d/%d], Step [%d/%d], Loss: %.5f'
                      % (epoch + 1, num_epochs, i + 1, trnum, running_loss))
                running_loss = 0
                print('Time elasped:', time.time() - start_time)
        correct = 0
        total = 0

        optimizer = lr_scheduler(optimizer, epoch, learning_rate, 40)


        for i in range(tsnum):
            images = ts_images_currset[:,:,:,i]
            
            
            lbls=torch.FloatTensor(10)
            lbls.zero_()
            lbls[ts_labels_currset[i]]=1.

            images = images.float().to(device)
            outputs = snn(images, time_window,1) #/ time_window
            
            _,preds[i]=torch.max(outputs.data,0)

            #_, predicted = torch.max(outputs.data, 0)
            #_, lbls = torch.max(lbls.data, 0)
            #total += lbls.size(0)
            #correct += (predicted.cpu() == lbls).sum()
            #correct += (predicted.cpu() == tr_labels).sum()
            #print(correct,total)


            #i +=1
        ct = (preds[0:tsnum]==ts_labels_currset[0:tsnum])
        #print(preds[0:shortiter])
        print('\n\n\n')
        #print(tr_labels[0:shortiter])
        #print('\n\n\n')
        #print(preds.shape,tr_labels.shape)
        #print(ct.shape)
        print('Iters:', epoch, '\n')
        correct=ct.sum()
        print(correct)
        print('Test Accuracy of the model: %.3f  ' % (100 * correct.float() / tsnum))
        
        if correct.float()>=bestresult:
            bestresult=correct.float()
            if kfold==0:
                sio.savemat('bestresults_preds_splitnet_k0.mat',{'preds':preds})
                print('saved')
            if kfold==1:
                sio.savemat('bestresults_preds_splitnet_k1.mat',{'preds':preds})
                print('saved')
            if kfold==2:
                sio.savemat('bestresults_preds_splitnet_k2.mat',{'preds':preds})
                print('saved')
            if kfold==3:
                sio.savemat('bestresults_preds_splitnet_k3.mat',{'preds':preds})
                print('saved')
            
        
        print(torch.sum(snn.conv0.weight.data))
        print(torch.max(snn.conv0.weight.data))







