# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import argparse
import os.path as osp
import time
from collections import OrderedDict

import numpy as np
# https://github.com/numpy/numpy/issues/21079
try:
    import numpy.distutils
    numpy.distutils.__config__.blas_opt_info = np.distutils.__config__.blas_ilp64_opt_info
except Exception:
    pass
from nlgeval import NLGEval

import torch
import torchvision.transforms as transforms
import torchvision.transforms._transforms_video as transforms_video

from lavila.data import datasets
from lavila.data.video_transforms import Permute, SpatialCrop, TemporalCrop
from lavila.models import models
from lavila.models.utils import inflate_positional_embeds
from lavila.utils import distributed as dist_utils
from lavila.utils.preprocess import generate_tokenizer


def decode_one(generated_ids, tokenizer):
    # get the index of <EOS>
    if tokenizer.eos_token_id == tokenizer.bos_token_id:
        if tokenizer.eos_token_id in generated_ids[1:].tolist():
            eos_id = generated_ids[1:].tolist().index(tokenizer.eos_token_id) + 1
        else:
            eos_id = len(generated_ids.tolist()) - 1
    elif tokenizer.eos_token_id in generated_ids.tolist():
        eos_id = generated_ids.tolist().index(tokenizer.eos_token_id)
    else:
        eos_id = len(generated_ids.tolist()) - 1
    generated_text_str = tokenizer.tokenizer.decode(generated_ids[1:eos_id].tolist())
    return generated_text_str


def get_args_parser():
    parser = argparse.ArgumentParser(description='LAVILA 0-shot evaluations', add_help=False)
    parser.add_argument('--dataset', default='ego4d', type=str,
                        choices=['ego4d'])
    parser.add_argument('--root',
                        default='datasets/Ego4D/video_5min_chunks_288px/',
                        type=str, help='path to dataset root')
    parser.add_argument('--metadata-val',
                        default='datasets/Ego4D/ego4d_val.pkl',
                        type=str, help='path to metadata file (val set)')
    parser.add_argument('--output-dir', default='./', type=str, help='output dir')
    parser.add_argument('--num-crops', default=1, type=int, help='number of crops in transforms')
    parser.add_argument('--num-clips', default=1, type=int, help='number of clips (for untrimmed videos, eg. Charades)')
    parser.add_argument('--clip-length', default=4, type=int, help='clip length')
    parser.add_argument('--clip-stride', default=16, type=int, help='clip stride')
    parser.add_argument('--sparse-sample', action='store_true', help='switch to sparse sampling')
    parser.add_argument('--batch-size', default=16, type=int, help='batch_size')
    # captioning options
    parser.add_argument('--caption-sample', default='multinomial_sample',
                        choices=['multinomial_sample', 'beam_sample', 'group_beam_search'])
    parser.add_argument('--caption-top-k', default=None, type=int, help='top-k sampling (predecessor of nucleus sampling)')
    parser.add_argument('--caption-top-p', default=0.95, type=float, help='top-p sampling sampling (aka nucleus sampling)')
    parser.add_argument('--caption-num-beams', default=3, type=int)
    parser.add_argument('--caption-num-beam-groups', default=1, type=int)
    parser.add_argument('--caption-temperature', default=0.7, type=float)
    parser.add_argument('--caption-length-penalty', default=1.0, type=float)
    parser.add_argument('--caption-num-return-sequences', default=1, type=int)
    parser.add_argument('--caption-max-len', default=77, type=int)
    parser.add_argument('--caption-disable-visual', action='store_true')
    parser.add_argument('--caption-early-stop', action='store_true', help='early stopping to save computation')
    parser.add_argument('--caption-output-filename', default='caption.txt', type=str)
    # others
    parser.add_argument('--eval-freq', default=1000, type=int,
                        help='percentage (1/eval_freq) of val data to evaluate (for fast prototyping)')
    parser.add_argument('--print-freq', default=10, type=int)
    parser.add_argument('-j', '--workers', default=10, type=int, metavar='N',
                        help='number of data loading workers per process')
    parser.add_argument('--resume', default='', type=str, help='path to latest checkpoint')
    parser.add_argument('--use-half', action='store_true')
    return parser


