
class DatasetConfig:

    ignore_label = -1
    datasets = ['gta5','synthia','cityscapes','idd','bdd100k','mapillary','shift']
    mean_bgr = {
        'gta5': [127.5, 127.5, 127.5],
        'synthia': [127.5, 127.5, 127.5],
        'cityscapes': [72.39239876, 82.90891754, 73.15835921],
        'idd': [127.5, 127.5, 127.5],
        'bdd100k': [127.5, 127.5, 127.5],
        'mapillary': [127.5, 127.5, 127.5],
        'shift': [127.5, 127.5, 127.5],
        'imagenet': [104.00698793, 116.66876762, 122.67891434],
        'uniform': [127.5, 127.5, 127.5]
    }

    # class names from original dataset
    id2name = {
        'gta5': {0: 'unlabeled', 1: 'ego vehicle', 2: 'rectification border', 3: 'out of roi', 4: 'static', 5: 'dynamic', 6: 'ground', 7: 'road',
                 8: 'sidewalk', 9: 'parking', 10: 'rail track', 11: 'building', 12: 'wall', 13: 'fence', 14: 'guard rail', 15: 'bridge', 16: 'tunnel',
                 17: 'pole', 18: 'polegroup', 19: 'traffic light', 20: 'traffic sign', 21: 'vegetation', 22: 'terrain', 23: 'sky', 24: 'person', 25: 'rider',
                 26: 'car', 27: 'truck', 28: 'bus', 29: 'caravan', 30: 'trailer', 31: 'train', 32: 'motorcycle', 33: 'bicycle', 34: 'license plate'},
        'synthia': {0: 'void', 1: 'sky', 2: 'building', 3: 'road', 4: 'sidewalk', 5: 'fence', 6: 'vegetation', 7: 'pole', 8: 'car', 9: 'traffic sign',
                    10: 'person',  11: 'bicycle', 12: 'motorcycle', 13: 'parking-slot', 14: 'road-work', 15: 'traffic light', 16: 'terrain',
                    17: 'rider', 18: 'truck', 19: 'bus',  20: 'train', 21: 'wall', 22: 'lanemarking'},
        'cityscapes': {0: 'unlabeled', 1: 'ego vehicle', 2: 'rectification border', 3: 'out of roi', 4: 'static', 5: 'dynamic', 6: 'ground', 7: 'road',
                       8: 'sidewalk', 9: 'parking', 10: 'rail track', 11: 'building', 12: 'wall', 13: 'fence', 14: 'guard rail', 15: 'bridge', 16: 'tunnel',
                       17: 'pole', 18: 'polegroup', 19: 'traffic light', 20: 'traffic sign', 21: 'vegetation', 22: 'terrain', 23: 'sky', 24: 'person',
                       25: 'rider', 26: 'car', 27: 'truck', 28: 'bus', 29: 'caravan', 30: 'trailer', 31: 'train', 32: 'motorcycle', 33: 'bicycle', -1: 'license plate'},
        'idd': {0: 'road', 1: 'drivable fallback', 2: 'sidewalk', 3: 'non-drivable fallback', 4: 'person', 5: 'rider', 6: 'motorcycle', 7: 'bicycle', 8: 'autorickshaw', 9: 'car',
                10: 'truck', 11: 'bus', 12: 'train', 13: 'curb', 14: 'wall', 15: 'fence', 16: 'guard rail', 17: 'billboard', 18: 'traffic sign', 19: 'traffic light',
                20: 'pole', 21: 'obs-str-bar-fallback', 22: 'building', 23: 'tunnel', 24: 'vegetation', 25: 'sky', 255: 'unlabeled'},
        # data already in cityscapes 19-class format
        'bdd100k': {0: 'road', 1: 'sidewalk', 2: 'building', 3: 'wall', 4: 'fence', 5: 'pole', 6: 'traffic light', 7: 'traffic sign', 8: 'vegetation', 9: 'terrain', 10: 'sky',
                    11: 'person', 12: 'rider', 13: 'car', 14: 'truck', 15: 'bus', 16: 'train', 17: 'motorcycle', 18: 'bicycle', 255: 'unlabeled'},
        'mapillary': {0: 'Bird', 1: 'Ground Animal', 2: 'Curb', 3: 'Fence', 4: 'Guard Rail', 5: 'Barrier', 6: 'Wall', 7: 'Bike Lane', 8: 'Crosswalk - Plain', 9: 'Curb Cut', 10: 'Parking',
                      11: 'Pedestrian Area', 12: 'Rail Track', 13: 'Road', 14: 'Service Lane', 15: 'Sidewalk', 16: 'Bridge', 17: 'Building', 18: 'Tunnel', 19: 'Person', 20: 'Bicyclist',
                      21: 'Motorcyclist', 22: 'Other Rider', 23: 'Lane Marking - Crosswalk', 24: 'Lane Marking - General', 25: 'Mountain', 26: 'Sand', 27: 'Sky', 28: 'Snow', 29: 'Terrain', 30: 'Vegetation',
                      31: 'Water', 32: 'Banner', 33: 'Bench', 34: 'Bike Rack', 35: 'Billboard', 36: 'Catch Basin', 37: 'CCTV Camera', 38: 'Fire Hydrant', 39: 'Junction Box', 40: 'Mailbox',
                      41: 'Manhole', 42: 'Phone Booth', 43: 'Pothole', 44: 'Street Light', 45: 'Pole', 46: 'Traffic Sign Frame', 47: 'Utility Pole', 48: 'Traffic Light', 49: 'Traffic Sign (Back)',
                      50: 'Traffic Sign (Front)', 51: 'Trash Can', 52: 'Bicycle', 53: 'Boat', 54: 'Bus', 55: 'Car', 56: 'Caravan', 57: 'Motorcycle', 58: 'On Rails', 59: 'Other Vehicle', 60: 'Trailer',
                      61: 'Truck', 62: 'Wheeled Slow', 63: 'Car Mount', 64: 'Ego Vehicle', 65: 'Unlabeled', 66: None},
        'shift': {0: 'unlabeled', 1: 'building', 2: 'fence', 3: 'other', 4: 'pedestrian', 5: 'pole', 6: 'road line', 7: 'road',
                  8: 'sidewalk', 9: 'vegetation', 10: 'vehicles', 11: 'wall', 12: 'traffic sign', 13: 'sky', 14: 'ground', 15: 'bridge', 16: 'rail track',
                  17: 'guard rail', 18: 'traffic light', 19: 'static', 20: 'dynamic', 21: 'water', 22: 'terrain'}
    }

    name2color = {
        'city19': {'road': [128, 64,128], 'sidewalk': [244, 35,232], 'building': [70, 70, 70], 'wall': [102,102,156], 'fence': [190,153,153], 'pole': [153,153,153],
                   'traffic light': [250,170, 30], 'traffic sign': [220,220,  0], 'vegetation': [107,142, 35], 'terrain': [152,251,152], #'non-drivable fallback': [152,251,152],  # 'non-drivable fallback' as 'terrain' to adapt IDD to 19 classes
                   'sky': [70,130,180], 'person': [220, 20, 60], 'rider': [255,  0,  0], 'car': [0, 0, 142], 'truck': [0,  0, 70], 'bus': [0, 60,100], 'train': [0, 80,100],
                   'motorcycle': [0,  0,230], 'bicycle': [119, 11, 32], 'unknown':[0, 0, 0]},
        'idd18': {'road': [128, 64,128], 'sidewalk': [244, 35,232], 'building': [70, 70, 70], 'wall': [102,102,156], 'fence': [190,153,153], 'pole': [153,153,153],
                  'traffic light': [250,170, 30], 'traffic sign': [220,220,  0], 'vegetation': [107,142, 35],
                  'sky': [70,130,180], 'person': [220, 20, 60], 'rider': [255,  0,  0], 'car': [0, 0, 142], 'truck': [0,  0, 70], 'bus': [0, 60,100], 'train': [0, 80,100],
                  'motorcycle': [0,  0,230], 'bicycle': [119, 11, 32], 'unknown':[0, 0, 0]},
        'synthia16': {'road': [128, 64,128], 'sidewalk': [244, 35,232], 'building': [70, 70, 70], 'wall': [102,102,156], 'fence': [190,153,153], 'pole': [153,153,153],
                      'traffic light': [250,170, 30], 'traffic sign': [220,220,  0], 'vegetation': [107,142, 35],
                      'sky': [70,130,180], 'person': [220, 20, 60], 'rider': [255,  0,  0], 'car': [0, 0, 142], 'bus': [0, 60,100],
                      'motorcycle': [0,  0,230], 'bicycle': [119, 11, 32], 'unknown':[0, 0, 0]},
        'mapillary': {'Bird': [165, 42, 42], 'Ground Animal': [0, 192, 0], 'Curb': [196, 196, 196], 'Fence': [190, 153, 153], 'Guard Rail': [180, 165, 180], 'Barrier': [90, 120, 150],
                      'Wall': [102, 102, 156], 'Bike Lane': [128, 64, 255], 'Crosswalk - Plain': [140, 140, 200], 'Curb Cut': [170, 170, 170], 'Parking': [250, 170, 160],
                      'Pedestrian Area': [96, 96, 96], 'Rail Track': [230, 150, 140], 'Road': [128, 64, 128], 'Service Lane': [110, 110, 110], 'Sidewalk': [244, 35, 232],
                      'Bridge': [150, 100, 100], 'Building': [70, 70, 70], 'Tunnel': [150, 120, 90], 'Person': [220, 20, 60], 'Bicyclist': [255, 0, 0], 'Motorcyclist': [255, 0, 100],
                      'Other Rider': [255, 0, 200], 'Lane Marking - Crosswalk': [200, 128, 128], 'Lane Marking - General': [255, 255, 255], 'Mountain': [64, 170, 64], 'Sand': [230, 160, 50],
                      'Sky': [70, 130, 180], 'Snow': [190, 255, 255], 'Terrain': [152, 251, 152], 'Vegetation': [107, 142, 35], 'Water': [0, 170, 30], 'Banner': [255, 255, 128],
                      'Bench': [250, 0, 30], 'Bike Rack': [100, 140, 180], 'Billboard': [220, 220, 220], 'Catch Basin': [220, 128, 128], 'CCTV Camera': [222, 40, 40],
                      'Fire Hydrant': [100, 170, 30], 'Junction Box': [40, 40, 40], 'Mailbox': [33, 33, 33], 'Manhole': [100, 128, 160], 'Phone Booth': [142, 0, 0], 'Pothole': [70, 100, 150],
                      'Street Light': [210, 170, 100], 'Pole': [153, 153, 153], 'Traffic Sign Frame': [128, 128, 128], 'Utility Pole': [0, 0, 80], 'Traffic Light': [250, 170, 30],
                      'Traffic Sign (Back)': [192, 192, 192], 'Traffic Sign (Front)': [220, 220, 0], 'Trash Can': [140, 140, 20], 'Bicycle': [119, 11, 32], 'Boat': [150, 0, 255],
                      'Bus': [0, 60, 100], 'Car': [0, 0, 142], 'Caravan': [0, 0, 90], 'Motorcycle': [0, 0, 230], 'On Rails': [0, 80, 100], 'Other Vehicle': [128, 64, 64],
                      'Trailer': [0, 0, 110], 'Truck': [0, 0, 70], 'Wheeled Slow': [0, 0, 192], 'Car Mount': [32, 32, 32], 'Ego Vehicle': [120, 10, 10], 'Unlabeled': [0, 0, 0], 'unknown':[0, 0, 0]},
        'shift': {'unlabeled': (0, 0, 0), 'building': (70, 70, 70), 'fence': (100, 40, 40), 'other': (55, 90, 80), 'pedestrian': (220, 20, 60), 'pole': (153, 153, 153),
                  'road line': (157, 234, 50), 'road': (128, 64, 128), 'sidewalk': (244, 35, 232), 'vegetation': (107, 142, 35), 'vehicles': (0, 0, 142), 'wall': (102, 102, 156),
                  'traffic sign': (220, 220, 0), 'sky': (70, 130, 180), 'ground': (81, 0, 81), 'bridge': (150, 100, 100), 'rail track': (230, 150, 140), 'guard rail': (180, 165, 180),
                  'traffic light': (250, 170, 30), 'static': (110, 190, 160), 'dynamic': (170, 120, 50), 'water': (45, 60, 150), 'terrain': (145, 170, 100), 'unknown':(0, 0, 0)}

    }
    
    @staticmethod
    def name2name():
        # Convert from inner-key dataset to outer-key dataset class sets: used to rename classes, or group sets of classes of first dataset to single classes in destination dataset
        # Final class set should perfectly match destination class set (no original classes not present in destination set still available)
        name2name = {
            'cityscapes': {
                'mapillary': {'Bird': None, 'Ground Animal': None, 'Curb': 'sidewalk', 'Fence': 'fence', 'Guard Rail': None, 'Barrier': None, 'Wall': 'wall', 'Bike Lane': 'sidewalk',   #!!
                              'Crosswalk - Plain': None, 'Curb Cut': 'sidewalk', 'Parking': None, 'Pedestrian Area': 'sidewalk', 'Rail Track': None, 'Road': 'road', 'Service Lane': 'road', 'Sidewalk': 'sidewalk',
                              'Bridge': None, 'Building': 'building', 'Tunnel': None, 'Person': 'person', 'Bicyclist': 'rider', 'Motorcyclist': 'rider', 'Other Rider': 'rider', 'Lane Marking - Crosswalk': None,
                              'Lane Marking - General': 'road', 'Mountain': None, 'Sand': 'terrain', 'Sky': 'sky', 'Snow': None, 'Terrain': 'terrain', 'Vegetation': 'vegetation', 'Water': None, 'Banner': None,
                              'Bench': None, 'Bike Rack': None, 'Billboard': None, 'Catch Basin': None, 'CCTV Camera': None, 'Fire Hydrant': None, 'Junction Box': None, 'Mailbox': None, 'Manhole': None,
                              'Phone Booth': None, 'Pothole': 'road', 'Street Light': 'pole', 'Pole': 'pole', 'Traffic Sign Frame': 'pole', 'Utility Pole': 'pole', 'Traffic Light': 'traffic light',
                              'Traffic Sign (Back)': None, 'Traffic Sign (Front)': 'traffic sign', 'Trash Can': None, 'Bicycle': 'bicycle', 'Boat': None, 'Bus': 'bus', 'Car': 'car', 'Caravan': None,
                              'Motorcycle': 'motorcycle', 'On Rails': 'train', 'Other Vehicle': None, 'Trailer': None, 'Truck': 'truck', 'Wheeled Slow': None, 'Car Mount': None, 'Ego Vehicle': None,
                              'Unlabeled': None, None: None},
                'idd': {**{v:v for v in DatasetConfig.id2name['idd'].values() if v in DatasetConfig.id2name['cityscapes'].values()},
                        **{v:None for v in DatasetConfig.id2name['idd'].values() if v not in DatasetConfig.id2name['cityscapes'].values()},
                        'non-drivable fallback':'terrain'}  # note: new key-value pairs added manually will overwrite entries from **{} with same keys
            }
        }
        return name2name
        
    @staticmethod
    def name2iid():
        # 'idd18': terrain missing, 'synthia16': terrain, truck, train missing
        name2iid = {
            'standard': {
                'city19': ({'road': 0, 'sidewalk': 1, 'building': 2, 'wall': 3, 'fence': 4, 'pole': 5, 'traffic light': 6, 'traffic sign': 7, 'vegetation': 8, 'terrain': 9,
                            'sky': 10, 'person': 11, 'rider': 12, 'car': 13, 'truck': 14, 'bus': 15, 'train': 16, 'motorcycle': 17, 'bicycle': 18},),
                'idd18': ({'road': 0, 'sidewalk': 1, 'building': 2, 'wall': 3, 'fence': 4, 'pole': 5, 'traffic light': 6, 'traffic sign': 7, 'vegetation': 8, 'sky': 9,  # terrain missing
                           'person': 10, 'rider': 11, 'car': 12, 'truck': 13, 'bus': 14, 'train': 15, 'motorcycle': 16, 'bicycle': 17},),
                'synthia16': ({'road': 0, 'sidewalk': 1, 'building': 2, 'wall': 3, 'fence': 4, 'pole': 5, 'traffic light': 6, 'traffic sign': 7, 'vegetation': 8, 'sky': 9,
                               'person': 10, 'rider': 11, 'car': 12, 'bus': 13, 'motorcycle': 14, 'bicycle': 15},),
                'mapillary': ({**{v:k for k,v in DatasetConfig.id2name['mapillary'].items() if v not in ('Unlabeled',None)}, 'Unlabeled': -1, None: -1},),
                'shift': ({**{v: k for v, k in zip([vv for vv in DatasetConfig.id2name['shift'].values() if vv not in ('unlabeled', 'other')],
                                                   range(len([vv for vv in DatasetConfig.id2name['shift'].values() if vv not in ('unlabeled', 'other')])))
                              }, **{v: -1 for v in ('unlabeled', 'other')}},)
            },
            'CIL': {
                'city19': ({'unknown': 0, 'road': 1, 'sidewalk': 2, 'vegetation': 3, 'terrain': 4, 'sky': 5},
                           {'unknown': 0, 'building': 6, 'wall': 7, 'fence': 8, 'pole': 9, 'traffic light': 10, 'traffic sign': 11},
                           {'unknown': 0, 'person': 12, 'rider': 13, 'car': 14, 'truck': 15, 'bus': 16, 'train': 17, 'motorcycle': 18, 'bicycle': 19}),
                'idd18': ({'unknown': 0, 'road': 1, 'sidewalk': 2, 'vegetation': 3, 'sky': 4},  # terrain missing
                          {'unknown': 0, 'building': 5, 'wall': 6, 'fence': 7, 'pole': 8, 'traffic light': 9, 'traffic sign': 10},
                          {'unknown': 0, 'person': 11, 'rider': 12, 'car': 13, 'truck': 14, 'bus': 15, 'train': 16, 'motorcycle': 17, 'bicycle': 18}),
                'synthia16': ({'unknown': 0, 'road': 1, 'sidewalk': 2, 'vegetation': 3, 'sky': 4},  # terrain missing
                              {'unknown': 0, 'building': 5, 'wall': 6, 'fence': 7, 'pole': 8, 'traffic light': 9, 'traffic sign': 10},
                              {'unknown': 0, 'person': 11, 'rider': 12, 'car': 13, 'bus': 14, 'motorcycle': 15, 'bicycle': 16}),  # truck, train missing
                'shift': ({'unknown': 0, 'road line': 1, 'road': 2, 'sidewalk': 3, 'vegetation': 4, 'sky': 5, 'ground': 6, 'water': 7, 'terrain': 8},
                          {'unknown': 0, 'building': 9, 'fence': 10, 'pole': 11, 'wall': 12, 'traffic sign': 13, 'bridge': 14, 'rail track': 15, 'guard rail':16, 'traffic light': 17, 'static':18},
                          {'unknown': 0, 'pedestrian': 19, 'vehicles': 20, 'dynamic': 21}),
            },
            'doubleCIL': {
                'city19': ({'unknown': 0, 'road': 1, 'sidewalk': 2},
                           {'unknown': 0, 'vegetation': 3, 'terrain': 4, 'sky': 5},
                           {'unknown': 0, 'building': 6, 'wall': 7, 'fence': 8},
                           {'unknown': 0, 'pole': 9, 'traffic light': 10, 'traffic sign': 11},
                           {'unknown': 0, 'person': 12, 'rider': 13, 'motorcycle': 14, 'bicycle': 15},
                           {'unknown': 0, 'car': 16, 'truck': 17, 'bus': 18, 'train': 19}),

            },
            'reverseCIL': {
                'city19': ({'unknown': 0, 'road': 1, 'sidewalk': 2, 'vegetation': 3, 'terrain': 4, 'sky': 5},
                           {'unknown': 0, 'person': 6, 'rider': 7, 'car': 8, 'truck': 9, 'bus': 10, 'train': 11, 'motorcycle': 12, 'bicycle': 13},
                           {'unknown': 0, 'building': 14, 'wall': 15, 'fence': 16, 'pole': 17, 'traffic light': 18, 'traffic sign': 19},),
            }
        }
        return name2iid


    # 1. id2name
    # 2. name2iid
    # 3. id2iid

    @staticmethod
    def get_id2name(dataset, convert_names_to=None):
        """
        Mapping from class GT id to class name
        Args:
            dataset: dataset name
            convert_names_to (optional): dataset to whose set convert class name
        Returns: dict where (key,value) pairs are (class GT id, class name) pairs
        """
        assert dataset in DatasetConfig.datasets, f'{dataset} not in set of available datasets: {DatasetConfig.datasets}'
        id2name = DatasetConfig.id2name[dataset]
        if convert_names_to is not None:
            assert convert_names_to in DatasetConfig.name2name(), f'{convert_names_to} not in set of available destination datasets: {DatasetConfig.name2name()}'
            assert dataset in DatasetConfig.name2name()[convert_names_to], f'{dataset} not in set of available source datasets with {convert_names_to} as destination: {DatasetConfig.name2name()[convert_names_to]}'
            name2name = DatasetConfig.name2name()[convert_names_to][dataset]
            id2name = {k: name2name[v] for k,v in id2name.items()}
        return id2name

    @staticmethod
    def get_name2iid(incr_class_split, class_set, incr_class_step, training):
        name2iid_all_steps = DatasetConfig.name2iid()[incr_class_split][class_set]
        if not training:
            name2iid = name2iid_all_steps[0]
            for i in range(1,incr_class_step+1):
                name2iid = {**name2iid, **{k:v for k,v in name2iid_all_steps[i].items() if v != DatasetConfig.ignore_label}}
            return name2iid
        return name2iid_all_steps[incr_class_step]

    @staticmethod
    def get_id2iid(dataset, incr_class_split, class_set, incr_class_step, training, convert_names_to=None):
        id2name = DatasetConfig.get_id2name(dataset, convert_names_to=convert_names_to)
        name2iid = DatasetConfig.get_name2iid(incr_class_split, class_set, incr_class_step, training)

        for i,n in id2name.items():
            if n not in name2iid.keys():
                name2iid[n] = DatasetConfig.ignore_label

        id2iid = {}
        for i,n in id2name.items():
            if n in name2iid.keys():
                id2iid[i] = name2iid[n]
            else: # it should never enter the else
                id2iid[i] = DatasetConfig.ignore_label
        id2iid[255] = DatasetConfig.ignore_label

        if 'unknown' in name2iid:
            for k,v in id2iid.items():
                if v == DatasetConfig.ignore_label:
                    id2iid[k] = name2iid['unknown']

        return id2iid