import streamlit as st
import cv2
import numpy as np
from io import BytesIO
from PIL import Image
import base64
from fpdf import FPDF
import os

# Define augmentation functions
def rotate_image(image, angle):
    (h, w) = image.shape[:2]
    center = (w // 2, h // 2)
    rotation_matrix = cv2.getRotationMatrix2D(center, angle, 1.0)
    rotated_image = cv2.warpAffine(image, rotation_matrix, (w, h))
    return rotated_image

def add_noise(image):
    s_vs_p = 0.5
    amount = 0.02
    noisy_image = image.copy()
    num_salt = np.ceil(amount * image.size * s_vs_p)
    num_pepper = np.ceil(amount * image.size * (1.0 - s_vs_p))
    coords = [np.random.randint(0, i - 1, int(num_salt)) for i in image.shape[:2]]
    noisy_image[coords[0], coords[1], :] = 255
    coords = [np.random.randint(0, i - 1, int(num_pepper)) for i in image.shape[:2]]
    noisy_image[coords[0], coords[1], :] = 0
    return noisy_image

def blur_image(image, kernel_size):
    blurred_image = cv2.GaussianBlur(image, (kernel_size, kernel_size), 0)
    return blurred_image

def translate_image(image, tx, ty):
    translation_matrix = np.float32([[1, 0, tx], [0, 1, ty]])
    height, width = image.shape[:2]
    translated_image = cv2.warpAffine(image, translation_matrix, (width, height))
    return translated_image

# Download functions
def download_image(image, file_format):
    if file_format.lower() == "jpg":
        file_format = "jpeg"
    pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
    buffer = BytesIO()
    pil_image.save(buffer, format=file_format.upper())
    buffer.seek(0)
    b64 = base64.b64encode(buffer.read()).decode()
    href = f'<a href="data:file/{file_format};base64,{b64}" download="augmented_image.{file_format.lower()}">Download {file_format.upper()} File</a>'
    return href

def download_pdf(image):
    pdf = FPDF()
    pdf.add_page()

    temp_image_path = "temp_image.jpg"
    pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
    pil_image.save(temp_image_path, format="JPEG")

    pdf.image(temp_image_path, x=10, y=10, w=190)
    os.remove(temp_image_path)

    pdf_output = BytesIO()
    pdf.output(pdf_output, "F")
    pdf_output.seek(0)

    b64 = base64.b64encode(pdf_output.read()).decode()
    href = f'<a href="data:application/octet-stream;base64,{b64}" download="augmented_image.pdf">Download PDF</a>'
    return href

# Streamlit app
st.title("Image Augmentation Tool")

# File upload
uploaded_file = st.file_uploader("Upload an Image", type=["jpg", "png", "jpeg"])
if uploaded_file:
    image = Image.open(uploaded_file)
    image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
    st.image(image, caption="Uploaded Image", use_column_width=True)

    st.write("### Select Transformations to Apply:")
    rotate = st.checkbox("Rotate (45 degrees)")
    noise = st.checkbox("Add Noise")
    blur = st.checkbox("Blur (5x5 Kernel)")
    translate = st.checkbox("Translate (50 pixels right and down)")

    transformed_image = None  # Initialize transformed_image

    if st.button("Apply Transformations"):
        transformed_image = image.copy()

        if rotate:
            transformed_image = rotate_image(transformed_image, 45)
        if noise:
            transformed_image = add_noise(transformed_image)
        if blur:
            transformed_image = blur_image(transformed_image, 5)
        if translate:
            transformed_image = translate_image(transformed_image, 50, 50)

        st.image(cv2.cvtColor(transformed_image, cv2.COLOR_BGR2RGB), caption="Transformed Image", use_column_width=True)

    # Provide download options after transformations
    if transformed_image is not None:
        st.markdown("### Download Transformed Image:")
        download_format = st.radio("Select Download Format:", ["JPG", "PNG", "PDF"], key="download_format")

        if download_format == "JPG" or download_format == "PNG":
            st.markdown(download_image(transformed_image, download_format), unsafe_allow_html=True)
        elif download_format == "PDF":
            st.markdown(download_pdf(transformed_image), unsafe_allow_html=True)