#!/usr/bin/env python3

"""
RBE/CS549 Spring 2023: Computer Vision

Author(s): 
Uday Sankar (usankar@wpi.edu)
Mechatronics, Robotics and Automation Engineering,
Worcester Polytechnic Institute
"""

# Dependencies:
# opencv, do (pip install opencv-python)
# skimage, do (apt install python-skimage)
# termcolor, do (pip install termcolor)


import torch
import torchvision
from torch.utils.tensorboard import SummaryWriter
from torchvision import datasets, transforms
from torch.optim import AdamW
from torchvision.datasets import CIFAR10
import cv2
import sys
import os
import numpy as np
import random
import skimage
import PIL
import os
import glob
import random
from skimage import data, exposure, img_as_float
import matplotlib.pyplot as plt
import time
from torchvision.transforms import ToTensor
import argparse
import shutil
import string
from termcolor import colored, cprint
import math as m
from tqdm.notebook import tqdm
# import Misc.ImageUtils as iu
from Network.Network import CIFAR10Model
from Misc.MiscUtils import *
from Misc.DataUtils import *

# setting the device as 'cuda'
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
print(f'Device being used: {device}')

# global variables for testing
ModelNumber = 1
CheckPointPath = 'Checkpoints/BasicNet/14model.ckpt'  # BasiNet by default

# Don't generate pyc codes
sys.dont_write_bytecode = True

    
def GenerateBatch(TrainSet, TrainLabels, ImageSize, MiniBatchSize):
    """
    Inputs: 
    TrainSet - Variable with Subfolder paths to train files
    NOTE that Train can be replaced by Val/Test for generating batch corresponding to validation (held-out testing in this case)/testing
    TrainLabels - Labels corresponding to Train
    NOTE that TrainLabels can be replaced by Val/TestLabels for generating batch corresponding to validation (held-out testing in this case)/testing
    ImageSize is the Size of the Image
    MiniBatchSize is the size of the MiniBatch
   
    Outputs:
    I1Batch - Batch of images
    LabelBatch - Batch of one-hot encoded labels 
    """
    I1Batch = []
    LabelBatch = []
    
    ImageNum = 0
    while ImageNum < MiniBatchSize:
        # Generate random image
        RandIdx = random.randint(0, len(TrainSet)-1)
        
        ImageNum += 1
        
          ##########################################################
          # Add any standardization or data augmentation here!
          ##########################################################

        I1, Label = TrainSet[RandIdx]

        # Append All Images and Mask
        I1Batch.append(I1)
        LabelBatch.append(torch.tensor(Label))
        
    return torch.stack(I1Batch).to(device), torch.stack(LabelBatch).to(device)


def PrettyPrint(NumEpochs, DivTrain, MiniBatchSize, NumTrainSamples, LatestFile):
    """
    Prints all stats with all arguments
    """
    print('Number of Epochs Training will run for ' + str(NumEpochs))
    print('Factor of reduction in training data is ' + str(DivTrain))
    print('Mini Batch Size ' + str(MiniBatchSize))
    print('Number of Training Images ' + str(NumTrainSamples))
    if LatestFile is not None:
        print('Loading latest checkpoint with the name ' + LatestFile)              

    

