import streamlit as st
import numpy as np
import pandas as pd
import os
import zipfile


import cv2

st.markdown("""
    <style>
    /* Set full-page width and remove Streamlit default padding */
    .stApp {
        background-color: #d0e8ff;
    }
    .main .block-container {
        padding-top: 2rem;
        padding-right: 2rem;
        padding-left: 2rem;
        padding-bottom: 2rem;
        max-width: 100%;
    }
    </style>
""", unsafe_allow_html=True)

# Main Heading

st.markdown('<h2 style="text-align: center;"><span style="color:#003366;">Generate Augmented Images </span></h2>', unsafe_allow_html=True)

st.markdown("<br><br>", unsafe_allow_html=True)

#st.markdown('<h4 style="text-align: center;"><span style="color:#003366;">Upload the image </span></h4>', unsafe_allow_html=True)

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

image_count=st.number_input("Number of Images to be generated",1,50,placeholder="Type a number...")

st.write("Select the Affine Transformation to be applied:")

trans = st.checkbox("Translation")
rot = st.checkbox("Rotation")
scal = st.checkbox("Scaling")
shear=st.checkbox("Shearing")
crop=st.checkbox("Crop")


transformation=[]   
btn=st.button("Preview Image")

#sub=st.button("Preview Image")

def clear_folder(folder_name="output_images"):
    if os.path.exists(folder_name):
        for file in os.listdir(folder_name):
            file_path=os.path.join(folder_name,file)
            if os.path.isfile(file_path):
                os.remove(file_path)
    else:
        os.makedirs(folder_name)

def save_image_and_show(image, folder_name="output_images"):
    if not os.path.exists(folder_name):
        os.makedirs(folder_name)

    image_name = f"image_{len(os.listdir(folder_name)) + 1}.jpg"
    image_path = os.path.join(folder_name, image_name)

    # Save the image
    cv2.imwrite(image_path, image)
    
    if btn:
        # Showing the success message
        st.success(f"Image saved to {image_path} (local path)")
        # Show the image to the user (not the path)
        st.image(image, caption="Saved Image", channels="BGR")



def zip_file(folder_name="output_images"):
    zip_filename = f"{folder_name}.zip"
    with zipfile.ZipFile(zip_filename, 'w') as zipf:
        for root, dirs, files in os.walk(folder_name):
            for file in files:
                filepath = os.path.join(root, file)
                arcname = os.path.relpath(filepath, folder_name)
                zipf.write(filepath, arcname)

    # Read zip for download
    with open(zip_filename, "rb") as f:
        st.download_button(
            label="Download All Images as ZIP",
            data=f,
            file_name=zip_filename,
            mime="application/zip"
        )



def translation(img,n=5):
    for i in range(n):
        tx=np.random.randint(2,50)
        ty=np.random.randint(1,50)
        tm=np.array([[1,0,tx],[0,1,ty]],dtype=np.float32)
        trans_img=cv2.warpAffine(img,tm,dsize=(img.shape[1],img.shape[0]),borderMode=cv2.BORDER_REFLECT)
        #img_trans_rgb = cv2.cvtColor(trans_img, cv2.COLOR_BGR2RGB)
        save_image_and_show(trans_img)
    #st.image(img_trans_rgb, caption="Translation Applied", use_container_width=True)

def rotation(img,n=5):
    for i in range(n):
        angle=np.random.choice([30,45,90, 180, -90,-30,-45])
        scale=np.round(np.random.uniform(1,1.6),2)
        r_x=img.shape[0]//2
        r_y=img.shape[1]//2
        rm=cv2.getRotationMatrix2D((r_y,r_x),angle,scale)
        r_img=cv2.warpAffine(img,rm,(img.shape[1],img.shape[0]))
        #img_rot_rgb = cv2.cvtColor(r_img, cv2.COLOR_BGR2RGB)
        save_image_and_show(r_img)
        #st.image(img_rot_rgb, caption="Rotation Applied", use_container_width=True)

