import numpy as np

from datasets.BaseDataset import BaseDataset


class Cityscapes(BaseDataset):
    
    # trans_labels=[]
    # # Standard Constructor
    # def __init__(self):
    #     BaseDataset.__init__(self)#call std constructor of base class
    #     if self.isCustomData==False:
    #        self.trans_labels = [7, 8, 11, 12, 13, 17, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 31, 32, 33]
    #     else:
    #         self.trans_labels = self.pickTransLabels()              


    @classmethod
    def get_class_colors(cls, *args):
        if cls.isCustomData==False:
            return [[128, 64, 128], [244, 35, 232], [70, 70, 70],
                    [102, 102, 156], [190, 153, 153], [153, 153, 153],
                    [250, 170, 30], [220, 220, 0], [107, 142, 35],
                    [152, 251, 152], [70, 130, 180], [220, 20, 60], [255, 0, 0],
                    [0, 0, 142], [0, 0, 70], [0, 60, 100], [0, 80, 100],
                    [0, 0, 230], [119, 11, 32]]
        else:
            return cls.pickColor()

    @classmethod
    def get_class_names(cls, *args):
        if cls.isCustomData==False:
            # class counting(gtFine)
            # 2953 2811 2934  970 1296 2949 1658 2808 2891 1654 2686 2343 1023 2832
            # 359  274  142  513 1646
            return ['road', 'sidewalk', 'building', 'wall', 'fence', 'pole',
                    'traffic light', 'traffic sign',
                    'vegetation', 'terrain', 'sky', 'person', 'rider', 'car',
                    'truck', 'bus', 'train', 'motorcycle', 'bicycle']
        else:
            return cls.pickNames()
            
    @classmethod
    def transform_label(cls, pred, name):
        label = np.zeros(pred.shape)
        ids = np.unique(pred)
        for id in ids:
            label[np.where(pred == id)] = cls.trans_labels[id]

        new_name = (name.split('.')[0]).split('_')[:-1]
        new_name = '_'.join(new_name) + '.png'

        print('Trans', name, 'to', new_name, '    ',
              np.unique(np.array(pred, np.uint8)), ' ---------> ',
              np.unique(np.array(label, np.uint8)))
        return label, new_name