import cv2
import numpy as np
import math
from sklearn.cluster import KMeans
from collections import Counter
from src.color_utils import rgb_to_hex, hex_to_bgr, hex_to_rgb
from src.hair_utils import HairColorPalette
from src.skin_utils import SkinTonePalette
from src.eyes_utils import EyeColorPalette
from sklearn.metrics import silhouette_score
import matplotlib.pyplot as plt


# Function to create a color bar
def create_color_bar(height, width, color):
    bar = np.zeros((height, width, 3), dtype=np.uint8)
    bar[:] = color
    return bar


# Function to get dominant colors and their percentages
def get_dominant_colors(image, mask, n_colors, debug=True):
    image_np = image[mask > 0]
    pixels = image_np.reshape((-1, 3))
    n_colors_elbow = optimal_clusters_elbow(pixels, max_clusters=15)
    # n_colors_silhouette = optimal_clusters_silhouette(pixels, max_clusters=5)

    # kmeans_silhouette = KMeans(n_clusters=n_colors_silhouette)
    # kmeans_silhouette.fit(pixels)
    kmeans_elbow = KMeans(n_clusters=n_colors_elbow)
    kmeans_elbow.fit(pixels)
    dominant_colors = kmeans_elbow.cluster_centers_
    counts = Counter(kmeans_elbow.labels_)
    total_count = sum(counts.values())
    dominant_colors = [dominant_colors[i] for i in counts.keys()]
    dominant_percentages = [counts[i] / total_count for i in counts.keys()]
    if debug:
        visualize_clusters(image, mask, kmeans_elbow, tag="elbow")
        # visualize_clusters(image, mask, kmeans_silhouette, tag="silhouette")
    return dominant_colors, dominant_percentages


def optimal_clusters_elbow(skin_pixels, max_clusters=10):
    distortions = []
    for i in range(1, max_clusters + 1):
        kmeans = KMeans(n_clusters=i, random_state=42)
        kmeans.fit(skin_pixels)
        distortions.append(kmeans.inertia_)

    # Compute the second derivative to find the "elbow" point
    second_derivative = np.diff(np.diff(distortions))
    optimal_k_elbow = (
        np.argmax(second_derivative) + 2
    )  # +2 because of the second derivative

    plt.figure(figsize=(10, 8))
    plt.plot(range(1, max_clusters + 1), distortions, marker="o")
    plt.xlabel("Number of clusters")
    plt.ylabel("Distortion (Inertia)")
    plt.title("Elbow Method For Optimal Clusters")
    plt.axvline(
        x=optimal_k_elbow,
        linestyle="--",
        color="r",
        label=f"Optimal k={optimal_k_elbow}",
    )
    plt.legend()
    plt.savefig("workspace/kmeans_elbow.png")

    return optimal_k_elbow


def optimal_clusters_silhouette(skin_pixels, max_clusters=10):
    silhouette_scores = []
    for i in range(2, max_clusters + 1):  # Silhouette score is undefined for k=1
        kmeans = KMeans(n_clusters=i, random_state=42)
        kmeans.fit(skin_pixels)
        score = silhouette_score(skin_pixels, kmeans.labels_)
        silhouette_scores.append(score)

    optimal_k_silhouette = (
        np.argmax(silhouette_scores) + 2
    )  # +2 because range starts at 2

    plt.figure(figsize=(10, 8))
    plt.plot(range(2, max_clusters + 1), silhouette_scores, marker="o")
    plt.xlabel("Number of clusters")
    plt.ylabel("Silhouette Score")
    plt.title("Silhouette Method For Optimal Clusters")
    plt.axvline(
        x=optimal_k_silhouette,
        linestyle="--",
        color="r",
        label=f"Optimal k={optimal_k_silhouette}",
    )
    plt.legend()
    plt.savefig("workspace/kmeans_sillouette.png")

    return optimal_k_silhouette


