import streamlit as st
from streamlit.runtime.scriptrunner.script_run_context import add_script_run_ctx
import threading
import cv2
from ultralytics import YOLO

model = YOLO('best.pt')
model.fuse()

# Configurar cámaras
camaras = {
    'camara1': 'IMG_0559.mov', #0,
    'camara2': 'IMG_0559.mov', #0,
}

def tracking(source, iou, conf):
    result = model.track(source=source, iou=iou, conf=conf, 
                        persist=True,
                        device='mps', 
                        tracker='botsort.yaml',
                        verbose=False)
    return result

def inferir_camara(index, source, iou, conf, st_frame, st_cantidad, stop_button_pressed):

    camera1 = cv2.VideoCapture(source)  # Open camera 0
    
    while camera1.isOpened() and not stop_button_pressed:
        ret, frame = camera1.read()
        if not ret:
            st.write(f"Video {str(index)} Capture Ended")
            break
        
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        
        resultado = tracking(frame, iou, conf)

        # Extract bounding boxes (replace this with the specific method from your library)
        objects_detected = resultado[0].boxes.cls  # Boxes object for bbox outputs
        object_count = len(objects_detected)
        st_cantidad.write(f'Cantidad detectados en cámara {str(index)}: {object_count}')
        
        res_plotted = resultado[0].plot()
        st_frame.image(res_plotted,
                        #caption='Detected Video',
                        channels="RGB",
                        use_column_width=True
                        )

        #if cv2.waitKey(1) & 0xFF == ord("q") or stop_button_pressed:
        if stop_button_pressed:
            break
    
    # Release cameras
    #camera1.release()
    #cv2.destroyAllWindows()

# Interfaz
st.set_page_config(page_title="Tracking YOLOv8")
st.title("Tracking YOLOv8")
iou = float(st.sidebar.slider("NMS IoU threshold", 30, 100, 80)) / 100 # Umbral de intersección sobre unión (IoU) para NMS
conf = float(st.sidebar.slider("Umbral o threshold", 30, 100, 80)) / 100 # Select Model Confidence

col1, col2 = st.columns(2)
play_button_pressed = st.empty()
with col1:
    stop_button_pressed_1 = st.button("Detener cámara 1")
    play_button_1 = st.empty()
    st_frame_camara_1 = st.empty()
    st_cantidad_1 = st.empty()

    tracker_thread1 = threading.Thread(target=inferir_camara, 
                            args=("1", camaras.get('camara1', 0),
                                iou, conf, 
                                st_frame_camara_1, st_cantidad_1, 
                                stop_button_pressed_1), 
                            daemon=True)
    add_script_run_ctx(tracker_thread1)
    play_button_1 = st.button("Iniciar cámara 1", on_click=tracker_thread1.start())

with col2:
    stop_button_pressed_2 = st.button("Detener cámara 2")
    st_frame_camara_2 = st.empty()
    st_cantidad_2 = st.empty()

    tracker_thread2 = threading.Thread(target=inferir_camara, 
                            args=("2", camaras.get('camara2', 0), 
                                iou, conf, 
                                st_frame_camara_2, st_cantidad_2, 
                                stop_button_pressed_2), 
                            daemon=True)
    add_script_run_ctx(tracker_thread2)
    play_button_2 = st.button("Iniciar cámara 2", on_click=tracker_thread2.start())

# Wait for the tracker threads to finish
tracker_thread1.join()
tracker_thread2.join()