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

# Custom CSS for enhanced styling
st.markdown(
    """
    <style>
    .stApp {
        background-color: #121212;
        color: #E0E0E0;
    }
    .stTitle {
        text-align: center;
        color: #4CAF50;
    }
    .css-1d391kg p, .css-1v0mbdj p {
        color: #B0BEC5;
    }
    .css-1aumxhk .stSlider {
        background-color: #4CAF50;
    }
    .css-2trqyj, .css-16idsys p {
        color: #BB86FC;
    }
    </style>
    """,
    unsafe_allow_html=True,
)

st.title("🌟 Image Augmentation Tool Hands-On Experience", anchor=False)

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

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

        st.sidebar.subheader("✨ Augmentation Options")
        
        augmentation_options = st.sidebar.multiselect(
            "Choose augmentations:",
            ["Cropping", "Flipping", "Rotation", "Zoom In", "Zoom Out", "Translation", "Shearing"]
        )
        
        augmented_image = image.copy()
        
        # Live preview toggle
        live_preview = st.sidebar.checkbox("🔄 Live Preview", value=True)

        for augmentation_option in augmentation_options:
            if augmentation_option == "Cropping":
                left = st.slider("Left Crop", 0, image.width // 2, 0)
                top = st.slider("Top Crop", 0, image.height // 2, 0)
                right = st.slider("Right Crop", 0, image.width // 2, 0)
                bottom = st.slider("Bottom Crop", 0, image.height // 2, 0)
                augmented_image = augmented_image.crop((left, top, image.width - right, image.height - bottom))

            elif augmentation_option == "Flipping":
                flip_option = st.selectbox("🔄 Flip Type", ["Horizontal", "Vertical"], key=f"flip_{augmentation_option}")
                if flip_option == "Horizontal":
                    augmented_image = ImageOps.mirror(augmented_image)
                elif flip_option == "Vertical":
                    augmented_image = ImageOps.flip(augmented_image)

            elif augmentation_option == "Rotation":
                rotation_angle = st.slider("⤾ Rotation (°)", 0, 360, 0)
                augmented_image = augmented_image.rotate(rotation_angle, expand=True)

            elif augmentation_option == "Zoom In":
                scale_factor = st.slider("🔍 Zoom Factor (%)", 100, 200, 100) / 100.0
                width, height = augmented_image.size
                new_width = int(width * scale_factor)
                new_height = int(height * scale_factor)
                zoomed_image = augmented_image.resize((new_width, new_height))
                augmented_image = zoomed_image.crop(((new_width - width) // 2, (new_height - height) // 2,
                                                    (new_width + width) // 2, (new_height + height) // 2))

            elif augmentation_option == "Zoom Out":
                scale_factor = st.slider("🔎 Zoom Out Factor (%)", 50, 100, 100) / 100.0
                width, height = augmented_image.size
                new_width = int(width * scale_factor)
                new_height = int(height * scale_factor)
                zoomed_image = augmented_image.resize((new_width, new_height))
                padded_image = Image.new("RGB", (width, height), (0, 0, 0))
                padded_image.paste(zoomed_image, ((width - new_width) // 2, (height - new_height) // 2))
                augmented_image = padded_image

            elif augmentation_option == "Translation":
                translation_x = st.slider("↔ Horizontal Shift (Pixels)", -100, 100, 0)
                translation_y = st.slider("↕ Vertical Shift (Pixels)", -100, 100, 0)
                translation_matrix = (1, 0, translation_x, 0, 1, translation_y)
                augmented_image = augmented_image.transform(
                    augmented_image.size, Image.AFFINE, translation_matrix, fillcolor=(0, 0, 0)
                )

            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 = augmented_image.transform(
                    (augmented_image.width + abs(int(augmented_image.height * np.tan(np.radians(shear_angle)))), augmented_image.height),
                    Image.AFFINE,
                    shear_matrix,
                    fillcolor=(0, 0, 0),
                )
        
        st.subheader("🖼️ Augmented Image")
        st.image(augmented_image, caption="Augmentation Applied", use_column_width=True)

        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}")