import gradio as gr
print("[STARTUP] Imports completed (gradio loaded)", flush=True)

import numpy as np
import onnxruntime as ort
from PIL import Image, ImageDraw

print("[STARTUP] onnxruntime, numpy, PIL loaded", flush=True)

# ── Configuration ─────────────────────────────────────────────────────────
MODEL_PATH = "best.onnx"
INPUT_SIZE = 640
CONF_THRESHOLD = 0.4
IOU_THRESHOLD = 0.45

CLASS_NAMES = [
    'corrosion', 'delamination', 'cover_detachment', 'efflorescence',
    'crack', 'rust', 'scaling', 'spalling', 'void',
    'tile_crack', 'tile_delamination', 'tile_loss', 'peeling'
]

CLASS_MAPPING = {
    'crack': 'Crack',
    'tile_crack': 'Crack',
    'spalling': 'Spalling',
    'corrosion': 'Damage',
    'delamination': 'Damage',
    'cover_detachment': 'Spalling',
    'efflorescence': 'Damage',
    'rust': 'Damage',
    'scaling': 'Damage',
    'void': 'Damage',
    'tile_delamination': 'Damage',
    'tile_loss': 'Damage',
    'peeling': 'Damage',
}

MERGED_COLORS = {
    'Crack': (255, 0, 0),
    'Spalling': (255, 165, 0),
    'Damage': (0, 0, 255),
}

print("[STARTUP] Loading ONNX model from best.onnx ...", flush=True)
session = ort.InferenceSession(MODEL_PATH, providers=["CPUExecutionProvider"])
input_name = session.get_inputs()[0].name
print("[STARTUP] Model loaded successfully.", flush=True)


# ── Preprocessing ─────────────────────────────────────────────────────────
def preprocess(image: Image.Image):
    original_w, original_h = image.size
    resized = image.resize((INPUT_SIZE, INPUT_SIZE))
    img_array = np.array(resized).astype(np.float32) / 255.0
    img_array = img_array.transpose(2, 0, 1)
    img_array = np.expand_dims(img_array, axis=0)
    return img_array, original_w, original_h


# ── IoU + NMS ──────────────────────────────────────────────────────────────
def compute_iou(box1, box2):
    x1 = max(box1[0], box2[0])
    y1 = max(box1[1], box2[1])
    x2 = min(box1[2], box2[2])
    y2 = min(box1[3], box2[3])
    inter_area = max(0, x2 - x1) * max(0, y2 - y1)
    box1_area = (box1[2] - box1[0]) * (box1[3] - box1[1])
    box2_area = (box2[2] - box2[0]) * (box2[3] - box2[1])
    union_area = box1_area + box2_area - inter_area
    return inter_area / union_area if union_area > 0 else 0


def non_max_suppression(detections, iou_threshold):
    detections = sorted(detections, key=lambda d: d["confidence"], reverse=True)
    keep = []
    while detections:
        best = detections.pop(0)
        keep.append(best)
        detections = [d for d in detections if compute_iou(best["bbox"], d["bbox"]) < iou_threshold]
    return keep


# ── Parse YOLO output: shape [1, 17, 8400] ────────────────────────────────
def parse_yolo_output(output, original_w, original_h):
    output = output[0]
    output = output.T

    scale_x = original_w / INPUT_SIZE
    scale_y = original_h / INPUT_SIZE

    detections = []
    for row in output:
        cx, cy, w, h = row[0], row[1], row[2], row[3]
        class_scores = row[4:4 + len(CLASS_NAMES)]
        class_id = int(np.argmax(class_scores))
        confidence = float(class_scores[class_id])

        if confidence < CONF_THRESHOLD:
            continue

        cx = cx * scale_x
        cy = cy * scale_y
        w = w * scale_x
        h = h * scale_y
        x1, y1 = cx - w / 2, cy - h / 2
        x2, y2 = cx + w / 2, cy + h / 2

        original_class = CLASS_NAMES[class_id]
        merged_class = CLASS_MAPPING[original_class]

        detections.append({
            "bbox": [x1, y1, x2, y2],
            "confidence": confidence,
            "class_id": class_id,
            "original_class": original_class,
            "class_name": merged_class,
        })

    return non_max_suppression(detections, IOU_THRESHOLD)


def draw_detections(image: Image.Image, detections):
    image = image.copy()
    draw = ImageDraw.Draw(image)
    for det in detections:
        x1, y1, x2, y2 = det["bbox"]
        color = MERGED_COLORS.get(det["class_name"], (255, 255, 255))
        label = f"{det['class_name']} {det['confidence']:.2f}"
        draw.rectangle([x1, y1, x2, y2], outline=color, width=3)
        text_bbox = draw.textbbox((x1, y1), label)
        draw.rectangle(text_bbox, fill=color)
        draw.text((x1, y1), label, fill=(255, 255, 255))
    return image


def detect_damage(input_image: Image.Image):
    print("[REQUEST] Received new image for detection", flush=True)

    if input_image is None:
        return None, "No image provided."

    input_image = input_image.convert("RGB")
    img_array, original_w, original_h = preprocess(input_image)

    outputs = session.run(None, {input_name: img_array})
    detections = parse_yolo_output(outputs[0], original_w, original_h)

    result_image = draw_detections(input_image, detections)

    if not detections:
        summary = "No damage detected."
    else:
        lines = [f"Found {len(detections)} detection(s):\n"]
        for d in detections:
            x1, y1, x2, y2 = d["bbox"]
            lines.append(
                f"- category: {d['class_name']}\n"
                f"  confidence: {d['confidence']:.2f}\n"
                f"  bbox: [{x1:.1f}, {y1:.1f}, {x2:.1f}, {y2:.1f}]\n"
            )
        summary = "\n".join(lines)

    print(f"[REQUEST] Done. {len(detections)} detections.", flush=True)
    return result_image, summary


print("[STARTUP] Building Gradio interface ...", flush=True)
demo = gr.Interface(
    fn=detect_damage,
    inputs=gr.Image(type="pil", label="Upload an image"),
    outputs=[
        gr.Image(type="pil", label="Detection result"),
        gr.Textbox(label="Summary"),
    ],
    title="Urban Infrastructure Defect Detection",
    description=(
        "Upload an image of a building, bridge, pavement, or tiled surface to detect structural defects. "
        "Detections are grouped into three categories: Crack (red), Spalling (orange), and Damage (blue). "
    ),
)

print("[STARTUP] Calling demo.launch() now ....", flush=True)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)