# Training DNC

import sys

import time

import torch

import random

import numpy as np

# torch.autograd.set_detect_anomaly(True) # Setting Anomaly Detection True for finding bad operations

####### Following function is adapted from the implemention by loudinthecloud on Github ##########

def random_seed():

    seed = int(time.time()*10000000)

    random.seed(seed)

    np.random.seed(int(seed/10000000))      # NumPy seed Range is 2**32 - 1 max

    torch.manual_seed(seed)

##########################################################################################################

def main():

    if len(sys.argv) > 4:

        if sys.argv[1] == "1":

            if sys.argv[2] == "GPU" and torch.cuda.is_available():  # Checking if GPU Request is given or not and availability of CUDA

                from tasks.babi_task_GPU import task_babi

            elif sys.argv[2] == "CPU":

                from tasks.babi_task import task_babi

            else:

                print("Please specify the run device (GPU/CPU)")

                exit()

            c_task = task_babi()                    # Initialization of the bAbI Task

            print("\nStarting bAbI Question Answering Task for DNC\n")

        elif sys.argv[1] == "2":

            if sys.argv[2] == "GPU" and torch.cuda.is_available():  # Checking if GPU Request is given or not and availability of CUDA

                from tasks.ner_task_bio_GPU import task_NER

            elif sys.argv[2] == "CPU":

                from tasks.ner_task_bio import task_NER

            else:

                print("Please specify the run device (GPU/CPU)")

                exit()

            c_task = task_NER()                    # Initialization of the NER Task (This is for BIO Tagging. Edit at line 33 and 35 for BIEOS tagging)

            print("\nStarting Medical NER Task for DNC\n")

        else:

            print("Unidentified task, please refer README file")

            exit()

    else:

        print("Incorrect Number of arguments")

        exit()

    epoch = sys.argv[3] # Last Epoch number till the model was trained (eg: 0)

    batch = sys.argv[4] # Last Batch Number till the model was trained (eg: 1000)

    batch_size = 1

    # Random Seed

    random_seed()

    c_task.init_dnc()

    c_task.init_loss()

    c_task.batch_size = batch_size

    c_task.load_model(2, epoch, batch)

    results = c_task.test_model()

    print(results)

if __name__ == '__main__':

    main()