def main(args):
    if args.resume:
        ckpt_path = args.resume
    elif osp.isfile(osp.join(args.output_dir, 'checkpoint_best.pt')):
        ckpt_path = osp.join(args.output_dir, 'checkpoint_best.pt')
    else:
        raise Exception('no checkpoint found')

    ckpt = torch.load(ckpt_path, map_location='cpu')

    # create model
    state_dict = OrderedDict()
    for k, v in ckpt['state_dict'].items():
        state_dict[k.replace('module.', '')] = v

    old_args = ckpt['args']
    print('=> creating model: {}'.format(old_args.model))
    model = getattr(models, old_args.model)(
        text_use_cls_token=old_args.use_cls_token,
        project_embed_dim=old_args.project_embed_dim,
        gated_xattn=False if 'gated_xattn' not in old_args else old_args.gated_xattn,
        timesformer_gated_xattn=False if 'timesformer_gated_xattn' not in old_args else old_args.timesformer_gated_xattn,
        timesformer_freeze_space=False if 'timesformer_freeze_space' not in old_args else old_args.timesformer_freeze_space,
        freeze_lm_vclm=False if 'freeze_lm_vclm' not in old_args else old_args.freeze_lm_vclm,
        freeze_visual_vclm=False if 'freeze_visual_vclm' not in old_args else old_args.freeze_visual_vclm,
        num_frames=args.clip_length,
        drop_path_rate=0,
    )
    model.cuda()
    if 'TIMESFORMER' in old_args.model or 'EGOVLP' in old_args.model:
        # inflate weight
        print('=> inflating PE in models due to different frame numbers')
        state_dict = inflate_positional_embeds(
            model.state_dict(), state_dict,
            num_frames=args.clip_length,
            load_temporal_fix='bilinear',
        )
    model.load_state_dict(state_dict, strict=True)
    print("=> loaded resume checkpoint '{}' (epoch {}, best_metric = {})".format(args.resume, ckpt['epoch'], ckpt['best_acc1']))

    torch.backends.cudnn.benchmark = True

    tokenizer = generate_tokenizer(old_args.model)
    crop_size = 224 if '336PX' not in old_args.model else 336
    if args.num_crops == 1 and args.num_clips == 1:
        val_transform = transforms.Compose([
            Permute([3, 0, 1, 2]),  # T H W C -> C T H W
            transforms.Resize(crop_size),
            transforms.CenterCrop(crop_size),
            (transforms_video.NormalizeVideo(mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375]) if ('OPENAI' not in old_args.model) else
             transforms_video.NormalizeVideo(mean=[108.3272985, 116.7460125, 104.09373615000001], std=[68.5005327, 66.6321579, 70.32316305])),
        ])
    else:
        val_transform = transforms.Compose([
            Permute([3, 0, 1, 2]),  # T H W C -> C T H W
            transforms.Resize(crop_size),
            (transforms_video.NormalizeVideo(mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375]) if ('OPENAI' not in old_args.model) else
             transforms_video.NormalizeVideo(mean=[108.3272985, 116.7460125, 104.09373615000001], std=[68.5005327, 66.6321579, 70.32316305])),
            TemporalCrop(frames_per_clip=args.clip_length, stride=args.clip_length),
            SpatialCrop(crop_size=crop_size, num_crops=args.num_crops),
            ])

    val_dataset = datasets.VideoCaptionDatasetCLIP(
            args.dataset,
            args.root,
            args.metadata_val,
            transform=val_transform,
            is_training=False,
            tokenizer=tokenizer,
            clip_length=args.clip_length,
            clip_stride=args.clip_stride,
            sparse_sample=False,
            subsample_stride=args.eval_freq,
        )

    val_loader = torch.utils.data.DataLoader(
        val_dataset, batch_size=args.batch_size, shuffle=False,
        num_workers=args.workers, pin_memory=True, drop_last=False)

    validate_caption(val_loader, model, tokenizer, args.caption_output_filename, use_half=args.use_half)