def scaling(img,n=5):
    for i in range(n):
        sx=np.random.uniform(1,2)
        sy=np.random.uniform(1,2)
        tx=np.random.randint(1,16)
        ty=np.random.randint(1,12)
        sm=np.array([[sx,0,tx],[0,sy,ty]],dtype=np.float32)
        s_img=cv2.warpAffine(img,sm,(img.shape[1],img.shape[0]))
        #img_scal_rgb=cv2.cvtColor(s_img,cv2.COLOR_BGR2RGB)
        save_image_and_show(s_img)
        #st.image(img_scal_rgb, caption="Scaling Applied", use_container_width=True)

def shearing(img,n=5):
    for i in range(n):
        sx=np.round(np.random.uniform(1,1.5),2)
        shx=np.round(np.random.uniform(0,0.5))
        tx=np.random.randint(1,10)
        shy=np.round(np.random.uniform(0,0.5))
        sy=np.round(np.random.uniform(1,1.2),2)
        ty=np.random.randint(1,15)
        shm=np.array([[sx,shx,tx],[shy,sy,ty]],dtype=np.float32)
        img_shear=cv2.warpAffine(img,shm,(img.shape[1],img.shape[0]),borderMode=cv2.BORDER_REFLECT)
        save_image_and_show(img_shear)
        #img_shear_rgb=cv2.cvtColor(img_all,cv2.COLOR_BGR2RGB)
        #st.image(img_shear_rgb, caption="Shearing Applied", use_container_width=True)

def cropping(img,n=5):
    for i in range(n):
        x_1=np.random.randint(20,60)
        y_1=np.random.randint(110,160)
        x_2=np.random.randint(20,60)
        y_2=np.random.randint(110,160)
        img_crop=img[x_1:y_1,x_2:y_2]
        save_image_and_show(img_crop)
        # img_crop_rgb=cv2.cvtColor(img_crop,cv2.COLOR_BGR2RGB)
        # st.image(img_crop_rgb, caption="Shearing Applied", use_container_width=True)

# def save_image(image, folder_name="output_images"):
#     if not os.path.exists(folder_name):
#         os.makedirs(folder_name)
    
#     image_name = f"image_{len(os.listdir(folder_name)) + 1}.jpg"
#     image_path = os.path.join(folder_name, image_name)
#     cv2.imwrite(image_path, image)
#     st.success(f"Image saved to {image_path}")



if uploaded_file is not None:
    # Converting the image into array
    img_array=np.asarray(bytearray(uploaded_file.read()),dtype=np.uint8)
    # Decode or converting the array back to image array
    img=cv2.imdecode(img_array,cv2.IMREAD_COLOR)
    #  # Convert BGR to RGB as streamlit expects RGB colorspace
    # img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    
    transformation = []

    if trans:
        transformation.append("Translation")
    if rot:
        transformation.append("Rotation")
    if scal:
        transformation.append("Scaling")
    if shear:
        transformation.append("Shearing")
    if crop:
        transformation.append("Crop")



    if transformation:
        clear_folder()
        # per_trans = max(1,(image_count // len(transformation)))
        # trans_count=[per_trans+1 if (per_trans * len(transformation))!= image_count else per_trans][0]
        num_trans = len(transformation)
        base = image_count // num_trans
        remainder = image_count % num_trans
    
    # Create a list like [base+1, base, base, ...] to ensure sum == image_count
        counts = [base + 1 if i < remainder else base for i in range(num_trans)]

        # if "Translation" in transformation:
        #     translation(img,trans_count)
        # if "Rotation" in transformation:
        #     rotation(img, per_trans)
        # if "Scaling" in transformation:
        #     scaling(img, per_trans)
        # if "Shearing" in transformation:
        #     shearing(img, per_trans)
        # if "Crop" in transformation:
        #     cropping(img, per_trans)

        for t_name, t_count in zip(transformation, counts):
            if t_name == "Translation":
                translation(img, t_count)
            elif t_name == "Rotation":
                rotation(img, t_count)
            elif t_name == "Scaling":
                scaling(img, t_count)
            elif t_name == "Shearing":
                shearing(img, t_count)
            elif t_name == "Crop":
                cropping(img, t_count)


        zip_file()

    #st.download_button("Download File",img_rgb)

    # # Displaying image in streamlit
    # st.image(img_rgb, caption="Uploaded Image (OpenCV)", use_container_width=True)