from os.path import expanduser
import torch
import json
import torchvision
from general_utils import get_from_repository
from general_utils import log
from torchvision import transforms

PASCAL_VOC_CLASSES_ZS = [['cattle.n.01', 'motorcycle.n.01'], ['aeroplane.n.01', 'sofa.n.01'], 
                         ['cat.n.01', 'television.n.03'], ['train.n.01', 'bottle.n.01'],
                          ['chair.n.01', 'pot_plant.n.01']]


class PascalZeroShot(object):

    def __init__(self, split, n_unseen, image_size=224) -> None:
        super().__init__()

        import sys
        sys.path.append('third_party/JoEm')
        from third_party.JoEm.data_loader.dataset import VOCSegmentation
        from third_party.JoEm.data_loader import get_seen_idx, get_unseen_idx, VOC

        self.pascal_classes = VOC
        self.image_size = image_size

        self.transform = transforms.Compose([
            transforms.Resize((image_size, image_size)),
        ])

        if split == 'train':
            self.voc = VOCSegmentation(get_unseen_idx(n_unseen), get_seen_idx(n_unseen), 
                                       split=split, transform=True, transform_args=dict(base_size=312, crop_size=312), 
                                       ignore_bg=False, ignore_unseen=False, remv_unseen_img=True)
        elif split == 'val':
            self.voc = VOCSegmentation(get_unseen_idx(n_unseen), get_seen_idx(n_unseen), 
                                       split=split, transform=False, 
                                       ignore_bg=False, ignore_unseen=False)

        self.unseen_idx = get_unseen_idx(n_unseen)

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

    def __getitem__(self, i):

        sample = self.voc[i]
        label = sample['label'].long()
        all_labels = [l for l in torch.where(torch.bincount(label.flatten())>0)[0].numpy().tolist() if l != 255]
        class_indices = [l for l in all_labels]
        class_names = [self.pascal_classes[l] for l in all_labels]

        image = self.transform(sample['image'])

        label = transforms.Resize((self.image_size, self.image_size), 
            interpolation=torchvision.transforms.InterpolationMode.NEAREST)(label.unsqueeze(0))[0]

        return (image,), (label, )