"""Hugging Face Space for street-scene detection and segmentation."""

from __future__ import annotations

import os
import time
import uuid
from functools import lru_cache
from pathlib import Path

import gradio as gr
import numpy as np
import torch
from PIL import Image
from transformers import AutoImageProcessor, SegformerForSemanticSegmentation

try:
    import spaces
except ImportError:
    # `spaces` is injected by the ZeroGPU runtime. Keep local/CPU execution valid.
    class _SpacesFallback:
        @staticmethod
        def GPU(*_args, **_kwargs):
            def decorator(function):
                return function

            return decorator

    spaces = _SpacesFallback()

from segmentation_utils import (
    build_class_table,
    render_segmentation,
    resize_for_output,
    write_class_csv,
)
from detection_utils import (
    build_detection_summary,
    build_detection_table,
    build_street_indicators,
    render_detection,
    write_detection_csv,
)

SEGMENTATION_MODEL_ID = "nvidia/segformer-b0-finetuned-cityscapes-1024-1024"
DETECTION_MODEL_ID = "yolo26s.pt"
OUTPUT_ROOT = Path("/tmp/street-scene-vision")
SAMPLE_ROOT = (
    "https://raw.githubusercontent.com/"
    "LabMingzeChen/HNIVision/main/space/examples"
)
SAMPLE_IMAGES = [f"{SAMPLE_ROOT}/ubc-campus-main-mall.jpeg"]


@lru_cache(maxsize=1)
def load_model():
    """Download once per container, then reuse the processor and model."""
    torch.set_num_threads(max(1, min(4, os.cpu_count() or 1)))
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    processor = AutoImageProcessor.from_pretrained(SEGMENTATION_MODEL_ID)
    model = (
        SegformerForSemanticSegmentation.from_pretrained(SEGMENTATION_MODEL_ID)
        .to(device)
        .eval()
    )
    id2label = {int(key): value for key, value in model.config.id2label.items()}
    return processor, model, id2label, device


@lru_cache(maxsize=1)
def load_detector():
    """Download YOLO26-s once per container and reuse it."""
    from ultralytics import YOLO

    return YOLO(DETECTION_MODEL_ID)


def _format_detection_indicators(
    detections: list[dict[str, object]],
) -> str:
    indicators = build_street_indicators(detections)
    average_confidence = (
        sum(float(item["confidence"]) for item in detections) / len(detections)
        if detections
        else 0.0
    )
    return f"""
### Detection-based street indicators

| Indicator | Visible count |
|---|---:|
| People | {indicators['people']} |
| Active-mobility objects (`person` + `bicycle`) | {indicators['active_mobility']} |
| Motor vehicles | {indicators['motor_vehicles']} |
| All transport objects | {indicators['all_transport']} |
| All detected objects | {len(detections)} |

Average detection confidence: **{average_confidence:.2f}**

> Counts describe visible COCO detections in this image. They are not traffic-flow,
> occupancy, accessibility, or safety measurements.
"""


@spaces.GPU(duration=90)
def detect_street_objects(
    image: Image.Image | None,
    confidence_threshold: float,
):
    """Run YOLO object detection and return visual and tabular outputs."""
    if image is None:
        raise gr.Error("Please upload a street-scene image first.")

    started_at = time.perf_counter()
    prepared_image = resize_for_output(image)
    device = "cuda" if torch.cuda.is_available() else "cpu"

    try:
        detector = load_detector()
        predictions = detector.predict(
            source=np.asarray(prepared_image),
            conf=float(confidence_threshold),
            imgsz=1024,
            device=device,
            max_det=100,
            verbose=False,
        )
        prediction = predictions[0]
        names = prediction.names
        detections: list[dict[str, object]] = []
        if prediction.boxes is not None:
            coordinates = prediction.boxes.xyxy.detach().cpu().tolist()
            confidences = prediction.boxes.conf.detach().cpu().tolist()
            class_ids = prediction.boxes.cls.detach().cpu().tolist()
            for coordinates_row, confidence, class_id_value in zip(
                coordinates,
                confidences,
                class_ids,
            ):
                class_id = int(class_id_value)
                detections.append(
                    {
                        "class_id": class_id,
                        "class_name": str(names[class_id]),
                        "confidence": float(confidence),
                        "x1": float(coordinates_row[0]),
                        "y1": float(coordinates_row[1]),
                        "x2": float(coordinates_row[2]),
                        "y2": float(coordinates_row[3]),
                    }
                )
    except Exception as exc:
        raise gr.Error(
            f"Object detection failed: {type(exc).__name__}: {exc}"
        ) from exc

    overlay = render_detection(prepared_image, detections)
    summary_rows = build_detection_summary(detections)
    detection_rows = build_detection_table(detections)

    output_dir = OUTPUT_ROOT / uuid.uuid4().hex
    output_dir.mkdir(parents=True, exist_ok=True)
    overlay_path = output_dir / "street_object_detection_overlay.png"
    csv_path = output_dir / "street_object_detections.csv"
    overlay.save(overlay_path)
    write_detection_csv(csv_path, detection_rows)

    elapsed = time.perf_counter() - started_at
    visible_classes = len(summary_rows)
    status = (
        f"Done · {prepared_image.width}×{prepared_image.height} · "
        f"{len(detections)} objects · {visible_classes} COCO classes · "
        f"{elapsed:.1f}s · device={device}"
    )
    return (
        overlay,
        summary_rows,
        detection_rows,
        [str(overlay_path), str(csv_path)],
        _format_detection_indicators(detections),
        status,
    )


