# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: CC-BY-NC-4.0
import os
import abc
import torch
import logging
import jsonlines
import numpy as np
import _pickle as cPickle
from torch.utils.data import Dataset


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


class GuessWhatDataset(Dataset):
    def __init__(
        self,
        dataroot,
        model,  # 'guesser', 'oracle', 'qgen'
        split,  # 'train', 'valid', 'test'
        image_features_reader,
        tokenizer,
        image_features_reader_gt=None,
        dataset_name='guesswhat',
        **kwargs
    ):
        super().__init__()
        model = model.lower()
        split = split.lower()
        # assert model in ['guesser', 'oracle', 'qgen']
        assert split in ['train', 'valid', 'test']
        self.dataroot = dataroot
        self.split = split
        self.model = model
        # self.name = 'guessWhatDataset-%s'.format(model)
        self.answer2id = {'Yes': 0, 'No': 1, 'N/A': 2}
        self.yes_id = tokenizer.yes_id
        self.no_id = tokenizer.no_id
        self.na_id = tokenizer.na_id
        self.eoq_id = tokenizer.eoq_id
        self.eod_id = tokenizer.eod_id
        self.sos_id = tokenizer.sos_id
        self.pad_id = tokenizer.pad_id
        self.answer2token = {
            'Yes': self.yes_id,
            'No': self.no_id,
            'N/A': self.na_id
            }
        self._image_features_reader = image_features_reader
        self._image_features_reader_gt = image_features_reader_gt
        self._tokenizer = tokenizer
        # cache_path = os.path.join(
        #     dataroot, 'cache',  self.name + '_' + split + '.pkl'
        # )
        
        self.data_path = os.path.join(self.dataroot, '%s.%s.jsonl' % (dataset_name, self.split))
        self.entries = self._load_dataset()
        self.tensorize()


    @abc.abstractmethod
    def _load_dataset(self):
        raise NotImplementedError

    @abc.abstractmethod
    def tensorize(self):
        raise NotImplementedError

    def __len__(self):
        return len(self.entries)

    @abc.abstractmethod
    def __getitem__(self, index):
        raise NotImplementedError



