import numpy as np

def create_ade20k_label_colormap_official():
  """Creates a label colormap used in ADE20K segmentation benchmark.

  Returns:
    A colormap for visualizing segmentation results.
  """
  return [
      [120, 120, 120],
      [180, 120, 120],
      [6, 230, 230],
      [80, 50, 50],
      [4, 200, 3],
      [120, 120, 80],
      [140, 140, 140],
      [204, 5, 255],
      [230, 230, 230],
      [4, 250, 7],
      [224, 5, 255],
      [235, 255, 7],
      [150, 5, 61],
      [120, 120, 70],
      [8, 255, 51],
      [255, 6, 82],
      [143, 255, 140],
      [204, 255, 4],
      [255, 51, 7],
      [204, 70, 3],
      [0, 102, 200],
      [61, 230, 250],
      [255, 6, 51],
      [11, 102, 255],
      [255, 7, 71],
      [255, 9, 224],
      [9, 7, 230],
      [220, 220, 220],
      [255, 9, 92],
      [112, 9, 255],
      [8, 255, 214],
      [7, 255, 224],
      [255, 184, 6],
      [10, 255, 71],
      [255, 41, 10],
      [7, 255, 255],
      [224, 255, 8],
      [102, 8, 255],
      [255, 61, 6],
      [255, 194, 7],
      [255, 122, 8],
      [0, 255, 20],
      [255, 8, 41],
      [255, 5, 153],
      [6, 51, 255],
      [235, 12, 255],
      [160, 150, 20],
      [0, 163, 255],
      [140, 140, 140],
      [250, 10, 15],
      [20, 255, 0],
      [31, 255, 0],
      [255, 31, 0],
      [255, 224, 0],
      [153, 255, 0],
      [0, 0, 255],
      [255, 71, 0],
      [0, 235, 255],
      [0, 173, 255],
      [31, 0, 255],
      [11, 200, 200],
      [255, 82, 0],
      [0, 255, 245],
      [0, 61, 255],
      [0, 255, 112],
      [0, 255, 133],
      [255, 0, 0],
      [255, 163, 0],
      [255, 102, 0],
      [194, 255, 0],
      [0, 143, 255],
      [51, 255, 0],
      [0, 82, 255],
      [0, 255, 41],
      [0, 255, 173],
      [10, 0, 255],
      [173, 255, 0],
      [0, 255, 153],
      [255, 92, 0],
      [255, 0, 255],
      [255, 0, 245],
      [255, 0, 102],
      [255, 173, 0],
      [255, 0, 20],
      [255, 184, 184],
      [0, 31, 255],
      [0, 255, 61],
      [0, 71, 255],
      [255, 0, 204],
      [0, 255, 194],
      [0, 255, 82],
      [0, 10, 255],
      [0, 112, 255],
      [51, 0, 255],
      [0, 194, 255],
      [0, 122, 255],
      [0, 255, 163],
      [255, 153, 0],
      [0, 255, 10],
      [255, 112, 0],
      [143, 255, 0],
      [82, 0, 255],
      [163, 255, 0],
      [255, 235, 0],
      [8, 184, 170],
      [133, 0, 255],
      [0, 255, 92],
      [184, 0, 255],
      [255, 0, 31],
      [0, 184, 255],
      [0, 214, 255],
      [255, 0, 112],
      [92, 255, 0],
      [0, 224, 255],
      [112, 224, 255],
      [70, 184, 160],
      [163, 0, 255],
      [153, 0, 255],
      [71, 255, 0],
      [255, 0, 163],
      [255, 204, 0],
      [255, 0, 143],
      [0, 255, 235],
      [133, 255, 0],
      [255, 0, 235],
      [245, 0, 255],
      [255, 0, 122],
      [255, 245, 0],
      [10, 190, 212],
      [214, 255, 0],
      [0, 204, 255],
      [20, 0, 255],
      [255, 255, 0],
      [0, 153, 255],
      [0, 41, 255],
      [0, 255, 204],
      [41, 0, 255],
      [41, 255, 0],
      [173, 0, 255],
      [0, 245, 255],
      [71, 0, 255],
      [122, 0, 255],
      [0, 255, 184],
      [0, 92, 255],
      [184, 255, 0],
      [0, 133, 255],
      [255, 214, 0],
      [25, 194, 194],
      [102, 255, 0],
      [92, 0, 255],
      [0, 0, 0]
  ]