@spaces.GPU(duration=90)
def segment_street_scene(
    image: Image.Image | None,
    opacity: float,
    min_share_percent: float,
):
    """Run semantic segmentation and return visual, tabular, and raw outputs."""
    if image is None:
        raise gr.Error("Please upload a street-scene image first.")

    started_at = time.perf_counter()
    prepared_image = resize_for_output(image)

    try:
        processor, model, id2label, device = load_model()
        inputs = processor(images=prepared_image, return_tensors="pt")
        inputs = {name: tensor.to(device) for name, tensor in inputs.items()}

        with torch.inference_mode():
            outputs = model(**inputs)

        target_size = (prepared_image.height, prepared_image.width)
        class_map_tensor = processor.post_process_semantic_segmentation(
            outputs,
            target_sizes=[target_size],
        )[0]
        class_map = class_map_tensor.cpu().numpy().astype(np.uint8)
    except Exception as exc:
        raise gr.Error(
            f"Segmentation failed: {type(exc).__name__}: {exc}"
        ) from exc

    overlay, color_mask = render_segmentation(
        prepared_image,
        class_map,
        id2label,
        float(opacity),
    )
    rows = build_class_table(class_map, id2label, float(min_share_percent))

    output_dir = OUTPUT_ROOT / uuid.uuid4().hex
    output_dir.mkdir(parents=True, exist_ok=True)
    overlay_path = output_dir / "street_segmentation_overlay.png"
    mask_path = output_dir / "street_segmentation_color_mask.png"
    class_ids_path = output_dir / "street_segmentation_class_ids.png"
    csv_path = output_dir / "street_segmentation_classes.csv"

    overlay.save(overlay_path)
    color_mask.save(mask_path)
    Image.fromarray(class_map).save(class_ids_path)
    write_class_csv(csv_path, rows)

    elapsed = time.perf_counter() - started_at
    visible_classes = len(np.unique(class_map))
    status = (
        f"Done · {prepared_image.width}×{prepared_image.height} · "
        f"{visible_classes} street-scene classes · {elapsed:.1f}s · "
        f"device={device.type}"
    )

    return (
        overlay,
        color_mask,
        rows,
        [str(overlay_path), str(mask_path), str(class_ids_path), str(csv_path)],
        status,
    )


CSS = """
.gradio-container {max-width: 1260px !important;}
.hero {text-align: center; margin: 0 auto 1rem;}
.hero h1 {font-size: 2.1rem; margin-bottom: .3rem;}
.muted {color: #64748b;}
.project-links {display: flex; justify-content: center; gap: .55rem; flex-wrap: wrap; margin-top: .75rem;}
.project-link {
  display: inline-block; padding: .42rem .78rem; border: 1px solid #d7deea;
  border-radius: 999px; color: inherit !important; text-decoration: none !important;
  background: white; font-size: .92rem; font-weight: 600;
}
.project-link:hover {border-color: #6366f1; box-shadow: 0 2px 8px rgba(99, 102, 241, .12);}
.guide-grid {display: grid; grid-template-columns: repeat(3, minmax(0, 1fr)); gap: .8rem; margin: .8rem 0;}
.guide-card {border: 1px solid #e2e8f0; border-radius: 12px; padding: .85rem 1rem; background: rgba(255,255,255,.55);}
.guide-card h3 {margin: 0 0 .35rem; font-size: 1rem;}
.guide-card p {margin: 0; color: #475569; font-size: .93rem; line-height: 1.45;}
@media (max-width: 760px) {.guide-grid {grid-template-columns: 1fr;}}
"""


