# 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

from tqdm import tqdm

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


class SelfPlayDataset(GuessWhatDataset):
    def __init__(
        self,
        dataroot,
        split,
        image_features_reader,
        tokenizer,
        **kwargs
    ):
        super().__init__(
            dataroot,
            'selfplay',
            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
                # dialog = []

                qs = []
                # q_len = []
                for qa in game.qas:
                    q = self._tokenizer.encode(qa['question'])[:-1] + [self.eoq_id]
                    qs.append(q)
                    # q_len.append(len(q))
                

                    # a_token = self.answer2token[qa['answer']]
                #     dialog.extend(q_tokens + [a_token])

                item = dict()
                item['game'] = game
                item['image_id'] = game.image_id
                item['target_index'] = game.target_index
                item['categories'] = game.categories
                item['bboxs'] = [
                    bbox2spatial_gw(box, game.image_width, game.image_height) for box in game.bboxs]

                item['qs'] = qs
                # item['q_len'] = q_len
                entries.append(item)
        return entries

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


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

        tgt_index = entry['target_index']
        cats = entry['categories']
        bboxs = entry['bboxs']
        
        tgt_cat = cats[tgt_index]
        tgt_bbox = bboxs[tgt_index]
        label = torch.LongTensor([tgt_index])

        qs = entry['qs']
        # q_len = entry['q_len']

        return (
            game,
            qgen_img_feats,
            qgen_bboxs,
            tgt_cat,
            tgt_bbox,
            cats,
            bboxs,
            label,
            qs,
            # q_len
        )


def collate_fn(batch, wrd_pad_id):
    batch_size = len(batch)
    # batch
    game, qgen_img_feats, qgen_bboxs, tgt_cat, tgt_bbox, cats, bboxs, label, qs = zip(*batch)

    # Dealing with ground truth questions
    # qs: [batch size, turns (not padded), seq len (not padded)]
    
    qs = list(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 = []
    for b in range(len(qs)):
        _q_len = []
        for t in range(len(qs[b])):
            _q_len.append(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])
        assert len(qs[b]) == max_num_turns
        assert len(q_len[b]) == max_num_turns
        # q_len: [batch size, turns (not padded)]            

    qs = torch.LongTensor(qs)
    q_len = torch.LongTensor(q_len)

    bboxs_mask = [torch.ones(len(xs)) for xs in bboxs]
    
    tgt_cat = torch.stack(tgt_cat).long()
    tgt_bbox = torch.stack(tgt_bbox).float()
    # (batch_size, padded_num_obj)
    cats = pad_sequence(cats, batch_first=True).long()

    # (batch_size, padded_num_bboxs, spatial dim)
    bboxs = pad_sequence(bboxs, batch_first=True).float()
    # (batch_size, padded_num_obj)
    bboxs_mask = pad_sequence(bboxs_mask, batch_first=True).bool()
    # (batch_size, padded_seq_length)
    label = torch.stack(label).view(-1)

    qgen_img_feats = torch.stack(qgen_img_feats).float()
    qgen_bboxs = torch.stack(qgen_bboxs).float()
    # obj_feats = torch.cat([qgen_img_feats, qgen_bboxs], dim=-1)

    return game, qgen_img_feats, qgen_bboxs, tgt_cat, tgt_bbox, cats, bboxs, bboxs_mask, label, qs, q_len

