import torch
from datasets import load_dataset
from torch.utils.data import Dataset


class BeansDataset(Dataset):

    def __init__(self, image_processor, tp_size=1, split='train'):

        super().__init__()
        self.image_processor = image_processor
        self.ds = load_dataset('beans')[split]
        self.label_names = self.ds.features['labels'].names
        while len(self.label_names) % tp_size != 0:
            # ensure that the number of labels is multiple of tp_size
            self.label_names.append(f"pad_label_{len(self.label_names)}")
        self.num_labels = len(self.label_names)
        self.inputs = []
        for example in self.ds:
            self.inputs.append(self.process_example(example))

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

    def __getitem__(self, idx):
        return self.inputs[idx]

    def process_example(self, example):
        input = self.image_processor(example['image'], return_tensors='pt')
        input['labels'] = example['labels']
        return input


def beans_collator(batch):
    return {
        'pixel_values': torch.cat([data['pixel_values'] for data in batch], dim=0),
        'labels': torch.tensor([data['labels'] for data in batch], dtype=torch.int64)
    }