from torchvision.transforms.functional import normalize
import torch.nn as nn
import numpy as np


def denormalize(tensor, mean, std):
    mean = np.array(mean)
    std = np.array(std)

    _mean = -mean/std
    _std = 1/std
    return normalize(tensor, _mean, _std)


class Denormalize(object):
    def __init__(self, mean, std):
        mean = np.array(mean)
        std = np.array(std)
        self._mean = -mean/std
        self._std = 1/std

    def __call__(self, tensor):
        if isinstance(tensor, np.ndarray):
            return (tensor - self._mean.reshape(-1,1,1)) / self._std.reshape(-1,1,1)
        return normalize(tensor, self._mean, self._std)


def fix_bn(model):
    for m in model.modules():
        if isinstance(m, nn.BatchNorm2d):
            m.eval()
            m.weight.requires_grad = False
            m.bias.requires_grad = False


def color_map(dataset):
    if dataset=='voc':
        return voc_cmap()
    elif dataset=='cityscapes':
        return cityscapes_cmap()
    elif dataset=='ade':
        return ade_cmap()
    elif dataset =='isaid':
        return isaid_cmap()
    elif dataset =='SAR2020':
        return SAR2020_cmap()
    elif dataset =='Unify_single':
        return unify_single_cmap()
    elif dataset == 'Unify_double':
        return unify_cmap()
    elif dataset == 'Unify_YIJISAR':
        return unify_YIJISAR_cmap()
    elif dataset == 'Unify_Vai':
        return unify_Vai_cmap()

def unify_Vai_cmap():
    cmap = np.zeros((255, 3), dtype=np.uint8)
    colors = [
        [0, 0, 0],#0
        [0, 255, 0],
        [255, 255, 255],
        [0, 0, 255],
        [0, 0, 0],
        [255,255,0],
        [0, 255, 255], #6
        [255,195,128],
        [0, 0, 0],
        [255,255,0],
        [0, 0, 0], #10
        [0, 0, 0],
        [0,0,0],#12
        [0, 0, 0],
        [0,0,255], #14

    ]
    for i in range(len(colors)):
        cmap[i] = colors[i]

    return cmap.astype(np.uint8)

def cityscapes_cmap():
    return np.array([(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), (  0,  0,  0)], 
                         dtype=np.uint8)

def SAR2020_cmap():
    cmap = np.zeros((256, 3), dtype=np.uint8)
    colors = [[255,255,255], [255,255,0], [0,0,255], [0, 255,0], [255,0,0], [0,255,255]]
    for i in range(len(colors)):
        cmap[i] = colors[i]

    return cmap.astype(np.uint8)


def isaid_cmap():
    cmap = np.zeros((255, 3), dtype=np.uint8)
    colors = [[0, 0, 0],
              [0,63,0],
              [0,63,191],
              [0, 127, 63],
              [0, 63, 255],
              [0,100,155],
              [0, 0, 191], [0, 127, 127],
              [0, 127, 255], [0, 191, 127], [0, 0, 63], [0, 191, 127], [0, 127, 191],[0, 63, 63], [0 ,0 ,255], [0,63,127]]
    for i in range(len(colors)):
        cmap[i] = colors[i]

    return cmap.astype(np.uint8)



def unify_cmap():
    cmap = np.zeros((255, 3), dtype=np.uint8)
    colors = [
    [0, 0, 0],
    [0, 127, 255],
    # [0, 0, 191],
    [0, 63, 0],
    [0, 127, 63],
    [0, 63, 255],
    [0, 127, 127],
    # [0, 0, 127],
    [0, 0, 63],
    [0, 63, 127],
    [0, 63, 191],
    [0, 63, 63],
    [0, 127, 191],
    [0, 191, 127],
    [0, 100, 155],
    [0, 0, 255],
        # [255, 255, 255], #0
        #       [0,63,0],  #1
        #       [0,127,63],#2
        #       [0, 63, 255],#3
        #       [0, 127, 127],#4
        #       [0,128,0], #5,
        #       [0, 0, 63],#6
        #       [0,63,127],#7
        #       [0, 63, 191],#8
        #       [0, 63, 63],#9
        #       [0, 127, 191], #10
        #       [0, 191, 127],#11
        #       [0,100,155], #12
        #       [0, 0, 255], #13
              [0,255,0], #14
              [0,153,204], #15
              [204,204,68], #16
              [255, 204, 51], #17
              [255, 255, 204], #18
              [0, 255, 255],
              [255,102,102], #20
              [0,255,0] ,#21
              [255,255,0],#22
              [255,0,0],#23
              [255,195,128],#24
              [153,102,153]
    ]
    for i in range(len(colors)):
        cmap[i] = colors[i]

    return cmap.astype(np.uint8)



