"""Florence-2 scene tagger (ZeroGPU): proposes class names from an image.

The first implementation used only the Florence ``<OD>`` task, which is too
sparse for street-level panoptic work because object detection misses broad
surface/stuff classes. This endpoint now combines object detection, dense
region captions, and detailed caption text into one open-vocabulary concept
list for SAM3.
"""
import os
from collections import Counter
import re

import gradio as gr
import spaces
import torch
from transformers import AutoProcessor, Florence2ForConditionalGeneration

MODEL_ID = "florence-community/Florence-2-large"

# Built at import on CPU; moved to CUDA inside the @spaces.GPU function.
processor = AutoProcessor.from_pretrained(MODEL_ID)
model = Florence2ForConditionalGeneration.from_pretrained(MODEL_ID)
model.eval()

TASKS = ("<OD>", "<DENSE_REGION_CAPTION>", "<MORE_DETAILED_CAPTION>")
DROP_WORDS = {
    "a", "an", "the", "this", "that", "these", "those", "there", "here",
    "image", "photo", "picture", "view", "scene", "background", "foreground",
    "left", "right", "top", "bottom", "front", "back", "side", "area", "part",
    "visible", "large", "small", "several", "multiple", "many", "some",
    "is", "are", "was", "were", "be", "being", "been", "appears", "appear",
    "of", "in", "it", "its", "to", "for", "by", "as", "at", "from", "into",
    "parked", "parking", "surrounded", "including", "taken", "shining",
    "brightly", "different", "models", "colors", "color", "angle", "high",
    "panoramic", "modern", "curved", "few", "blue", "red", "white",
}
KEEP_PHRASES = {
    "parking lot",
}
DROP_LABELS = {
    "roof", "floor", "floors", "balcony", "balconies", "sun", "cloud", "clouds",
    "sky", "brightly", "angle", "high angle",
}
CAPTION_SPLIT = re.compile(
    r"[,.;:]|\\bwith\\b|\\band\\b|\\bnext to\\b|\\bin front of\\b|\\bbehind\\b|\\bon\\b|\\balong\\b|\\bnear\\b|\\bsurrounded by\\b",
    flags=re.IGNORECASE,
)


def _clean_label(value: str) -> str:
    text = str(value).strip().lower()
    text = re.sub(r"[_/\\-]+", " ", text)
    text = re.sub(r"[^a-z0-9\\s]+", " ", text)
    for phrase in KEEP_PHRASES:
        if phrase in text:
            return phrase
    words = [w for w in text.split() if w and w not in DROP_WORDS]
    if not words:
        return ""
    if len(words) == 1:
        word = words[0]
        if len(word) > 3 and word.endswith("ies"):
            word = word[:-3] + "y"
        elif len(word) > 3 and word.endswith("s") and not word.endswith("ss"):
            word = word[:-1]
        words = [word]
    if len(words) > 5:
        words = words[-5:]
    label = " ".join(words)
    return "" if label in DROP_LABELS else label


def _labels_from_caption(text: str) -> list[str]:
    labels: list[str] = []
    for chunk in CAPTION_SPLIT.split(str(text)):
        clean = _clean_label(chunk)
        if not clean:
            continue
        words = clean.split()
        candidates = [clean]
        if len(words) >= 3:
            candidates.append(" ".join(words[-2:]))
        candidates.append(words[-1])
        for candidate in candidates:
            candidate = _clean_label(candidate)
            if candidate and len(candidate) > 2 and candidate not in labels:
                labels.append(candidate)
                break
    return labels


def _extract_labels(parsed) -> list[str]:
    labels: list[str] = []

    def walk(obj, key_hint: str = ""):
        if isinstance(obj, dict):
            for key, value in obj.items():
                k = str(key).lower()
                if k in {"label", "labels", "caption", "captions", "text", "description", "descriptions"}:
                    walk(value, k)
                else:
                    walk(value, key_hint)
        elif isinstance(obj, (list, tuple)):
            for item in obj:
                walk(item, key_hint)
        elif isinstance(obj, str):
            if key_hint in {"label", "labels"}:
                clean = _clean_label(obj)
                if clean:
                    labels.append(clean)
            else:
                labels.extend(_labels_from_caption(obj))

    walk(parsed)
    return labels


def _run_florence_task(image, task: str, num_beams: int) -> dict:
    device = "cuda"
    model.to(device)
    inputs = processor(text=task, images=image, return_tensors="pt").to(device)
    with torch.no_grad():
        gen = model.generate(**inputs, max_new_tokens=1536, num_beams=int(num_beams))
    text = processor.batch_decode(gen, skip_special_tokens=False)[0]
    parsed = processor.post_process_generation(text, task=task, image_size=image.size)
    return {"task": task, "text": text, "parsed": parsed, "labels": _extract_labels(parsed)}


@spaces.GPU(duration=120)
def api_autotag(image, max_tags, num_beams=3):
    if image is None:
        return {"error": "no image provided"}
    image = image.convert("RGB")
    task_results = [_run_florence_task(image, task, int(num_beams)) for task in TASKS]
    labels = []
    for result in task_results:
        labels.extend(result["labels"])
    counts = Counter(label for label in labels if label)
    tags = [{"name": n, "count": c} for n, c in counts.most_common(int(max_tags))]
    return {"model": MODEL_ID, "tasks": task_results, "tags": tags, "labels": [t["name"] for t in tags]}


with gr.Blocks(title="SAM3 AutoTag") as demo:
    gr.Markdown("# Florence-2 AutoTag\nUpload an image; returns multi-task scene class names for SAM3.")
    with gr.Row():
        inp = gr.Image(type="pil", label="Image")
        out = gr.JSON(label="Tags")
    mt = gr.Slider(1, 50, value=20, step=1, label="Max tags")
    nb = gr.Slider(1, 8, value=3, step=1, label="Beams (higher = more precise)")
    gr.Button("Tag").click(api_autotag, [inp, mt, nb], out, api_name="api_autotag")

if __name__ == "__main__":
    demo.queue().launch(show_error=True)