from PIL import Image
import numpy as np
from sklearn.cluster import KMeans
from collections import Counter
from src.color_utils import calculate_color_distance_lab, rgb_to_lab
import cv2


class SkinTonePalette:
    def __init__(self):
        self.palette = {
            "Ebony": ((55, 48, 40), "#373028"),
            "Deep Mocha": ((66, 40, 17), "#422811"),
            "Dark Chocolate": ((81, 59, 46), "#513b2e"),
            "Warm Almond": ((111, 80, 60), "#6f503c"),
            "Golden Bronze": ((129, 101, 79), "#81654f"),
            "Honey": ((157, 122, 84), "#9d7a54"),
            "Caramel": ((190, 160, 126), "#bea07e"),
            "Light Tan": ((229, 200, 166), "#e5c8a6"),
            "Peach": ((231, 193, 184), "#e7c1b8"),
            "Fair": ((243, 218, 214), "#f3dad6"),
            "Ivory": ((251, 242, 243), "#fbf2f3"),
        }

    def get_dominant_colors(self, image_np, n_colors):
        pixels = image_np.reshape((-1, 3))
        kmeans = KMeans(n_clusters=n_colors)
        kmeans.fit(pixels)
        dominant_colors = kmeans.cluster_centers_
        counts = Counter(kmeans.labels_)
        dominant_colors = [dominant_colors[i] for i in counts.keys()]
        return dominant_colors

    def get_closest_color(self, image, mask, n_colors=3):
        image_np = np.array(image)
        mask_np = np.array(mask)
        if image_np.shape[:2] != mask_np.shape[:2]:
            raise ValueError("Image and mask must have the same dimensions")

        skin_pixels = image_np[mask_np > 0]
        dominant_colors = self.get_dominant_colors(skin_pixels, n_colors)

        closest_color = None
        closest_hex = None
        min_distance = float("inf")
        for dom_color in dominant_colors:
            for color_name, (color_value, color_hex) in self.palette.items():
                distance = calculate_color_distance_lab(dom_color, color_value)
                if distance < min_distance:
                    min_distance = distance
                    closest_color = color_name
                    closest_hex = color_hex

        return closest_color, closest_hex

    def calculate_ita(self, rgb_color):
        lab_color = rgb_to_lab(rgb_color)
        L = lab_color.lab_l
        b = lab_color.lab_b
        ita = np.arctan2(L - 50, b) * (180 / np.pi)
        return ita

    def is_within_vectorscope_skin_tone_line(self, rgb_color):
        ycbcr_color = cv2.cvtColor(np.uint8([[rgb_color]]), cv2.COLOR_RGB2YCrCb)[0][0]
        cb, cr = ycbcr_color[1], ycbcr_color[2]
        return 80 <= cb <= 120 and 133 <= cr <= 173