# 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.data.utils import Game, bbox2spatial_gw
from src.data.dataset import GuessWhatDataset


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


class OracleDataset(GuessWhatDataset):
    def __init__(
        self,
        dataroot,
        split,
        image_features_reader,
        tokenizer,
        **kwargs
    ):
        super().__init__(
            dataroot,
            'oracle',
            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 >= 1000: break
                game = Game.from_annotation(annotation)
                if game.status != 'success':
                    continue
                for qa in game.qas:
                    item = dict()
                    item['game'] = game
                    item['image_id'] = game.image_id
                    item['target_index'] = game.target_index
                    item['bboxs'] = [
                        bbox2spatial_gw(box, game.image_width, game.image_height) for box in game.bboxs]
                    item['categories'] = game.categories
                    item['question_id'] = qa['id']
                    item['question'] = qa['question']
                    # item['q_tokens'] = self._tokenizer.encode(qa['question'].replace('?', '')) + [self.eoq_id]
                    item['q_tokens'] = self._tokenizer.encode(qa['question'])[:-1] + [self.eoq_id]
                    item['answer'] = [int(self.answer2id[qa['answer']])]
                    entries.append(item)
        return entries

    def tensorize(self):
        for entry in self.entries:
            entry['categories'] = torch.from_numpy(np.array(entry['categories']))
            entry['bboxs'] = torch.from_numpy(np.array(entry['bboxs']))
            entry['q_tokens'] = torch.from_numpy(np.array(entry['q_tokens']))
            entry['answer'] = torch.from_numpy(np.array(entry['answer']))
            entry['target_index'] = torch.from_numpy(np.array(entry['target_index']))

    def __getitem__(self, index):
        entry = self.entries[index]
        #image_id = entry['image_id']
        game = entry['game']
        q_tokens = entry['q_tokens']
        tgt_index = entry['target_index']
        cats = entry['categories']
        bboxs = entry['bboxs']
        # bboxs = torch.from_numpy(self._image_features_reader[image_id])
        tgt_cat = cats[tgt_index]
        tgt_bbox = bboxs[tgt_index]
        answer = torch.LongTensor(entry['answer'])

        return (
            game,
            tgt_cat,
            tgt_bbox,
            cats,
            bboxs,
            q_tokens,
            answer
        )


def collate_fn(batch, wrd_pad_id):
    batch_size = len(batch)
    # batch
    game, tgt_cat, tgt_bbox, cats, bboxs, q_tokens, answer = zip(*batch)

    q_len = torch.LongTensor([len(q) for q in q_tokens])

    tgt_cat = torch.stack(tgt_cat).long()
    tgt_bbox = torch.stack(tgt_bbox).float()
    # (batch_size, padded_seq_length)
    q_tokens = pad_sequence(q_tokens, batch_first=True, padding_value=wrd_pad_id).long()
    # (batch_size)
    answer = torch.stack(answer).view(batch_size).long()

    return game, tgt_cat, tgt_bbox, q_tokens, q_len, answer

