import streamlit as st
import cv2
from PIL import Image
from io import BytesIO
from ultralytics import YOLO
import numpy as np
from streamlit_option_menu import option_menu
from markup import real_estate_app, real_estate_app_hf

model = YOLO('yolov8n.onnx')

PASSWORD = 'Ethan101'

def authenticate(password):
    return password == PASSWORD

def tab1():
    st.header("Human and Vehicle Recognition Demo")  
    col1, col2 = st.columns([1, 2])
    with col1:
        st.image("image.jpg", use_column_width=True)
    with col2:
        st.markdown(real_estate_app(), unsafe_allow_html=True)
    st.markdown(real_estate_app_hf(),unsafe_allow_html=True) 


    github_link = '[<img src="https://badgen.net/badge/icon/github?icon=github&label">](https://github.com/ethanrom)'
    #huggingface_link = '[<img src="https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue">](https://huggingface.co/ethanrom)'

    st.write(github_link + '&nbsp;&nbsp;&nbsp;', unsafe_allow_html=True)

def tab2():
    st.header("Test Detection Algorithm")
    uploaded_file = st.file_uploader('Choose an image', type=['jpg', 'jpeg', 'png'])
    if uploaded_file is not None:
        image = Image.open(uploaded_file)
        image_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)

        col1, col2 = st.columns([2,1])
        with col2:    
            iou_threshold = st.slider('IoU Threshold', min_value=0.0, max_value=1.0, value=0.7)
            conf_threshold = st.slider('Confidence Threshold', min_value=0.0, max_value=1.0, value=0.65)
            show_labels = st.checkbox('Show Labels', value=False)
            show_conf = st.checkbox('Show Confidence Scores', value=False)
            boxes = st.checkbox('Show Boxes', value=True)

        with col1:
            st.image(image, caption='Input Image', use_column_width=True)
            
            if st.button('Apply and Predict'):
                results = model(
                    image_cv,
                    classes=[0,2,7,3,5],
                    iou=iou_threshold,
                    conf=conf_threshold,
                    show_labels=show_labels,
                    show_conf=show_conf,
                    boxes=boxes,        
                    )

                annotated_frame = results[0].plot()
                annotated_image = Image.fromarray(cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB))
                st.image(annotated_image, caption='Annotated Image', use_column_width=True)



def tab3():
    st.header("Configure and save script")
    st.write("Please send me a DM to get the password")

    password_input = st.text_input('Enter Password', type='password')
    if authenticate(password_input):

        source_folder = st.text_input('Source Folder Location')
        results_folder = st.text_input('Destination Folder Location for Object Detection Results')

        script = f"""
import os
import cv2
from ultralytics import YOLO
from tqdm import tqdm

model = YOLO('yolov8n.pt')

def detect_cars_humans(image_path):
    image_cv = cv2.imread(image_path)
    
    # Perform object detection
    results = model(
        image_cv,
        classes=[0, 2, 7, 3, 5],
        iou=0.7,
        conf=0.65,
        show_labels=False,
        show_conf=False,
        boxes=True
    )

    if len(results[0].boxes.xyxy) == 0:
        return
    
    # Create the destination folder if it doesn't exist
    os.makedirs(r"{results_folder}", exist_ok=True)
    
    # Save the annotated image in the results folder
    annotated_image_path = os.path.join(r"{results_folder}", os.path.basename(image_path))
    cv2.imwrite(annotated_image_path, results[0].plot())

source_folder = r"{source_folder}"
image_files = [f for f in os.listdir(source_folder) if f.endswith(".png") or f.endswith(".jpg")]

with tqdm(total=len(image_files), desc='Processing Images') as pbar:
    for filename in image_files:
        image_path = os.path.join(source_folder, filename)
        detect_cars_humans(image_path)
        pbar.update(1)
"""
        st.code(script, language='python')

        if st.button('Download Script'):
            script_filename = 'object_detection_script.py'
            with open(script_filename, 'w') as file:
                file.write(script)
            st.download_button(
                label='Download Script',
                data=script_filename,
                file_name=script_filename
            )

        if st.button('Download Requirements'):
            requirements_filename = 'requirements.txt'
            with open(requirements_filename, 'r') as file:
                requirements_content = file.read()
            st.download_button(
                label='Download Requirements',
                data=requirements_content,
                file_name=requirements_filename
            )

        st.subheader("Instructions:")
        st.write("1. Set the source folder and destination folder locations.")
        st.write("2. Click on the 'Download Script' button to download the object detection script.")
        st.write("3. Click on the 'Download Requirements' button to download the requirements.txt file.")
        st.write("4. Open a terminal or command prompt and navigate to the project directory.")
        st.write("5. Run the following command to install the required packages:")
        st.code("pip install -r requirements.txt")
        st.write("6. Finally, run the object detection script using the following command:")
        st.code("python object_detection_script.py")

    else:
        # Password is incorrect, show an error message
        st.error('Invalid password. Access denied.')

def main():
    st.set_page_config(page_title="Human and vehicle recognition", page_icon=":memo:", layout="wide")
    tabs = ["Intro", "Test", "Download Script"]

    with st.sidebar:

        current_tab = option_menu("Select a Tab", tabs, menu_icon="cast")

    tab_functions = {
    "Intro": tab1,
    "Test": tab2,
    "Download Script": tab3,
    }

    if current_tab in tab_functions:
        tab_functions[current_tab]()

if __name__ == "__main__":
    main()