import streamlit as st
import os
import shutil
import zipfile
from PIL import Image, ImageOps

st.set_page_config(page_title="🖼️ Sequential Image Transformer", layout="wide")
st.title("🧪 Sequential Image Augmentation Tool")

uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])

# Sidebar: Choose transformations in order
st.sidebar.header("🛠️ Select Transformations (applied in order)")
options = st.sidebar.multiselect(
    "Choose transformations in order:",
    ["Grayscale", "Horizontal Flip", "Vertical Flip", "Rotate 45°", "Zoom", "Crop 10px"],
)

output_dir = "augmented_sequential"
zip_file = "augmented_sequential.zip"

if uploaded_file:
    # Prepare fresh folder
    if os.path.exists(output_dir):
        shutil.rmtree(output_dir)
    os.makedirs(output_dir)

    # Load & resize image
    original = Image.open(uploaded_file).convert("RGB").resize((224, 224))
    transformed = original.copy()

    st.subheader("📷 Original Image")
    st.image(original, width=200)

    # Apply transformations in order
    for i, op in enumerate(options):
        if op == "Grayscale":
            transformed = ImageOps.grayscale(transformed).convert("RGB")
        elif op == "Horizontal Flip":
            transformed = ImageOps.mirror(transformed)
        elif op == "Vertical Flip":
            transformed = ImageOps.flip(transformed)
        elif op == "Rotate 45°":
            transformed = transformed.rotate(45)
        elif op == "Zoom":
            crop = 40
            w, h = transformed.size
            transformed = transformed.crop((crop, crop, w - crop, h - crop)).resize((224, 224))
        elif op == "Crop 10px":
            transformed = transformed.crop((10, 10, 214, 214)).resize((224, 224))

        # Save intermediate result
        transformed.save(os.path.join(output_dir, f"{i+1}_{op.replace(' ', '_')}.jpg"))

    # Display final result
    st.subheader("🧪 Final Transformed Image")
    st.image(transformed, width=200)

    # Show transformation steps
    st.subheader("🧬 Steps Preview")
    num_cols = max(1, min(3, len(options)))  # fix for columns error
    cols = st.columns(num_cols)
    for i, file in enumerate(sorted(os.listdir(output_dir))):
        img = Image.open(os.path.join(output_dir, file))
        with cols[i % num_cols]:
            st.image(img, caption=file[:-4], width=180)

    # Create ZIP
    with zipfile.ZipFile(zip_file, 'w') as zipf:
        for file in os.listdir(output_dir):
            zipf.write(os.path.join(output_dir, file), arcname=file)

    with open(zip_file, "rb") as f:
        st.download_button(
            label="📦 Download All Transformations as ZIP",
            data=f,
            file_name="sequential_transformed_images.zip",
            mime="application/zip"
        )
else:
    st.info("Please upload an image to begin.")