import cv2
import numpy as np
from PIL import Image

def reinhards_color_transfer(source, target):
    # Convert the images from the RGB to the Lab color space
    source_lab = cv2.cvtColor(source, cv2.COLOR_BGR2Lab).astype(np.float64)
    target_lab = cv2.cvtColor(target, cv2.COLOR_BGR2Lab).astype(np.float64)
    
    # Compute mean and standard deviation for each channel in both images
    l_mean_src, l_std_src = np.mean(source_lab[:, :, 0]), np.std(source_lab[:, :, 0])
    a_mean_src, a_std_src = np.mean(source_lab[:, :, 1]), np.std(source_lab[:, :, 1])
    b_mean_src, b_std_src = np.mean(source_lab[:, :, 2]), np.std(source_lab[:, :, 2])

    l_mean_tar, l_std_tar = np.mean(target_lab[:, :, 0]), np.std(target_lab[:, :, 0])
    a_mean_tar, a_std_tar = np.mean(target_lab[:, :, 1]), np.std(target_lab[:, :, 1])
    b_mean_tar, b_std_tar = np.mean(target_lab[:, :, 2]), np.std(target_lab[:, :, 2])
    
    # Subtract the means from the source image
    source_lab[:, :, 0] -= l_mean_src
    source_lab[:, :, 1] -= a_mean_src
    source_lab[:, :, 2] -= b_mean_src

    # Scale by the standard deviations
    source_lab[:, :, 0] = (l_std_tar / l_std_src) * source_lab[:, :, 0]
    source_lab[:, :, 1] = (a_std_tar / a_std_src) * source_lab[:, :, 1]
    source_lab[:, :, 2] = (b_std_tar / b_std_src) * source_lab[:, :, 2]

    # Add the target means
    source_lab[:, :, 0] += l_mean_tar
    source_lab[:, :, 1] += a_mean_tar
    source_lab[:, :, 2] += b_mean_tar
    
    # Clip pixel values to ensure they fall within the valid Lab range
    source_lab[:, :, 0] = np.clip(source_lab[:, :, 0], 0, 255)
    source_lab[:, :, 1] = np.clip(source_lab[:, :, 1], 0, 255)
    source_lab[:, :, 2] = np.clip(source_lab[:, :, 2], 0, 255)
    
    # Convert back to RGB
    transferred_rgb = cv2.cvtColor(source_lab.astype(np.uint8), cv2.COLOR_Lab2BGR)
    return transferred_rgb

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


# if __name__ == "__main__":
#     source = cv2.imread("estampa.jpg")  
#     colors = ['#6b3d68', '#6d2055', '#695977', '#6b7988', '#6f9b9b']
#     # Generate palette image
#     palette_img = create_color_palette(colors)
#     # Convert the palette image to BGR format
#     palette_bgr = cv2.cvtColor(np.array(palette_img), cv2.COLOR_RGB2BGR)
    
#     # Save the palette image
#     # cv2.imwrite("palette.jpg", palette_bgr)
    
#     target = palette_bgr#cv2.imread("palette.jpg")

#     transferred = reinhards_color_transfer(source, target)
#     cv2.imwrite("transferred_reinhard.jpg", transferred) 


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

    recolored = reinhards_color_transfer(source, palette_bgr)
    recoloredFile = cv2.imwrite("result.jpg", recolored) 
    return recolored