# -*- coding: utf-8 -*-
"""Assignment2CV.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1qq-TVQGZgr6KqMs6ZlxBIl7QhkvADzB9
"""

import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
import numpy as np
import torchvision
from torchvision import datasets, models, transforms
import matplotlib.pyplot as plt
import time
import os
import copy
import itertools
from sklearn.metrics import confusion_matrix
from sklearn.svm import LinearSVC
from sklearn.multiclass import OneVsRestClassifier

from google.colab import drive
drive.mount('/content/drive')

data_transforms = {
    'train': transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
    'validation': transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
    'test': transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ]),
}

def plot_confusion_matrix(cm, classes,
                          cmap=plt.cm.Blues):
   

    plt.imshow(cm, interpolation='nearest', cmap=cmap)
    plt.colorbar()
    tick_marks = np.arange(len(classes))
    plt.xticks(tick_marks, classes, rotation=45)
    plt.yticks(tick_marks, classes)
    
    threshold = cm.max() / 2.
    for i, j in itertools.product(range(cm.shape[0]), range(cm.shape[1])):
        plt.text(j, i, format(cm[i, j]),horizontalalignment="center",color="white" if cm[i, j] > threshold else "black")

    plt.ylabel('True label')
    plt.xlabel('Predicted label')

class Identity(nn.Module):
    def __init__(self):
        super(Identity, self).__init__()
        
    def forward(self, x):
        return x

def retrievefeatures(model,phase):
    model.eval()
    model.to(device)
    features = []
    feature_classes = []
    with torch.no_grad():
        for i, (inputs, classes) in enumerate(dataloaders[phase]):
            inputs = inputs.to(device)
            outputs = model(inputs)
            feature_classes.extend(classes.cpu().numpy())
            features.extend(outputs.cpu().numpy())
        

    return features,feature_classes

def getdataloader_sizes(batchsize): 
  dataset_directory = 'drive/My Drive/dataset'
  image_datasets = {x: datasets.ImageFolder(os.path.join(dataset_directory, x),data_transforms[x])for x in ['train', 'validation', 'test']}
  #Batch size is set as 64 
  dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=batchsize,
                                               shuffle=True, num_workers=8)
                for x in ['train', 'validation' ,'test']}
  dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'validation','test']}

  print({x: len(image_datasets[x]) for x in ['train', 'validation','test']})
  class_names = image_datasets['train'].classes
  return dataloaders,dataset_sizes,class_names

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

model= models.vgg16(pretrained=True)
for param in model.features.parameters():
    param.requires_grad = False

def calculateTestAcc(trained_model,dataloaders,dataset_sizes):
  confusion_matrixx = torch.zeros(10, 10)
  np.set_printoptions(precision=2)
  current_phase_correct_outputnumber = 0
  topk = 0
  with torch.no_grad():
    for i, (inputs, classes) in enumerate(dataloaders['test']):
        inputs = inputs.to(device)
        classes = classes.to(device)
        outputs = trained_model(inputs)
        _, preds = torch.max(outputs, 1)
        current_phase_correct_outputnumber += torch.sum(preds == classes.data)
        
        probabilities,labels = outputs.topk(5,dim=1)
        
        classes_size = labels.size(0)
        for p in range(classes_size):
          if classes[p] in labels[p]:
            topk+=1
          
        for t, p in zip(classes.view(-1), preds.view(-1)):
            confusion_matrixx[t.long(), p.long()] += 1
    #### Top 1 score
    test_acc = 100*current_phase_correct_outputnumber.double() / dataset_sizes['test']
    #### top5 score
    top5_score = 100*topk/dataset_sizes['test']
    
    #Top 1 and Top 5 accuracies printed
    print('Test Acc: {:4f}'.format(test_acc))
    print('Top5 Acc: {:4f}'.format(top5_score))
  #Plot size is set
  plt.figure(figsize = (10,10))
  plot_confusion_matrix(confusion_matrixx,classes=class_names)
  plt.show()

def plot_graph(plotlist1,plotlist2,ylabel):
   #Plot accuracy graph 
    plt.xlabel("Training Epochs")
    plt.ylabel(ylabel)
    plt.plot(plotlist1, color="green")
    plt.plot(plotlist2, color="yellow")
    
    plt.gca().legend(('Train', 'Validation'))
    plt.show()

