import streamlit as st
from PIL import Image, ImageEnhance, ImageOps
import random
import io

# --- Helper Functions for Image Transformations ---
def translate_image(image, x_offset, y_offset):
    """Translate the image by x and y offsets."""
    return image.transform(
        image.size,
        Image.AFFINE,
        (1, 0, x_offset, 0, 1, y_offset),
        resample=Image.BICUBIC
    )

def rotate_image(image, angle):
    """Rotate the image by a given angle."""
    return image.rotate(angle, expand=True)

def scale_image(image, scale_factor):
    """Scale the image by a given factor."""
    new_size = (int(image.width * scale_factor), int(image.height * scale_factor))
    return image.resize(new_size)

def crop_image(image, left, upper, right, lower):
    """Crop the image."""
    return image.crop((left, upper, right, lower))

def flip_image(image, flip_type):
    """Flip the image horizontally or vertically."""
    if flip_type == 'Horizontal':
        return ImageOps.mirror(image)
    elif flip_type == 'Vertical':
        return ImageOps.flip(image)
    return image

# --- Main App Function ---
def main():
    st.set_page_config(
        page_title="Interactive Image Augmentation Tool",
        layout="wide"
    )
    
    # --- Styling ---
    st.markdown(
        """
        <style>
        .stTitle {
            color: #3E2723;
            text-align: center;
            font-family: 'Arial', sans-serif;
        }
        .stButton>button {
            background-color: #C19A6B;
            color: white;
            border-radius: 8px;
        }
        .stDownloadButton>button {
            background-color: #8B5E3C;
            color: white;
            border-radius: 8px;
        }
        </style>
        """,
        unsafe_allow_html=True
    )

    # --- Title and Introduction ---
    st.title("🖼️ Interactive Image Augmentation Tool")
    st.write("""
    Welcome to the **Interactive Image Augmentation Tool**! This app allows you to transform and enhance your images 
    using various augmentation techniques like **Translation**, **Rotation**, **Scaling**, **Cropping**, and **Flipping**.
    """)

    st.write("""
    ### 📝 Instructions for Use:
    - Upload an image using the uploader below.  
    - Select a transformation technique from the dropdown menu.  
    - Adjust the parameters using sliders and dropdowns.  
    - Preview the transformed image.  
    - Download the transformed image if you're satisfied.  
    """)

    st.write("""
    ### 🌟 Who Can Use This App?
    - **Data Scientists:** Generate augmented datasets for image classification models.  
    - **Graphic Designers:** Experiment with creative visual effects.  
    - **Researchers:** Prepare data for analysis and experiments.  
    - **Students:** Learn image processing techniques interactively.  
    """)

    # --- File Uploader ---
    uploaded_file = st.file_uploader("📤 Upload an Image", type=['jpg', 'jpeg', 'png'])
    if uploaded_file:
        image = Image.open(uploaded_file)
        st.image(image, caption="Original Image", use_column_width=True)
        
        st.subheader("🔄 Apply Transformations")
        transformation = st.selectbox(
            "Choose a Transformation",
            ["Translate", "Rotate", "Scale", "Crop", "Flip"]
        )
        
        if transformation == "Translate":
            x_offset = st.slider("X Offset", -100, 100, 0)
            y_offset = st.slider("Y Offset", -100, 100, 0)
            transformed_image = translate_image(image, x_offset, y_offset)
        
        elif transformation == "Rotate":
            angle = st.slider("Rotation Angle", 0, 360, 0)
            transformed_image = rotate_image(image, angle)
        
        elif transformation == "Scale":
            scale = st.slider("Scale Factor", 0.5, 2.0, 1.0)
            transformed_image = scale_image(image, scale)
        
        elif transformation == "Crop":
            left = st.slider("Left", 0, image.width, 0)
            upper = st.slider("Upper", 0, image.height, 0)
            right = st.slider("Right", left, image.width, image.width)
            lower = st.slider("Lower", upper, image.height, image.height)
            transformed_image = crop_image(image, left, upper, right, lower)
        
        elif transformation == "Flip":
            flip_type = st.selectbox("Flip Type", ["None", "Horizontal", "Vertical"])
            transformed_image = flip_image(image, flip_type)
        
        else:
            st.warning("Select a valid transformation!")
            return
        
        st.image(transformed_image, caption="Transformed Image", use_column_width=True)
        
        # --- Download Image ---
        img_buffer = io.BytesIO()
        transformed_image.save(img_buffer, format="PNG")
        st.download_button(
            label="📥 Download Transformed Image",
            data=img_buffer.getvalue(),
            file_name="transformed_image.png",
            mime="image/png"
        )

    st.markdown("---")
    st.write("🎯 **Experiment with transformations and download your enhanced images today!**")

if __name__ == "__main__":
    main()