import numpy as np
from PIL import Image, ImageOps, ImageEnhance
import colorsys


##############################################################################
# Visual
def visual_robustness(tests, noise_level=0.3, gray=True, contrast=True, inv=True, temp=True, color=True, s_and_p=True, gaus=True, rot=True, flip=True, crop=True):
    noises = []
    if gray:
        noises.append(grayscale)
    if contrast:
        noises.append(low_contrast)
    if inv:
        noises.append(inversion)
    if temp:
        noises.append(WB)
    if color:
        noises.append(colorize)
    if s_and_p:
        noises.append(salt_and_pepper)
    if gaus:
        noises.append(gaussian)
    if rot:
        noises.append(rotate)
    if flip:
        noises.append(horizontal_flip)
    if crop:
        noises.append(random_crop)
    robustness_tests = []
    for i in range(len(tests)):
        img = Image.fromarray(tests[i])
        mode = img.mode
        img = img.convert('RGB')
        for noise in noises:
            img = noise(img, noise_level)
        img = img.convert(mode)
        robustness_tests.append(np.array(img))
    return robustness_tests


def grayscale(img, p):
    if np.random.sample() <= p:
        return ImageOps.grayscale(img)
    else:
        return img


def low_contrast(img, p):
    if np.random.sample() <= p:
        enhancer = ImageEnhance.Contrast(img)
        return enhancer.enhance(0.5)
    else:
        return img


def inversion(img, p):
    if np.random.sample() <= p:
        return ImageOps.invert(img)
    else:
        return img


def WB(img, p):
    if np.random.sample() <= p and img.mode == 'RGB':
        kelvin_table = {1000: (255,56,0), 1500: (255,109,0), 2000: (255,137,18), 2500: (255,161,72), 3000: (255,180,107), 3500: (255,196,137), 4000: (255,209,163), 4500: (255,219,186), 5000: (255,228,206), 5500: (255,236,224), 6000: (255,243,239), 6500: (255,249,253), 7000: (245,243,255), 7500: (235,238,255), 8000: (227,233,255), 8500: (220,229,255), 9000: (214,225,255), 9500: (208,222,255), 10000: (204,219,255)}
        temps = list(kelvin_table.keys())
        temp = temps[np.random.choice(len(temps))]
        r, g, b = kelvin_table[temp]
        matrix = ( r / 255.0, 0.0, 0.0, 0.0,
                0.0, g / 255.0, 0.0, 0.0,
                0.0, 0.0, b / 255.0, 0.0 )
        return img.convert('RGB', matrix)
    else:
        return img


def colorize(img, p):
    if np.random.sample() <= p and img.mode == 'RGB':
        color = np.random.choice(['red', 'blue', 'green'])
        layer = Image.new('RGB', img.size, color)
        return Image.blend(img, layer, 0.3)
    else:
        return img


def salt_and_pepper(img, p):
    if np.random.sample() <= p:
        img = ImageOps.grayscale(img)
        output = np.copy(np.array(img))
        nb_salt = np.ceil(p*output.size*0.5)
        coords = [np.random.randint(0, i-1, int(nb_salt)) for i in output.shape]
        for i in range(int(nb_salt)):
            output[coords[0][i]][coords[1][i]] = 1
        nb_pepper = np.ceil(p*output.size*0.5)
        coords = [np.random.randint(0, i-1, int(nb_pepper)) for i in output.shape]
        for i in range(int(nb_pepper)):
            output[coords[0][i]][coords[1][i]] = 0
        return Image.fromarray(output)
    else:
        return img


def gaussian(img, p):
    if np.random.sample() <= p:
        dim = np.array(img).shape
        gauss = np.random.normal(0, p, (dim[0], dim[1]))
        return Image.fromarray((np.array(ImageOps.grayscale(img))+gauss).astype('uint8'))
    else:
        return img


def rotate(img, p):
    if np.random.sample() <= p:
        angle = np.random.random_sample()*40-20
        return img.rotate(angle, Image.BILINEAR)
    else:
        return img


def horizontal_flip(img, p):
    if np.random.sample() <= p:
        return img.transpose(Image.FLIP_LEFT_RIGHT)
    else:
        return img


def random_crop(img, p):
    if np.random.sample() <= p:
        dim = np.array(img).shape
        height = dim[0]
        width = dim[1]
        cropped_height = height / 5
        cropped_width = width / 5
        init_height = np.random.random_sample() * cropped_height
        init_width = np.random.random_sample() * cropped_width
        end_height = height - cropped_height + init_height
        end_width = width - cropped_width + init_width
        return img.crop((init_width, init_height, end_width, end_height)).resize((height, width))
    else:
        return img


def periodic(img, periodic_noise_filename="periodic_noise"):
    height = img.height
    width = img.width
    output = []
    for i in range(6):
        noise = Image.open("{}_{}.png".format(periodic_noise_filename, i+1)).convert("RGBA")
        noise = random_crop(rotate(noise.resize((width*2, height*2)), np.random.random_sample()*360, 'white'), height, width)
        output.append(Image.blend(img.convert("RGBA"), noise, 0.3))
    return output