def visualize_clusters(image, mask, kmeans, tag="_none"):
    clustered_image = np.zeros_like(image)
    mask = mask > 0
    labels = kmeans.labels_
    cluster_centers = kmeans.cluster_centers_
    skin_coords = np.where(mask)

    for label, (x, y) in zip(labels, zip(skin_coords[0], skin_coords[1])):
        clustered_image[x, y] = cluster_centers[label]

    clustered_image = clustered_image.astype(np.uint8)

    plt.figure(figsize=(12, 6))
    plt.subplot(1, 2, 1)
    plt.title("Original Image")
    plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
    plt.subplot(1, 2, 2)
    plt.title("Clustered Image")
    plt.imshow(cv2.cvtColor(clustered_image, cv2.COLOR_BGR2RGB))
    plt.savefig(f"workspace/kmeans_visual_{tag}.png")


# Function to get the closest color from the palette
def get_closest_color(dominant_colors, palette):
    min_distance = float("inf")
    for dom_color in dominant_colors:
        for color_name, (color_value, color_hex) in palette.items():
            distance = np.linalg.norm(dom_color - np.array(color_value))
            if distance < min_distance:
                min_distance = distance
                closest_color = color_name
                closest_hex = color_hex
    return closest_color, closest_hex, min_distance


# Function to create the dominant color bar
def create_dominant_color_bar(
    report_image, dominant_colors, dominant_percentages, bar_width
):
    color_bars = []
    total_height = 0
    for color, pct in zip(dominant_colors, dominant_percentages):
        bar_height = int(math.floor(report_image.shape[0] * pct))
        total_height += bar_height
        bar = create_color_bar(bar_height, bar_width, color)
        color_bars.append(bar)
    padding_height = report_image.shape[0] - total_height
    if padding_height > 0:
        padding = create_color_bar(padding_height, bar_width, (255, 255, 255))
        color_bars.append(padding)
    return np.vstack(color_bars)


# Function to create the tone palette bar
def create_tone_palette_bar(report_image, tone_id, skin_tone_palette, bar_width):
    palette_bars = []
    tone_height = report_image.shape[0] // len(skin_tone_palette)
    tone_bgrs = []
    for tone in skin_tone_palette.values():
        color_bgr = hex_to_bgr(tone[1])
        tone_bgrs.append(color_bgr)
        bar = create_color_bar(tone_height, bar_width, color_bgr)
        palette_bars.append(bar)
    padding_height = report_image.shape[0] - tone_height * len(skin_tone_palette)
    if padding_height > 0:
        padding = create_color_bar(padding_height, bar_width, (255, 255, 255))
        palette_bars.append(padding)
    bar = np.vstack(palette_bars)

    padding = 1
    start_point = (padding, tone_id * tone_height + padding)
    end_point = (bar_width - padding, (tone_id + 1) * tone_height)
    bar = cv2.rectangle(bar, start_point, end_point, (255, 0, 0), 2)
    return bar


# Function to create the message bar
def create_message_bar(
    dominant_colors, dominant_percentages, tone_hex, distance, img_shape
):
    bar_width = img_shape[1]
    bar_height = img_shape[0] // 30
    msg_bar = create_color_bar(
        height=bar_height, width=bar_width, color=(243, 239, 214)
    )
    b, g, r = np.around(dominant_colors[0]).astype(int)
    dominant_color_hex = "#%02X%02X%02X" % (r, g, b)
    pct = f"{dominant_percentages[0] * 100:.2f}%"

    font, font_scale, txt_color, thickness, line_type = (
        cv2.FONT_HERSHEY_SIMPLEX,
        1,
        (0, 0, 0),
        1,
        cv2.LINE_AA,
    )
    x, y = 2, 15
    msg = f"- Dominant color: {dominant_color_hex}, percent: {pct}"
    cv2.putText(msg_bar, msg, (x, y), font, font_scale, txt_color, thickness, line_type)

    text_size, _ = cv2.getTextSize(msg, font, font_scale, thickness)
    line_height = text_size[1] + 10
    accuracy = round(100 - distance, 2)
    cv2.putText(
        msg_bar,
        f"- Skin tone: {tone_hex}, accuracy: {accuracy}",
        (x, y + line_height),
        font,
        font_scale,
        txt_color,
        thickness,
        cv2.LINE_AA,
    )

    return msg_bar