def unify_YIJISAR_cmap():
    cmap = np.zeros((255, 3), dtype=np.uint8)
    colors = [
    [255, 255, 0],#0
    [0, 255, 0],
    [0, 63, 0],
    [0, 0, 255],
    [0, 63, 255],
    [255, 0, 0],
    [0, 0, 63],
    [0, 63, 127],
    [0, 63, 191],
    [0,255, 255],
    [0, 127, 191],
    [0, 191, 127],
    [0,0,0],#12
    [0, 0, 255],
              [0,255,0], #14
              [0,153,204], #15
              [204,204,68], #16
              [255, 204, 51], #17
              [255, 255, 204], #18
              [0, 255, 255],
              [255,102,102], #20
              [0,255,0] ,#21
              [255,255,0],#22
              [255,0,0],#23
              [255,195,128],#24
              [153,102,153]
    ]
    for i in range(len(colors)):
        cmap[i] = colors[i]

    return cmap.astype(np.uint8)



def unify_single_cmap():
    cmap = np.zeros((255, 3), dtype=np.uint8)
    colors = [
    [0, 127, 255],#0
    [0, 0, 0],
    [0, 0, 0],
    [0, 0, 0], #3
    [0, 0, 0], #4
    [0,255,0], #5
    [0, 0, 0],
    [0, 0, 0],  # 7
    [0, 0, 0],  # 8
    [0, 0, 0],
    [0, 0, 0],  # 10
    [0, 0, 0],  # 11
    [0, 0, 0],
    [0, 0, 0],  # 13
    [0, 0, 0],  # 14
    [0, 0, 0], #15
      [159,129,183], #16
      [0, 0, 0], #17
      [255, 195, 128], #18
      [0, 0, 0],
    [255, 0, 0],#20
    [255,255,0],
    [0,0,255], #22
      [0, 0, 0],
      [0, 0, 0],
    [0,0,0]
    ]
    for i in range(len(colors)):
        cmap[i] = colors[i]

    return cmap.astype(np.uint8)

def ade_cmap():
    cmap = np.zeros((256, 3), dtype=np.uint8)
    colors = [
        [0, 0, 0],
        [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]
    ]

    for i in range(len(colors)):
        cmap[i] = colors[i]

    return cmap.astype(np.uint8)


def voc_cmap(N=256, normalized=False):
    def bitget(byteval, idx):
        return ((byteval & (1 << idx)) != 0)

    dtype = 'float32' if normalized else 'uint8'
    cmap = np.zeros((N, 3), dtype=dtype)
    for i in range(N):
        r = g = b = 0
        c = i
        for j in range(8):
            r = r | (bitget(c, 0) << 7-j)
            g = g | (bitget(c, 1) << 7-j)
            b = b | (bitget(c, 2) << 7-j)
            c = c >> 3

        cmap[i] = np.array([r, g, b])

    cmap = cmap/255 if normalized else cmap
    return cmap


class Label2Color(object):
    def __init__(self, cmap):
        self.cmap = cmap

    def __call__(self, lbls):
        return self.cmap[lbls]


def convert_bn2gn(module):
    mod = module
    if isinstance(module, nn.modules.batchnorm._BatchNorm):
        num_features = module.num_features
        num_groups = num_features//16
        mod = nn.GroupNorm(num_groups=num_groups, num_channels=num_features)
    for name, child in module.named_children():
        mod.add_module(name, convert_bn2gn(child))
    del module
    return mod