import streamlit as st
from ultralytics import YOLO
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from datetime import datetime
import random
from sklearn.linear_model import LinearRegression

st.set_page_config(page_title="Parking Traffic Predictor", layout="wide")

st.title("🚗 AI-Based Parking Space Traffic Predictor (Simulation Only)")
st.markdown("Upload a **top-view parking image**, and the system will automatically detect cars, count them, and predict parking trends — no sensors needed!")

# Load YOLOv8 model (pretrained)
@st.cache_resource
def load_model():
    return YOLO("yolov8n.pt")  # small model for fast inference

model = load_model()

# Upload image
uploaded_image = st.file_uploader("Upload a parking lot image", type=["jpg", "jpeg", "png"])

if uploaded_image:
    st.image(uploaded_image, caption="Uploaded Image", use_container_width=True)

    with st.spinner("Detecting cars..."):
        results = model.predict(uploaded_image, conf=0.4)
        boxes = results[0].boxes
        detections = 0

        # Count cars only (class 2 = car in COCO dataset)
        for c in boxes.cls:
            if int(c) == 2:
                detections += 1

        st.success(f"✅ Cars Detected: {detections}")

        # Display YOLO annotated image
        annotated_img = results[0].plot()
        st.image(annotated_img, caption="Detected Cars", use_container_width=True)

        # Simulated data for prediction
        st.markdown("### 📈 Parking Traffic Prediction")

        # Create fake historical data (simulating last 10 time intervals)
        timestamps = pd.date_range(datetime.now().replace(minute=0, second=0), periods=10, freq='-1H')
        car_counts = np.clip([detections + random.randint(-5, 5) for _ in range(10)], 0, None)
        df = pd.DataFrame({"Time": timestamps, "Cars": car_counts})
        df.sort_values(by="Time", inplace=True)

        # Train a simple model
        X = np.arange(len(df)).reshape(-1, 1)
        y = df["Cars"].values
        model_lr = LinearRegression().fit(X, y)
        next_X = np.array([[len(df)]])
        predicted_cars = int(model_lr.predict(next_X))

        # Plot trend
        fig, ax = plt.subplots()
        ax.plot(df["Time"], df["Cars"], marker="o", label="Past Traffic")
        ax.plot([df["Time"].iloc[-1] + pd.Timedelta(hours=1)], [predicted_cars],
                marker="x", color="red", label="Predicted Next Hour")
        ax.set_xlabel("Time")
        ax.set_ylabel("Number of Cars")
        ax.set_title("Parking Traffic Trend")
        ax.legend()
        st.pyplot(fig)

        # Prediction message
        if predicted_cars < detections:
            st.info(f"🟢 Parking availability is **likely to increase** next hour. (Predicted {predicted_cars} cars)")
        else:
            st.warning(f"🔴 Parking might get **busier** next hour. (Predicted {predicted_cars} cars)")
else:
    st.info("👆 Upload an image of a parking lot (top view) to get started.")