def color_analysis(skin, hair, eyes):
    analysis = {}

    # Determine Season
    if (
        skin == "light"
        and (hair in ["golden blonde", "light brown"])
        and (eyes in ["blue", "green"])
    ):
        analysis["season"] = "Spring"
    elif (
        skin == "light"
        and (hair in ["ash blonde", "light brown"])
        and (eyes in ["blue", "green"])
    ):
        analysis["season"] = "Summer"
    elif (
        skin in ["medium", "dark"]
        and (hair in ["red", "brown"])
        and (eyes in ["green", "hazel"])
    ):
        analysis["season"] = "Autumn"
    elif (
        skin in ["medium", "dark"]
        and (hair in ["dark brown", "black"])
        and (eyes in ["blue", "brown"])
    ):
        analysis["season"] = "Winter"

    # Determine Warm/Cool
    if skin in ["light", "medium", "dark"] and hair in ["golden", "red", "caramel"]:
        analysis["warm/cool"] = "Warm"
    else:
        analysis["warm/cool"] = "Cool"

    # Determine Intensity
    if eyes in ["bright blue", "bright green"] or hair in ["black", "vivid red"]:
        analysis["intensity"] = "High"
    else:
        analysis["intensity"] = "Low"

    # Determine Value
    if (
        skin == "light"
        and hair in ["blonde", "light brown"]
        and eyes in ["blue", "green"]
    ):
        analysis["value"] = "Light"
    elif skin == "medium" and hair in ["brown", "red"] and eyes in ["green", "hazel"]:
        analysis["value"] = "Medium"
    else:
        analysis["value"] = "Dark"

    # Determine Tone
    if analysis["value"] == "Light" and analysis["intensity"] == "Low":
        analysis["tone"] = "Light and Soft"
    elif analysis["value"] == "Light" and analysis["intensity"] == "High":
        analysis["tone"] = "Light and Bright"
    elif analysis["value"] == "Dark" and analysis["intensity"] == "Low":
        analysis["tone"] = "Dark and Soft"
    else:
        analysis["tone"] = "Dark and Bright"

    # Determine Saturation
    if hair in ["bright red", "black"] or eyes in ["clear blue"]:
        analysis["saturation"] = "High"
    else:
        analysis["saturation"] = "Low"

    return analysis


# # Example usage
# result = color_analysis("light", "golden blonde", "blue")
# print(result)
def create_combined_overlay(image, hair_mask, skin_mask, eye_mask):
    """
    Create an overlay image by combining the original image with the hair, skin, and eye masks.
    :param image: Original image as a numpy array.
    :param hair_mask: Hair mask as a numpy array.
    :param skin_mask: Skin mask as a numpy array.
    :param eye_mask: Eye mask as a numpy array.
    :return: Combined overlay image as a numpy array.
    """
    # Create overlays for different parts
    hair_overlay = np.zeros_like(image)
    skin_overlay = np.zeros_like(image)
    eye_overlay = np.zeros_like(image)

    # Color the masks
    hair_overlay[hair_mask > 0] = [0, 255, 0]  # Green for hair
    skin_overlay[skin_mask > 0] = [255, 0, 0]  # Red for skin
    eye_overlay[eye_mask > 0] = [0, 0, 255]  # Blue for eyes

    # Combine the overlays with the original image
    combined_overlay = cv2.addWeighted(image, 0.8, hair_overlay, 0.2, 0)
    combined_overlay = cv2.addWeighted(combined_overlay, 0.8, skin_overlay, 0.2, 0)
    combined_overlay = cv2.addWeighted(combined_overlay, 0.8, eye_overlay, 0.2, 0)

    return combined_overlay


