import numpy as np
import random
import streamlit as st
import cv2
from PIL import Image
import base64
from io import BytesIO
import zipfile

st.markdown("""
    <style>
        .stButton>button {
            color: white;
            border-radius: 8px;
            padding: 10px 20px;
            font-weight: bold;
            transition: background-color 0.3s ease;
        }
        .stButton>button:hover {
            background-color: Black;
            border-color: white;
            color: white;
        }
    </style>
""", unsafe_allow_html=True)

def encode_image(image_path):
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode()

def add_bg_from_local(image_file):
    encoded_string = encode_image(image_file)
    st.markdown(
        f"""
        <style>
        .stApp {{
            background-image: url(data:image/{"png"};base64,{encoded_string});
            background-size: cover;
            background-repeat: no-repeat;
            background-attachment: fixed;
        }}
        </style>
        """,
        unsafe_allow_html=True
    )

add_bg_from_local("Images/Gemini_Generated_Image_pnb6j1pnb6j1pnb6.jpg")

st.sidebar.header("Augmentation Options")
apply_translation = st.sidebar.checkbox("Apply Translation", value=True)
apply_rotation = st.sidebar.checkbox("Apply Rotation", value=True)
apply_brightness = st.sidebar.checkbox("Apply Brightness Adjustment", value=True)

if apply_brightness:
    brightness_mode = st.sidebar.radio(
        "Brightness Adjustment Mode",
        options=["Increase Brightness", "Decrease Brightness", "Random Adjustment"],
        index=2
    )
    if brightness_mode == "Increase Brightness":
        brightness_range = (1.0, 2.0)

    elif brightness_mode == "Decrease Brightness":
        brightness_range = (0.1, 1.0)
    
    elif brightness_mode == "Random Adjustment":
        st.sidebar.header("Random Adjustment Settings")
        brightness_min = st.sidebar.slider(
            "Minimum brightness factor",
            min_value=0.1, max_value=1.0, value=0.7, key="brightness_min_random"
        )
        brightness_max = st.sidebar.slider(
            "Maximum brightness factor",
            min_value=1.0, max_value=2.0, value=1.3, key="brightness_max_random"
        )
        if brightness_min >= brightness_max:
            st.sidebar.error("Minimum brightness factor must be less than maximum brightness factor.")
        brightness_range = (brightness_min, brightness_max)


def augment_image(image, tx_range=30, angle_range=15, brightness_range=(0.7, 1.3)):
    rows, cols, _ = image.shape
    augmented_image = image.copy()

    if apply_translation:
        tx = random.randint(-tx_range, tx_range)
        ty = random.randint(-tx_range, tx_range)
        translation_matrix = np.float32([[1, 0, tx], [0, 1, ty]])
        augmented_image = cv2.warpAffine(augmented_image, translation_matrix, (cols, rows))

    if apply_rotation:
        angle = random.randint(-angle_range, angle_range)
        rotation_matrix = cv2.getRotationMatrix2D((cols / 2, rows / 2), angle, 1)
        augmented_image = cv2.warpAffine(augmented_image, rotation_matrix, (cols, rows))

    if apply_brightness:
        factor = random.uniform(*brightness_range)
        augmented_image = cv2.convertScaleAbs(augmented_image, alpha=factor, beta=0)

    return augmented_image


def display_images_in_grid(images, cols=3):
    rows = (len(images) + cols - 1) // cols
    for row in range(rows):
        cols_to_display = images[row * cols : (row + 1) * cols]
        columns = st.columns(len(cols_to_display))
        for col, img in zip(columns, cols_to_display):
            with col:
                st.image(img, channels="RGB", use_container_width=True)


st.title("Image Augmentation App")
st.write("Upload an image and generate augmented images. For more control, use Advanced Settings.")

uploaded_file = st.file_uploader("Choose an image", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    image = Image.open(uploaded_file)
    image_array = np.array(image)
    st.image(image, caption="Original Image", use_container_width=True)

    num_augmentations = st.sidebar.number_input("Number of augmented images", min_value=1, max_value=200, value=5, step=1)

    advanced_settings = st.sidebar.checkbox("Advanced Settings")

    tx_range = 30
    angle_range = 15
    brightness_range = (0.7, 1.3)

    if advanced_settings:
        st.sidebar.header("Augmentation Settings")
        tx_range = st.sidebar.slider("Translation range (pixels)", min_value=0, max_value=100, value=30, key="tx_range")
        angle_range = st.sidebar.slider("Rotation range (degrees)", min_value=0, max_value=90, value=15, key="angle_range")

    if not (apply_brightness and brightness_mode == "Random Adjustment"):
        if advanced_settings:
            brightness_min = st.sidebar.slider(
                "Minimum brightness factor",
                min_value=0.1, max_value=1.0, value=0.7, key="brightness_min_advanced"
            )
            brightness_max = st.sidebar.slider(
                "Maximum brightness factor",
                min_value=1.0, max_value=2.0, value=1.3, key="brightness_max_advanced"
            )
            if brightness_min >= brightness_max:
                st.sidebar.error("Minimum brightness factor must be less than maximum brightness factor.")
            brightness_range = (brightness_min, brightness_max)

    if apply_translation or apply_rotation or apply_brightness:
        if st.sidebar.button("Generate Augmented Images"):
            st.write("Augmented Images:")
            augmented_images = set()
            unique_images = []

            while len(unique_images) < num_augmentations:
                augmented_image = augment_image(image_array.copy(), tx_range, angle_range, brightness_range)
                image_hash = hash(augmented_image.tobytes())

                if image_hash not in augmented_images:
                    augmented_images.add(image_hash)
                    unique_images.append(augmented_image)
            
            st.session_state.unique_images = unique_images

            display_images_in_grid(unique_images, cols=3)

            if unique_images:
                zip_buffer = BytesIO()
                with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zip_file:
                    for i, aug_img in enumerate(unique_images):
                        aug_img = cv2.cvtColor(aug_img, cv2.COLOR_BGR2RGB)
                        _, img_encoded = cv2.imencode(".jpg", aug_img)
                        zip_file.writestr(f"augmented_image_{i + 1}.jpg", img_encoded.tobytes())
                    
                zip_buffer.seek(0)
                st.download_button(
                    label="Download Augmented Images as ZIP",
                    data=zip_buffer.getvalue(),
                    file_name="augmented_images.zip",
                    mime="application/zip",
                )
    else:
        if st.sidebar.button("Generate Augmented Images"):
            st.warning("No augmentation options selected. Please select at least one option (Translation, Rotation, or Brightness).")