import streamlit as st
from PIL import Image, ImageEnhance, ImageOps
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="Image Augmentation Tool",
        layout="wide"
    )
    
    # --- Nude Theme Styling ---
    st.markdown(
        """
        <style>
        body {
            background-color: #F2D1B3;  /* Soft nude beige */
        }
        .stTitle {
            color: #5F3B26;  /* Warm brown */
            text-align: center;
            font-family: 'Arial', sans-serif;
        }
        .stButton>button {
            background-color: #D1A39A;  /* Muted pinkish nude */
            color: white;
            border-radius: 8px;
        }
        .stDownloadButton>button {
            background-color: #A87C61;  /* Warm sand color */
            color: white;
            border-radius: 8px;
        }
        .stSlider {
            color: #5F3B26;  /* Warm brown */
        }
        </style>
        """,
        unsafe_allow_html=True
    )

    # --- Title and Introduction ---
    st.title("🖼️ Image Augmentation Tool with Nude Theme")
    st.write("""
    Welcome to the **Interactive Image Augmentation Tool**! This tool allows you to apply transformations such as 
    **Translation**, **Rotation**, **Scaling**, **Cropping**, and **Flipping** either **all at once** or **individually**.
    """)

    # --- Instructions ---
    st.write("""
    ### 📝 Instructions:
    - Upload an image using the uploader below.  
    - Apply **All Transformations** by adjusting default sliders.  
    - Or select a **Specific Transformation** to refine changes.  
    - Preview the transformed image.  
    - Download the result if you're satisfied.  
    """)

    # --- 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_container_width=True)
        
        st.subheader("🎛️ Transformation Controls")

        # --- Transformation Mode Selection ---
        transformation_mode = st.radio(
            "Choose Transformation Mode:",
            ["All Transformations", "Specific Transformation"]
        )
        
        # --- Default Transformations ---
        transformed_image = image.copy()  # Keep a copy for showing all intermediate transformations

        if transformation_mode == "All Transformations":
            st.write("### 🌟 Adjust All Transformations")

            # Translation
            x_offset = st.slider("X Offset", -100, 100, 0)
            y_offset = st.slider("Y Offset", -100, 100, 0)
            translated_image = translate_image(transformed_image, x_offset, y_offset)
            st.image(translated_image, caption="After Translation", use_container_width=True)

            # Rotation
            angle = st.slider("Rotation Angle", 0, 360, 0)
            rotated_image = rotate_image(transformed_image, angle)
            st.image(rotated_image, caption="After Rotation", use_container_width=True)

            # Scaling
            scale = st.slider("Scale Factor", 0.5, 2.0, 1.0)
            scaled_image = scale_image(transformed_image, scale)
            st.image(scaled_image, caption="After Scaling", use_container_width=True)

            # Cropping
            left = st.slider("Crop Left", 0, image.width, 0)
            upper = st.slider("Crop Upper", 0, image.height, 0)
            right = st.slider("Crop Right", left, image.width, image.width)
            lower = st.slider("Crop Lower", upper, image.height, image.height)
            cropped_image = crop_image(transformed_image, left, upper, right, lower)
            st.image(cropped_image, caption="After Cropping", use_container_width=True)

            # Flipping
            flip_type = st.selectbox("Flip Image", ["None", "Horizontal", "Vertical"])
            flipped_image = flip_image(transformed_image, flip_type)
            st.image(flipped_image, caption="After Flipping", use_container_width=True)

        # --- Specific Transformation ---
        elif transformation_mode == "Specific Transformation":
            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)
                image = translate_image(image, x_offset, y_offset)

            elif transformation == "Rotate":
                angle = st.slider("Rotation Angle", 0, 360, 0)
                image = rotate_image(image, angle)

            elif transformation == "Scale":
                scale = st.slider("Scale Factor", 0.5, 2.0, 1.0)
                image = scale_image(image, scale)

            elif transformation == "Crop":
                left = st.slider("Crop Left", 0, image.width, 0)
                upper = st.slider("Crop Upper", 0, image.height, 0)
                right = st.slider("Crop Right", left, image.width, image.width)
                lower = st.slider("Crop Lower", upper, image.height, image.height)
                image = crop_image(image, left, upper, right, lower)

            elif transformation == "Flip":
                flip_type = st.selectbox("Flip Type", ["None", "Horizontal", "Vertical"])
                image = flip_image(image, flip_type)

        # --- Display Final Transformed Image ---
        st.image(image, caption="Transformed Image", use_container_width=True)

        # --- Download Transformed Image ---
        img_buffer = io.BytesIO()
        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("🎯 **Enjoy experimenting with image transformations in style!** 🖼️✨")

if __name__ == "__main__":
    main()