def train_model(model, criterion, optimizer, epoch_number,device,earlystopping):
   
    model.to(device)
    best_model_wts = copy.deepcopy(model.state_dict())
    best_train_acc = 0.0
    best_val_acc = 0.0
    best_test_acc = 0.0
    train_acc_history = list()
    train_loss_history =list()
    val_acc_history = list()
    val_loss_history =list()
    
    counter = 0
    stop =False
    best_loss = None
    
    #early stopping
    n_epochs_stop = 1
    min_val_loss = np.Inf
    epochs_no_improve = 0
    
    for epoch in range(epoch_number):
        if stop:
          break
        print('Epoch {}/{}'.format(epoch, epoch_number - 1))
        
        # Each epoch has a training and validation phase
        for part in ['train', 'validation']:
            if part == 'train':
                
                model.train()  
            else:
                model.eval()  

            current_loss = 0.0
            current_phase_correct_outputnumber = 0
            # For each phase in datasets are iterated
            for inputs, labels in dataloaders[part]:
                inputs = inputs.to(device)
                labels = labels.to(device)
                # zero the parameter gradients
                optimizer.zero_grad()
                # forward
                with torch.set_grad_enabled(part == 'train'):
                    outputs = model(inputs)
                    _, preds = torch.max(outputs, 1)
                    loss = criterion(outputs, labels)
                    # Backpropagate and opitimize Training part
                    if part == 'train':
                        loss.backward()
                        optimizer.step()

                # statistics
                current_loss += loss.item() * inputs.size(0)
                current_phase_correct_outputnumber += torch.sum(preds == labels.data)

            current_loss = current_loss / dataset_sizes[part]
            epoch_acc = 100*current_phase_correct_outputnumber.double() / dataset_sizes[part]

            if part == 'validation':
                val_acc_history.append(epoch_acc)
                val_loss_history.append(current_loss)
                if earlystopping:
                  # If the validation loss is at a minimum
                  if current_loss < min_val_loss:
                    # Save the model
                    epochs_no_improve = 0
                    min_val_loss = current_loss

                  else:
                    epochs_no_improve += 1
                    # Check early stopping condition
                    if epochs_no_improve == n_epochs_stop:
                      print('Early stopping!')
                      
                      #Printed best accuracies
                      print('Best train Acc: {:4f}'.format(best_train_acc))
                      print('Best validation Acc: {:4f}'.format(best_val_acc))

                      print()

                      #Printed best accuracies
                      print('Best train Acc: {:4f}'.format(best_train_acc))
                      print('Best validation Acc: {:4f}'.format(best_val_acc))

                      # load best model weights
                      model.load_state_dict(best_model_wts)
                      #Plot accuracy graph 
                      plot_graph(train_acc_history,val_acc_history,"Accuracy")
                      plot_graph(train_loss_history,val_loss_history,"Loss")
                      
                      return model                  
            else:
                train_acc_history.append(epoch_acc)
                train_loss_history.append(current_loss)

            print('{} Loss: {:.4f} Acc: {:.4f}'.format(
                part, current_loss, epoch_acc))

            # deep copy the model
            if part == 'train' and epoch_acc > best_train_acc:
                  best_train_acc = epoch_acc
                
            if part == 'validation' and epoch_acc > best_val_acc:             
                best_val_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())
              
        print() 
    
    print('Best train Acc: {:4f}'.format(best_train_acc))
    print('Best validation Acc: {:4f}'.format(best_val_acc))
              
    print()
    #Printed best accuracies
    print('Best train Acc: {:4f}'.format(best_train_acc))
    print('Best validation Acc: {:4f}'.format(best_val_acc))
    
    # load best model weights
    model.load_state_dict(best_model_wts)
    #Plot accuracy graph 
    plot_graph(train_acc_history,val_acc_history,"Accuracy")
    plot_graph(train_loss_history,val_loss_history,"Loss")
  
    return model

trainingmodel= models.vgg16(pretrained=True)
for param in trainingmodel.features.parameters():
    param.requires_grad = False
    
learning_rate = 0.001
epoch = 20
batchsize = 32
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(trainingmodel.parameters(), lr=learning_rate)

earlystoping = True

num_ftrs = trainingmodel.classifier[6].in_features

#Last layer is removed and modified with a linear layer
trainingmodel.classifier[6] = nn.Linear(num_ftrs,10)

dataloaders,dataset_sizes,class_names = getdataloader_sizes(batchsize)


trained_model = train_model(trainingmodel, criterion, optimizer,epoch,device,earlystoping)

calculateTestAcc(trained_model,dataloaders,dataset_sizes)

#PART 3
#Removed effect of last 2 layer by putting identity layer
trained_model.classifier[5] = Identity()
trained_model.classifier[6] = Identity()  

x_test ,y_test =retrievefeatures(trained_model,'test')
x_train,y_train = retrievefeatures(trained_model,'train')

clf = LinearSVC(max_iter = 100000)

classifier = clf.fit(x_train, y_train)
y_preds = clf.predict(x_test)

###Overall accuracy printed
print('Part3 Acc: {:.2f}'.format(100*clf.score(x_test,y_test)))

###Confusion matrix is created and plotted
conf_matrix = confusion_matrix(y_preds,y_test)
plt.figure(figsize = (10,10))
plot_confusion_matrix(conf_matrix,class_names)
for clas,i in zip(class_names,conf_matrix.diagonal()/conf_matrix.sum(axis=1)):
    print("Accuracy of "+clas+" : ", 100*i)

##########   PART1   ######
#Removed effect of last 2 layer by putting identity layer
model.classifier[5] = Identity()
model.classifier[6] = Identity()  

x_test ,y_test =retrievefeatures(model,'test')
x_train,y_train = retrievefeatures(model,'train')


clf = LinearSVC(max_iter = 100000)
classifier = clf.fit(x_train, y_train)
y_preds = clf.predict(x_test)
print('Part1 Acc: {:.2f}'.format(100*clf.score(x_test,y_test)))
#Created confusion matrix from predictions and test values
conf_matrix = confusion_matrix(y_preds,y_test)
#Plots confusion matrix
plt.figure(figsize = (10,10))
plot_confusion_matrix(conf_matrix,class_names)
#prints for each class accuracy
for clas,i in zip(class_names,conf_matrix.diagonal()/conf_matrix.sum(axis=1)):
    print("Accuracy of "+clas+" : ", 100*i)