import os
import cv2
import glob
import torch
import streamlit as st
import numpy as np
import base64
from random import randrange
from PIL import Image, ImageDraw
from transformers import pipeline
from tempfile import NamedTemporaryFile
from st_clickable_images import clickable_images

def clearEphimeral():
    for filename in glob.glob('/tmp/cropped_image_*'):
        os.remove(filename)


st.set_page_config(layout="wide")

clearEphimeral()

if 'wip' not in st.session_state:
    st.session_state['wip'] = False
    st.session_state['image'] = None

detector = pipeline("object-detection", model="sergioska/yolo-finetuned-flyers")

col1, col2, col3 = st.columns([0.5, 0.3, 0.2])

gibsOk = False
totalGibs = 0

with col1:

    clearEphimeral()

    st.header('flyer')
    uploaded_image_zero_file = st.file_uploader("Choose an image file (zero)")

    if not st.session_state['wip']:
        imageLocation = st.empty()
    else:
        imageLocation = st.image(st.session_state['image'])
    
    if uploaded_image_zero_file is not None:
        image = Image.open(uploaded_image_zero_file)
        outputImage = np.array(image)
        predictions = detector(
            image,
            candidate_labels=['product']
        )
        imageLocation.image(outputImage)
    
    if st.button('detect') or st.session_state['wip']:
    
        draw = ImageDraw.Draw(image)

        arrImg = np.array(image)
        imageCv = cv2.cvtColor(arrImg, cv2.COLOR_RGB2BGR)

        i = 0
        for prediction in predictions:
            box = prediction["box"]
            label = prediction["label"]
            score = prediction["score"]
    
            xmin, ymin, xmax, ymax = box.values()
            draw.rectangle((xmin, ymin, xmax, ymax), outline="red", width=5)
            draw.text((xmin, ymin), f"{label}: {round(score,5)}", fill="white")

            crop_img = imageCv[ymin:ymax, xmin:xmax]
            if crop_img.size > 0:
                cv2.imwrite("/tmp/cropped_image_{}.jpg".format(i), crop_img)
                i += 1

        totalGibs = i
        gibsOk = True
        imageLocation.image(image)
        st.session_state['wip'] = True
        st.session_state['image'] = image
        st.write(predictions)
        
with col2:
    st.header('gibs')
    images = []
    imagesForOpen = []
    for cropped in glob.glob('/tmp/cropped_image_*'):
        with open(cropped, "rb") as image:
            encoded = base64.b64encode(image.read()).decode()
            images.append(f"data:image/jpeg;base64,{encoded}")
            imagesForOpen.append(cropped)
    clicked = clickable_images(
        images,
        titles=[f"Image #{str(i)}" for i in range(totalGibs)],
        div_style={"display": "flex", "justify-content": "center", "flex-wrap": "wrap"},
        img_style={"margin": "5px", "height": "150px"},
        key="cropped_images"
    )

with col3:
    st.header('details')
    colors = ["red", "yellow", "purple", "blue", "green"]
    if gibsOk:
        if clicked == -1:
            clicked = 0
        detectorForContents = pipeline("object-detection", model="sergioska/yolo-finetuned-flyers-content")
        imageForDetails = Image.open(imagesForOpen[clicked])
        outputImageForDetails = np.array(imageForDetails)
        predictionsForDetails = detectorForContents(
            imageForDetails,
            candidate_labels=['title', 'price']
        )
        drawForDetails = ImageDraw.Draw(imageForDetails)
        c = 0
        for predictionsForDetail in predictionsForDetails:
            box = predictionsForDetail["box"]
            label = predictionsForDetail["label"]
            score = predictionsForDetail["score"]
    
            xmin, ymin, xmax, ymax = box.values()
            drawForDetails.rectangle((xmin, ymin, xmax, ymax), outline=colors[c], width=5)
            drawForDetails.text((xmin, ymin), f"{label}: {round(score,5)}", fill="white")
            c += 1

        st.image(imageForDetails)
        st.write(predictionsForDetails)
        

    st.markdown(f"Image #{clicked} clicked" if clicked > -1 else "No image clicked")