''' Training Scropt for V2C captioning task. '''

__author__ = 'Jacob Zhiyuan Fang'

import os
import numpy as np
from opts import *
from utils.utils import *
import torch.optim as optim
from model.Model import Model
from torch.utils.data import DataLoader
from utils.dataloader import VideoDataset
from model.transformer.Optim import ScheduledOptim


def train(loader, model, optimizer, opt, cap_vocab, cms_vocab):

    model.train()

    for epoch in range(opt['epochs']):
        iteration = 0

        for data in loader:
            torch.cuda.synchronize()

            if opt['cms'] == 'int':
                cms_labels = data['int_labels']
            elif opt['cms'] == 'eff':
                cms_labels = data['eff_labels']
            else:
                cms_labels = data['att_labels']

            if opt['cuda']:
                fc_feats = data['fc_feats'].cuda()
                cap_labels = data['cap_labels'].cuda()
                cms_labels = cms_labels.cuda()

            optimizer.zero_grad()

            # cap_probs, cms_probs = model(fc_feats, cap_labels, cap_pos, cms_labels, cms_pos)
            cap_probs, _, cms_probs, _ = model(fc_feats, cap_labels, cms_labels)

            # note: currently we just used most naive cross-entropy as training objective,
            # advanced loss func. like SELF-CRIT, different loss weights or stronger video feature
            # may lead performance boost, however is not the goal of this work.
            cap_loss, cap_n_correct = cal_performance(cap_probs.view(-1, cap_probs.shape[-1]),
                                                      cap_labels[:, 1:], smoothing=True)
            cms_loss, cms_n_correct = cal_performance(cms_probs.view(-1, cms_probs.shape[-1]),
                                                      cms_labels[:, 1:], smoothing=True)

            # compute the token prediction Acc.
            non_pad_mask = cap_labels[:, 1:].ne(Constants.PAD)
            n_word = non_pad_mask.sum().item()
            cms_non_pad_mask = cms_labels[:, 1:].ne(Constants.PAD)
            cms_n_word = cms_non_pad_mask.sum().item()
            cap_loss /= n_word
            cms_loss /= n_word

            loss = cms_loss + cap_loss

            loss.backward()
            optimizer.step_and_update_lr()
            torch.nn.utils.clip_grad_norm_(filter(lambda p: p.requires_grad, model.parameters()), 1)

            # update parameters
            cap_train_loss = cap_loss.item()
            cms_train_loss = cms_loss.item()

            # multi-gpu case, not necessary in newer PyTorch version or on single GPU.
            if opt['cuda']: torch.cuda.synchronize()

            iteration += 1

            if iteration % opt['print_loss_every'] ==0:
                print('iter %d (epoch %d), cap_train_loss = %.6f, cms_train_loss = %.6f,'
                      ' current step = %d, current lr = %.3E, cap_acc = %.3f, cms_acc = %.3f'
                      % (iteration, epoch, cap_train_loss, cms_train_loss, optimizer.n_current_steps,
                         optimizer._optimizer.param_groups[0]['lr'],
                         cap_n_correct/n_word, cms_n_correct/cms_n_word))

                # show the intermediate generations
                if opt['show_predict']:
                    cap_pr, cap_gt = show_prediction(cap_probs, cap_labels[:, :-1], cap_vocab, caption=True)
                    cms_pr, cms_gt = show_prediction(cms_probs, cms_labels[:, :-1], cms_vocab, caption=False)
                    print(' \n')

                    with open(opt['info_path'], 'a') as f:
                        f.write('model_%d, cap_loss: %.6f, cms_loss: %.6f\n'
                                % (epoch, cap_train_loss, cms_train_loss))
                        f.write('\n %s \n %s' % (cap_pr, cap_gt))
                        f.write('\n %s \n %s' % (cms_pr, cms_gt))
                        f.write('\n')

        if epoch % opt['save_checkpoint_every'] == 0:

            # save the checkpoint
            model_path = os.path.join(opt['output_dir'],
                                      'CMS_CAP_MODEL_{}_lr_{}_BS_{}_Layer_{}_ATTHEAD_{}_HID_{}_RNNLayer_{}_epoch_{}.pth'
                                      .format(opt['cms'], opt['init_lr'], opt['batch_size'], opt['num_layer'],
                                              opt['num_head'], opt['dim_model'], opt['rnn_layer'], epoch))

            torch.save(model.state_dict(), model_path)

            print('model saved to %s' % model_path)
            with open(opt['model_info_path'], 'a') as f:
                f.write('model_%d, cap_loss: %.6f, cms_loss: %.6f\n'
                        % (epoch, cap_train_loss/n_word, cms_train_loss/n_word))


