import streamlit as st
from PIL import Image, ImageEnhance, ImageOps
import numpy as np
import io
import zipfile

def apply_basic_augmentations(image):
    """Applies basic augmentations such as rotation and color jitter."""
    image = image.rotate(np.random.uniform(-30, 30))
    enhancer = ImageEnhance.Color(image)
    image = enhancer.enhance(np.random.uniform(0.75, 1.25))
    if np.random.rand() > 0.5:
        image = ImageOps.mirror(image)
    return image

def simulate_latent_space_noising(image, noise_scale=25):
    """Simulates latent space manipulation by adding noise."""
    image_array = np.array(image)
    noise = np.random.normal(0, noise_scale, image_array.shape)
    noised_image_array = np.clip(image_array + noise, 0, 255).astype(np.uint8)
    return Image.fromarray(noised_image_array)

def augment_image(image, augmentations_count):
    """Generates augmented versions of a single image."""
    augmented_images = []
    for _ in range(augmentations_count):
        augmented_image = apply_basic_augmentations(image)
        augmented_image = simulate_latent_space_noising(augmented_image)
        augmented_images.append(augmented_image)
    return augmented_images

def create_downloadable_zip(augmented_images):
    """Creates a ZIP file in memory for downloading."""
    zip_buffer = io.BytesIO()
    with zipfile.ZipFile(zip_buffer, "a", zipfile.ZIP_DEFLATED, False) as zip_file:
        for idx, image in enumerate(augmented_images):
            img_byte_arr = io.BytesIO()
            image.save(img_byte_arr, format="JPEG")
            zip_file.writestr(f"augmented_image_{idx+1}.jpg", img_byte_arr.getvalue())
    zip_buffer.seek(0)
    return zip_buffer

st.title("Ready-To-Use Synthetic Image Dataset Generation with Few-shots")

st.write("""
1. Easily prepare your dataset by uploading up to 10 images and specifying desired augmentations.
2. Utilize advanced image processing techniques to automatically generate multiple variations per image.
3. Reduce data preprocessing time with automated, customizable enhancements.
4. Quickly download your augmented images in a ZIP file for immediate use in your projects.
""")

uploaded_files = st.file_uploader("Choose images (1-10)", accept_multiple_files=True, type=["jpg", "jpeg", "png"])
augmentations_count = st.number_input("Number of augmented samples per image", min_value=1, max_value=10, value=3)

if uploaded_files:
    all_augmented_images = []
    for uploaded_file in uploaded_files:
        image = Image.open(uploaded_file).convert("RGB")
        augmented_images = augment_image(image, augmentations_count)
        all_augmented_images.extend(augmented_images)

    if st.button("Generate Synthetic Dataset") and all_augmented_images:
        zip_buffer = create_downloadable_zip(all_augmented_images)
        st.download_button(
            label="Download ZIP",
            data=zip_buffer,
            file_name="augmented_images.zip",
            mime="application/zip"
        )