from torch.utils.data import Dataset
import pandas as pd
import os
from PIL import Image
from torchvision import transforms
from collections import defaultdict
import torch
import pickle

class ImageTextContrastiveCollator:
    def __init__(self):
        return
    def __call__(self, batch):
        inputs = defaultdict(list)
        for data in batch:
            inputs['image'].append(data['image'])
            inputs['question'].append(data['question'])
            inputs['answer'].append(data['answer'])
            

        # inputs['image'] = torch.stack(inputs['image'])

        return inputs
pkl_path = '../PathVQA/pvqa/qas/test_vqa.pkl'
class PVQAdataset(Dataset):
    def __init__(self):
        # self.df = pd.read_csv(csv_path)
        with open(pkl_path, 'rb') as f:
            self.data = pickle.load(f)
       
        normalize = transforms.Normalize(
            (0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)
        )

        self.transform = transforms.Compose(
            [
                transforms.RandomResizedCrop(224, scale=(0.2, 1.0)),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                normalize,
            ]
        )
        
    def __len__(self):
        return len(self.data)
    def __getitem__(self, index):
        question = self.data[index]['sent']
        answer = list(self.data[index]['label'].keys())[0]
        img_path = os.path.join('../PathVQA/pvqa/images', 'test', self.data[index]['img_id'])+".jpg"
        return {
            "image": img_path,
            "question": question,
            "answer": answer,
        }
        # return {
        #     "image": img_path,
        #     "text_input": caption,
        #     "text_output": caption,
        # }
         
        
        
if __name__ == '__main__':
    test = PVQAdataset()
    print(test.__len__())
    print(test.__getitem__(0))
    print(test.__getitem__(1))