"""
Segmentação Panóptica de Cenas Urbanas — Araguatins/TO
Projeto PDI — Antonio Agacy & Maria Divina · IFTO · 2026
HuggingFace Spaces · Gradio · OneFormer (fine-tuned)
"""

import io
import time

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

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches

from transformers import (
    OneFormerProcessor,
    OneFormerForUniversalSegmentation,
)

# ─── Configuração ────────────────────────────────────────────
MODEL_NAME = "Agacy/PDI"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# ─── Carrega modelo uma vez ──────────────────────────────────
print(f"Carregando OneFormer ({MODEL_NAME}) no dispositivo: {DEVICE}")
processor = OneFormerProcessor.from_pretrained(MODEL_NAME)
model = OneFormerForUniversalSegmentation.from_pretrained(
    MODEL_NAME, is_training=False
)
model = model.to(DEVICE).eval()
ID2LABEL = {int(k): v for k, v in model.config.id2label.items()}
NUM_CLASSES = len(ID2LABEL)
print(f"✅ Modelo pronto! {NUM_CLASSES} classes: {list(ID2LABEL.values())}")

# ─── Paleta FIXA POR CLASSE (cor consistente entre imagens) ──
PALETTE = [
    [228, 26, 28], [55, 126, 184], [77, 175, 74], [152, 78, 163],
    [255, 127, 0], [255, 255, 51], [166, 86, 40], [247, 129, 191],
    [153, 153, 153], [0, 206, 209], [124, 252, 0], [138, 43, 226],
    [255, 99, 71], [70, 130, 180], [240, 230, 140], [219, 112, 147],
]


def cor_da_classe(label_id: int) -> np.ndarray:
    return np.array(PALETTE[label_id % len(PALETTE)], dtype=float)


# Nomes em português para exibição
PT_LABELS = {
    "building": "edifício", "bus": "ônibus", "car": "carro",
    "motorcycle": "moto", "person": "pessoa", "river": "rio",
    "road": "rua", "sidewalk": "calçada", "sky": "céu",
    "terrain": "terreno", "truck": "caminhão", "vegetation": "vegetação",
}


def nome_pt(label: str) -> str:
    return PT_LABELS.get(label, label)


# ─── Função principal ────────────────────────────────────────
def segmentar(imagem_pil, tarefa: str, opacidade: float, progress=gr.Progress()):
    if imagem_pil is None:
        return None, "⚠️ Envie uma imagem primeiro."

    t0 = time.time()
    progress(0.1, desc="Pré-processando…")

    task_map = {
        "🔀 Panóptico (things + stuff)": "panoptic",
        "🎨 Semântico (regiões)": "semantic",
        "📦 Instâncias (objetos)": "instance",
    }
    task_token = task_map.get(tarefa, "panoptic")

    imagem_pil = imagem_pil.convert("RGB")

    # Reduz imagens muito grandes (CPU do Space agradece)
    MAX_SIDE = 1024
    if max(imagem_pil.size) > MAX_SIDE:
        ratio = MAX_SIDE / max(imagem_pil.size)
        novo = (int(imagem_pil.width * ratio), int(imagem_pil.height * ratio))
        imagem_pil = imagem_pil.resize(novo, Image.BILINEAR)

    inputs = processor(
        images=imagem_pil, task_inputs=[task_token], return_tensors="pt"
    )
    inputs = {
        k: v.to(DEVICE) if isinstance(v, torch.Tensor) else v
        for k, v in inputs.items()
    }

    progress(0.3, desc="Rodando o OneFormer… (CPU pode levar ~1 min)")
    with torch.no_grad():
        outputs = model(**inputs)

    progress(0.8, desc="Gerando máscaras…")
    img_size = [imagem_pil.size[::-1]]  # (H, W)
    img_np = np.array(imagem_pil)
    overlay = img_np.copy().astype(float)
    total_px = img_np.shape[0] * img_np.shape[1]

    handles_por_classe = {}
    linhas_resumo = []

    if task_token == "panoptic":
        resultado = processor.post_process_panoptic_segmentation(
            outputs, target_sizes=img_size
        )[0]
        seg_map = resultado["segmentation"].cpu().numpy()
        seg_info = resultado["segments_info"]

        for seg in seg_info:
            cls = seg.get("label_id", 0)
            cor = cor_da_classe(cls)
            mask = seg_map == seg["id"]
            overlay[mask] = overlay[mask] * (1 - opacidade) + cor * opacidade

            label = nome_pt(ID2LABEL.get(cls, "?"))
            pct = 100 * mask.sum() / total_px
            score = seg.get("score", 0)
            linhas_resumo.append(
                f"| {label} | {pct:.1f}% | {score:.0%} |"
            )
            if cls not in handles_por_classe:
                handles_por_classe[cls] = mpatches.Patch(
                    color=cor / 255.0, label=label
                )
        n_seg = len(seg_info)

    elif task_token == "semantic":
        resultado = processor.post_process_semantic_segmentation(
            outputs, target_sizes=img_size
        )[0]
        seg_map = resultado.cpu().numpy()
        ids_presentes = np.unique(seg_map)

        for cid in ids_presentes:
            cls = int(cid)
            cor = cor_da_classe(cls)
            mask = seg_map == cid
            overlay[mask] = overlay[mask] * (1 - opacidade) + cor * opacidade

            label = nome_pt(ID2LABEL.get(cls, "?"))
            pct = 100 * mask.sum() / total_px
            linhas_resumo.append(f"| {label} | {pct:.1f}% | — |")
            handles_por_classe[cls] = mpatches.Patch(
                color=cor / 255.0, label=label
            )
        n_seg = len(ids_presentes)

    else:  # instance
        resultado = processor.post_process_instance_segmentation(
            outputs, target_sizes=img_size
        )[0]
        seg_map = resultado["segmentation"].cpu().numpy()
        seg_info = resultado["segments_info"]

        for i, seg in enumerate(seg_info):
            cls = seg.get("label_id", 0)
            cor = cor_da_classe(cls)
            mask = seg_map == seg["id"]
            overlay[mask] = overlay[mask] * (1 - opacidade) + cor * opacidade

            label = nome_pt(ID2LABEL.get(cls, "?"))
            pct = 100 * mask.sum() / total_px
            score = seg.get("score", 0)
            linhas_resumo.append(
                f"| {label} #{i + 1} | {pct:.1f}% | {score:.0%} |"
            )
            if cls not in handles_por_classe:
                handles_por_classe[cls] = mpatches.Patch(
                    color=cor / 255.0, label=label
                )
        n_seg = len(seg_info)

    # ─── Figura lado a lado ──────────────────────────────────
    fig, axes = plt.subplots(1, 2, figsize=(14, 6))
    axes[0].imshow(img_np)
    axes[0].set_title("Original", fontsize=13)
    axes[0].axis("off")

    axes[1].imshow(overlay.astype(np.uint8))
    axes[1].set_title(
        f"Segmentação — modo {task_token}", fontsize=13
    )
    axes[1].axis("off")

    handles = list(handles_por_classe.values())
    if handles:
        axes[1].legend(
            handles=handles,
            loc="upper left",
            bbox_to_anchor=(1.01, 1.0),
            fontsize=9,
            framealpha=0.9,
        )

    plt.tight_layout()
    buf = io.BytesIO()
    plt.savefig(buf, format="png", dpi=120, bbox_inches="tight")
    plt.close(fig)
    buf.seek(0)
    img_resultado = Image.open(buf).copy()
    buf.close()

    # ─── Resumo em tabela Markdown ──────────────────────────
    dt = time.time() - t0
    resumo = (
        f"### {n_seg} segmento(s) · modo **{task_token}** · {dt:.1f}s\n\n"
        "| Classe | Área da imagem | Confiança |\n"
        "|---|---|---|\n"
        + "\n".join(linhas_resumo[:25])
    )
    if len(linhas_resumo) > 25:
        resumo += f"\n\n*…e mais {len(linhas_resumo) - 25} segmento(s).*"

    return img_resultado, resumo


