# %% [code]
"""
Author: Anish Aritakula
Title: Implementation of Convolutional Neural Network from scratch(almost).

"""

import torch
import pandas as pd
from itertools import product

import torch.nn.functional as F
import time
torch.set_printoptions(precision=4, sci_mode=False)
import numpy as np

#Logging my run results with key parameters
import wandb
wandb.login()


eps=1e-7

#Function to calculate cross entropy loss for multi-class classificaton
def _cross_entropy_loss(pred,act,n_classes):
    pred=torch.clamp(pred,min=0)
    one_hot=F.one_hot(act,num_classes=n_classes)
    log_pred=torch.log2(pred+eps)
    #print(pred,log_pred,one_hot)
    return torch.sum(one_hot*log_pred)*-1

#Functuon to calculate model accuracy
def _accuracy(pred,act):
    preds_idx=torch.argmax(pred,dim=1)
    return (preds_idx==act).sum()/act.shape[0]


class Convolution:
    def __init__(self,filter_shape,n_filters):
        self.filter_length,self.filter_width=filter_shape
        self.n_filters=n_filters
        self.filter=torch.tensor(np.random.uniform(-0.5,0.5,(self.filter_length,self.filter_width,self.n_filters)),dtype=torch.float32,requires_grad=True)
    
    #This function is invoked when the Convolution object is invoked
    def __call__(self,Images):
        num_images,len_images,wid_images=Images.shape
        conv_output=torch.zeros((Images.shape[0],len_images - self.filter_length + 1,wid_images - self.filter_width + 1,self.n_filters))
        
        
        for b in range(num_images):
            for k in range(self.n_filters):
                for i,j in product(range(len_images),range(wid_images)):
                    if i+self.filter_length<len_images and j+self.filter_width<wid_images:
                        conv_output[b,i,j,k]=max(0,torch.sum(Images[b,i:i+self.filter_length,j:j+self.filter_width]*self.filter[:,:,k]))
        
        return conv_output