with gr.Blocks(title="Street Scene Vision Toolkit", theme=gr.themes.Soft(), css=CSS) as demo:
    gr.Markdown(
        """
        <div class="hero">
          <h1>🚦 Street Scene Vision Toolkit</h1>
          <p>Map every pixel with semantic segmentation, then detect individual objects with bounding boxes.</p>
          <p class="muted">YOLO26-s · COCO 80 objects · SegFormer-B0 · 19 Cityscapes classes · no API key required</p>
          <div class="project-links">
            <a class="project-link" href="https://huggingface.co/spaces/Mingze/StreetSceneSegmentation" target="_blank">🤗 Hugging Face Space</a>
            <a class="project-link" href="https://docs.ultralytics.com/models/yolo26/" target="_blank">📦 YOLO26</a>
            <a class="project-link" href="https://huggingface.co/nvidia/segformer-b0-finetuned-cityscapes-1024-1024" target="_blank">🎨 SegFormer</a>
            <a class="project-link" href="https://github.com/LabMingzeChen/StreetSceneSegmentation" target="_blank">⭐ GitHub source</a>
          </div>
        </div>
        """
    )

    with gr.Row(equal_height=False):
        with gr.Column(scale=5):
            image_input = gr.Image(
                type="pil",
                label="Upload a street-scene image",
                height=470,
                sources=["upload", "clipboard", "webcam"],
            )
            if SAMPLE_IMAGES:
                gr.Examples(
                    examples=SAMPLE_IMAGES,
                    inputs=image_input,
                    label="Try the UBC campus street example",
                    examples_per_page=1,
                )
            with gr.Accordion("Segmentation settings", open=True):
                opacity_input = gr.Slider(
                    0.15,
                    0.85,
                    value=0.55,
                    step=0.05,
                    label="Overlay opacity",
                )
                min_share_input = gr.Slider(
                    0.0,
                    5.0,
                    value=0.1,
                    step=0.1,
                    label="Minimum class area shown in table (%)",
                )
            with gr.Accordion("Object-detection settings", open=False):
                confidence_input = gr.Slider(
                    0.05,
                    0.90,
                    value=0.25,
                    step=0.05,
                    label="Minimum detection confidence",
                )
            with gr.Row():
                segmentation_button = gr.Button(
                    "Segment pixels",
                    variant="primary",
                    size="lg",
                )
                detection_button = gr.Button("Detect objects", size="lg")
            clear_button = gr.ClearButton(value="Clear image", components=[image_input])

        with gr.Column(scale=7):
            with gr.Tabs(selected="segmentation") as visual_tabs:
                with gr.Tab("Segmentation overlay", id="segmentation"):
                    overlay_output = gr.Image(label="Segmentation overlay", height=470)
                    segmentation_status = gr.Markdown()
                with gr.Tab("Color mask", id="mask"):
                    mask_output = gr.Image(label="Cityscapes color mask", height=470)
                with gr.Tab("Object detection", id="detection"):
                    detection_output = gr.Image(
                        label="YOLO26-s bounding boxes",
                        height=470,
                    )
                    detection_status = gr.Markdown()

    with gr.Tabs(selected="segmentation-results") as result_tabs:
        with gr.Tab("Segmentation results", id="segmentation-results"):
            table_output = gr.Dataframe(
                headers=["Class ID", "Class", "Pixels", "Area share (%)", "Color"],
                datatype=["number", "str", "number", "number", "str"],
                label="Detected street-scene classes",
                interactive=False,
                wrap=True,
            )
            files_output = gr.File(
                label="Download segmentation overlay, color mask, class IDs, and CSV",
                file_count="multiple",
            )

        with gr.Tab("Detection results", id="detection-results"):
            detection_indicators = gr.Markdown()
            detection_summary = gr.Dataframe(
                headers=["Class", "Count", "Average confidence", "Maximum confidence"],
                datatype=["str", "number", "number", "number"],
                label="Detected object classes",
                interactive=False,
                wrap=True,
            )
            with gr.Accordion("Detailed bounding-box coordinates", open=False):
                detection_table = gr.Dataframe(
                    headers=["Object ID", "Class", "Confidence", "x1", "y1", "x2", "y2"],
                    datatype=["number", "str", "number", "number", "number", "number", "number"],
                    label="Individual detections",
                    interactive=False,
                    wrap=True,
                )
            detection_files = gr.File(
                label="Download detection overlay and bounding-box CSV",
                file_count="multiple",
            )

    gr.Markdown(
        """
        ## How to use the app

        <div class="guide-grid">
          <div class="guide-card"><h3>1 · Choose an image</h3><p>Upload, paste, use a webcam, or select the UBC campus example.</p></div>
          <div class="guide-card"><h3>2 · Segment, then detect</h3><p>Start with SegFormer semantic segmentation, then optionally run YOLO object detection on the same image.</p></div>
          <div class="guide-card"><h3>3 · Explore and download</h3><p>Compare boxes, overlays, masks, counts, pixel shares, coordinates, and reusable CSV outputs.</p></div>
        </div>

        ## What the results mean

        - **Object detection** finds separate COCO objects, draws bounding boxes, and reports a confidence score for each detection.
        - **Detection indicators** summarize visible people, active-mobility objects, and transport objects. They are transparent image counts, not traffic-flow estimates.
        - **Segmentation overlay** blends the Cityscapes prediction with the original photograph. White lines mark class boundaries.
        - **Color mask and area share** show pixel-level scene composition. Area share describes visual coverage, not physical land area.
        - **Downloadable data** include bounding-box coordinates, class-ID pixels, overlays, masks, and CSV summaries.

        | Scene layer | Segmentation classes |
        |---|---|
        | Travel surfaces | road, sidewalk |
        | Built environment | building, wall, fence, pole, traffic light, traffic sign |
        | Nature and sky | vegetation, terrain, sky |
        | People | person, rider |
        | Transport | car, truck, bus, train, motorcycle, bicycle |

        ## Classroom and research ideas

        - Compare detected people, bicycles, and motor vehicles across several street images.
        - Compare what bounding boxes reveal with what pixel-level segmentation reveals.
        - Discuss missed objects, false positives, confidence thresholds, and segmentation boundary errors.
        - Export both CSV files and build object-count and class-coverage charts.
        - Compare the same location across seasons, weather conditions, or camera viewpoints.

        > **Important:** predictions are model estimates, not ground truth. COCO detection is limited to its trained object vocabulary, while Cityscapes segmentation is specialized for road-driving imagery. Do not use either output for safety-critical decisions, surveillance, or identifying individuals.

        [Read the YOLO26 documentation](https://docs.ultralytics.com/models/yolo26/) ·
        [Explore the COCO dataset](https://cocodataset.org/) ·
        [Read the SegFormer paper](https://arxiv.org/abs/2105.15203) ·
        [Explore the Cityscapes dataset](https://www.cityscapes-dataset.com/) ·
        [View the source on GitHub](https://github.com/LabMingzeChen/StreetSceneSegmentation)
        """
    )

    detection_event = detection_button.click(
        fn=detect_street_objects,
        inputs=[image_input, confidence_input],
        outputs=[
            detection_output,
            detection_summary,
            detection_table,
            detection_files,
            detection_indicators,
            detection_status,
        ],
        api_name="detect",
        scroll_to_output=True,
    )
    detection_event.then(
        fn=lambda: (
            gr.Tabs(selected="detection"),
            gr.Tabs(selected="detection-results"),
        ),
        outputs=[visual_tabs, result_tabs],
        queue=False,
        api_name=False,
    )

    segmentation_event = segmentation_button.click(
        fn=segment_street_scene,
        inputs=[image_input, opacity_input, min_share_input],
        outputs=[
            overlay_output,
            mask_output,
            table_output,
            files_output,
            segmentation_status,
        ],
        api_name="segment",
        scroll_to_output=True,
    )
    segmentation_event.then(
        fn=lambda: (
            gr.Tabs(selected="segmentation"),
            gr.Tabs(selected="segmentation-results"),
        ),
        outputs=[visual_tabs, result_tabs],
        queue=False,
        api_name=False,
    )


if __name__ == "__main__":
    demo.queue(max_size=8, default_concurrency_limit=1).launch()