# ─── Interface Gradio ────────────────────────────────────────
DESCRICAO = """
# 🏙️ Segmentação Panóptica de Cenas Urbanas — Araguatins/TO

**Projeto PDI — Antonio Agacy Silva Lima Júnior & Maria Divina** · Instituto Federal do Tocantins (IFTO) · 2026

Modelo [OneFormer](https://arxiv.org/abs/2211.06220) (Jain et al., CVPR 2023), pré-treinado no **Cityscapes**
e **ajustado (fine-tuned) com fotos reais de Araguatins-TO** anotadas pela equipe — 12 classes,
incluindo cenas do rio Araguaia, ruas de bloquete e vegetação do cerrado.

⏱️ *Rodando em CPU gratuita: a inferência leva de 30 a 90 segundos.*
"""

RODAPE = """
---
**Classes do modelo:** edifício · ônibus · carro · moto · pessoa · rio · rua · calçada · céu · terreno · caminhão · vegetação

**Pipeline:** anotação no Roboflow → fine-tuning no Kaggle (T4) → publicação no 🤗 Hub → demo neste Space

**Referências:** Kirillov et al. (2019) — *Panoptic Segmentation* · Jain et al. (2023) — *OneFormer* ·
Ren et al. (2024) · Ravi et al. (2024)
"""

with gr.Blocks(title="Segmentação Panóptica — PDI IFTO") as demo:
    gr.Markdown(DESCRICAO)

    with gr.Row():
        with gr.Column(scale=1):
            img_input = gr.Image(
                type="pil",
                label="📷 Imagem de entrada",
                sources=["upload", "clipboard", "webcam"],
            )
            tarefa = gr.Radio(
                choices=[
                    "🔀 Panóptico (things + stuff)",
                    "🎨 Semântico (regiões)",
                    "📦 Instâncias (objetos)",
                ],
                value="🔀 Panóptico (things + stuff)",
                label="Modo de segmentação",
            )
            opacidade = gr.Slider(
                minimum=0.3, maximum=0.9, value=0.55, step=0.05,
                label="Opacidade das máscaras",
            )
            btn = gr.Button("🚀 Segmentar", variant="primary", size="lg")

            gr.Markdown(
                "**Como usar:** envie uma foto de rua/trânsito, "
                "escolha o modo e clique em Segmentar."
            )

        with gr.Column(scale=2):
            img_output = gr.Image(type="pil", label="Resultado")
            texto_output = gr.Markdown()

    # Exemplos prontos (adicione os arquivos na pasta examples/ do Space)
    gr.Examples(
        examples=[
            ["examples/ifto_estacionamento.jpg"],
            ["examples/rio_araguaia.jpg"],
            ["examples/rua_centro.jpg"],
        ],
        inputs=[img_input],
        label="📁 Exemplos de Araguatins (clique para carregar)",
    )

    btn.click(
        fn=segmentar,
        inputs=[img_input, tarefa, opacidade],
        outputs=[img_output, texto_output],
    )

    gr.Markdown(RODAPE)

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Soft(primary_hue="violet"))