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

# Custom CSS for styling
st.markdown("""
    <style>
    /* General body styling */
    .stApp {
        background-color: #1e1e2f;
        color: #e0e0e0;
    }
    /* Title styling */
    .stTitle {
        color: #4caf50;
        font-weight: bold;
        text-align: center;
        margin-bottom: 25px;
    }
    /* Sidebar styling */
    .stSidebar {
        background-color: #28293d;
        color: #e0e0e0;
        border-radius: 10px;
    }
    .stSidebar h2, .stSidebar h3 {
        color: #81d4fa;
    }
    /* Sliders and buttons */
    .stSlider {
        background-color: #4caf50;
    }
    .stButton button {
        background-color: #4caf50;
        color: white;
        border-radius: 10px;
        border: none;
    }
    .stButton button:hover {
        background-color: #81d4fa;
    }
    </style>
""", unsafe_allow_html=True)

# App title
st.title("🌟 Interactive Image Augmentation Tool")

# Sidebar: Upload image
uploaded_image = st.sidebar.file_uploader("📤 Upload an Image", type=["jpg", "png", "jpeg"])
if uploaded_image:
    # Open and display the uploaded image
    image = Image.open(uploaded_image)
    st.sidebar.image(image, caption="📸 Uploaded Image",width=400)

    # Sidebar: Augmentation options
    st.sidebar.header("✨ Augmentation Options")
    augmentation_options = st.sidebar.multiselect(
        "Choose augmentations:",
        ["Cropping", "Flipping", "Rotation", "Zoom In", "Zoom Out", "Translation", "Shearing"]
    )

    # Sidebar: Live preview toggle
    live_preview = st.sidebar.checkbox("🔄 Enable Live Preview", value=True)

    # Make a copy of the image for augmentation
    augmented_image = image.copy()

    # Apply augmentations based on user selections
    for option in augmentation_options:
        if option == "Cropping":
            left = st.sidebar.slider("Left Crop", 0, image.width // 2, 0)
            top = st.sidebar.slider("Top Crop", 0, image.height // 2, 0)
            right = st.sidebar.slider("Right Crop", 0, image.width // 2, 0)
            bottom = st.sidebar.slider("Bottom Crop", 0, image.height // 2, 0)
            augmented_image = augmented_image.crop((left, top, image.width - right, image.height - bottom))

        elif option == "Flipping":
            flip_type = st.sidebar.selectbox("🔄 Flip Type", ["Horizontal", "Vertical"])
            if flip_type == "Horizontal":
                augmented_image = ImageOps.mirror(augmented_image)
            elif flip_type == "Vertical":
                augmented_image = ImageOps.flip(augmented_image)

        elif option == "Rotation":
            angle = st.sidebar.slider("⤾ Rotation Angle (°)", 0, 360, 0)
            augmented_image = augmented_image.rotate(angle, expand=True)

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

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

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

        elif option == "Shearing":
            shear_angle = st.sidebar.slider("🌀 Shear 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)
            )

    # Display the augmented image
    st.subheader("🖼️ Augmented Image")
    st.image(augmented_image, caption="Augmented Output",width=400)

    # Download button for the augmented image
    buffer = io.BytesIO()
    augmented_image.save(buffer, format="PNG")
    buffer.seek(0)
    st.download_button("📥 Download Augmented Image", buffer, file_name="augmented_image.png")
else:
    st.sidebar.warning("📂 Please upload an image to get started!")