import gradio as gr
import numpy as np
import onnxruntime as ort
from PIL import Image, ImageDraw, ImageFont
import cv2

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

CLASS_NAMES = ["crack", "spalling", "pothole"]
CLASS_COLORS = {
    "crack": (255, 0, 0),       # red
    "spalling": (255, 165, 0),  # orange
    "pothole": (255, 255, 0),   # yellow
}

# ── Load ONNX model once at startup ─────────────────────────────────────────
session = ort.InferenceSession(MODEL_PATH, providers=["CPUExecutionProvider"])
input_name = session.get_inputs()[0].name


# ── 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
    # HWC -> CHW
    img_array = img_array.transpose(2, 0, 1)
    # Add batch dimension
    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 YOLOv8 raw output: shape [1, 7, 8400] ────────────────────────────
# 7 = 4 (cx, cy, w, h) + 3 (class scores, no separate objectness column)
# 8400 = number of candidate detections across all scales
def parse_yolo_output(output, original_w, original_h):
    output = output[0]  # shape: [7, 8400]
    output = output.T   # transpose -> shape: [8400, 7]

    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, cy, w, h = cx * scale_x, cy * scale_y, w * scale_x, h * scale_y
        x1, y1 = cx - w / 2, cy - h / 2
        x2, y2 = cx + w / 2, cy + h / 2

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

    return non_max_suppression(detections, IOU_THRESHOLD)


# ── Draw detections on the image ────────────────────────────────────────
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 = CLASS_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=(0, 0, 0))

    return image


# ── Main inference function (called by Gradio) ─────────────────────────────
def detect_damage(input_image: Image.Image):
    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):"]
        for d in detections:
            lines.append(
                f"- {d['class_name']} (confidence: {d['confidence']:.2f})"
            )
        summary = "\n".join(lines)

    return result_image, summary


# ── Gradio interface ────────────────────────────────────────────────────
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="Structural Damage Detection (Crack / Spalling / Pothole)",
    description="Upload an image of a road or concrete structure to detect cracks, spalling, and potholes using a YOLOv8 model.",
)

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