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

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

# Load YOLOv11 model
@st.cache_resource
def load_model():
    return YOLO("yolo11l .pt")  # Ensure model filename matches

model = load_model()

# ------------------ Improved Action Classification Logic ------------------
def classify_action(detections):
    """
    Classify activity based on object types, count, and confidence.
    """
    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
    has_bag = any(obj in objects for obj in ['handbag', 'backpack'])
    few_objects = len(set(objects)) <= 2
    mostly_person = objects.count('person') >= len(objects) * 0.6 if objects else False
    max_conf = max(confidences) if confidences else 0.0

    # Decision tree for classification
    if has_person:
        if has_bag and len(objects) >= 3:
            action_scores['Stealing'] += 1.0
        elif max_conf < 0.55 and few_objects:
            action_scores['Sneaking'] += 1.0
        elif mostly_person and few_objects and max_conf >= 0.55:
            action_scores['Peaking'] += 1.0
        else:
            action_scores['Normal'] += 1.0
    else:
        action_scores['Normal'] += 1.0

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

    return action_scores

# ------------------ Detection Function ------------------
def detect_action(image_path):
    results = model.predict(source=image_path, conf=0.35, 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)
    ]

    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 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)

    st.image(detected_image, caption="🔍 Detection Results", use_column_width=True)

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

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