import numpy as np
import cv2
from PIL import Image

def read_file(source_image, target_image):
    s = cv2.cvtColor(source_image, cv2.COLOR_RGB2LAB)
    t = cv2.cvtColor(np.array(target_image), cv2.COLOR_RGB2LAB)
    return s, t

def get_mean_and_std(x):
    x_mean, x_std = cv2.meanStdDev(x)
    x_mean = np.hstack(np.around(x_mean,2))
    x_std = np.hstack(np.around(x_std,2))
    return x_mean, x_std

def color_transfer(source, target_image):
    """
    The primary goal of this algorithm is to adjust the colors of the source image
    such that its color distribution looks like the distribution of the target image.
    """
    s, t = read_file(source, target_image)
    s_mean, s_std = get_mean_and_std(s)
    t_mean, t_std = get_mean_and_std(t)
    height, width, channel = s.shape
    for i in range(0,height):
        for j in range(0,width):
            for k in range(0,channel):
                x = s[i,j,k]
                x = ((x-s_mean[k])*(t_std[k]/s_std[k]))+t_mean[k]
                # round or +0.5
                x = round(x)
                # boundary check
                x = 0 if x<0 else x
                x = 255 if x>255 else x
                s[i,j,k] = x
    s = cv2.cvtColor(s,cv2.COLOR_LAB2BGR)
    cv2.imwrite('result.jpg',s)

def rgb_to_hex(rgb):
    return '#{:02x}{:02x}{:02x}'.format(rgb[0], rgb[1], rgb[2])

def hex_to_rgb(hex):
    hex = hex.lstrip('#')
    return tuple(int(hex[i:i+2], 16) for i in (0, 2, 4))

def create_color_palette(colors, palette_width=800, palette_height=200):
    """
    Receives a list of colors in hex format and creates a palette image
    """
    pixels = []
    n_colors = len(colors)
    for i in range(n_colors):
        color = hex_to_rgb(colors[i])
        for j in range(palette_width//n_colors * palette_height):
            pixels.append(color)
    img = Image.new('RGB', (palette_height, palette_width))
    img.putdata(pixels) 
    # img.show()
    return img

# colors = ['#d54b18', '#555015', '#98680e', '#1f3919', '#da8007', '#a50637', '#b50406', '#d52350', '#e2090f', '#f8545c']
# img = create_color_palette(colors)
# color_transfer('s2', img)

def recolor(source, colors):
    palette_img = create_color_palette(colors)
    palette_bgr = cv2.cvtColor(np.array(palette_img), cv2.COLOR_RGB2BGR)
    source_bgr = cv2.cvtColor(source, cv2.COLOR_RGB2BGR)
    recolored = color_transfer(source_bgr, palette_bgr)
    return recolored