class PoolingLayer:
    def __init__(self,pooling_shape):
        self.pool_len,self.pool_wid=pooling_shape
    
    def __call__(self,Input):
        batch_size,max_i,max_j,n_filters=Input.shape
        new1=torch.zeros((batch_size,max_i//2,max_j//2,n_filters))

        for b in range(batch_size):
            for k in range(n_filters):
                for i,j in product(range(new1.shape[0]),range(new1.shape[1])):
                    if self.pool_len*i<max_i and self.pool_wid*j<max_j:
                        new1[b,i,j,k]=torch.max(Input[b,i*self.pool_len:i*self.pool_len+self.pool_len,
                                                       j*self.pool_wid:j*self.pool_wid+self.pool_wid,k])
        return new1

class NeuralNetLayer:
    def __init__(self,network_shape):
        self.n_neurons,self.n_activations=network_shape
        self.weights=torch.tensor(np.random.uniform(-0.5,0.5,(self.n_neurons,self.n_activations)),dtype=torch.float32,requires_grad=True)
        self.bias=torch.tensor(np.random.uniform(-0.5,0.5,self.n_activations),dtype=torch.float32,requires_grad=True)
    
    #Data is first flattened and then a neuaral network is created
    def __call__(self,Input):
        flattend_data=Input.reshape(Input.shape[0],-1)
        raw_val_matrix=torch.matmul(flattend_data,self.weights)+self.bias
        return raw_val_matrix

class SoftmaxLayer:
    def __init__(self,num_classes):
        self.num_classes=num_classes
    
    def __call__(self,Input):
        #Input normalization first to provide numerical stability
        Input_norm=Input-torch.max(Input,dim=1,keepdim=True).values
        return torch.exp(Input_norm)/torch.sum(torch.exp(Input_norm),dim=1,keepdim=True)

class CNN:
    def __init__(self):
        self.conv=Convolution((3,3),6)
        self.max_pool=PoolingLayer((2,2))
        self.nn=NeuralNetLayer((169*6,10))
        self.sm=SoftmaxLayer(10)
    
    def __call__(self,Images):
        x=self.conv(Images)
        x=self.max_pool(x)
        x=self.nn(x)
        x=self.sm(x)
        return x

def main():
    st=time.time()
    data=pd.read_csv('/kaggle/input/digit-recognizer/train.csv')

    Y=data.iloc[:1000,0].values
    X=data.iloc[:1000,1:].values

    all_idx=np.array(range(X.shape[0]))
    train_idx=np.argsort(np.random.rand(int(X.shape[0]*.75)))
    test_idx=np.setdiff1d(all_idx,train_idx)

    X_train,Y_train=torch.tensor(X[train_idx].reshape(-1,28,28)/255),torch.tensor(Y[train_idx])
    X_test,Y_test=torch.tensor(X[test_idx].reshape(-1,28,28)/255),torch.tensor(Y[test_idx])
    print(X_train.shape,Y_train.shape,X_test.shape,Y_test.shape)
    

    n_epochs=10
    batch_size=32
    learning_rate=.01
    n_classes=10

    cnn_net=CNN()
    wandb.init(
    # Set the project where this run will be logged
    project="CNN-train_test", 
    # We pass a run name (otherwise it’ll be randomly assigned, like sunshine-lollypop-10)
    name=f"experiment_2", 
    # Track hyperparameters and run metadata
    config={
      "learning_rate": learning_rate,
      "architecture": "CNN",
      "dataset": "MNIST-kaggle",
      "epochs": n_epochs,
      "batch_size":batch_size,
      "records":X.shape[0],
      "params_on":"Weights+Biases+n_filters"})
    

    #Training process here
    """Mini batch gradient descent"""
    for epoch in range(n_epochs):
        #Shuffling data before each epoch
        shuffled_idx=torch.argsort(torch.rand(X_train.shape[0]))
        shuffled_X=X_train[shuffled_idx]
        shuffled_Y=Y_train[shuffled_idx]
        epoch_loss=0

        for batch_num in range(X_train.shape[0]//batch_size + 1):
                
                batch_start=batch_num*batch_size
                batch_end=batch_start + batch_size

                #Forward pass
                pred_probs=cnn_net(shuffled_X[batch_start:batch_end])

                #Backward pass
                loss=_cross_entropy_loss(pred_probs,shuffled_Y[batch_start:batch_end],n_classes)
                loss.backward()
                epoch_loss+=loss

                #Parameters update
                with torch.no_grad():
                    cnn_net.conv.filter-=cnn_net.conv.filter.grad*learning_rate
                    cnn_net.nn.weights-=cnn_net.nn.weights.grad*learning_rate
                    cnn_net.nn.bias-=cnn_net.nn.bias.grad*learning_rate
                
                #Zeroing out gradients after each mini batch
                cnn_net.conv.filter.grad.zero_()
                cnn_net.nn.weights.grad.zero_()
                cnn_net.nn.bias.grad.zero_()
            
        print(f"The total train loss after epoch {epoch} is {epoch_loss}",flush=True)
        with torch.no_grad():
            pred_probs_epoch=cnn_net(X_train)
            accuracy=_accuracy(pred_probs_epoch,Y_train)

            pred_probs_test=cnn_net(X_test)
            test_loss=_cross_entropy_loss(pred_probs_test,Y_test,n_classes)
            accuracy_test=_accuracy(pred_probs_test,Y_test)
        print(f"The total val loss after epoch {epoch} is {test_loss}",flush=True)
        print(f"The total train accuracy after epoch {epoch} is {accuracy}",flush=True)
        print(f"The total test accuracy after epoch {epoch} is {accuracy_test}",flush=True)
        print("\n")
        wandb.log({"train_acc": accuracy, "train_loss": epoch_loss,"filter":
                   cnn_net.conv.filter,"bias":cnn_net.nn.bias,"weights":cnn_net.nn.weights,
                   "test_loss":test_loss,'test_accuracy':accuracy_test})
    et=time.time()
    print(f"Execution time is {et-st}")
    wandb.log({'execution_time':et-st})
    wandb.finish()


if __name__=="__main__":
    main()