def analyze_and_visualize(image, hair_mask, skin_mask, eye_mask, n_colors=3):
    image_np = np.array(image)
    hair_mask_np = np.array(hair_mask)
    skin_mask_np = np.array(skin_mask)
    eye_mask_np = np.array(eye_mask)

    if not (
        image_np.shape[:2]
        == hair_mask_np.shape[:2]
        == skin_mask_np.shape[:2]
        == eye_mask_np.shape[:2]
    ):
        raise ValueError("Image and all masks must have the same dimensions")

    hair_palette = HairColorPalette()
    hair_dominant_colors, hair_dominant_percentages = get_dominant_colors(
        image_np, hair_mask_np, n_colors, debug=True
    )
    hair_color, hair_hex, hair_distance = get_closest_color(
        hair_dominant_colors, hair_palette.palette
    )

    skin_palette = SkinTonePalette()
    skin_dominant_colors, skin_dominant_percentages = get_dominant_colors(
        image_np, skin_mask_np, n_colors, debug=True
    )
    skin_color, skin_hex, skin_distance = get_closest_color(
        skin_dominant_colors, skin_palette.palette
    )

    # Calculate ITA for the dominant skin color
    dominant_skin_color = skin_dominant_colors[0]
    ita = skin_palette.calculate_ita(dominant_skin_color)
    vectorscope_check = skin_palette.is_within_vectorscope_skin_tone_line(
        dominant_skin_color
    )

    eye_palette = EyeColorPalette()
    eye_dominant_colors, eye_dominant_percentages = get_dominant_colors(
        image_np, eye_mask_np, n_colors, debug=True
    )
    eye_color, eye_hex, eye_distance = get_closest_color(
        eye_dominant_colors, eye_palette.palette
    )

    combined_overlay = create_combined_overlay(
        image_np, hair_mask_np, skin_mask_np, eye_mask_np
    )

    bar_width = 50
    hair_color_bar = create_dominant_color_bar(
        image_np, hair_dominant_colors, hair_dominant_percentages, bar_width
    )
    skin_color_bar = create_dominant_color_bar(
        image_np, skin_dominant_colors, skin_dominant_percentages, bar_width
    )
    eye_color_bar = create_dominant_color_bar(
        image_np, eye_dominant_colors, eye_dominant_percentages, bar_width
    )

    hair_palette_bar = create_tone_palette_bar(
        image_np,
        list(hair_palette.palette.keys()).index(hair_color),
        hair_palette.palette,
        bar_width,
    )
    skin_palette_bar = create_tone_palette_bar(
        image_np,
        list(skin_palette.palette.keys()).index(skin_color),
        skin_palette.palette,
        bar_width,
    )
    eye_palette_bar = create_tone_palette_bar(
        image_np,
        list(eye_palette.palette.keys()).index(eye_color),
        eye_palette.palette,
        bar_width,
    )

    output_image = np.hstack(
        [
            combined_overlay,
            hair_color_bar,
            hair_palette_bar,
            skin_color_bar,
            skin_palette_bar,
            eye_color_bar,
            eye_palette_bar,
        ]
    )
    img_shape = output_image.shape
    msg_bar_hair = create_message_bar(
        hair_dominant_colors,
        hair_dominant_percentages,
        hair_hex,
        hair_distance,
        img_shape,
    )
    msg_bar_skin = create_message_bar(
        skin_dominant_colors,
        skin_dominant_percentages,
        skin_hex,
        skin_distance,
        img_shape,
    )
    msg_bar_eye = create_message_bar(
        eye_dominant_colors, eye_dominant_percentages, eye_hex, eye_distance, img_shape
    )

    output_image = np.vstack([output_image, msg_bar_hair, msg_bar_skin, msg_bar_eye])

    analysis_record = {
        "hair": {
            "dominant_colors": [rgb_to_hex(color) for color in hair_dominant_colors],
            "dominant_percentages": hair_dominant_percentages,
            "closest_color": hair_color,
            "closest_color_hex": hair_hex,
            "distance": hair_distance,
        },
        "skin": {
            "dominant_colors": [rgb_to_hex(color) for color in skin_dominant_colors],
            "dominant_percentages": skin_dominant_percentages,
            "closest_color": skin_color,
            "closest_color_hex": skin_hex,
            "distance": skin_distance,
            "ita": ita,
            # "vectorscope_check": vectorscope_check,
        },
        "eyes": {
            "dominant_colors": [rgb_to_hex(color) for color in eye_dominant_colors],
            "dominant_percentages": eye_dominant_percentages,
            "closest_color": eye_color,
            "closest_color_hex": eye_hex,
            "distance": eye_distance,
        },
    }

    return output_image, analysis_record