def create_coco_stuff_label_colormap():
    return {
        0: [128, 0, 0],
        1: [0, 128, 0],
        2: [128, 128, 0],
        3: [0, 0, 128],
        4: [128, 0, 128],
        5: [0, 128, 128],
        6: [128, 128, 128],
        7: [64, 0, 0],
        8: [192, 0, 0],
        9: [64, 128, 0],
        10: [192, 128, 0],
        11: [0, 0, 0],
        12: [192, 0, 128],
        13: [64, 128, 128],
        14: [192, 128, 128],
        15: [0, 64, 0],
        16: [128, 64, 0],
        17: [0, 192, 0],
        18: [128, 192, 0],
        19: [0, 64, 128],
        20: [128, 64, 128],
        21: [0, 192, 128],
        22: [128, 192, 128],
        23: [64, 64, 0],
        24: [192, 64, 0],
        25: [0, 0, 0],
        26: [192, 192, 0],
        27: [64, 64, 128],
        28: [0, 0, 0],
        29: [0, 0, 0],
        30: [192, 192, 128],
        31: [0, 0, 64],
        32: [128, 0, 64],
        33: [0, 128, 64],
        34: [128, 128, 64],
        35: [0, 0, 192],
        36: [128, 0, 192],
        37: [0, 128, 192],
        38: [128, 128, 192],
        39: [64, 0, 64],
        40: [192, 0, 64],
        41: [64, 128, 64],
        42: [192, 128, 64],
        43: [64, 0, 192],
        44: [0, 0, 0],
        45: [64, 128, 192],
        46: [192, 128, 192],
        47: [0, 64, 64],
        48: [128, 64, 64],
        49: [0, 192, 64],
        50: [128, 192, 64],
        51: [0, 64, 192],
        52: [128, 64, 192],
        53: [0, 192, 192],
        54: [128, 192, 192],
        55: [64, 64, 64],
        56: [192, 64, 64],
        57: [64, 192, 64],
        58: [192, 192, 64],
        59: [64, 64, 192],
        60: [192, 64, 192],
        61: [64, 192, 192],
        62: [192, 192, 192],
        63: [32, 0, 0],
        64: [160, 0, 0],
        65: [0, 0, 0],
        66: [160, 128, 0],
        67: [0, 0, 0],
        68: [0, 0, 0],
        69: [32, 128, 128],
        70: [0, 0, 0],
        71: [96, 0, 0],
        72: [224, 0, 0],
        73: [96, 128, 0],
        74: [224, 128, 0],
        75: [96, 0, 128],
        76: [224, 0, 128],
        77: [96, 128, 128],
        78: [224, 128, 128],
        79: [32, 64, 0],
        80: [160, 64, 0],
        81: [32, 192, 0],
        82: [0, 0, 0],
        83: [32, 64, 128],
        84: [160, 64, 128],
        85: [32, 192, 128],
        86: [160, 192, 128],
        87: [96, 64, 0],
        88: [224, 64, 0],
        89: [96, 192, 0],
        90: [0, 0, 0],
        91: [96, 64, 128],
        92: [224, 64, 128],
        93: [96, 192, 128],
        94: [224, 192, 128],
        95: [32, 0, 64],
        96: [160, 0, 64],
        97: [32, 128, 64],
        98: [160, 128, 64],
        99: [32, 0, 192],
        100: [160, 0, 192],
        101: [32, 128, 192],
        102: [160, 128, 192],
        103: [96, 0, 64],
        104: [224, 0, 64],
        105: [96, 128, 64],
        106: [224, 128, 64],
        107: [96, 0, 192],
        108: [224, 0, 192],
        109: [96, 128, 192],
        110: [224, 128, 192],
        111: [32, 64, 64],
        112: [160, 64, 64],
        113: [32, 192, 64],
        114: [160, 192, 64],
        115: [32, 64, 192],
        116: [160, 64, 192],
        117: [32, 192, 192],
        118: [160, 192, 192],
        119: [96, 64, 64],
        120: [224, 64, 64],
        121: [96, 192, 64],
        122: [224, 192, 64],
        123: [96, 64, 192],
        124: [224, 64, 192],
        125: [96, 192, 192],
        126: [224, 192, 192],
        127: [0, 32, 0],
        128: [128, 32, 0],
        129: [0, 160, 0],
        130: [128, 160, 0],
        131: [0, 32, 128],
        132: [128, 32, 128],
        133: [0, 160, 128],
        134: [128, 160, 128],
        135: [64, 32, 0],
        136: [192, 32, 0],
        137: [64, 160, 0],
        138: [192, 160, 0],
        139: [64, 32, 128],
        140: [192, 32, 128],
        141: [64, 160, 128],
        142: [192, 160, 128],
        143: [0, 96, 0],
        144: [128, 96, 0],
        145: [0, 224, 0],
        146: [128, 224, 0],
        147: [0, 96, 128],
        148: [128, 96, 128],
        149: [0, 224, 128],
        150: [128, 224, 128],
        151: [64, 96, 0],
        152: [192, 96, 0],
        153: [64, 224, 0],
        154: [192, 224, 0],
        155: [64, 96, 128],
        156: [192, 96, 128],
        157: [64, 224, 128],
        158: [192, 224, 128],
        159: [0, 32, 64],
        160: [128, 32, 64],
        161: [0, 160, 64],
        162: [128, 160, 64],
        163: [0, 32, 192],
        164: [128, 32, 192],
        165: [0, 160, 192],
        166: [128, 160, 192],
        167: [64, 32, 64],
        168: [192, 32, 64],
        169: [64, 160, 64],
        170: [192, 160, 64],
        171: [64, 32, 192],
        172: [192, 32, 192],
        173: [64, 160, 192],
        174: [192, 160, 192],
        175: [0, 96, 64],
        176: [128, 96, 64],
        177: [0, 224, 64],
        178: [128, 224, 64],
        179: [0, 96, 192],
        180: [128, 96, 192],
        181: [0, 224, 192],
        255: [0, 0, 0]
    }

