import itertools
from collections import defaultdict

import torch
from torch.distributed import get_rank
from tqdm import tqdm

import masker
import utils
from masker import gen_pointing_text_mask_locs


def accuracy(output, target, topk=(1,), tok_groups=None, tok_group_labels=None, as_list=False):
    """Computes the accuracy over the k top predictions for the specified values of k"""
    mask = target != -1

    output = output[mask]
    target = target[mask]

    maxk = max(topk)

    _, pred = output.topk(maxk, 1, True, True)
    pred = pred.t()
    correct = pred.eq(target[None])

    res = {}
    if tok_groups:
        try:
            correct = correct.split(tok_groups, dim=1)
        except:
            return res
        try:
            correct = torch.stack([grp.sum(dim=1).cumsum(dim=0) == torch.ones_like(grp).sum(dim=1) for grp in correct],
                                  dim=-1)
        except:
            print('wrong', get_rank())
    for k in topk:
        try:
            if as_list:
                res[f'top{k}'] = correct[:k].max(dim=0).values
            else:
                correct_k = correct[:k].max(dim=0).values.sum(dtype=torch.float32)
                res[f'top{k}'] = ((correct_k * (100.0 / correct.shape[1])).item(), correct.shape[1])
        except:  # if no items are masked in batch, max operation doesn't work since correct is empty tensor
            if as_list:
                res[f'top{k}'] = [0]
            else:
                res[f'top{k}'] = (0, 0)

    if tok_group_labels:
        label_set = set(tok_group_labels)
        for label in label_set:
            for k in topk:
                correct_ = correct[:, [_ == label for _ in tok_group_labels]]
                correct_k = correct_[:k].max(dim=0).values.sum(dtype=torch.float32)
                res[f'{label}_top{k}'] = ((correct_k * (100.0 / correct_.shape[1])).item(), correct_.shape[1])

    return res


def test_accuracy(trainer, masking_policy):
    trainer.masker.test_masking_policy = masking_policy
    trainer.test_loader.dataset.test_masking_policy = masking_policy
    # behavior only differs in testmasker and accuracy function call.
    avg_meters = defaultdict(lambda: utils.AverageMeter())

    # Switch to eval mode
    trainer.model.eval()

    with torch.no_grad(), tqdm(trainer.test_loader, desc=f'Testing {masking_policy}',
                               disable=trainer.args.local_rank > 0) as t:
        for batch_idx, data in enumerate(t):
            imgs = data['imgs'].cuda()
            if trainer.args.pointing:
                text_mask_locs, text_no_mask_locs = gen_pointing_text_mask_locs(data, only_mask_tgt_toks=True)
                mask_kwargs = {
                    'mask_locs': text_mask_locs,
                    'no_mask_locs': text_no_mask_locs,
                }
                text, lm_labels, input_pointing_labels, *_ = \
                    trainer.masker.mask_text(data['text'].cuda(), input_pointing=trainer.args.input_pointing,
                                             **mask_kwargs, **data)
                if sum(map(len, data['tok_groups'])) > 0:
                    tok_groups = data['tok_groups']
                    tok_group_labels = data['tok_group_labels']
                else:
                    tok_groups = None
                    tok_group_labels = None
            else:
                try:
                    mask_kwargs = {
                        'actions': [trainer.test_loader.dataset.actions[v][a] for v, a in
                                    zip(data['vid_ids'], data['act_ids'])]
                    }
                except:
                    mask_kwargs = {
                        'actions': [trainer.test_loader.dataset.actions[trainer.test_loader.dataset.indices[i]] for i in
                                    data['indices'].squeeze().tolist()]
                    }
                text, lm_labels, input_pointing_labels, tok_groups, tok_group_labels = \
                    trainer.masker.mask_text(data['text'].cuda(), input_pointing=trainer.args.input_pointing,
                                             prob_mask=1, **mask_kwargs, **data)

            img_bboxes = data['img_bboxes'].cuda()
            imgs_len = data['imgs_len'].cuda()
            text_len = data['text_len'].cuda()

            img_locs = txt_locs = None

            if trainer.args.pointing:
                attn_mask, img_locs, txt_locs = masker.attn_mask_pointing(imgs_len, text_len, data['seq_type'],
                                                                          data['num_seqs'].cuda(),
                                                                          trainer.args.attn_masking)

            else:
                img_attn_mask = \
                    torch.arange(trainer.args.max_img_seq_len, device=imgs.device)[None, :] < imgs_len[:, None]
                text_attn_mask = \
                    torch.arange(trainer.args.max_txt_seq_len, device=imgs.device)[None, :] < text_len[:, None]
                attn_mask = torch.cat((text_attn_mask[:, :1], img_attn_mask, text_attn_mask[:, 1:]), dim=1)

            # text starts with [IMG] token that gets moved to beginning of input in forward pass

            lm_preds, vm_preds, input_pointing_pred, hidden_states, *_ = \
                trainer.model(imgs, text, img_bboxes, attention_mask=attn_mask, img_lens=imgs_len,
                              txt_lens=text_len, img_locs=img_locs, txt_locs=txt_locs)

            if trainer.args.pointing:
                pointing_loss, pointing_acc, chance_pointing_acc, pointing_scores = \
                    trainer.losses.pointing_loss(data, hidden_states, lm_labels, text, text_len, txt_locs, tok_groups,
                                                 tok_group_labels, log=True)
                avg_meters['pointing acc'].update(*pointing_acc)
                avg_meters['chance pointing acc'].update(*chance_pointing_acc)
                non_padding_text = (torch.arange(text.shape[1], device=text.device)[None, :] <
                                    text_len.cumsum(dim=1)[:, -1][:, None])
                if trainer.args.input_pointing:
                    input_pointing_loss, input_pointing_acc, chance_input_pointing_acc = \
                        trainer.losses.input_pointing_pointing_loss(input_pointing_pred[0], input_pointing_pred[1], input_pointing_labels,
                                                              txt_locs, lm_labels, text, tok_groups, tok_group_labels,
                                                              log=True, data=data)
                    avg_meters['input_pointing acc'].update(*input_pointing_acc)
                    avg_meters['chance input_pointing acc'].update(*chance_input_pointing_acc)
                if tok_groups is not None:
                    tok_groups = list(itertools.chain.from_iterable(tok_groups))
                    tok_group_labels = list(itertools.chain.from_iterable(tok_group_labels))
                lm_labels = lm_labels[non_padding_text]

            results = accuracy(lm_preds, lm_labels, topk=(1, 5), tok_groups=tok_groups,
                               tok_group_labels=tok_group_labels)
            for k in results:
                avg_meters[k].update(*results[k])

    for k, v in avg_meters.items():
        out = f'{trainer.masker.test_masking_policy} {k}: {utils.gather_score(v.avg, v.count)} on ' \
              f'{utils.gather_score(v.count, 1) * trainer.args.n_gpu} examples'
        if trainer.args.local_rank <= 0:
            print(out, flush=True)
