from ..data_aug import imagenet_like_image_train_aug, imagenet_like_image_test_aug
from ..ab_dataset import ABDataset
from ..dataset_split import train_val_split, train_val_test_split
from torchvision.datasets import ImageFolder
import os
from typing import Dict, List, Optional
from torchvision.transforms import Compose

from ..registery import dataset_register

# with open(os.path.join(os.path.dirname(__file__), 'fruits360_classes.txt'), 'r') as f:
#     classes = [line.split(':')[0].strip('"') for line in f.readlines()]
#     assert len(classes) == 131


# gta_classes = [
#     'road', 'sidewalk', 'building', 'wall',
#     'fence', 'pole', 'light', 'sign',
#     'vegetation', 'terrain', 'sky', 'people', # person
#     'rider', 'car', 'truck', 'bus', 'train',
#     'motocycle', 'bicycle'
# ]
# cityscapes_classes = []

# ignore_label = 255
# m = {-1: ignore_label, 0: ignore_label, 1: ignore_label, 2: ignore_label,
#                 3: ignore_label, 4: ignore_label, 5: ignore_label, 6: ignore_label,
#                 7: 0, 8: 1, 9: ignore_label, 10: ignore_label, 11: 2, 12: 3, 13: 4,
#                 14: ignore_label, 15: ignore_label, 16: ignore_label, 17: 5,
#                 18: ignore_label, 19: 6, 20: 7, 21: 8, 22: 9, 23: 10, 24: 11, 25: 12, 26: 13, 27: 14,
#                 28: 15, 29: ignore_label, 30: ignore_label, 31: 16, 32: 17, 33: 18}

# for ci, c in enumerate(gta_classes):
#     for k, v in m.items():
#         if v == ci:
#             cityscapes_classes += [c]
# print(cityscapes_classes)
# exit()

@dataset_register(
    name='CityscapesCls', 
    classes=[
        'road', 'sidewalk', 'building', 'wall',
        'fence', 'pole', 'light', 'sign',
        'vegetation', 'terrain', 'sky', 'people', # person
        'rider', 'car', 'truck', 'bus', 'train',
        'motocycle', 'bicycle'
    ],
    task_type='Image Classification',
    object_type='Autonomous Driving',
    class_aliases=[],
    shift_type=None
)

class CityscapesCls(ABDataset):
    def create_dataset(self, root_dir: str, split: str, transform: Optional[Compose],
                       classes: List[str], ignore_classes: List[str], idx_map: Optional[Dict[int, int]]):
        if transform is None:
            transform = imagenet_like_image_train_aug() if split == 'train' else imagenet_like_image_test_aug()
            self.transform = transform
        #root_dir = os.path.join(root_dir, 'train' if split != 'test' else 'val')
        dataset = ImageFolder(root_dir, transform=transform)

        if len(ignore_classes) > 0:
            ignore_classes_idx = [classes.index(c) for c in ignore_classes]
            dataset.samples = [s for s in dataset.samples if s[1] not in ignore_classes_idx]

        if idx_map is not None:
            dataset.samples = [(s[0], idx_map[s[1]]) if s[1] in idx_map.keys() else s for s in dataset.samples]

        dataset = train_val_test_split(dataset, split)
        return dataset