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

from tqdm import tqdm


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

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

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

    def _load_dataset(self):
        # LENS = []
        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 >= 400: break
                game = Game.from_annotation(annotation)
                if game.status != 'success':
                    continue
                
                questions = []
                answers = []
                answers_id = []
                # qgen_in = [self.sos_id]
                # qgen_tgt = []
                # LENS.append(0)
                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 = 'N/A' if qa['answer'] == '<na>' else qa['answer']
                    a_token = self.answer2token[a]
                    a_id = int(self.answer2id[a])
                    questions.append([self._tokenizer.cls_id] + q_tokens)
                    answers.append(a_token)
                    answers_id.append(a_id)
                # Add an end-of-dialog token
                questions.append([self._tokenizer.cls_id, self.eod_id])

                item = dict()
                item['game'] = game
                item['end_turn'] = min(MAX_TURNS, len(game.qas)) - 1
                item['image_id'] = game.image_id
                # item['categories'] = game.categories
                item['questions'] = questions
                item['answers'] = answers
                item['answers_id'] = answers_id

                item['image_width'] = game.image_width
                item['image_height'] = game.image_height
                item['target_index'] = game.target_index

                # feats, _, bboxs, _ = self._image_features_reader[game.image_id]
                # item['image_features'] = feats
                # item['bboxs'] = bboxs
                
                entries.append(item)
        return entries

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

    def __getitem__(self, index):
        entry = self.entries[index]
        game = entry['game']
        # bboxs = entry['bboxs']
        # img_feats = entry['image_features']
        img_id = entry['image_id']
        feats, _, bboxs, _ = self._image_features_reader[img_id]
        img_feats = torch.from_numpy(np.array(feats))
        bboxs = torch.from_numpy(np.array(bboxs))
        # img_feats = torch.from_numpy(np.array(entry['image_features']))
        # bboxs = torch.from_numpy(np.array(entry['bboxs']))

        questions = entry['questions']
        answers = entry['answers']
        answers_id = entry['answers_id']
        # label = torch.LongTensor([entry['target_index']])
        end_turn = torch.LongTensor([entry['end_turn']])

        return (
            game,
            questions,
            answers,
            answers_id,
            end_turn,
            img_feats,
            bboxs,
            # label
        )


def collate_fn(batch, wrd_pad_id):
    batch_size = len(batch)
    # batch
    game, questions, answers, answers_id, end_turn, img_feats, bboxs = zip(*batch)

    # Pad questions
    qs = list(questions)
    qs_tf_in = copy.deepcopy(qs)
    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]))
            n_pad = max_q_seq_len - len(qs[b][t])
            _txt_attn_mask.append([1] * len(qs[b][t]) + [0] * n_pad)
            # Pad each sentence to max_q_seq_len
            qs[b][t].extend(n_pad * [wrd_pad_id])
            # print(max_q_seq_len, ":", qs[b][t])
            qs_tf_in[b][t] = qs_tf_in[b][t][:-1] + n_pad * [wrd_pad_id]
            assert len(qs[b][t]) == max_q_seq_len
        # Pad each batch to max_num_turns
        n_pad_turns = max_num_turns - len(qs[b])
        qs[b] += n_pad_turns * [max_q_seq_len * [wrd_pad_id]]
        qs_tf_in[b] += n_pad_turns * [(max_q_seq_len-1) * [wrd_pad_id]]
        q_len.append(_q_len + n_pad_turns * [0])
        txt_attn_mask.append(_txt_attn_mask + n_pad_turns * [[0] * max_q_seq_len])
        assert len(qs[b]) == max_num_turns
        assert len(q_len[b]) == max_num_turns
    qs = torch.LongTensor(qs)
    qs_tf_in = torch.LongTensor(qs_tf_in)
    # print(qs[0, :])
    # print(qs_tf_in[0, :])
    # aasasa()
    q_len = torch.LongTensor(q_len)
    txt_attn_mask = torch.from_numpy(np.array(txt_attn_mask))
    
    # (batch_size, num_max_turns)
    answers = pad_sequence(answers, batch_first=True).long()
    answers_id = pad_sequence(answers_id, batch_first=True).long()
    img_feats = torch.stack(img_feats).float()
    bboxs = torch.stack(bboxs).float()
    end_turn = torch.stack(end_turn).view(-1).long()
    # label = torch.stack(label).view(-1)
    # obj_feats = torch.cat([img_feats, bboxs], dim=-1)

    return game, qs, qs_tf_in, answers, answers_id, q_len, img_feats, bboxs, txt_attn_mask, end_turn

