# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: CC-BY-NC-4.0
import os
import torch
import logging
import jsonlines
import numpy as np
from tqdm import tqdm
import _pickle as cPickle
from functools import partial
from torch.utils.data import Dataset
from torch.nn.utils.rnn import pad_sequence
from src.tools.utils import (
    Game, 
    bbox2spatial_gw, 
    bbox2spatial_vilbert, 
    add_global_vilbert_feats
)
from src.data.dataset import GuessWhatDataset


logger = logging.getLogger(__name__)  # pylint: disable=invalid-name

MAX_Q_LEN = 15
MAX_TURNS = 12 # 12: 95%, 10: 90%

class GuesserDataset(GuessWhatDataset):
    def __init__(
        self,
        dataroot,
        split,
        image_features_reader,
        tokenizer,
        **kwargs
    ):
        super().__init__(
            dataroot,
            'guesser',
            split,
            image_features_reader,
            tokenizer,
            **kwargs)

    def _load_dataset(self):
        with jsonlines.open(self.data_path) as reader:
            # Build an index which maps image id with a list of qa annotations.
            entries = []
            for cur, annotation in tqdm(enumerate(reader)):
                # if cur >= 100: break
                game = Game.from_annotation(annotation)
                if game.status != 'success':
                    continue
                questions = []
                answers = []

                for turn, qa in enumerate(game.qas):
                    if self.split == 'train' and turn == MAX_TURNS: break
                    q_tokens = self._tokenizer.encode(qa['question'])
                    if self.split == 'train' and len(q_tokens) > MAX_Q_LEN:
                        q_tokens = q_tokens[:MAX_Q_LEN-1] + [self.eoq_id]
                    assert q_tokens[-1] == self.eoq_id,\
                        "Game %d has question which is not ended with `?`." % game.id
                    # a_token = self.answer2token[qa['answer']]
                    a_token = int(self.answer2id[qa['answer']])
                    questions.append([self._tokenizer.cls_id] + q_tokens)
                    answers.append(a_token)

                item = dict()
                item['game'] = game
                # 0-based
                item['end_turn'] = min(MAX_TURNS, len(game.qas)) - 1
                item['target_index'] = game.target_index
                # 99: global
                item['categories'] = [99] + game.categories
                
                feats, _, _ = self._image_features_reader[game.image_id]
                bboxs = np.array([
                    bbox2spatial_vilbert(box, game.image_width, game.image_height, mode='xyxy') 
                    for box in game.bboxs])
                # Insert a global feat in front of feats & bboxs
                feats, bboxs = add_global_vilbert_feats(feats, bboxs)
                item['bboxs'] = bboxs
                item['image_features'] = feats
                item['questions'] = questions
                item['answers'] = answers
                entries.append(item)
            
        return entries

    def tensorize(self):
        for entry in self.entries:
            # entry['end_turn'] = torch.from_numpy(np.array(entry['end_turn']))
            entry['target_index'] = torch.from_numpy(np.array(entry['target_index']))
            entry['categories'] = torch.from_numpy(np.array(entry['categories']))
            # entry['questions'] = torch.from_numpy(np.array(entry['questions']))
            entry['answers'] = torch.from_numpy(np.array(entry['answers']))
            entry['image_features'] = torch.from_numpy(np.array(entry['image_features']))
            entry['bboxs'] = torch.from_numpy(np.array(entry['bboxs']))


    def __getitem__(self, index):
        entry = self.entries[index]
        game = entry['game']
        #image_id = entry['image_id']
        questions = entry['questions']
        answers = entry['answers']
        image_features = entry['image_features']
        bboxs = entry['bboxs']

        end_turn = entry['end_turn']
        tgt_index = entry['target_index']
        cats = entry['categories']
        # bboxs = torch.from_numpy(self._image_features_reader[image_id])

        label = torch.LongTensor([tgt_index])
        end_turn = torch.LongTensor([end_turn])
        
        return (
            game,
            questions,
            answers,
            end_turn,
            image_features,
            bboxs,
            cats,
            label
        )


def collate_fn(batch, wrd_pad_id):
    batch_size = len(batch)
    # batch
    game, questions, answers, end_turn, image_features, bboxs, cats, label = zip(*batch)

    # Pad questions
    qs = list(questions)
    max_q_seq_len = max([len(q) for _qs in qs for q in _qs])
    max_num_turns = max([len(_qs) for _qs in qs])
    q_len = []
    txt_attn_mask = []
    for b in range(len(qs)):
        _q_len = []
        _txt_attn_mask = []
        for t in range(len(qs[b])):
            _q_len.append(len(qs[b][t]))
            _txt_attn_mask.append([1] * len(qs[b][t]) + [0] * (max_q_seq_len - len(qs[b][t])))
            # Pad each sentence to max_q_seq_len
            qs[b][t].extend((max_q_seq_len - len(qs[b][t])) * [wrd_pad_id])
            assert len(qs[b][t]) == max_q_seq_len
        # Pad each batch to max_num_turns
        qs[b] += (max_num_turns - len(qs[b])) * [max_q_seq_len * [wrd_pad_id]]
        q_len.append(_q_len + (max_num_turns - len(_q_len)) * [0])
        txt_attn_mask.append(_txt_attn_mask + (max_num_turns - len(_q_len)) * [[0] * max_q_seq_len])
        assert len(qs[b]) == max_num_turns
        assert len(q_len[b]) == max_num_turns
    qs = torch.LongTensor(qs)
    q_len = torch.LongTensor(q_len)
    txt_attn_mask = torch.from_numpy(np.array(txt_attn_mask))

    answers = pad_sequence(answers, batch_first=True, padding_value=wrd_pad_id).long()
    # (batch_size, padded_num_obj)
    cats = pad_sequence(cats, batch_first=True).long()
    # (batch_size, padded_num_bboxs, image dim)
    image_features = pad_sequence(image_features, batch_first=True).float()
    # (batch_size, padded_num_bboxs, spatial dim)
    bboxs = pad_sequence(bboxs, batch_first=True).float()
    # (batch_size, padded_num_obj)
    bboxs_mask = pad_sequence(
        [torch.ones(len(xs)) for xs in bboxs], batch_first=True)
    bboxs_mask_vb = bboxs_mask.long() # for vilbert
    bboxs_mask = bboxs_mask.bool()
    end_turn = torch.stack(end_turn).view(-1).long()
    label = torch.stack(label).view(-1)

    return game, qs, q_len, answers, end_turn, cats, image_features, bboxs, bboxs_mask, bboxs_mask_vb, txt_attn_mask, label

