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

st.set_page_config(page_title="Image Augmentation Tool 🧰", page_icon="🖼️", layout="centered")


def get_base64_image(image_path):
    with open(image_path, "rb") as f:
        data = f.read()
    return base64.b64encode(data).decode()

image_base64 = get_base64_image("innomatics-footer-logo.png")

st.markdown(
    f"""
    <div style='text-align: center;'>
        <img src="data:image/png;base64,{image_base64}" width="600" height="100">
    </div>
    """,
    unsafe_allow_html=True
)

st.markdown("<h1 style='text-align: center;'>🖼️ Image Augmentation Tool 🎨</h1>", unsafe_allow_html=True)
st.markdown("<h4 style='text-align: center;'>✨ Apply real-time transformations to your images</h4>", unsafe_allow_html=True)
st.markdown("---")

# Upload section
st.subheader("📤 Upload Your Image")
uploaded_file = st.file_uploader("Choose an image file (jpg, jpeg, png)", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)
    image = cv2.imdecode(file_bytes, 1)
    original_image = image.copy()

    # Sidebar options
    st.sidebar.title("🛠️ Transformations Panel")
    st.sidebar.markdown("🔄 Choose your desired transformations below:")

    grayscale = st.sidebar.checkbox("🖤 Convert to Grayscale")
    rotate_angle = st.sidebar.slider("🌀 Rotate", -180, 180, 0)
    flip_horizontal = st.sidebar.checkbox("↔️ Flip Horizontally")
    flip_vertical = st.sidebar.checkbox("↕️ Flip Vertically")
    shear_factor = st.sidebar.slider("📐 Shearing", -0.5, 0.5, 0.0, step=0.01)
    translate_x = st.sidebar.slider("➡️ Translate X", -100, 100, 0)
    translate_y = st.sidebar.slider("⬇️ Translate Y", -100, 100, 0)
    crop = st.sidebar.checkbox("✂️ Crop Image")

    if crop:
        crop_x = st.sidebar.slider("🔲 Crop Start X", 0, image.shape[1] - 1, 0)
        crop_y = st.sidebar.slider("🔳 Crop Start Y", 0, image.shape[0] - 1, 0)
        crop_w = st.sidebar.slider("🔲 Crop Width", 1, image.shape[1] - crop_x, image.shape[1] - crop_x)
        crop_h = st.sidebar.slider("🔳 Crop Height", 1, image.shape[0] - crop_y, image.shape[0] - crop_y)

    # Apply transformations
    try:
        if grayscale:
            image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
            image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)

        if rotate_angle != 0:
            (h, w) = image.shape[:2]
            M = cv2.getRotationMatrix2D((w // 2, h // 2), rotate_angle, 1.0)
            image = cv2.warpAffine(image, M, (w, h))

        if flip_horizontal:
            image = cv2.flip(image, 1)

        if flip_vertical:
            image = cv2.flip(image, 0)

        if shear_factor != 0.0:
            rows, cols, _ = image.shape
            M = np.float32([
                [1, shear_factor, 0],
                [0, 1, 0]
            ])
            image = cv2.warpAffine(image, M, (cols, rows))

        if translate_x != 0 or translate_y != 0:
            M = np.float32([
                [1, 0, translate_x],
                [0, 1, translate_y]
            ])
            image = cv2.warpAffine(image, M, (image.shape[1], image.shape[0]))

        if crop:
            image = image[crop_y:crop_y + crop_h, crop_x:crop_x + crop_w]

        # Display results
        st.subheader("🔍 Preview")
        col1, col2 = st.columns(2)
        with col1:
            st.image(original_image, caption="🖼️ Original Image", use_container_width=True, channels="BGR")
        with col2:
            st.image(image, caption="✨ Transformed Image", use_container_width=True, channels="BGR")

        # Download button
        img_pil = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
        buf = io.BytesIO()
        img_pil.save(buf, format="PNG")
        byte_im = buf.getvalue()

        st.download_button(
            label="📥 Download Transformed Image",
            data=byte_im,
            file_name="transformed_image.png",
            mime="image/png"
        )

    except Exception as e:
        st.error(f"🚨 Error: {e}")

else:
    st.info("📎 Please upload an image to begin.")