import streamlit as st
from PIL import Image
import numpy as np
from ultralytics import YOLO

# Set Streamlit page config first
st.set_page_config(page_title="Suspicious Activity Detection", layout="centered")

# Load the YOLOv11 model
@st.cache_resource
def load_model():
    return YOLO("yolo11l.pt")  # Ensure your model is uploaded to the app folder

model = load_model()

# ------------------ Intelligent Action Classification Logic ------------------
def classify_action(detections):
    """
    Classify based on detected object types and confidence scores.
    """
    action_scores = {
        'Stealing': 0,
        'Sneaking': 0,
        'Peaking': 0,
        'Normal': 0
    }

    objects = [d[0] for d in detections]
    confidences = [d[1] for d in detections]

    has_person = 'person' in objects
    num_objects = len(objects)

    # Confidence-boosted rules
    if has_person:
        if 'handbag' in objects or 'backpack' in objects:
            action_scores['Stealing'] += 0.6
        if 'refrigerator' in objects or 'microwave' in objects:
            action_scores['Stealing'] += 0.4
        if max(confidences) < 0.55:
            action_scores['Sneaking'] += 0.6
        if num_objects <= 2:
            action_scores['Peaking'] += 0.5

    # Fallback if nothing suspicious detected
    if all(score == 0 for score in action_scores.values()):
        action_scores['Normal'] = 0.8

    # Normalize scores (to sum to 1.0)
    total = sum(action_scores.values())
    if total > 0:
        for key in action_scores:
            action_scores[key] /= total

    return action_scores

# ------------------ Detection Function ------------------
def detect_action(image_path):
    results = model.predict(source=image_path, conf=0.4, iou=0.5, save=False, verbose=False)
    result = results[0]

    detections = [
        (model.names[int(cls)], float(conf))
        for cls, conf in zip(result.boxes.cls, result.boxes.conf)
    ]

    # Generate visualization
    annotated_image = result.plot()
    action_scores = classify_action(detections)

    return annotated_image, action_scores

# ------------------ Streamlit UI ------------------
st.title("🔍 Suspicious Activity Detection")
st.markdown("Upload an image to detect if someone is **Stealing**, **Sneaking**, **Peaking**, or just acting **Normal**.")

uploaded_file = st.file_uploader("📤 Upload an image", type=["jpg", "jpeg", "png"])

if uploaded_file:
    image = Image.open(uploaded_file).convert("RGB")
    st.image(image, caption="Uploaded Image", use_column_width=True)

    temp_path = "/tmp/uploaded.jpg"
    image.save(temp_path)

    with st.spinner("🕵️ Detecting suspicious activity..."):
        detected_image, action_scores = detect_action(temp_path)

    # Show results
    st.image(detected_image, caption="🔍 Detected Image", use_column_width=True)

    st.subheader("📊 Action Confidence Scores")
    for action, score in action_scores.items():
        st.write(f"**{action}**: {score:.2%}")

    most_likely_action = max(action_scores.items(), key=lambda x: x[1])
    st.success(f"🎯 **Predicted Action**: {most_likely_action[0]} ({most_likely_action[1]:.2%} confidence)")