def convert_gray_to_color_official(color_map, gray_image):
    gray_image[gray_image == 255] = 150
    color_image_c1 = []
    color_image_c2 = []
    color_image_c3 = []
    h, w = gray_image.shape
    for i in range(gray_image.shape[0]):
        for j in range(gray_image.shape[1]):
            color = gray_image[i, j]
            color_image_c1.append(color_map[color][0])
            color_image_c2.append(color_map[color][1])
            color_image_c3.append(color_map[color][2])
    return np.stack([np.asarray(color_image_c1, dtype=np.uint8),
                    np.asarray(color_image_c2, dtype=np.uint8),
                    np.asarray(color_image_c3, dtype=np.uint8)], axis=-1).reshape(h, w, 3)


def gen_ade20k_color_map():
    res = []
    for i in range(7):
        for j in range(7):
            for k in range(7):
                res.append((i*30, j*30, k*30))
    return res

def convert_gray_to_color(color_map, gray_image):
    color_image_c1 = []
    color_image_c2 = []
    color_image_c3 = []
    h, w = gray_image.shape
    for i in range(gray_image.shape[0]):
        for j in range(gray_image.shape[1]):
            color = gray_image[i, j]
            color_image_c1.append(color_map[color][0])
            color_image_c2.append(color_map[color][1])
            color_image_c3.append(color_map[color][2])
    return np.stack([np.asarray(color_image_c1, dtype=np.uint8),
                    np.asarray(color_image_c2, dtype=np.uint8),
                    np.asarray(color_image_c3, dtype=np.uint8)], axis=-1).reshape(h, w, 3)
