# --------------------------------------------------------
# mcan-vqa (Deep Modular Co-Attention Networks)
# Licensed under The MIT License [see LICENSE for details]
# Written by Yuhao Cui https://github.com/cuiyuhao1996
# --------------------------------------------------------

#from core.data.load_data_swapmix import DataSet
#from core.model.net import Net
from core.model.optim import get_optim, adjust_lr
from core.data.data_utils import shuffle_list
from core.dataloader.frcnn_modify import FRCNN

import os, json, torch, datetime, pickle, copy, shutil, time
import numpy as np
import torch.nn as nn
import torch.utils.data as Data
import pdb
from numpyencoder import NumpyEncoder
from collections import OrderedDict


class Execution:
    def __init__(self, __C):
        self.__C = __C
        if self.__C.features == 'frcnn' :
            from core.data.load_data_frcnn_swapmix import DataSet
            from core.model.net import Net

            ## Loading fastrcnn featres
            self.frcnn = FRCNN(__C)

        elif self.__C.features == 'scene_graph' : 
            from core.data.load_data_scene_swapmix import DataSet
            from core.model.net_scene import Net


        self.Net = Net
        print('Loading training set ........')
        self.dataset = DataSet(__C)


        self.dataset_eval = None
        if __C.EVAL_EVERY_EPOCH:
            __C_eval = copy.deepcopy(__C)
            setattr(__C_eval, 'RUN_MODE', 'val')

            print('Loading validation set for per-epoch evaluation ........')
            self.dataset_eval = DataSet(__C_eval)


    def train(self, dataset, dataset_eval=None):

        # Obtain needed information
        data_size = dataset.data_size
        token_size = dataset.token_size
        ans_size = dataset.ans_size
        pretrained_emb = dataset.pretrained_emb

        # Define the MCAN model
        net = self.Net(
            self.__C,
            pretrained_emb,
            token_size,
            ans_size
        )
        net.cuda()
        net.train()

        # Define the multi-gpu training if needed
        if self.__C.N_GPU > 1:
            net = nn.DataParallel(net, device_ids=self.__C.DEVICES)

        # Define the binary cross entropy loss
        # loss_fn = torch.nn.BCELoss(size_average=False).cuda()
        loss_fn = torch.nn.BCELoss(reduction='sum').cuda()

        # Load checkpoint if resume training
        if self.__C.RESUME:
            print(' ========== Resume training')

            if self.__C.CKPT_PATH is not None:
                print('Warning: you are now using CKPT_PATH args, '
                      'CKPT_VERSION and CKPT_EPOCH will not work')

                path = self.__C.CKPT_PATH
            else:
                path = self.__C.CKPTS_PATH + \
                       'ckpt_' + self.__C.CKPT_VERSION + \
                       '/epoch' + str(self.__C.CKPT_EPOCH) + '.pkl'

            # Load the network parameters
            print('Loading ckpt {}'.format(path))
            ckpt = torch.load(path)
            print('Finish!')
            net.load_state_dict(ckpt['state_dict'])

            # Load the optimizer paramters
            optim = get_optim(self.__C, net, data_size, ckpt['lr_base'])
            optim._step = int(data_size / self.__C.BATCH_SIZE * self.__C.CKPT_EPOCH)
            optim.optimizer.load_state_dict(ckpt['optimizer'])

            start_epoch = self.__C.CKPT_EPOCH

        else:
            if ('ckpt_' + self.__C.VERSION) in os.listdir(self.__C.CKPTS_PATH):
                shutil.rmtree(self.__C.CKPTS_PATH + 'ckpt_' + self.__C.VERSION)

            os.mkdir(self.__C.CKPTS_PATH + 'ckpt_' + self.__C.VERSION)

            optim = get_optim(self.__C, net, data_size)
            start_epoch = 0

        loss_sum = 0
        named_params = list(net.named_parameters())
        grad_norm = np.zeros(len(named_params))

        # Define multi-thread dataloader
        if self.__C.SHUFFLE_MODE in ['external']:
            dataloader = Data.DataLoader(
                dataset,
                #batch_size=4,
                batch_size=self.__C.BATCH_SIZE,
                shuffle=False,
                num_workers=self.__C.NUM_WORKERS,
                pin_memory=self.__C.PIN_MEM,
                drop_last=True
            )
        else:
            dataloader = Data.DataLoader(
                dataset,
                batch_size=self.__C.BATCH_SIZE,
                shuffle=True,
                num_workers=self.__C.NUM_WORKERS,
                pin_memory=self.__C.PIN_MEM,
                drop_last=True
            )

        # Training script
        for epoch in range(start_epoch, self.__C.MAX_EPOCH):

            # Save log information
            logfile = open(
                self.__C.LOG_PATH +
                'log_run_' + self.__C.VERSION + '.txt',
                'a+'
            )
            logfile.write(
                'nowTime: ' +
                datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S') +
                '\n'
            )
            logfile.close()

            # Learning Rate Decay
            if epoch in self.__C.LR_DECAY_LIST:
                adjust_lr(optim, self.__C.LR_DECAY_R)

            # Externally shuffle
            if self.__C.SHUFFLE_MODE == 'external':
                shuffle_list(dataset.ans_list)

            time_start = time.time()

            if self.__C.features == 'frcnn' :

                # Iteration
                for step, (
                        feature_iter,
                        ques_ix_iter,
                        ans_iter
                ) in enumerate(dataloader):
                    #feature_iter, ques_ix_iter, ans_iter = dataset.extract(imp)
                    optim.zero_grad()

                    feature_iter = feature_iter.cuda()
                    ques_ix_iter = ques_ix_iter.cuda()
                    ans_iter = ans_iter.cuda()

                    for accu_step in range(self.__C.GRAD_ACCU_STEPS):

                        sub_feature_iter = \
                            feature_iter[accu_step * self.__C.SUB_BATCH_SIZE:
                                          (accu_step + 1) * self.__C.SUB_BATCH_SIZE]
                        sub_ques_ix_iter = \
                            ques_ix_iter[accu_step * self.__C.SUB_BATCH_SIZE:
                                         (accu_step + 1) * self.__C.SUB_BATCH_SIZE]
                        sub_ans_iter = \
                            ans_iter[accu_step * self.__C.SUB_BATCH_SIZE:
                                     (accu_step + 1) * self.__C.SUB_BATCH_SIZE]

                        pred = net(
                            sub_feature_iter,
                            sub_ques_ix_iter
                        )

                        loss = loss_fn(pred, sub_ans_iter)
                        loss.backward()
                        loss_sum += loss.cpu().data.numpy() * self.__C.GRAD_ACCU_STEPS

                        if self.__C.VERBOSE:
                            if dataset_eval is not None:
                                mode_str = self.__C.SPLIT['train'] + '->' + self.__C.SPLIT['val']
                            else:
                                mode_str = self.__C.SPLIT['train'] + '->' + self.__C.SPLIT['test']

                            print("\r[version %s][epoch %2d][step %4d/%4d][%s] loss: %.4f, lr: %.2e" % (
                                self.__C.VERSION,
                                epoch + 1,
                                step,
                                int(data_size / self.__C.BATCH_SIZE),
                                mode_str,
                                loss.cpu().data.numpy() / self.__C.SUB_BATCH_SIZE,
                                optim._rate
                            ), end='          ')

                    # Gradient norm clipping
                    if self.__C.GRAD_NORM_CLIP > 0:
                        nn.utils.clip_grad_norm_(
                            net.parameters(),
                            self.__C.GRAD_NORM_CLIP
                        )

                    # Save the gradient information
                    for name in range(len(named_params)):
                        norm_v = torch.norm(named_params[name][1].grad).cpu().data.numpy() \
                            if named_params[name][1].grad is not None else 0
                        grad_norm[name] += norm_v * self.__C.GRAD_ACCU_STEPS

                    optim.step()

            if self.__C.features == "scene_graph" :

                for step, (
                        img_feat_obj_iter,
                        img_feat_attr_iter,
                        ques_ix_iter,
                        ans_iter,
                        img_feat_bbox
                ) in enumerate(dataloader):

                #for imp in range(100) :
                #    img_feat_obj_iter, img_feat_attr_iter, ques_ix_iter, ans_iter, img_feat_bbox = dataset.extract(imp)

                    optim.zero_grad()

                    img_feat_obj_iter = torch.tensor(img_feat_obj_iter, dtype=torch.float32)
                    img_feat_attr_iter = torch.tensor(img_feat_attr_iter, dtype=torch.float32)
                    img_feat_attr_iter = img_feat_attr_iter.cuda()
                    img_feat_obj_iter = img_feat_obj_iter.cuda()
                    ques_ix_iter = ques_ix_iter.cuda()
                    img_feat_bbox = torch.tensor(img_feat_bbox, dtype=torch.float32)
                    img_feat_bbox = img_feat_bbox.cuda()
                    ans_iter = ans_iter.cuda()
                    for accu_step in range(self.__C.GRAD_ACCU_STEPS):

                        sub_img_feat_obj_iter = \
                            img_feat_obj_iter[accu_step * self.__C.SUB_BATCH_SIZE:
                                          (accu_step + 1) * self.__C.SUB_BATCH_SIZE]
                        sub_img_feat_attr_iter = \
                            img_feat_attr_iter[accu_step * self.__C.SUB_BATCH_SIZE:
                                          (accu_step + 1) * self.__C.SUB_BATCH_SIZE]
                        sub_ques_ix_iter = \
                            ques_ix_iter[accu_step * self.__C.SUB_BATCH_SIZE:
                                         (accu_step + 1) * self.__C.SUB_BATCH_SIZE]
                        sub_ans_iter = \
                            ans_iter[accu_step * self.__C.SUB_BATCH_SIZE:
                                     (accu_step + 1) * self.__C.SUB_BATCH_SIZE]
                        sub_img_feat_bbox = \
                            img_feat_bbox[accu_step * self.__C.SUB_BATCH_SIZE:
                                     (accu_step + 1) * self.__C.SUB_BATCH_SIZE]

                        #net = net.float()
                        pred = net(
                            sub_img_feat_obj_iter,
                            sub_img_feat_attr_iter,
                            sub_ques_ix_iter,
                            sub_img_feat_bbox
                        )

                        loss = loss_fn(pred, sub_ans_iter)
                        loss.backward()
                        loss_sum += loss.cpu().data.numpy() * self.__C.GRAD_ACCU_STEPS

                        if self.__C.VERBOSE:
                            if dataset_eval is not None:
                                mode_str = self.__C.SPLIT['train'] + '->' + self.__C.SPLIT['val']
                            else:
                                mode_str = self.__C.SPLIT['train'] + '->' + self.__C.SPLIT['test']

                            print("\r[version %s][epoch %2d][step %4d/%4d][%s] loss: %.4f, lr: %.2e" % (
                                self.__C.VERSION,
                                epoch + 1,
                                step,
                                int(data_size / self.__C.BATCH_SIZE),
                                mode_str,
                                loss.cpu().data.numpy() / self.__C.SUB_BATCH_SIZE,
                                optim._rate
                            ), end='          ')

                    # Gradient norm clipping
                    if self.__C.GRAD_NORM_CLIP > 0:
                        nn.utils.clip_grad_norm_(
                            net.parameters(),
                            self.__C.GRAD_NORM_CLIP
                        )

                    # Save the gradient information
                    for name in range(len(named_params)):
                        norm_v = torch.norm(named_params[name][1].grad).cpu().data.numpy() \
                            if named_params[name][1].grad is not None else 0
                        grad_norm[name] += norm_v * self.__C.GRAD_ACCU_STEPS

                    optim.step()


            time_end = time.time()
            print('Finished in {}s'.format(int(time_end-time_start)))

            # print('')
            epoch_finish = epoch + 1

            if epoch_finish % 2 == 0 :

            # Save checkpoint
                state = {
                    'state_dict': net.state_dict(),
                    'optimizer': optim.optimizer.state_dict(),
                    'lr_base': optim.lr_base
                }
                if not os.path.exists(self.__C.CKPTS_PATH +'ckpt_' + self.__C.VERSION) :
                    os.mkdir(self.__C.CKPTS_PATH +'ckpt_' + self.__C.VERSION)

                torch.save(
                    state,
                    self.__C.CKPTS_PATH +
                    'ckpt_' + self.__C.VERSION +
                    '/epoch' + str(epoch_finish) +
                    '.pkl'
                )

            # Logging
            logfile = open(
                self.__C.LOG_PATH +
                'log_run_' + self.__C.VERSION + '.txt',
                'a+'
            )
            logfile.write(
                'epoch = ' + str(epoch_finish) +
                '  loss = ' + str(loss_sum / data_size) +
                '\n' +
                'lr = ' + str(optim._rate) +
                '\n\n'
            )
            logfile.close()

            # Eval after every epoch
            if dataset_eval is not None:
                self.eval(
                    dataset_eval,
                    state_dict=net.state_dict(),
                    valid=True
                )

            # if self.__C.VERBOSE:
            #     logfile = open(
            #         self.__C.LOG_PATH +
            #         'log_run_' + self.__C.VERSION + '.txt',
            #         'a+'
            #     )
            #     for name in range(len(named_params)):
            #         logfile.write(
            #             'Param %-3s Name %-80s Grad_Norm %-25s\n' % (
            #                 str(name),
            #                 named_params[name][0],
            #                 str(grad_norm[name] / data_size * self.__C.BATCH_SIZE)
            #             )
            #         )
            #     logfile.write('\n')
            #     logfile.close()

            loss_sum = 0
            grad_norm = np.zeros(len(named_params))

    # Evaluation
    def eval(self, dataset, state_dict=None, valid=False):

        # Load parameters
        if self.__C.CKPT_PATH is not None:
            print('Warning: you are now using CKPT_PATH args, '
                  'CKPT_VERSION and CKPT_EPOCH will not work')

            path = self.__C.CKPT_PATH
        else:
            path = self.__C.CKPTS_PATH + \
                   'ckpt_' + self.__C.CKPT_VERSION + \
                   '/epoch' + str(self.__C.CKPT_EPOCH) + '.pkl'

        val_ckpt_flag = False
        if state_dict is None:
            val_ckpt_flag = True
            print('Loading ckpt {}'.format(path))
            state_dict = torch.load(path)['state_dict']
            print('Finish!')

        # Store the prediction list
        qid_list = [ques['question_id'] for ques in dataset.ques_list]
        ans_ix_list = []
        ans_index_list = []
        pred_list = []

        data_size = dataset.data_size
        token_size = dataset.token_size
        ans_size = dataset.ans_size
        pretrained_emb = dataset.pretrained_emb

        net = self.Net(
            self.__C,
            pretrained_emb,
            token_size,
            ans_size
        )
        net.cuda()
        net.eval()

        if self.__C.N_GPU > 1:
            net = nn.DataParallel(net, device_ids=self.__C.DEVICES)

        net.load_state_dict(state_dict)
        
        dataloader = Data.DataLoader(
            dataset,
            batch_size=self.__C.EVAL_BATCH_SIZE,
            shuffle=False,
            num_workers=self.__C.NUM_WORKERS,
            pin_memory=True,
        )

        ques_list = dataset.val_ques_list
    
        correct = 0
        total = 0
        for step, (
                feature_iter,
                ques_ix_iter,
                ans_iter
        ) in enumerate(dataloader):

            feature_iter = feature_iter.cuda()
            ques_ix_iter = ques_ix_iter.cuda()

            pred = net(
                feature_iter,
                ques_ix_iter
            )

            pred_np = pred.cpu().data.numpy()
            pred_argmax = np.argmax(pred_np, axis=1)

            for i in range(ans_iter.shape[0]) :
                try :
                    ans_index = ans_iter[i].nonzero(as_tuple=True)[0].item()
                except :
                    ans_index = 0
                total +=1
                if ans_index == pred_argmax[i] :
                    correct += 1

        accuracy = float(correct)/total
        # Logging
        logfile = open(
            self.__C.LOG_PATH +
            'log_run_' + self.__C.VERSION + '.txt',
            'a+'
        )
        logfile.write(
            'Evaluation : ' + '\n' +
            'correct = ' + str(correct) +
            '\n' +
            'accuracy = ' + str(accuracy) +
            '\n\n'
        )
        logfile.close()


    def run(self, run_mode):
        if run_mode == 'train':
            self.empty_log(self.__C.VERSION)

            if self.__C.features == 'frcnn' : 
                self.train(self.dataset, self.dataset_eval)

            if self.__C.features == 'scene_graph' :
                self.train(self.dataset)

        elif run_mode == 'val':
            self.eval(self.dataset, valid=True)

        elif run_mode == 'test':
            self.eval(self.dataset)

        else:
            exit(-1)


    def empty_log(self, version):
        print('Initializing log file ........')
        if (os.path.exists(self.__C.LOG_PATH + 'log_run_' + version + '.txt')):
            os.remove(self.__C.LOG_PATH + 'log_run_' + version + '.txt')
        print('Finished!')
        print('')




