import streamlit as st
import numpy as np
import cv2
from PIL import Image
from ultralytics import YOLO
from huggingface_hub import hf_hub_download

# Streamlit App
st.set_page_config(layout="wide", initial_sidebar_state="collapsed")

# Define static class names
CLASS_NAMES = {0: 'dengue-regions', 1: 'wet_surface'}
CONFIDENCE_VALUE = 0.3

# Load YOLOv8 (.pt) model
@st.cache_resource
def load_pt_model():
    # Update the path to your .pt model file
    model_path = hf_hub_download(repo_id="DhominickJ/DengueTect", filename="best.pt")
    return YOLO(model_path)

# Draw bounding boxes on the image and return the image with detections drawn
def draw_boxes_pt(image, results):
    image_cv = np.array(image)
    image_cv = cv2.cvtColor(image_cv, cv2.COLOR_RGB2BGR)
    for result in results:
        if result.boxes is not None:
            # Get bounding boxes, confidences, and class IDs
            boxes = result.boxes.xyxy.cpu().numpy()  # (n,4)
            confidences = result.boxes.conf.cpu().numpy()  # (n,)
            classes = result.boxes.cls.cpu().numpy()  # (n,)
            for box, conf, cls in zip(boxes, confidences, classes):
                x1, y1, x2, y2 = box.astype(int)
                label = f"{CLASS_NAMES.get(int(cls), 'Unknown')}: {conf:.2f}"
                # Set color based on class
                color = (0, 0, 255) if CLASS_NAMES[int(cls)] == 'dengue-regions' else (255, 0, 0)
                cv2.rectangle(image_cv, (x1, y1), (x2, y2), color, 3)
                cv2.putText(image_cv, label, (x1, max(y1-10, 0)), 
                            cv2.FONT_HERSHEY_SIMPLEX, 2.0, color, 3)
    # Convert back to RGB
    image_cv = cv2.cvtColor(image_cv, cv2.COLOR_BGR2RGB)
    return Image.fromarray(image_cv)

def process_image(model, image):
    image = Image.open(image).convert("RGB")
    # st.image(image, caption="Input Image", use_column_width=True)
    results = model.predict(source=np.array(image), conf=CONFIDENCE_VALUE, imgsz=640)
    detections = []
    for result in results:
        if result.boxes is not None:
            boxes = result.boxes.xyxy.cpu().numpy()
            confidences = result.boxes.conf.cpu().numpy()
            classes = result.boxes.cls.cpu().numpy()
            for box, conf, cls in zip(boxes, confidences, classes):
                x1, y1, x2, y2 = box.astype(int)
                detections.append((int(cls), x1, y1, x2, y2, conf))
                
    if detections:
        st.success(f"‼️Detected {len(detections)} objects!")
        # for det in detections:    
            # st.text(f"Detected: {CLASS_NAMES.get(det[0], 'Unknown')}, "
            #         f"BBox: ({det[1]}, {det[2]}), ({det[3]}, {det[4]}), "
            #         f"Confidence: {det[5]:.2f}")
            # st.text(f"Confidence: {det[5]:.2f}")
        result_image = draw_boxes_pt(image, results)
        st.image(result_image, caption="Detection Results")
    else:
        st.warning("⚠ No objects risk detected!")

# Load model once
model = load_pt_model()

# Set page background and styling
st.markdown("""
    <style>
        .stApp {
            background-color: #0A4D68;
            color: white;
        }
        .stButton button {
            background-color: #088395;
            color: white;
            border: 2px solid white;
            border-radius: 10px;
            transition: all 0.3s;
            padding: 1rem;
        }
        .stButton button:hover {
            background-color: #05BFDB;
            transform: scale(1.02);
        }
        .uploadedFile {
            background-color: #088395 !important;
            color: white !important;
            border-radius: 10px !important;
            padding: 10px !important;
        }
        .stAlert {
            background-color: #088395;
            color: white;
        }
    </style>
""", unsafe_allow_html=True)

# App title with animation
st.markdown("<h1 style='text-align: center; color: white; text-shadow: 2px 2px 4px #000000;'>🦟 DengueTect</h1>", unsafe_allow_html=True)
st.markdown("<p style='text-align: center; color: #05BFDB; font-size: 20px;'>Detect mosquito breeding grounds with a click!</p>", unsafe_allow_html=True)

# Create two columns for the input methods
col1, col2 = st.columns(2)

# with col1:
if st.button("📷 Take a Photo!", use_container_width=True):
    st.session_state.input_method = "camera"

with col2:
    if st.button("📤 Upload Image", use_container_width=True):
        st.session_state.input_method = "upload"

# Initialize session state
if 'input_method' not in st.session_state:
    st.session_state.input_method = None

# Handle different input methods
elif st.session_state.input_method == "camera":  
    with st.container():
        try:
            camera_image = st.camera_input("Take a photo", key="camera")
            if camera_image is not None:
                process_image(model, camera_image)
        except Exception as e:
            st.error("Error accessing camera. Please make sure your camera is connected and you have given permission to access it.")

elif st.session_state.input_method == "upload":
    with st.container():
        col1, col2, col3 = st.columns([1,3,1])
        with col2:
            st.markdown("""
                <style>
                    .uploadedFile {
                        text-align: center;
                        background-color: #05BFDB !important;
                        color: white !important;
                        padding: 20px !important;
                        border-radius: 10px !important;
                    }
                    [data-testid="stFileUploader"] {
                        background-color: #05BFDB !important;
                        color: white !important;
                        border: 2px dashed white !important;
                        border-radius: 10px;
                        padding: 20px !important;
                    }
                </style>
            """, unsafe_allow_html=True)
            uploaded_file = st.file_uploader("Choose an image file", type=["jpg", "jpeg", "png"])
            if uploaded_file is not None:
                process_image(model, uploaded_file)

if st.session_state.input_method is not None:
    if st.button("↩ Back", type="secondary"):
        st.session_state.input_method = None
        st.experimental_rerun()