import streamlit as st
import numpy as np
from sklearn.cluster import KMeans
from io import BytesIO
from PIL import Image
import base64

def compress_image_kmeans(image, k):
    # Convert PIL image to NumPy array
    img = np.array(image)
    h, w, c = img.shape
    
    # Reshape the image into a 2D array of pixels
    img_reshaped = img.reshape(-1, c)
    
    # Apply K-Means clustering
    kmeans = KMeans(n_clusters=k, random_state=42, n_init=10)
    labels = kmeans.fit_predict(img_reshaped)
    
    # Replace pixel values with cluster centers
    compressed_img = kmeans.cluster_centers_[labels].reshape(h, w, c).astype(np.uint8)
    
    return Image.fromarray(compressed_img)

def get_image_download_link(img, filename="compressed_image.png"):
    buffered = BytesIO()
    img.save(buffered, format="PNG")
    encoded = base64.b64encode(buffered.getvalue()).decode()
    href = f'<a href="data:file/png;base64,{encoded}" download="{filename}">Download Compressed Image</a>'
    return href

def main():
    st.title("Image Compression using K-Means Clustering")
    
    # Instructions Section
    st.header("About the App")
    st.write("This application allows you to compress images using K-Means clustering. By reducing the number of colors in the image, the app creates a simplified version while maintaining its visual integrity.")
    
    st.header("How to Use")
    st.write("1. Upload one or multiple images using the file uploader below.")
    st.write("2. Adjust the slider to select the number of colors (clusters) for compression.")
    st.write("3. Click the 'Compress' button to process the image.")
    st.write("4. Download the compressed image using the provided link.")
    
    uploaded_files = st.file_uploader("Choose images", type=["png", "jpg", "jpeg"], accept_multiple_files=True)
    
    if uploaded_files:
        for uploaded_file in uploaded_files:
            image = Image.open(uploaded_file)
            st.image(image, caption=f"Original Image - {uploaded_file.name}", use_container_width=True)
            
            k = st.slider(f"Select number of colors (clusters) for {uploaded_file.name}", min_value=2, max_value=64, value=16)
            
            if st.button(f"Compress {uploaded_file.name}"):
                compressed_img = compress_image_kmeans(image, k)
                st.image(compressed_img, caption=f"Compressed Image - {uploaded_file.name}", use_container_width=True)
                
                st.markdown(get_image_download_link(compressed_img, filename=f"compressed_{uploaded_file.name}"), unsafe_allow_html=True)

if __name__ == "__main__":
    main()