def TrainOperation(TrainLabels, NumTrainSamples, ImageSize, ModelNumber,
                   NumEpochs, MiniBatchSize, SaveCheckPoint, CheckPointPath,
                   DivTrain, LatestFile, TrainSet, LogsPath, TestSet, TestLabels):
    """
    Inputs: 
    TrainLabels - Labels corresponding to Train/Test
    NumTrainSamples - length(Train)
    ImageSize - Size of the image
    NumEpochs - Number of passes through the Train data
    MiniBatchSize is the size of the MiniBatch
    SaveCheckPoint - Save checkpoint every SaveCheckPoint iteration in every epoch, checkpoint saved automatically after every epoch
    CheckPointPath - Path to save checkpoints/model
    DivTrain - Divide the data by this number for Epoch calculation, use if you have a lot of dataor for debugging code
    LatestFile - Latest checkpointfile to continue training
    TrainSet - The training dataset
    LogsPath - Path to save Tensorboard Logs
    Outputs:
    Saves Trained network in CheckPointPath and Logs to LogsPath
    """
    # Initialize the model
    model = CIFAR10Model(ModelNumber)
    model = model.to(device)
    ###############################################
    # Fill your optimizer of choice here!
    ###############################################
    Optimizer = AdamW(model.parameters(), lr=1e-3)

    # Tensorboard
    # Create a summary to monitor loss tensor
    Writer = SummaryWriter(LogsPath)
    # adding the graph of the model to tensorboard for visualization
    Writer.add_graph(model, GenerateBatch(TrainSet, TrainLabels, ImageSize, MiniBatchSize)[0])

    if LatestFile is not None:
        CheckPoint = torch.load(CheckPointPath + LatestFile + '.ckpt')
        # Extract only numbers from the name
        StartEpoch = int(''.join(c for c in LatestFile.split('a')[0] if c.isdigit()))
        model.load_state_dict(CheckPoint['model_state_dict'])
        print('Loaded latest checkpoint with the name ' + LatestFile + '....')
    else:
        StartEpoch = 0
        print('New model initialized....')


    for Epochs in tqdm(range(StartEpoch, NumEpochs)):
        train_acc_plot = 0
        train_loss_plot = 0
        NumIterationsPerEpoch = int(NumTrainSamples/MiniBatchSize/DivTrain)
        for PerEpochCounter in tqdm(range(NumIterationsPerEpoch)):
            Batch = GenerateBatch(TrainSet, TrainLabels, ImageSize, MiniBatchSize)
            
            # Predict output with forward pass
            LossThisBatch = model.training_step(Batch)

            Optimizer.zero_grad()
            LossThisBatch.backward()
            Optimizer.step()
            
            # Save checkpoint every some SaveCheckPoint's iterations
            if PerEpochCounter % SaveCheckPoint == 0:
                # Save the Model learnt in this epoch
                SaveName =  CheckPointPath + str(Epochs) + 'a' + str(PerEpochCounter) + 'model.ckpt'
                
                torch.save({'epoch': Epochs,'model_state_dict': model.state_dict(),'optimizer_state_dict': Optimizer.state_dict(),'loss': LossThisBatch}, SaveName)
                print('\n' + SaveName + ' Model Saved...')

            result = model.validation_step(Batch)
            train_acc_plot += result['acc']
            train_loss_plot += result['loss']
            model.epoch_end(Epochs*NumIterationsPerEpoch + PerEpochCounter, result)
            # Tensorboard
            # Writer.add_scalar('LossEveryIter', result["loss"], Epochs*NumIterationsPerEpoch + PerEpochCounter)
            # Writer.add_scalar('Accuracy', result["acc"], Epochs*NumIterationsPerEpoch + PerEpochCounter)
            
            # # If you don't flush the tensorboard doesn't update until a lot of iterations!
            # Writer.flush()
        
        

        model.eval()
        test_loss_plot = 0
        test_acc_plot = 0
        with torch.no_grad():
            Batch_ = GenerateBatch(TestSet, TestLabels, ImageSize, MiniBatchSize)
            test_result = model.validation_step(Batch_)
            test_loss_plot += test_result["loss"]
            test_acc_plot += test_result["acc"]


        Writer.add_scalars('Training/Testing Loss', {'TrainLossPerEpoch': train_loss_plot/NumIterationsPerEpoch, 'TestLossPerEpoch': test_loss_plot}, Epochs)  
        Writer.add_scalars('Training/Testing Accuracy', {'TrainAccuracyPerEpoch': train_acc_plot/NumIterationsPerEpoch, 'TestAccuracyPerEpoch': test_acc_plot}, Epochs) 
        Writer.flush() 


        # Save model every epoch
        SaveName = CheckPointPath + str(Epochs) + 'model.ckpt'
        torch.save({'epoch': Epochs,'model_state_dict': model.state_dict(),'optimizer_state_dict': Optimizer.state_dict(),'loss': LossThisBatch}, SaveName)
        print('\n' + SaveName + ' Model Saved...')
        

