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

class ImageTextContrastiveCollator:

    def __init__(self):

        return

    def __call__(self, batch):

        inputs = defaultdict(list)

        for data in batch:

            inputs['image'].append(data['image'])

            inputs['text_input'].append(data['text_input'])

            inputs['text_output'].append(data['text_output'])

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

        return inputs

class Quiltdataset(Dataset):

    def __init__(self):

        # self.df = pd.read_csv(csv_path)

        self.df = pd.read_csv('../BLIP/LAVIS-main/quilt.csv')

        self.df = self.df.dropna(axis=0, subset=['pathology'])[400000:]

        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.df)

    def __getitem__(self, index):

        caption = self.df.iloc[index]['caption']

        if type(caption) == float:

            caption = "This is a image about the pathology."

        img_path = self.df.iloc[index]['image_path']

        # img_path = os.path.join("../", img_path)

        # image = Image.open(img_path).convert('RGB')

        # image = self.transform(image)

        # caption = self.text_processor(caption)

        # img = self.transform(img)

        caption = caption.split()

        prefix = caption[:int(len(caption) * 0.2)]

        subfix = caption[int(len(caption) * 0.2):]

        prefix = " ".join(prefix)

        subfix = " ".join(subfix)

        return {

            "image": img_path,

            "text_input": prefix,

            "text_output": subfix,

        }

        # return {

        #     "image": img_path,

        #     "text_input": caption,

        #     "text_output": caption,

        # }

if __name__ == '__main__':

    test = Quiltdataset()

    print(test.__len__())

    print(test.__getitem__(0))

    print(test.__getitem__(1))