app.py

"""
Streamlit dashboard for the Smart Parking & Vehicle Threat Detection System.

Tabs:
  1. Live Detection   – run YOLO on webcam or uploaded video
  2. Incident Log     – filterable table of all logged events
  3. Security Chat    – RAG-powered Q&A over incident history
  4. Analytics        – charts: status breakdown, zone activity
"""

import streamlit as st
import cv2
import numpy as np
import pandas as pd
import plotly.express as px
from PIL import Image

from yolo_module.detector import VehicleDetector
from yolo_module.ocr_reader import read_plate
from utils.incident_logger import log_incident, load_incidents, is_plate_flagged
from utils.anomaly_detector import analyse
from rag_module.rag_pipeline import ask
from rag_module.vector_store import ingest_incident

# ─── Page config ─────────────────────────────────────────────────────────────
st.set_page_config(
    page_title="Smart Parking Security",
    page_icon="🅿",
    layout="wide",
)

# ─── Session state ────────────────────────────────────────────────────────────
if "detector" not in st.session_state:
    st.session_state.detector = None
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []

# ─── Sidebar ─────────────────────────────────────────────────────────────────
st.sidebar.title("🅿 Smart Parking")
st.sidebar.markdown("YOLOv8 + RAG Security System")
zone = st.sidebar.selectbox("Active Zone", ["Gate_A", "Gate_B", "Gate_C", "VIP_LOT", "COMPACT_ONLY"])
model_size = st.sidebar.selectbox("YOLO Model", ["yolov8n.pt", "yolov8s.pt", "yolov8m.pt"])

if st.sidebar.button("Load / Reload Model"):
    with st.spinner("Loading YOLOv8..."):
        st.session_state.detector = VehicleDetector(model_size)
    st.sidebar.success("Model loaded!")

# ─── Tabs ─────────────────────────────────────────────────────────────────────
tab1, tab2, tab3, tab4 = st.tabs(["📷 Live Detection", "📋 Incident Log", "💬 Security Chat", "📊 Analytics"])

# ═══════════════════════════════════════════════════════════════════════════════
# TAB 1 – Live Detection
# ═══════════════════════════════════════════════════════════════════════════════
with tab1:
    st.header("Live Vehicle Detection")
    source_type = st.radio("Input source", ["Upload image", "Upload video"], horizontal=True)

    if source_type == "Upload image":
        uploaded = st.file_uploader("Choose an image", type=["jpg", "jpeg", "png"])
        if uploaded:
            img = Image.open(uploaded).convert("RGB")
            frame = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)

            if st.session_state.detector is None:
                st.warning("Load a model from the sidebar first.")
            else:
                with st.spinner("Detecting..."):
                    result = st.session_state.detector.detect_frame(frame)

                col_img, col_info = st.columns([2, 1])
                annotated_rgb = cv2.cvtColor(result["annotated_frame"], cv2.COLOR_BGR2RGB)
                col_img.image(annotated_rgb, width=700)

                col_info.subheader(f"Found {len(result['vehicles'])} vehicle(s)")
                for i, v in enumerate(result["vehicles"]):
                    with col_info.expander(f"Vehicle {i+1}: {v['class']}"):
                        plate = read_plate(v["plate_crop"])
                        status, notes = analyse(plate, v["class"], zone)
                        flagged = is_plate_flagged(plate)

                        st.metric("Plate", plate or "—")
                        st.metric("Confidence", f"{v['confidence']:.0%}")
                        color = {"normal": "🟢", "flagged": "🟡", "unauthorized": "🔴", "anomaly": "🟠"}
                        st.write(f"Status: {color.get(status, '⚪')} {status.upper()}")
                        if flagged:
                            st.error("⚠️ This plate is in the flagged list!")
                        if notes:
                            st.caption(notes)

                        if st.button(f"Log incident #{i+1}"):
                            row = log_incident(plate, v["class"], zone, status, notes)
                            ingest_incident(row)
                            st.success(f"Logged as ID {row['id']}")

    else:  # Video upload
        uploaded_vid = st.file_uploader("Choose a video", type=["mp4", "avi", "mov"])
        if uploaded_vid:
            import tempfile, os
            with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmp:
                tmp.write(uploaded_vid.read())
                tmp_path = tmp.name

            if st.session_state.detector is None:
                st.warning("Load a model from the sidebar first.")
            elif st.button("Start Processing"):
                frame_ph = st.empty()
                info_ph = st.empty()
                for result in st.session_state.detector.detect_video(tmp_path):
                    rgb = cv2.cvtColor(result["annotated_frame"], cv2.COLOR_BGR2RGB)
                    frame_ph.image(rgb, width=700)
                    info_ph.caption(f"Vehicles detected this frame: {len(result['vehicles'])}")
            os.unlink(tmp_path)