def validate_caption(val_loader, model, tokenizer, output_filename='caption.txt', use_half=False):
    model.eval()
    if args.use_half:
        model = model.half()
    nlgeval = NLGEval()
    f = open(output_filename, 'w')
    ppls_all = []
    ppls_with_teacher_all = []
    reference = []
    hypothesis = []
    end_time = time.time()
    id_offset = 0
    print('=> start forwarding')
    with torch.no_grad():
        for i, inputs in enumerate(val_loader):
            if i % args.print_freq == 0:
                print('finish batch {}/{} in {} sec'.format(i, len(val_loader), time.time() - end_time))
                end_time = time.time()
            images = inputs[0].cuda(non_blocking=True)
            if use_half:
                images = images.half()
            target = inputs[1].cuda(non_blocking=True)

            # encode images
            image_features = dist_utils.get_model(model).encode_image(images)

            # teacher forcing (to get standard ppl metric)
            generated_text_ids_with_teacher, ppls_with_teacher = dist_utils.get_model(model).generate(
                image_features,
                tokenizer,
                target=target,
                max_text_length=args.caption_max_len,
                top_k=args.caption_top_k,
                top_p=args.caption_top_p,
                teacher_forcing=True,
                early_stopping=args.caption_early_stop,
            )

            if args.caption_sample == 'multinomial_sample':
                assert args.caption_num_beam_groups == 1
                generated_text_ids, ppls = dist_utils.get_model(model).generate(
                    image_features,
                    tokenizer,
                    target=target.repeat_interleave(args.caption_num_return_sequences, dim=0),
                    max_text_length=args.caption_max_len,
                    top_k=args.caption_top_k,
                    top_p=args.caption_top_p,
                    num_return_sequences=args.caption_num_return_sequences,
                    temperature=args.caption_temperature,
                    early_stopping=args.caption_early_stop,
                )
            elif args.caption_sample == 'beam_sample':
                assert args.caption_num_beam_groups == 1
                generated_text_ids, ppls = dist_utils.get_model(model).beam_sample(
                    image_features,
                    tokenizer,
                    target=target,
                    max_text_length=args.caption_max_len,
                    top_k=args.caption_top_k,
                    top_p=args.caption_top_p,
                    temperature=args.caption_temperature,
                    length_penalty=args.caption_length_penalty,
                    num_beams=args.caption_num_beams,
                    num_return_sequences=args.caption_num_return_sequences,
                    early_stopping=args.caption_early_stop,
                )
            elif args.caption_sample == 'group_beam_search':
                assert args.caption_num_beam_groups > 1 and args.caption_num_beams % args.caption_num_beam_groups == 0
                generated_text_ids, ppls = dist_utils.get_model(model).group_beam_search(
                    image_features,
                    tokenizer,
                    target=target if not args.caption_no_gt else None,
                    max_text_length=args.caption_max_len,
                    top_k=args.caption_top_k,
                    top_p=args.caption_top_p,
                    temperature=args.caption_temperature,
                    length_penalty=args.caption_length_penalty,
                    num_beams=args.caption_num_beams,
                    num_beam_groups=args.caption_num_beam_groups,
                    num_return_sequences=args.caption_num_return_sequences,
                    early_stopping=args.caption_early_stop,
                )
            else:
                raise NotImplementedError
            ppls_all.append(ppls.reshape(-1, args.caption_num_return_sequences).mean(1))
            ppls_with_teacher_all.append(ppls_with_teacher)

            for j in range(generated_text_ids.shape[0] // args.caption_num_return_sequences):
                for k in range(args.caption_num_return_sequences):
                    jj = j * args.caption_num_return_sequences + k

                    generated_text_str = decode_one(generated_text_ids[jj], tokenizer)
                    gt_text = decode_one(target[j], tokenizer)
                    generated_text_str_with_teacher = decode_one(generated_text_ids_with_teacher[j], tokenizer)

                    from transformers import BertTokenizer
                    bert_tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
                    gt_text = bert_tokenizer.decode(bert_tokenizer(gt_text)['input_ids'][1:-1])
                    generated_text_str = bert_tokenizer.decode(bert_tokenizer(generated_text_str)['input_ids'][1:-1])
                    generated_text_str_with_teacher = bert_tokenizer.decode(bert_tokenizer(generated_text_str_with_teacher)['input_ids'][1:-1])
                    reference.append(gt_text)
                    hypothesis.append(generated_text_str)
                    s1 = '[{:6d}] Groundtruth              |                 | {}'.format(id_offset + j, gt_text)
                    s2 = '[{:6d}] Generated                | PPL : {:9.3f} | {}'.format(id_offset + j, ppls[jj], generated_text_str)
                    s3 = '[{:6d}] Generated (w/. teacher)  | PPL : {:9.3f} | {}'.format(id_offset + j, ppls_with_teacher[j], generated_text_str_with_teacher)
                    for s in [s1, s2, s3]:
                        # if i % args.print_freq == 0:
                        #     print(s)
                        f.write('{} \n'.format(s))
            id_offset += generated_text_ids.shape[0] // args.caption_num_return_sequences

    ppls_with_teacher_all = torch.cat(ppls_with_teacher_all, dim=0)
    ppls_all = torch.cat(ppls_all, dim=0)

    print('PPL (w/.  teacher) = {:9.3f}'.format(ppls_with_teacher_all.mean().item()))
    print('PPL (w/o. teacher) = {:9.3f}'.format(ppls_all.mean().item()))
    f.write('PPL (w/.  teacher) = {:9.3f} \n'.format(ppls_with_teacher_all.mean().item()))
    f.write('PPL (w/o. teacher) = {:9.3f} \n'.format(ppls_all.mean().item()))

    print('Avg length for reference:  {:9.3f}'.format(sum(map(lambda sentence: len(sentence.split(' ')), reference)) / len(reference)))
    print('Avg length for hypothesis: {:9.3f}'.format(sum(map(lambda sentence: len(sentence.split(' ')), hypothesis)) / len(hypothesis)))
    f.write('Avg length for reference:  {:9.3f} \n'.format(sum(map(lambda sentence: len(sentence.split(' ')), reference)) / len(reference)))
    f.write('Avg length for hypothesis: {:9.3f} \n'.format(sum(map(lambda sentence: len(sentence.split(' ')), hypothesis)) / len(hypothesis)))

    print('=> Calling NLGEval')
    f.write('=> Calling NLGEval\n')
    metrics_dict = nlgeval.compute_metrics([reference], hypothesis)
    for k in metrics_dict:
        print('{:16s} = {:9.3f}'.format(k, metrics_dict[k]))
        f.write('{:16s} = {:9.3f} \n'.format(k, metrics_dict[k]))
    f.close()


if __name__ == '__main__':
    parser = argparse.ArgumentParser('lavila 0-shot evaluations', parents=[get_args_parser()])
    args = parser.parse_args()
    main(args)