import os
import numpy as np
import torch
import torch.nn.functional as F
import gradio as gr
from PIL import Image
from transformers import SegformerForSemanticSegmentation, SegformerImageProcessor

# ============================================================
# Config
# ============================================================
MODEL_ID = "Dinusharg/segformer_environment_1"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# Class colors for your labels
PALETTE = {
    0: (0, 0, 0),        # background
    1: (255, 0, 0),      # footpath
    2: (0, 255, 0),      # grass
    3: (128, 128, 128),  # road
    4: (0, 0, 255),      # water
}

# Optional nicer class names for display
CLASS_NAMES = {
    0: "background",
    1: "footpath",
    2: "grass",
    3: "road",
    4: "water",
}

# ============================================================
# Load model once at startup
# ============================================================
processor = SegformerImageProcessor.from_pretrained(MODEL_ID)
model = SegformerForSemanticSegmentation.from_pretrained(MODEL_ID).to(DEVICE)
model.eval()


def make_color_mask(pred: np.ndarray) -> np.ndarray:
    """
    Convert class-id mask [H, W] into RGB color mask [H, W, 3].
    """
    color_mask = np.zeros((pred.shape[0], pred.shape[1], 3), dtype=np.uint8)
    for class_id, color in PALETTE.items():
        color_mask[pred == class_id] = color
    return color_mask


def make_overlay(image_np: np.ndarray, color_mask: np.ndarray, alpha: float = 0.5) -> np.ndarray:
    """
    Blend original image with color mask.
    """
    image_f = image_np.astype(np.float32)
    mask_f = color_mask.astype(np.float32)
    overlay = (image_f * (1.0 - alpha) + mask_f * alpha).clip(0, 255).astype(np.uint8)
    return overlay


def make_class_summary(pred: np.ndarray) -> str:
    """
    Summarize predicted pixel counts and percentages.
    """
    total_pixels = pred.size
    unique, counts = np.unique(pred, return_counts=True)

    lines = []
    for class_id, count in zip(unique.tolist(), counts.tolist()):
        pct = 100.0 * count / total_pixels
        class_name = CLASS_NAMES.get(class_id, f"class_{class_id}")
        lines.append(f"- {class_id}: {class_name} — {count} pixels ({pct:.2f}%)")

    return "\n".join(lines)


def predict(image: Image.Image):
    """
    Gradio callback.
    Input: PIL image
    Outputs:
      - segmentation mask
      - overlay
      - class summary
    """
    if image is None:
        return None, None, "No image provided."

    image = image.convert("RGB")
    original_np = np.array(image)

    inputs = processor(images=image, return_tensors="pt")
    inputs = {k: v.to(DEVICE) for k, v in inputs.items()}

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

    upsampled_logits = F.interpolate(
        outputs.logits,
        size=image.size[::-1],   # PIL gives (W, H), torch expects (H, W)
        mode="bilinear",
        align_corners=False
    )

    pred = upsampled_logits.argmax(dim=1)[0].detach().cpu().numpy()

    color_mask = make_color_mask(pred)
    overlay = make_overlay(original_np, color_mask, alpha=0.5)
    summary = make_class_summary(pred)

    return Image.fromarray(color_mask), Image.fromarray(overlay), summary


title = "SegFormer Semantic Segmentation Demo"
description = (
    "Upload an outdoor scene image to generate a semantic segmentation mask and overlay. "
    "Classes: background, footpath, grass, road, water."
)

article = """
### Notes
- The model is loaded from the Hugging Face Hub using `MODEL_ID`.
- To use your own model in this Space, set `MODEL_ID` in the Space settings or replace the default value in `app.py`.
- Later you can add your public Space link to the model README.
"""

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload Image"),
    outputs=[
        gr.Image(type="pil", label="Segmentation Mask"),
        gr.Image(type="pil", label="Overlay"),
        gr.Markdown(label="Class Summary"),
    ],
    title=title,
    description=description,
    examples=[
        # Add your sample images later, for example:
        # ["assets/sample1.jpg"],
        # ["assets/sample2.jpg"],
    ],
)

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