import numpy as np
import cv2
import ot
from PIL import Image

def transfer_channel(source_channel, target_channel):
    source_hist, _ = np.histogram(source_channel, bins=256, range=(0, 256))
    target_hist, _ = np.histogram(target_channel, bins=256, range=(0, 256))

    source_hist = source_hist.astype(np.float64) / source_hist.sum()
    target_hist = target_hist.astype(np.float64) / target_hist.sum()

    r = np.arange(256).reshape((-1, 1))
    c = np.arange(256).reshape((1, -1))
    M = (r - c) ** 2
    M = M / M.max()

    P = ot.emd(source_hist, target_hist, M)
    
    transferred_channel = np.zeros_like(source_channel)
    for i in range(256):
        transferred_channel[source_channel == i] = P[i].argmax()

    return transferred_channel

def optimal_transport_color_transfer(source, target):
    source_lab = cv2.cvtColor(source, cv2.COLOR_BGR2Lab).astype(np.float64)
    target_lab = cv2.cvtColor(target, cv2.COLOR_BGR2Lab).astype(np.float64)

    # Transfer all channels L, a, and b
    for ch in range(3):
        source_lab[:, :, ch] = transfer_channel(source_lab[:, :, ch], target_lab[:, :, ch])

    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-test.png")
#     target = cv2.imread("color.png")
#     transferred = optimal_transport_color_transfer(source, target)

#     smooth = False#test with true to have a different result.
#     if smooth:
#         # Apply bilateral filtering
#         diameter = 30  # diameter of each pixel neighborhood, adjust based on your image size
#         sigma_color = 25  # larger value means colors farther to each other will mix together
#         sigma_space = 25  # larger values means farther pixels will influence each other if their colors are close enough
#         smoothed = cv2.bilateralFilter(transferred, diameter, sigma_color, sigma_space)
#         cv2.imwrite("result_OTA.jpg", smoothed)
#     else:
#         cv2.imwrite("result_OTA.jpg", transferred)


def recolor(source, colors):
    pallete_img = create_color_palette(colors)
    palette_bgr = cv2.cvtColor(np.array(pallete_img), cv2.COLOR_RGB2BGR)
    recolored = optimal_transport_color_transfer(source, palette_bgr)
    smooth = True#test with true for different results.
    if smooth:
        # Apply bilateral filtering
        diameter = 10  # diameter of each pixel neighborhood, adjust based on your image size
        sigma_color = 25  # larger value means colors farther to each other will mix together
        sigma_space = 15  # larger values means farther pixels will influence each other if their colors are close enough
        smoothed = cv2.bilateralFilter(recolored, diameter, sigma_color, sigma_space)
        recoloredFile = cv2.imwrite("result.jpg", smoothed, [cv2.IMWRITE_JPEG_QUALITY, 100])
        return recoloredFile
    else:
        recoloredFile = cv2.imwrite("result.jpg", recolored)
        return recoloredFile