# 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
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


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


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):
        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 enumerate(reader):
                # if cur >= 1000: break
                game = Game.from_annotation(annotation)
                if game.status != 'success':
                    continue
                
                qgen_in = [self.sos_id]
                qgen_tgt = []
                for qa in game.qas:
                    q_tokens = self._tokenizer.encode(qa['question'])
                    a_token = self.answer2token[qa['answer']]
                    # Question mark will not be fed into the model
                    assert q_tokens[-1] == self._tokenizer.eoq_id,\
                        "There is question not ended with question mark in game-%d." % game.id
                    qgen_in.extend(q_tokens[:-1] + [a_token])
                    qgen_tgt.extend(q_tokens)
                qgen_tgt.extend([self.eod_id])

                item = dict()
                item['game'] = game
                item['categories'] = game.categories
                item['bboxs'] = [
                    bbox2spatial_gw(box, game.image_width, game.image_height) for box in game.bboxs]
                item['qgen_input'] = qgen_in
                item['qgen_target'] = qgen_tgt
                item['image_feature'] = self._image_features_reader[game.image_id]
                entries.append(item)
        return entries

    def tensorize(self):
        for entry in self.entries:
            entry['qgen_input'] = torch.from_numpy(np.array(entry['qgen_input']))
            entry['qgen_target'] = torch.from_numpy(np.array(entry['qgen_target']))
            entry['categories'] = torch.from_numpy(np.array(entry['categories']))
            entry['bboxs'] = torch.from_numpy(np.array(entry['bboxs']))
            entry['image_feature'] = torch.from_numpy(np.array(entry['image_feature']))
            


    def __getitem__(self, index):
        entry = self.entries[index]
        game = entry['game']
        qgen_in = entry['qgen_input']
        qgen_tgt = entry['qgen_target']
        cats = entry['categories']
        bboxs = entry['bboxs']
        # bboxs = torch.from_numpy(self._image_features_reader[image_id])
        # label = torch.LongTensor([tgt_index])
        img_feat = entry['image_feature']

        return (
            game,
            qgen_in,
            qgen_tgt,
            img_feat,
            cats,
            bboxs
        )


def collate_fn(batch, wrd_pad_id):
    batch_size = len(batch)
    # batch
    game, qgen_in, qgen_tgt, img_feat, _, _ = zip(*batch)

    # bboxs_mask = [torch.ones(len(xs)) for xs in bboxs]
    qgen_in_len = torch.LongTensor([len(d) for d in qgen_in])
    qgen_tgt_len = torch.LongTensor([len(d) for d in qgen_tgt])
    out_mask = [torch.ones(len(d)) for d in qgen_tgt]



    # (batch_size, padded_seq_length)
    qgen_in = pad_sequence(qgen_in, batch_first=True, padding_value=wrd_pad_id).long()
    qgen_tgt = pad_sequence(qgen_tgt, batch_first=True, padding_value=wrd_pad_id).long()
    # (batch_size, img_feat_size)
    img_feat = torch.stack(img_feat)
    # (batch_size, seq_len)
    out_mask = pad_sequence(out_mask, batch_first=True).float()
    return game, qgen_in, qgen_in_len, qgen_tgt, qgen_tgt_len, img_feat, out_mask