def main():
    """
    Inputs: 
    None
    Outputs:
    Runs the Training and testing code based on the Flag
    """
    # Parse Command Line arguments
    Parser = argparse.ArgumentParser()
    Parser.add_argument('--CheckPointPath', default='./Checkpoints/', help='Path to save Checkpoints, Default: ../Checkpoints/')
    Parser.add_argument('--NumEpochs', type=int, default=15, help='Number of Epochs to Train for, Default:15')
    Parser.add_argument('--DivTrain', type=int, default=1, help='Factor to reduce Train data by per epoch, Default:1')
    Parser.add_argument('--MiniBatchSize', type=int, default=32, help='Size of the MiniBatch to use, Default:32')
    Parser.add_argument('--LoadCheckPoint', type=int, default=0, help='Load Model from latest Checkpoint from CheckPointsPath?, Default:0')
    Parser.add_argument('--LogsPath', default='./Logs/', help='Path to save Logs for Tensorboard, Default=Logs/')
    Parser.add_argument('--ModelNumber', default=1, help='Model Type: 1-BasicNet, 2-BatchNormNet, 3-ResNet, 4-ResNeXt, 5-DenseNet, Default:1')

    Args = Parser.parse_args()
    NumEpochs = Args.NumEpochs
    DivTrain = float(Args.DivTrain)
    MiniBatchSize = Args.MiniBatchSize
    LoadCheckPoint = Args.LoadCheckPoint
    CheckPointPath = Args.CheckPointPath
    LogsPath = Args.LogsPath

    model_number_dict = {1: 'BasicNet', 2: 'BatchNormNet', 3: 'ResNet', 4: 'ResNeXt', 5: 'DenseNet'}
    global ModelNumber
    ModelNumber = int(Args.ModelNumber)
    # ModelNumber = int(input("Select model type (1-BasicNet, 2-BatchNormNet, 3-ResNet, 4-ResNeXt, 5-Densenet): "))
    # print(f"Selected Model Type: {model_number_dict[ModelNumber]}")

    # # setting the checkpoint and logs path based on the model selected
    CheckPointPath = CheckPointPath + model_number_dict[ModelNumber] + '/'
    LogsPath = LogsPath + model_number_dict[ModelNumber] + '/'
    
    # Default Hyperparameters
    NumEpochs = 15
    TrainSet = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=ToTensor())
    TestSet = CIFAR10(root='data/', train=False, transform=ToTensor())

    TestLabelsPath = "./TxtFiles/LabelsTest.txt"

    TestLabelsPath = "./TxtFiles/LabelsTest.txt"
    TestLabels = ReadLabels(TestLabelsPath)

    # Setup all needed parameters including file reading
    # SaveCheckPoint, ImageSize, NumTrainSamples, TrainLabels, NumClasses = SetupAll(TestLabelsPath, CheckPointPath)
    SaveCheckPoint, ImageSize, NumTrainSamples, TrainLabels, NumClasses = SetupAll(TestLabelsPath, CheckPointPath)


    # Find Latest Checkpoint File
    if LoadCheckPoint==1:
        LatestFile = FindLatestModel(CheckPointPath)
    else:
        LatestFile = None
    
    # Pretty print stats
    PrettyPrint(NumEpochs, DivTrain, MiniBatchSize, NumTrainSamples, LatestFile)

    TrainOperation(TrainLabels, NumTrainSamples, ImageSize, ModelNumber,
                NumEpochs, MiniBatchSize, SaveCheckPoint, CheckPointPath,
                DivTrain, LatestFile, TrainSet, LogsPath, TestSet, TestLabels)

    
if __name__ == '__main__':
    main()
 
