import sys
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

import numpy as np

import time
import os
from collections import defaultdict
import json

import captioning.utils.opts as opts
import captioning.models as models
from captioning.data.pth_loader import CaptionDataset
import captioning.utils.eval_utils as eval_utils
# import captioning.utils.vizwiz_eval_utils as vizwiz_eval_utils
import captioning.utils.misc as utils
from captioning.utils.rewards import init_scorer, get_self_critical_reward
from captioning.modules.loss_wrapper import LossWrapper

import pytorch_lightning as pl


class ModelCheckpoint(pl.callbacks.ModelCheckpoint):

    def on_keyboard_interrupt(self, trainer, pl_module):
        # Save model when keyboard interrupt
        filepath = os.path.join(self.dirpath, self.prefix + 'interrupt.ckpt')
        self._save_model(filepath)


if __name__ == '__main__':

    device = 'cuda'

    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--reward', type=str, default='mle')
    args = parser.parse_args()

    if args.reward == 'mle':
        cfg = f'configs/phase1/fg_clipRN50_{args.reward}.yml'
    else:
        cfg = f'configs/phase2/fg_clipRN50_{args.reward}.yml'

    print("Loading cfg from", cfg)

    opt = opts.parse_opt(parse=False, cfg=cfg)

    dataset = CaptionDataset(opt)

    opt.vocab_size = dataset.vocab_size
    opt.seq_length = dataset.seq_length

    opt.batch_size = 40

    opt.vocab = dataset.get_vocab()

    model = models.setup(opt)
    del opt.vocab

    ckpt_path = opt.checkpoint_path + '-last.ckpt'

    print("Loading checkpoint from", ckpt_path)
    raw_state_dict = torch.load(
        ckpt_path,
        map_location=device)

    strict = True

    state_dict = raw_state_dict['state_dict']

    if '_vocab' in state_dict:
        model.vocab = utils.deserialize(state_dict['_vocab'])
        del state_dict['_vocab']
    elif strict:
        raise KeyError
    if '_opt' in state_dict:
        saved_model_opt = utils.deserialize(state_dict['_opt'])
        del state_dict['_opt']
        # Make sure the saved opt is compatible with the curren topt
        need_be_same = ["caption_model",
                        "rnn_type", "rnn_size", "num_layers"]
        for checkme in need_be_same:
            if getattr(saved_model_opt, checkme) in ['updown', 'topdown'] and \
                    getattr(opt, checkme) in ['updown', 'topdown']:
                continue
            assert getattr(saved_model_opt, checkme) == getattr(
                opt, checkme), "Command line argument and saved model disagree on '%s' " % checkme
    elif strict:
        raise KeyError
    res = model.load_state_dict(state_dict, strict)
    print(res)

    opt.use_grammar = False

    lw_model = LossWrapper(model, opt)

    split = 'test'

    print("Building dataloader...")

    test_dataset = torch.utils.data.Subset(
        dataset,
        dataset.split_ix[split]
    )
    test_loader = torch.utils.data.DataLoader(
        test_dataset,
        batch_size=opt.batch_size,
        shuffle=False,
        num_workers=4,
        drop_last=False,
        collate_fn=dataset.collate_func
    )

    eval_kwargs = {'dataset': opt.input_json}
    eval_kwargs.update(vars(opt))

    verbose = eval_kwargs.get('verbose', True)
    verbose_beam = eval_kwargs.get('verbose_beam', 0)
    verbose_loss = eval_kwargs.get('verbose_loss', 1)
    # num_images = eval_kwargs.get('num_images', eval_kwargs.get('val_images_use', -1))
    # lang_eval = eval_kwargs.get('language_eval', 0)
    dataset = eval_kwargs.get('dataset', 'coco')
    beam_size = eval_kwargs.get('beam_size', 1)
    sample_n = eval_kwargs.get('sample_n', 1)
    remove_bad_endings = eval_kwargs.get('remove_bad_endings', 0)

    crit = lw_model.crit

    model = model.to(device)

    from tqdm import tqdm

    test_id2sent = {}

    model.eval()

    print("running inference...")

    for data in tqdm(test_loader):
        with torch.no_grad():
            # forward the model to get loss
            tmp = [data['fc_feats'], data['att_feats'],
                data['labels'], data['masks'], data['att_masks']]
            tmp = [d.to(device) if isinstance(d, torch.Tensor) else d for d in tmp]

            fc_feats, att_feats, labels, masks, att_masks = tmp

            loss = crit(model(fc_feats, att_feats,
                            labels[..., :-1], att_masks), labels[..., 1:], masks[..., 1:])

            # forward the model to also get generated samples for each image
            # Only leave one feature for each image, in case duplicate sample
            tmp_eval_kwargs = eval_kwargs.copy()
            tmp_eval_kwargs.update({'sample_n': 1})
            seq, seq_logprobs = model(
                fc_feats, att_feats, att_masks, opt=tmp_eval_kwargs, mode='sample')
            seq = seq.data
            entropy = - (F.softmax(seq_logprobs, dim=2) *
                        seq_logprobs).sum(2).sum(1) / ((seq > 0).to(seq_logprobs).sum(1)+1)
            perplexity = - \
                seq_logprobs.gather(2, seq.unsqueeze(2)).squeeze(
                    2).sum(1) / ((seq > 0).to(seq_logprobs).sum(1)+1)

            # Print beam search
            if beam_size > 1 and verbose_beam:
                for i in range(fc_feats.shape[0]):
                    print('\n'.join([utils.decode_sequence(model.vocab, _[
                        'seq'].unsqueeze(0))[0] for _ in model.done_beams[i]]))
                    print('--' * 10)
            sents = utils.decode_sequence(model.vocab, seq)

        for d, sent in zip(data['infos'], sents):
            test_id2sent[d['id']] = sent

    res_path = f'FineCapEval_results/clipRN50_{args.reward}.json'

    print("Results save at {}".format(res_path))

    with open(res_path, 'w') as f:
        json.dump(test_id2sent, f)