import streamlit as st
import cv2
import numpy as np
import io
import zipfile
import random

# Page configuration
st.set_page_config(page_title="✨ Image Augmenter", layout="centered")

# Custom Styling
st.markdown("""
    <style>
    html, body {
        background-color: #eaf3fc;
        font-family: 'Segoe UI', sans-serif;
    }
    .title {
        font-size: 2.8rem;
        font-weight: 800;
        color: #005bb5;
        text-align: center;
        margin-bottom: 0.5rem;
        letter-spacing: -1px;
    }
    .subtitle {
        text-align: center;
        font-size: 1.2rem;
        color: #333;
        margin-bottom: 2.5rem;
    }
    .stButton > button {
        background-color: #005bb5;
        color: white;
        font-weight: bold;
        border-radius: 10px;
        padding: 0.75rem 2rem;
        border: none;
        font-size: 1rem;
        transition: all 0.3s ease-in-out;
        box-shadow: 0 4px 10px rgba(0, 91, 181, 0.2);
    }
    .stButton > button:hover {
        background-color: #004494;
        transform: scale(1.05);
    }
    .stDownloadButton > button {
        background-color: #28a745;
        color: white;
        font-weight: bold;
        border-radius: 10px;
        padding: 0.75rem 2rem;
        border: none;
        font-size: 1rem;
        transition: all 0.3s ease-in-out;
        box-shadow: 0 4px 10px rgba(40, 167, 69, 0.2);
    }
    .stDownloadButton > button:hover {
        background-color: #1e7e34;
        transform: scale(1.05);
    }
    </style>
""", unsafe_allow_html=True)

# Titles
st.markdown('<div class="title">📸 Image Augmentation Studio</div>', unsafe_allow_html=True)
st.markdown('<div class="subtitle">Upload an image, choose transformations, preview, and download your dataset!</div>', unsafe_allow_html=True)

# File Upload
uploaded_file = st.file_uploader("📤 Upload an image", type=["jpg", "jpeg", "png"])

# Select transformations
transform_options = st.multiselect(
    "🔄 Choose one or more transformations", 
    ["Translation", "Cropping", "Shearing", "Rotation", "Scaling", "Grayscale", "Flip Horizontally", "Flip Vertically"],
    default=["Rotation"]
)

st.write("you selected : " , transform_options)

# Number of images
count = st.number_input("📈 Number of images to generate", min_value=10, max_value=200, value=100, step=10)

# Transformation logic
def apply_transformation(image, options, count):
    rows, cols = image.shape[:2]
    transformed_images = []

    for _ in range(count):
        img_out = image.copy()  # Always initialize

        if not options:
            transformed_images.append(img_out)
            continue

        for option in options:
            if option == "Translation":
                tx, ty = random.randint(-50, 50), random.randint(-50, 50)
                M = np.float32([[1, 0, tx], [0, 1, ty]])
                img_out = cv2.warpAffine(img_out, M, (cols, rows), borderMode=cv2.BORDER_REFLECT)

            elif option == "Cropping":
                x1 = random.randint(0, cols // 4)
                y1 = random.randint(0, rows // 4)
                x2 = random.randint(3 * cols // 4, cols)
                y2 = random.randint(3 * rows // 4, rows)
                cropped = img_out[y1:y2, x1:x2]
                img_out = cv2.resize(cropped, (cols, rows))

            elif option == "Shearing":
                shear_factor = random.uniform(-0.3, 0.3)
                M = np.float32([[1, shear_factor, 0], [shear_factor, 1, 0]])
                img_out = cv2.warpAffine(img_out, M, (cols, rows), borderMode=cv2.BORDER_REFLECT)

            elif option == "Rotation":
                angle = random.randint(-45, 45)
                center = (cols // 2, rows // 2)
                M = cv2.getRotationMatrix2D(center, angle, 1.0)
                img_out = cv2.warpAffine(img_out, M, (cols, rows), borderMode=cv2.BORDER_REFLECT)

            elif option == "Scaling":
                scale = random.uniform(0.5, 1.5)
                resized = cv2.resize(img_out, None, fx=scale, fy=scale, interpolation=cv2.INTER_LINEAR)
                img_out = cv2.resize(resized, (cols, rows))

            elif option == "Grayscale":
                gray = cv2.cvtColor(img_out, cv2.COLOR_BGR2GRAY)
                img_out = cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)

            elif option == "Flip Horizontally":
                img_out = cv2.flip(img_out, 1)

            elif option == "Flip Vertically":
                img_out = cv2.flip(img_out, 0)

    
        transformed_images.append(img_out)

    return transformed_images

# Main logic
if uploaded_file:
    file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)
    image = cv2.imdecode(file_bytes, 1)
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    st.image(image_rgb, caption="📷 Original Image", use_container_width=True)

    if st.button("✨ Generate Augmented Images"):
        with st.spinner("Generating images..."):
            augmented_images = apply_transformation(image, transform_options, count)

            st.markdown("### 🔍 Preview of Augmented Images")
            cols_preview = st.columns(3)
            for i in range(min(3, len(augmented_images))):
                preview = cv2.cvtColor(augmented_images[i], cv2.COLOR_BGR2RGB)
                cols_preview[i].image(preview, use_container_width=True)

            zip_buffer = io.BytesIO()
            with zipfile.ZipFile(zip_buffer, "w") as zip_file:
                for idx, img in enumerate(augmented_images):
                    _, buffer = cv2.imencode(".png", img)
                    zip_file.writestr(f"augmented_{idx+1}.png", buffer.tobytes())
            zip_buffer.seek(0)

        st.success(f"✅ {count} Augmented Images Generated!")
        st.download_button("📁 Download ZIP", zip_buffer, "augmented_images.zip", "application/zip")