# app.py
import gradio as gr
from ultralytics import YOLO
import numpy as np
from PIL import Image
import os
import csv
from datetime import datetime

# -------- CONFIG --------
MODEL_URL = "https://huggingface.co/santhosh1305/pcb-defect-detector/resolve/main/best.pt"
EXAMPLE_IMAGE = "/mnt/data/ce841e1a-f28d-44be-8053-9801c0531e6d.png"  # example image in the Space filesystem
LEADERBOARD_CSV = "leaderboard.csv"  # stored in the Space repo (persisted in the Space)
# ------------------------

# Load model (cached by HF/Ultralytics)
model = YOLO(MODEL_URL)

# Predefined class names (same as training)
DEFECT_CLASSES = {
    0: "missing_hole",
    1: "mouse_bite",
    2: "open_circuit",
    3: "short",
    4: "spur",
    5: "spurious_copper"
}

# Ensure leaderboard exists
if not os.path.exists(LEADERBOARD_CSV):
    with open(LEADERBOARD_CSV, "w", newline="") as f:
        writer = csv.writer(f)
        writer.writerow(["timestamp", "image_name", "total_defects", "per_class"])

def append_leaderboard(image_name, total_defects, per_class_dict):
    row = [
        datetime.utcnow().isoformat(),
        image_name,
        total_defects,
        ";".join([f"{k}:{v}" for k, v in per_class_dict.items()])
    ]
    with open(LEADERBOARD_CSV, "a", newline="") as f:
        writer = csv.writer(f)
        writer.writerow(row)

def read_leaderboard(limit=20):
    rows = []
    if os.path.exists(LEADERBOARD_CSV):
        with open(LEADERBOARD_CSV, "r") as f:
            reader = csv.reader(f)
            next(reader, None)  # skip header
            for r in list(reader)[-limit:][::-1]:
                rows.append({
                    "timestamp": r[0],
                    "image": r[1],
                    "total_defects": r[2],
                    "per_class": r[3]
                })
    return rows

def analyze(image):
    """Main detection function used by the UI."""
    if isinstance(image, Image.Image):
        img = np.array(image)
        img_name = getattr(image, "filename", "uploaded_image")
    else:
        # If given a path or numpy array
        img = np.array(Image.open(image)) if isinstance(image, (str,)) else image
        img_name = image if isinstance(image, str) else "uploaded_image"

    # Run inference
    results = model(img)[0]
    annotated = results.plot()[..., ::-1]  # BGR -> RGB for PIL/Gradio
    out_pil = Image.fromarray(annotated)

    # Gather defect stats
    per_class = {}
    for b in results.boxes:
        cls = int(b.cls[0])
        name = DEFECT_CLASSES.get(cls, f"class_{cls}")
        per_class[name] = per_class.get(name, 0) + 1
    total_defects = sum(per_class.values())

    # Append to leaderboard
    append_leaderboard(img_name, total_defects, per_class)

    # Create a human-friendly report
    if total_defects == 0:
        report = "✔ No defects detected. PCB looks good."
    else:
        lines = [f"❌ Total defects detected: {total_defects}"]
        for k, v in per_class.items():
            lines.append(f"- {k}: {v}")
        report = "\n".join(lines)

    return out_pil, report

def get_leaderboard_table():
    rows = read_leaderboard(50)
    if not rows:
        return "No entries yet."
    table = "timestamp | image | total | per_class\n---|---|---|---\n"
    for r in rows:
        table += f"{r['timestamp']} | {r['image']} | {r['total_defects']} | {r['per_class']}\n"
    return table

# ------- Build Gradio UI -------
with gr.Blocks(title="PCB Defect Detector (Pro Template)") as demo:
    gr.Markdown("# 🔍 PCB Defect Detector — Professional Demo")
    gr.Markdown("Upload a PCB image and the YOLOv8 model (hosted on HuggingFace) will detect defects. "
                "A lightweight leaderboard logs recent runs (timestamp, image, counts).")

    with gr.Row():
        with gr.Column(scale=2):
            img_input = gr.Image(type="pil", label="Upload PCB image")
            run_btn = gr.Button("Analyze")
            example_btn = gr.Button("Load example image")
            output_img = gr.Image(label="Detection result")
            report = gr.Textbox(label="Report", interactive=False, lines=6)
        with gr.Column(scale=1):
            gr.Markdown("### Leaderboard (recent runs)")
            leaderboard_md = gr.Markdown(get_leaderboard_table())
            refresh_btn = gr.Button("Refresh Leaderboard")

    # Example handling: load the sample image from the Space filesystem
    def load_example():
        if os.path.exists(EXAMPLE_IMAGE):
            return EXAMPLE_IMAGE
        return None

    example_btn.click(fn=load_example, inputs=None, outputs=img_input)

    def run_and_refresh(inp):
        out_img, rpt = analyze(inp)
        # update leaderboard markdown
        return out_img, rpt, get_leaderboard_table()

    run_btn.click(fn=run_and_refresh, inputs=[img_input], outputs=[output_img, report, leaderboard_md])
    refresh_btn.click(fn=lambda: get_leaderboard_table(), inputs=None, outputs=leaderboard_md)

    gr.Examples(examples=[EXAMPLE_IMAGE], inputs=img_input)

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