def main(opt):

    # load and define dataloader
    dataset = VideoDataset(opt, 'train')
    dataloader = DataLoader(dataset, batch_size=opt['batch_size'], shuffle=True)

    opt['cms_vocab_size'] = dataset.get_cms_vocab_size()
    opt['cap_vocab_size'] = dataset.get_cap_vocab_size()

    if opt['cms'] == 'int':
        cms_text_length = opt['int_max_len']
    elif opt['cms'] == 'eff':
        cms_text_length = opt['eff_max_len']
    else:
        cms_text_length = opt['att_max_len']

    # model initialization.
    from model.S2VTModel import S2VTModel
    model = S2VTModel(
                        dataset.get_cap_vocab_size(),
                        dataset.get_cms_vocab_size(),
                        opt['cap_max_len'],
                        cms_text_length,
                        opt["dim_model"],
                        opt["dim_word"],
                        opt['dim_vis_feat'],
                        n_layers=opt['rnn_layer'])

    # number of parameters
    model_parameters = filter(lambda p: p.requires_grad, model.parameters())
    params = sum([np.prod(p.size()) for p in model_parameters])
    print('number of learnable parameters are {}'.format(params))

    if opt['cuda']: model = model.cuda()

    # resume from previous checkpoint if indicated
    if opt['load_checkpoint'] and opt['resume']:
        cap_state_dict = torch.load(opt['load_checkpoint'])
        model_dict = model.state_dict()
        model_dict.update(cap_state_dict)
        model.load_state_dict(model_dict)

    optimizer = ScheduledOptim(optim.Adam(filter(lambda x: x.requires_grad, model.parameters()),
                                          betas=(0.9, 0.98), eps=1e-09), 512, opt['warm_up_steps'])

    # note: though we set the init learning rate as np.power(d_model, -0.5),
    # grid search indicates different LR may improve the results.
    opt['init_lr'] = round(optimizer.init_lr, 3)

    # create checkpoint output directory
    dir = os.path.join(opt['checkpoint_path'], 'S2VT_CMS_CAP_MODEL_{}_lr_{}_BS_{}_Layer_{}_ATTHEAD_{}_HID_{}_RNNLayer_{}'
                       .format(opt['cms'], opt['init_lr'], opt['batch_size'], opt['num_layer'],
                               opt['num_head'], opt['dim_model'], opt['rnn_layer']))

    if not os.path.exists(dir): os.makedirs(dir)

    # save the model snapshot to local
    info_path = os.path.join(dir, 'iteration_info_log.log')
    print('model architecture saved to {} \n {}'.format(info_path, str(model)))
    with open(info_path, 'a') as f:
        f.write(str(model))
        f.write('\n')
        f.write(str(params))
        f.write('\n')

    # log file directory
    opt['output_dir'] = dir
    opt['info_path'] = info_path
    opt['model_info_path'] = os.path.join(opt['output_dir'], 'checkpoint_loss_log.log')

    train(dataloader, model, optimizer, opt, dataset.get_cap_vocab(), dataset.get_cms_vocab())

if __name__ == '__main__':
    opt = parse_opt()
    opt = vars(opt)
    main(opt)
[(113, 'utils.dataloader.VideoDataset', 'VideoDataset', False, 'from utils.dataloader import VideoDataset\n'), (114, 'torch.utils.data.DataLoader', 'DataLoader', False, 'from torch.utils.data import DataLoader\n'), (167, 'os.path.join', 'os.path.join', False, 'import os\n'), (178, 'os.path.join', 'os.path.join', False, 'import os\n'), (164, 'os.path.exists', 'os.path.exists', False, 'import os\n'), (164, 'os.makedirs', 'os.makedirs', False, 'import os\n')]