import numpy as np
import cv2
from sklearn.cluster import KMeans
from PIL import Image

from logic.validation import ValidationError, ensure_positive_int

IMAGE_SIZE = 400 
SAMPLE_SIZE = 50000 


def resize_img(image: Image.Image, size: int = IMAGE_SIZE) -> Image.Image:
    """Resize an image to a square of the requested size.

    Args:
        image (PIL.Image.Image): Image to resize.
        size (int, optional): Target side length in pixels. Defaults to ``IMAGE_SIZE``.

    Returns:
        PIL.Image.Image: Cropped and resized copy in RGB.

    Raises:
        ValidationError: If ``size`` is not positive.

    Example:
        >>> resized = resize_img(Image.open('input.png'), 512)
    """
    ensure_positive_int(size, "Image size", minimum=32)
    width, height = image.size
    if width != height:
        min_size = min(width, height)
        left = (width - min_size) // 2
        top = (height - min_size) // 2
        image = image.crop((left, top, left + min_size, top + min_size))
    resized_img = image.resize((size, size), Image.Resampling.LANCZOS)
    return resized_img


def color_quantize(image: Image.Image, n_colors: int = 16):
    """Reduce the number of colors with KMeans clustering.

    Args:
        image (PIL.Image.Image): Image to quantize.
        n_colors (int, optional): Number of clusters. Defaults to 16.

    Returns:
        tuple[PIL.Image.Image, np.ndarray]: Quantized image and cluster centers in BGR order.

    Raises:
        ValidationError: If ``n_colors`` is outside ``[2, 64]``.
    """
    ensure_positive_int(n_colors, "Color palette size", minimum=2, maximum=64)
    img_np = np.array(image)
    original_shape = img_np.shape
    img_flat = img_np.reshape(-1, 3)
    
    sample_size = min(SAMPLE_SIZE, img_flat.shape[0])
    indices = np.random.choice(img_flat.shape[0], sample_size, replace=False)
    img_sample = img_flat[indices]
    
    kmeans = KMeans(
        n_clusters=n_colors, 
        random_state=42, 
        init='k-means++',
        n_init=5,
        max_iter=100,
        tol=1e-3,
        algorithm='lloyd'
    )
    
    kmeans.fit(img_sample)
    labels = kmeans.predict(img_flat)
    
    quantized_flat = kmeans.cluster_centers_[labels].astype(np.uint8)
    quantized_img_np = quantized_flat.reshape(original_shape)
    
    quantized_img = Image.fromarray(quantized_img_np)
    
    color_centers = kmeans.cluster_centers_.astype(np.uint8)
    color_centers = color_centers[:, [2, 1, 0]]  # Swap R and B channels
    return quantized_img, color_centers