"""UD Campus Parking Lab — peak-time lot occupancy. Campus operations only."""

from __future__ import annotations

from functools import lru_cache

import gradio as gr
import numpy as np
from PIL import Image

BANNER = (
    "**UD Campus Parking Lab.** Peak-time lot occupancy from a photo you took. "
    "Counts vehicles. Estimates how full the lot is. "
    "Does **not** read plates, identify people, or store a watchlist."
)
FOOTER_IDLE = "waiting for a lot photo"
YOLO_ID = "yolov8n.pt"
VLM_ID = "HuggingFaceTB/SmolVLM-500M-Instruct"
VEHICLE = {"car", "truck", "bus", "motorcycle"}
LOTS = [
    "Academic City — student lot",
    "Academic City — staff lot",
    "Visitor / short stay",
    "Overflow / event parking",
    "Other campus lot",
]
WINDOWS = [
    "Morning arrival (07:30–09:00)",
    "Class change",
    "Lunch",
    "Evening leave",
    "Weekend / off-peak",
]


def device() -> str:
    try:
        import torch

        return "cuda" if torch.cuda.is_available() else "cpu"
    except Exception:
        return "cpu"


@lru_cache(maxsize=1)
def yolo():
    from ultralytics import YOLO

    return YOLO(YOLO_ID)


@lru_cache(maxsize=1)
def vlm():
    import torch
    from transformers import AutoModelForVision2Seq, AutoProcessor

    proc = AutoProcessor.from_pretrained(VLM_ID)
    model = AutoModelForVision2Seq.from_pretrained(
        VLM_ID,
        torch_dtype=torch.float16 if device() == "cuda" else torch.float32,
    ).to(device())
    model.eval()
    return proc, model


def count_vehicles(image: Image.Image):
    result = yolo().predict(image, verbose=False)[0]
    plotted = result.plot()[:, :, ::-1]
    counts: dict[str, int] = {}
    if result.boxes is not None:
        names = result.names
        for b in result.boxes:
            name = names.get(int(b.cls.item()), "")
            if name in VEHICLE:
                counts[name] = counts.get(name, 0) + 1
    total = sum(counts.values())
    return Image.fromarray(plotted), counts, total


def lot_note(image: Image.Image, lot: str, window: str) -> str:
    import torch

    proc, model = vlm()
    prompt = (
        f"This is a photo of {lot} at {window}. "
        "In two sentences, say whether the lot looks empty, mixed, or packed. "
        "Mention parked cars and open stalls if you see them. "
        "Do not read license plates. Do not name or describe faces."
    )
    messages = [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": prompt}]}]
    text = proc.apply_chat_template(messages, add_generation_prompt=True)
    inputs = proc(text=text, images=[image], return_tensors="pt")
    inputs = {k: v.to(model.device) if hasattr(v, "to") else v for k, v in inputs.items()}
    with torch.inference_mode():
        out = model.generate(**inputs, max_new_tokens=80)
    return proc.batch_decode(out, skip_special_tokens=True)[0]


def peak_label(pct: float, window: str) -> str:
    if pct >= 90:
        return "STRESS — treat as full; send overflow"
    if pct >= 75:
        return "PEAK — few stalls left"
    if pct >= 40:
        return "MIXED"
    if "off-peak" in window.lower() or "Weekend" in window:
        return "QUIET (matches off-peak window)"
    return "QUIET — not peak yet"


def run(image, lot, window, capacity):
    if image is None:
        return BANNER, None, "Take or upload a photo of a campus lot.", FOOTER_IDLE
    if not isinstance(image, Image.Image):
        image = Image.fromarray(np.asarray(image)).convert("RGB")
    else:
        image = image.convert("RGB")
    cap = max(1, int(capacity or 40))
    vis, counts, total = count_vehicles(image)
    pct = min(100.0, 100.0 * total / cap)
    try:
        note = lot_note(image, lot, window)
    except Exception as e:
        note = f"Caption skipped ({type(e).__name__}). Counts still ran."

    breakdown = ", ".join(f"{k} {v}" for k, v in sorted(counts.items())) or "none"
    status = peak_label(pct, window)
    notes = (
        f"Lot: {lot}\n"
        f"Window: {window}\n"
        f"Vehicles counted: {total} ({breakdown})\n"
        f"Capacity you set: {cap} stalls\n"
        f"Estimated occupancy: {pct:.0f}%\n"
        f"Peak flag: {status}\n\n"
        f"Scene note:\n{note}\n\n"
        "This is a still-photo estimate, not a live sensor feed. "
        "Double-count and occlusion will skew the number."
    )
    footer = (
        f"yolo: {YOLO_ID} | vlm: {VLM_ID} | device: {device()} | "
        f"vehicles {total}/{cap} | {status}"
    )
    return BANNER, vis, notes, footer


with gr.Blocks(title="UD Campus Parking Lab") as demo:
    gr.Markdown("# UD Campus Parking Lab")
    banner = gr.Markdown(BANNER)
    gr.Markdown(
        "University of Dubai campus operations. "
        "Photograph a lot at peak arrival or class change. "
        "Set how many stalls that lot has. The app counts vehicles and flags peak stress."
    )
    with gr.Row():
        lot = gr.Dropdown(LOTS, value=LOTS[0], label="Campus lot")
        window = gr.Dropdown(WINDOWS, value=WINDOWS[0], label="Time window")
        capacity = gr.Number(value=40, precision=0, label="Stall capacity")
    with gr.Row():
        inp = gr.Image(type="pil", label="Lot photo or webcam", sources=["upload", "webcam", "clipboard"])
        out = gr.Image(type="pil", label="Detections (vehicles)")
    notes = gr.Textbox(label="Occupancy report", lines=16)
    footer = gr.Textbox(label="Footer", value=FOOTER_IDLE)
    gr.Button("Estimate occupancy").click(
        fn=run, inputs=[inp, lot, window, capacity], outputs=[banner, out, notes, footer]
    )

if __name__ == "__main__":
    demo.launch()