import numpy as np
import matplotlib.pyplot as plt
import streamlit as st
from skimage import io, color
from numpy.linalg import norm

def svd_compress(image, k):
    """Compress the image using SVD by keeping only the top k singular values."""
    U, S, Vt = np.linalg.svd(image, full_matrices=False)
    compressed_image = np.dot(U[:, :k], np.dot(np.diag(S[:k]), Vt[:k, :]))
    return compressed_image

def compute_norms(original, compressed):
    """Compute different norms to compare image quality."""
    frobenius_norm = norm(original - compressed, 'fro')
    l2_norm = norm(original - compressed)
    max_norm = norm(original - compressed, np.inf)
    return frobenius_norm, l2_norm, max_norm

def plot_images(original, compressed, k):
    """Plot original and compressed images side by side."""
    fig, axes = plt.subplots(1, 2, figsize=(12, 6))
    axes[0].imshow(original, cmap='gray')
    axes[0].set_title("Original Image")
    axes[0].axis('off')

    axes[1].imshow(compressed, cmap='gray')
    axes[1].set_title(f"Compressed Image (Rank {k})")
    axes[1].axis('off')

    st.pyplot(fig)

# Streamlit app
st.title("Image Compression using SVD")

# Upload an image
uploaded_file = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])
if uploaded_file is not None:
    # Load the image
    image = io.imread(uploaded_file)
    gray_image = color.rgb2gray(image)

    # Select compression rank
    k = st.slider("Select the rank for compression", min_value=1, max_value=min(gray_image.shape), value=50)

    # Compress the image
    compressed_image = svd_compress(gray_image, k)

    # Compute norms
    frobenius_norm, l2_norm, max_norm = compute_norms(gray_image, compressed_image)

    # Display norms
    st.write(f"Frobenius Norm: {frobenius_norm}")
    st.write(f"L2 Norm: {l2_norm}")
    st.write(f"Max Norm: {max_norm}")

    # Plot images
    plot_images(gray_image, compressed_image, k)