# ═══════════════════════════════════════════════════════════════════════════════
# TAB 2 – Incident Log
# ═══════════════════════════════════════════════════════════════════════════════
with tab2:
    st.header("Incident Log")
    df = load_incidents()

    if df.empty:
        st.info("No incidents logged yet. Run detection or seed sample data.")
    else:
        col_f1, col_f2, col_f3 = st.columns(3)
        status_filter = col_f1.multiselect("Status", df["status"].unique(), default=list(df["status"].unique()))
        zone_filter   = col_f2.multiselect("Zone",   df["zone"].unique(),   default=list(df["zone"].unique()))
        plate_search  = col_f3.text_input("Search plate")

        filtered = df[df["status"].isin(status_filter) & df["zone"].isin(zone_filter)]
        if plate_search:
            filtered = filtered[filtered["plate"].str.contains(plate_search.upper(), na=False)]

        st.dataframe(filtered, use_container_width=True, hide_index=True)
        st.caption(f"Showing {len(filtered)} of {len(df)} incidents")

        csv_bytes = filtered.to_csv(index=False).encode()
        st.download_button("Download CSV", csv_bytes, "incidents_export.csv", "text/csv")

# ═══════════════════════════════════════════════════════════════════════════════
# TAB 3 – Security Chat (RAG)
# ═══════════════════════════════════════════════════════════════════════════════
with tab3:
    st.header("Security Assistant")
    st.caption("Ask questions about the incident history using natural language.")

    example_queries = [
        "Was plate KA01XY9999 flagged recently?",
        "How many unauthorized vehicles were detected today?",
        "Which zone had the most incidents?",
        "Show me all anomalies in the last hour",
    ]
    st.markdown("**Example queries:**")
    cols = st.columns(2)
    for i, q in enumerate(example_queries):
        if cols[i % 2].button(q, key=f"eq{i}"):
            st.session_state.chat_history.append({"role": "user", "content": q})

    # Chat input
    user_input = st.chat_input("Ask the security assistant...")
    if user_input:
        st.session_state.chat_history.append({"role": "user", "content": user_input})

    # Render history
    for msg in st.session_state.chat_history:
        with st.chat_message(msg["role"]):
            st.write(msg["content"])
            if "retrieved" in msg:
                with st.expander("Retrieved incidents"):
                    st.dataframe(pd.DataFrame(msg["retrieved"]), use_container_width=True)

    # Generate response for latest user message
    if st.session_state.chat_history and st.session_state.chat_history[-1]["role"] == "user":
        query = st.session_state.chat_history[-1]["content"]
        with st.spinner("Searching incident logs..."):
            result = ask(query)
        response = {"role": "assistant", "content": result["answer"], "retrieved": result["retrieved_docs"]}
        st.session_state.chat_history.append(response)
        with st.chat_message("assistant"):
            st.write(result["answer"])
            if result["retrieved_docs"]:
                with st.expander("Retrieved incidents"):
                    st.dataframe(pd.DataFrame(result["retrieved_docs"]), use_container_width=True)

# ═══════════════════════════════════════════════════════════════════════════════
# TAB 4 – Analytics
# ═══════════════════════════════════════════════════════════════════════════════
with tab4:
    st.header("Analytics")
    df = load_incidents()

    if df.empty:
        st.info("No data yet.")
    else:
        col1, col2, col3, col4 = st.columns(4)
        col1.metric("Total Incidents", len(df))
        col2.metric("Flagged",    len(df[df["status"] == "flagged"]))
        col3.metric("Unauthorized", len(df[df["status"] == "unauthorized"]))
        col4.metric("Anomalies",   len(df[df["status"] == "anomaly"]))

        col_a, col_b = st.columns(2)
        fig1 = px.pie(df, names="status", title="Incidents by Status", hole=0.4)
        col_a.plotly_chart(fig1, use_container_width=True)

        fig2 = px.bar(df.groupby("zone").size().reset_index(name="count"),
                      x="zone", y="count", title="Incidents by Zone", color="count",
                      color_continuous_scale="Blues")
        col_b.plotly_chart(fig2, use_container_width=True)

        df["hour"] = df["timestamp"].dt.floor("H")
        timeline = df.groupby(["hour", "status"]).size().reset_index(name="count")
        fig3 = px.line(timeline, x="hour", y="count", color="status", title="Incident Timeline")
        st.plotly_chart(fig3, use_container_width=True)