# Python In-built packages
from pathlib import Path
import PIL
import time

# External packages
import streamlit as st
import cv2

# Local Modules
import settings
import helper

# Setting page layout
st.set_page_config(
    page_title="Object Detection using YOLOv9",
    page_icon="🤖",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Main page heading
st.title("Object Detection using YOLOv9")

# Sidebar
st.sidebar.header("ML Model Config")

# Model Options
model_type = st.sidebar.radio(
    "Select Task", ['Detection', 'Segmentation'])
confidence = float(st.sidebar.slider(
    "Select Model Confidence", 25, 100, 40)) / 100

# Selecting Detection Or Segmentation
if model_type == 'Detection':
    model_path = Path(settings.DETECTION_MODEL)
elif model_type == 'Segmentation':
    model_path = Path(settings.SEGMENTATION_MODEL)

# Load Pre-trained ML Model
try:
    model = helper.load_model(model_path)
except Exception as ex:
    st.error(f"Unable to load model. Check the specified path: {model_path}")
    st.error(ex)

st.sidebar.header("Image/Video Config")
source_radio = st.sidebar.radio(
    "Select Source", settings.SOURCES_LIST)
source_img = None

# If image is selected
if source_radio == settings.IMAGE:
    source_img = st.sidebar.file_uploader(
        "Choose an image...", type=("jpg", "jpeg", "png", 'bmp', 'webp'))
    col1, col2 = st.columns(2)
    with col1:
        try:
            if source_img is None:
                default_image_path = str(settings.DEFAULT_IMAGE)
                default_image = PIL.Image.open(default_image_path)
                st.image(default_image_path, caption="Default Image",
                         use_column_width=True)
            else:
                uploaded_image = PIL.Image.open(source_img)
                st.image(source_img, caption="Uploaded Image",
                         use_column_width=True)
        except Exception as ex:
            st.error("Error occurred while opening the image.")
            st.error(ex)
    with col2:
        if source_img is None:
            default_detected_image_path = str(settings.DEFAULT_DETECT_IMAGE)
            default_detected_image = PIL.Image.open(
                default_detected_image_path)
            st.image(default_detected_image_path, caption='Detected Image',
                     use_column_width=True)
        else:
            if st.sidebar.button('Detect Objects'):
                res = model.predict(uploaded_image,
                                    conf=confidence
                                    )
                boxes = res[0].boxes
                res_plotted = res[0].plot()[:, :, ::-1]
                st.image(res_plotted, caption='Detected Image',
                         use_column_width=True)
                try:
                    with st.expander("Detection Results"):
                        for box in boxes:
                            st.write(box.data)
                except Exception as ex:
                    # st.write(ex)
                    st.write("No image is uploaded yet!")

elif source_radio == settings.VIDEO:
    helper.play_stored_video(confidence, model)

elif source_radio == settings.YOUTUBE:
    helper.play_youtube_video(confidence, model)

elif source_radio == settings.WEBCAM:  # Adding webcam option
    st.sidebar.markdown("**Webcam**")
    webcam = cv2.VideoCapture(1)
    if not webcam.isOpened():
        st.error("Unable to open webcam. Please make sure it's connected.")
    else:
        webcam_status = st.sidebar.empty()
        start_stop_button = st.sidebar.button("Start Webcam")

        if start_stop_button:
            webcam_started = True
            webcam_status.text("Webcam Started!")

            while webcam_started:
                ret, frame = webcam.read()
                if not ret:
                    st.error("Failed to capture frame from webcam")
                    break
                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

                # Perform object detection on the frame
                res = model.predict(frame, conf=confidence)
                boxes = res[0].boxes
                frame = res[0].plot()[:, :, ::-1]

                # Display the frame
                st.image(frame, channels="RGB", use_column_width=True)

                stop_button = st.sidebar.button("Stop Webcam" + str(hash(time.time())))  # Unique key
                if stop_button:
                    webcam_started = False
                    webcam_status.text("Webcam Stopped!")

            webcam.release()