import numpy as np
from PIL import Image
import gradio as gr
import tempfile
import os

def initialize_centroids(X, k):
    np.random.seed(42)
    indices = np.random.choice(X.shape[0], k, replace=False)
    return X[indices]

def compute_centroids(X, labels, k):
    centroids = []
    for i in range(k):
        cluster_points = X[labels == i]
        if len(cluster_points) > 0:
            centroids.append(cluster_points.mean(axis=0))
        else:
            centroids.append(X[np.random.choice(len(X))])
    return np.array(centroids)

def assign_clusters(X, centroids):
    distances = np.linalg.norm(X[:, np.newaxis] - centroids, axis=2)
    return np.argmin(distances, axis=1)

def kmeans(X, k, max_iters=100):
    centroids = initialize_centroids(X, k)
    for _ in range(max_iters):
        labels = assign_clusters(X, centroids)
        new_centroids = compute_centroids(X, labels, k)
        if np.allclose(centroids, new_centroids):
            break
        centroids = new_centroids
    return centroids, labels

def compress_image(img, clusters=16):
    if clusters < 2:
        return None, "Minimum 3 clusters required."

    # Resize for speed & consistency
    max_dim = 512
    w, h = img.size
    scale = max_dim / max(w, h)
    img = img.resize((int(w * scale), int(h * scale)))

    img_np = np.array(img)
    shape = img_np.shape
    img_flat = img_np.reshape(-1, 3).astype(float)

    centroids, labels = kmeans(img_flat, clusters, max_iters=30)
    compressed_flat = centroids[labels].astype(np.uint8)
    compressed_img = compressed_flat.reshape(shape)

    # Save original and compressed to compare file sizes (use PNG to avoid JPEG inconsistency)
    with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as orig:
        Image.fromarray(img_np).save(orig.name, format='PNG')
        orig_size = os.path.getsize(orig.name)

    with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as comp:
        Image.fromarray(compressed_img).save(comp.name, format='PNG')
        comp_size = os.path.getsize(comp.name)

    os.remove(orig.name)
    os.remove(comp.name)

    percent = 100 * (orig_size - comp_size) / orig_size
    if percent < 0:
        status = f"⚠️ Larger by {abs(percent):.2f}% (due to more color detail)"
    else:
        status = f"✅ Compressed by {percent:.2f}%"

    return Image.fromarray(compressed_img), status

iface = gr.Interface(
    fn=compress_image,
    inputs=[
        gr.Image(type="pil", label="Upload Image"),
        gr.Slider(2, 64, value=16, step=1, label="Number of Clusters")
    ],
    outputs=[
        gr.Image(type="pil", label="Compressed Image"),
        gr.Label(label="Compression Info")
    ],
    title="Fixed K-Means Image Compressor",
    description="Uses KMeans clustering on color pixels. PNG is used to give accurate compression size comparison."
)

if __name__ == "__main__":
    iface.launch()