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

# Background color for the app
st.markdown(
    """
    <style>
    .stApp {
        background-color: #1E1E1E;
        color: #FFFFFF;
    }
    </style>
    """,
    unsafe_allow_html=True,
)

st.title("🔧 Image Data Augmentation Tool")
st.write("Enhance your images with various augmentation techniques.")

# Image uploader
uploaded_image = st.file_uploader("📤 Upload an Image (jpg, png, jpeg)", type=["jpg", "png", "jpeg"])

if uploaded_image:
    try:
        image = Image.open(uploaded_image)
        st.image(image, caption="Uploaded Image", use_column_width=True)

        st.subheader("🎨 Choose Augmentation Options")

        # Augmentation option
        augmentation_option = st.selectbox(
            "Select an Augmentation Technique:",
            [
                "None",
                "Cropping",
                "Flipping",
                "Rotation",
                "Zooming In",
                "Zooming Out",
                "Translation",
                "Shearing",
                "Grayscale",
                "Invert Colors",
            ],
        )

        augmented_image = image.copy()

        # Cropping
        if augmentation_option == "Cropping":
            col1, col2 = st.columns(2)
            with col1:
                left = st.slider("Left Crop", 0, image.width // 2, 0)
                top = st.slider("Top Crop", 0, image.height // 2, 0)
            with col2:
                right = st.slider("Right Crop", 0, image.width // 2, 0)
                bottom = st.slider("Bottom Crop", 0, image.height // 2, 0)
            augmented_image = image.crop((left, top, image.width - right, image.height - bottom))

        # Flipping
        elif augmentation_option == "Flipping":
            flip_option = st.radio("Flip Type", ["Horizontal", "Vertical"])
            augmented_image = ImageOps.mirror(image) if flip_option == "Horizontal" else ImageOps.flip(image)

        # Rotation
        elif augmentation_option == "Rotation":
            rotation_angle = st.slider("Rotation Angle (°)", -180, 180, 0)
            augmented_image = image.rotate(rotation_angle, expand=True)

        # Zooming In
        elif augmentation_option == "Zooming In":
            scale_factor = st.slider("Zoom Factor (%)", 100, 200, 100) / 100.0
            if scale_factor != 1.0:
                new_size = (int(image.width * scale_factor), int(image.height * scale_factor))
                zoomed_image = image.resize(new_size)
                augmented_image = zoomed_image.crop(((new_size[0] - image.width) // 2, 
                                                     (new_size[1] - image.height) // 2, 
                                                     (new_size[0] + image.width) // 2, 
                                                     (new_size[1] + image.height) // 2))

        # Zooming Out
        elif augmentation_option == "Zooming Out":
            scale_factor = st.slider("Zoom Factor (%)", 50, 100, 100) / 100.0
            new_size = (int(image.width * scale_factor), int(image.height * scale_factor))
            zoomed_image = image.resize(new_size)
            padded_image = Image.new("RGB", (image.width, image.height), (255, 255, 255))
            padded_image.paste(zoomed_image, ((image.width - new_size[0]) // 2, (image.height - new_size[1]) // 2))
            augmented_image = padded_image

        # Translation
        elif augmentation_option == "Translation":
            translation_x = st.slider("Horizontal Shift (px)", -100, 100, 0)
            translation_y = st.slider("Vertical Shift (px)", -100, 100, 0)
            augmented_image = image.transform(
                image.size, Image.AFFINE, (1, 0, translation_x, 0, 1, translation_y), fillcolor=(255, 255, 255)
            )

        # Shearing
        elif augmentation_option == "Shearing":
            shear_angle = st.slider("Shearing Angle (°)", -45, 45, 0)
            shear_matrix = (1, np.tan(np.radians(shear_angle)), 0, 0, 1, 0)
            augmented_image = image.transform(
                (image.width + abs(int(image.height * np.tan(np.radians(shear_angle)))), image.height),
                Image.AFFINE,
                shear_matrix,
                fillcolor=(255, 255, 255),
            )

        # Grayscale
        elif augmentation_option == "Grayscale":
            augmented_image = ImageOps.grayscale(image)

        # Invert Colors
        elif augmentation_option == "Invert Colors":
            augmented_image = ImageOps.invert(image)

        # Display augmented image
        st.subheader("🖼️ Augmented Image Preview")
        st.image(augmented_image, caption=f"{augmentation_option} Applied", use_column_width=True)

        # Download button
        buffer = io.BytesIO()
        augmented_image.save(buffer, format="PNG")
        buffer.seek(0)
        st.download_button("💾 Download Augmented Image", buffer, file_name="augmented_image.png")

    except Exception as e:
        st.error(f"❌ An error occurred: {e}")