"""
Author: Juan Pablo Triana Martinez
Gradio app — LinkNet Semantic Text Segmentation (12-class DocLayNet).
Entry-point for HuggingFace Spaces.
"""
import gradio as gr
import matplotlib.pyplot as plt
import numpy as np
import torch
from PIL import Image
from torchvision import transforms
from typing import Dict

from model import create_semantic_model

# ── Constants ──────────────────────────────────────────────────────────────────
MEAN        = [0.9329, 0.9343, 0.9341]
STD         = [0.1651, 0.1592, 0.1623]
IMG_SIZE    = 512
NUM_CLASSES = 12
MODEL_PATH  = "linknet_semantic_doclaynet_20_percent_seed_7_ce_0.25_dice_1.0_.pth"

CLASS_ID_TO_NAME: Dict[int, str] = {
    0:  "background",
    1:  "Caption",
    2:  "Footnote",
    3:  "Formula",
    4:  "List-item",
    5:  "Page-footer",
    6:  "Page-header",
    7:  "Picture",
    8:  "Section-header",
    9:  "Table",
    10: "Text",
    11: "Title",
}

# ── Transforms ─────────────────────────────────────────────────────────────────
inference_transform = transforms.Compose([
    transforms.Resize((IMG_SIZE, IMG_SIZE)),
    transforms.ToTensor(),
    transforms.Normalize(mean=MEAN, std=STD),
])

api_transform = transforms.Compose([
    transforms.Resize((IMG_SIZE, IMG_SIZE)),
    transforms.ToTensor(),
])

# ── Colormap ───────────────────────────────────────────────────────────────────
def _build_colormap(num_classes: int) -> np.ndarray:
    cmap = plt.get_cmap("tab20", num_classes)
    return np.array([cmap(i)[:3] for i in range(num_classes)])


def _apply_colormap(mask_hw: np.ndarray, lut: np.ndarray) -> np.ndarray:
    H, W = mask_hw.shape
    rgb  = np.zeros((H, W, 3), dtype=np.float32)
    for c in range(lut.shape[0]):
        rgb[mask_hw == c] = lut[c]
    return rgb


COLOR_LUT = _build_colormap(NUM_CLASSES)

# ── Model ──────────────────────────────────────────────────────────────────────
device = "cpu"
model  = create_semantic_model().to(device)
model.load_state_dict(torch.load(MODEL_PATH, map_location=device))
model.eval()

# ── Helpers ────────────────────────────────────────────────────────────────────
def _logits_to_mask(logits: torch.Tensor) -> torch.Tensor:
    return torch.argmax(logits, dim=1)  # (1, H, W)


def _mask_to_rgb(mask: torch.Tensor) -> np.ndarray:
    mask_hw = mask.squeeze().cpu().numpy().astype(np.int32)
    return (_apply_colormap(mask_hw, COLOR_LUT) * 255).clip(0, 255).astype(np.uint8)


def _overlay_mask(image: Image.Image, mask: torch.Tensor, alpha: float = 0.45) -> np.ndarray:
    img_np   = api_transform(image).permute(1, 2, 0).numpy()  # (H, W, 3) in [0,1]
    mask_rgb = _mask_to_rgb(mask).astype(np.float32) / 255.0
    blended  = (1 - alpha) * img_np + alpha * mask_rgb
    return (blended.clip(0, 1) * 255).astype(np.uint8)


# ── Inference ──────────────────────────────────────────────────────────────────
def inference(image: Image.Image, gt_mask: Image.Image = None):
    img_tensor = inference_transform(image).unsqueeze(0)  # (1, 3, H, W)

    with torch.inference_mode():
        logits    = model(img_tensor)          # (1, 12, H, W)
        pred_mask = _logits_to_mask(logits)    # (1, H, W)

    pred_mask_rgb = _mask_to_rgb(pred_mask)
    overlay       = _overlay_mask(image, pred_mask)

    if gt_mask is None:
        return pred_mask_rgb, overlay, "No ground truth provided → metrics unavailable"

    # Decode ground-truth mask (colour-coded RGB or greyscale class-index PNG)
    gt_pil = gt_mask.resize((IMG_SIZE, IMG_SIZE))
    if gt_pil.mode in ("L", "P"):
        gt_np = np.array(gt_pil.convert("L"), dtype=np.int64)
    else:
        gt_rgb = np.array(gt_pil.convert("RGB"), dtype=np.float32) / 255.0
        dists  = np.linalg.norm(
            gt_rgb[:, :, np.newaxis, :] - COLOR_LUT[np.newaxis, np.newaxis, :, :],
            axis=-1,
        )
        gt_np = np.argmin(dists, axis=-1).astype(np.int64)

    gt_tensor = torch.from_numpy(gt_np).unsqueeze(0)  # (1, H, W)
    pred_flat = pred_mask.squeeze().cpu()              # (H, W)
    gt_flat   = gt_tensor.squeeze()                   # (H, W)

    lines = ["=== Per-class IoU (background excluded) ==="]
    ious  = []
    for c in range(1, NUM_CLASSES):
        pred_c = (pred_flat == c)
        gt_c   = (gt_flat   == c)
        inter  = (pred_c & gt_c).sum().float()
        union  = (pred_c | gt_c).sum().float()
        iou    = (inter / (union + 1e-7)).item()
        ious.append(iou)
        name = CLASS_ID_TO_NAME.get(c, str(c))
        lines.append(f"  [{c:2d}] {name:<16s}: {iou:.4f}")

    import math
    macro_iou = float(sum(ious) / len(ious))
    lines.append(f" Macro IoU (excl. bg): {macro_iou:.4f}")

    return pred_mask_rgb, overlay, "".join(lines)


# ── Gradio Interface ───────────────────────────────────────────────────────────
demo = gr.Interface(
    fn=inference,
    inputs=[
        gr.Image(type="pil", label="Input Image"),
        gr.Image(type="pil", label="Ground Truth Semantic Mask (optional)"),
    ],
    outputs=[
        gr.Image(label="Predicted Semantic Mask"),
        gr.Image(label="Overlay"),
        gr.Textbox(label="Metrics"),
    ],
    title="LinkNet Semantic Segmentation — DocLayNet ChestNut 🌰",
    description=(
        "Upload a document image to obtain its 12-class semantic segmentation mask. "
        "Classes: background, Caption, Footnote, Formula, List-item, Page-footer, "
        "Page-header, Picture, Section-header, Table, Text, Title. "
        "Optionally upload a ground-truth mask (colour-coded PNG) to compute per-class IoU."
    ),
    examples=[
        ["examples/semantic_1_img.png", "examples/semantic_1_semantic_mask.png"],
        ["examples/semantic_2_img.png", "examples/semantic_2_semantic_mask.png"],
        ["examples/semantic_3_img.png", "examples/semantic_3_semantic_mask.png"],
    ],
)

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