import gradio as gr
import cv2
import numpy as np
from PIL import Image
import time
import tempfile
import os
from ultralytics import YOLO
from paddleocr import PaddleOCR

# ─────────────────────────────────────────────
# MODEL LOADING
# ─────────────────────────────────────────────
print("Loading models...")
model = YOLO("yolov8n.pt")   # swap with your fine-tuned weights if you have them
# PaddleOCR v3+ new API — removed use_gpu, show_log, use_angle_cls
ocr_engine = PaddleOCR(use_textline_orientation=True, lang="en")
print("Models ready.")


# ─────────────────────────────────────────────
# CORE PIPELINE
# ─────────────────────────────────────────────
def detect_plates(img_bgr, conf_threshold):
    results = model(img_bgr, conf=conf_threshold, verbose=False)
    boxes = []
    for r in results:
        for box in r.boxes:
            x1, y1, x2, y2 = map(int, box.xyxy[0])
            conf = float(box.conf[0])
            boxes.append((x1, y1, x2, y2, conf))
    return boxes


def read_plate(crop_bgr):
    """PaddleOCR v3+ returns a list of OCRResult objects, not raw nested lists."""
    texts = []
    try:
        results = ocr_engine.ocr(crop_bgr)
        if not results:
            return texts
        for res in results:
            # v3 API: res is an OCRResult with a .boxes attribute, each box has .text / .score
            if hasattr(res, 'boxes'):
                for box in res.boxes:
                    texts.append((box.text.strip().upper(), round(float(box.score), 3)))
            else:
                # Fallback: old-style nested list [[pts, (text, score)], ...]
                for item in res:
                    if isinstance(item, (list, tuple)) and len(item) == 2:
                        text_conf = item[1]
                        if isinstance(text_conf, (list, tuple)) and len(text_conf) == 2:
                            text, confidence = text_conf
                            texts.append((str(text).strip().upper(), round(float(confidence), 3)))
    except Exception as e:
        print(f"OCR error: {e}")
    return texts


def draw_annotations(img_bgr, detections, ocr_map):
    img = img_bgr.copy()
    for i, (x1, y1, x2, y2, conf) in enumerate(detections):
        cv2.rectangle(img, (x1, y1), (x2, y2), (0, 229, 255), 2)
        label_texts = ocr_map.get(i, [])
        plate_str = " ".join([t for t, _ in label_texts]) if label_texts else "PLATE"
        label = f"{plate_str} [{conf:.0%}]"
        (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.65, 2)
        cv2.rectangle(img, (x1, y1 - th - 12), (x1 + tw + 8, y1), (0, 229, 255), -1)
        cv2.putText(img, label, (x1 + 4, y1 - 4),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.65, (0, 0, 0), 2)
    return img


def full_pipeline(pil_image, conf_threshold):
    img_np  = np.array(pil_image.convert("RGB"))
    img_bgr = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)

    detections = detect_plates(img_bgr, conf_threshold)

    ocr_map    = {}
    plate_rows = []
    for i, (x1, y1, x2, y2, det_conf) in enumerate(detections):
        crop = img_bgr[y1:y2, x1:x2]
        if crop.size == 0:
            continue
        texts = read_plate(crop)
        ocr_map[i] = texts
        plate_str = " | ".join([t for t, _ in texts]) if texts else "—"
        ocr_conf  = f"{texts[0][1]:.1%}" if texts else "—"
        plate_rows.append([i + 1, plate_str, f"{det_conf:.1%}", ocr_conf])

    annotated_bgr = draw_annotations(img_bgr, detections, ocr_map)
    annotated_rgb = cv2.cvtColor(annotated_bgr, cv2.COLOR_BGR2RGB)
    annotated_pil = Image.fromarray(annotated_rgb)

    summary = f"✅ {len(detections)} plate(s) detected." if detections else "⚠️ No plates found. Try lowering the confidence threshold."
    return annotated_pil, plate_rows, summary


# ─────────────────────────────────────────────
# IMAGE TAB HANDLER
# ─────────────────────────────────────────────
def process_image(image, conf_threshold):
    if image is None:
        return None, [], "⚠️ Please upload an image."
    t0 = time.time()
    annotated, rows, summary = full_pipeline(image, conf_threshold)
    elapsed = time.time() - t0
    summary += f"  |  ⏱ {elapsed:.2f}s"
    return annotated, rows, summary


# ─────────────────────────────────────────────
# VIDEO TAB HANDLER
# ─────────────────────────────────────────────
def process_video(video_path, conf_threshold, frame_skip):
    if video_path is None:
        return None, [], "⚠️ Please upload a video."

    cap = cv2.VideoCapture(video_path)
    fps    = cap.get(cv2.CAP_PROP_FPS) or 25
    width  = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    out_path = tempfile.mktemp(suffix="_anpr.mp4")
    fourcc   = cv2.VideoWriter_fourcc(*"mp4v")
    writer   = cv2.VideoWriter(out_path, fourcc, fps, (width, height))

    all_plates   = {}
    frame_idx    = 0
    last_ocr_map = {}
    last_dets    = []

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        if frame_idx % int(frame_skip) == 0:
            dets = detect_plates(frame, conf_threshold)
            ocr_map = {}
            for i, (x1, y1, x2, y2, _) in enumerate(dets):
                crop = frame[y1:y2, x1:x2]
                if crop.size == 0:
                    continue
                texts = read_plate(crop)
                ocr_map[i] = texts
                for txt, conf in texts:
                    if txt not in all_plates or conf > all_plates[txt]:
                        all_plates[txt] = conf
            last_dets    = dets
            last_ocr_map = ocr_map

        annotated = draw_annotations(frame, last_dets, last_ocr_map)
        writer.write(annotated)
        frame_idx += 1

    cap.release()
    writer.release()

    rows = [[i + 1, plate, f"{conf:.1%}"] for i, (plate, conf) in enumerate(all_plates.items())]
    summary = f"✅ {len(all_plates)} unique plate(s) across {frame_idx} frames."
    return out_path, rows, summary


# ─────────────────────────────────────────────
# CUSTOM CSS
# ─────────────────────────────────────────────
css = """
@import url('https://fonts.googleapis.com/css2?family=Space+Mono:wght@400;700&family=Syne:wght@400;700;800&display=swap');

body, .gradio-container {
    background: #0a0a0f !important;
    font-family: 'Syne', sans-serif !important;
    color: #e8e8f0 !important;
}

.gradio-container {
    max-width: 1100px !important;
    margin: 0 auto !important;
}

#hero { padding: 2.5rem 0 1rem 0; }
#hero h1 {
    font-family: 'Syne', sans-serif;
    font-size: 3rem;
    font-weight: 800;
    letter-spacing: -0.04em;
    line-height: 1.1;
    background: linear-gradient(135deg, #ffffff 0%, #00e5ff 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
    margin: 0;
}
#hero p {
    font-family: 'Space Mono', monospace;
    color: #6b6b80;
    font-size: 0.85rem;
    margin-top: 8px;
    letter-spacing: 0.06em;
}
.tag {
    display: inline-block;
    background: rgba(0,229,255,0.08);
    border: 1px solid rgba(0,229,255,0.25);
    color: #00e5ff;
    font-family: 'Space Mono', monospace;
    font-size: 0.68rem;
    padding: 4px 12px;
    border-radius: 20px;
    letter-spacing: 0.08em;
    margin-right: 6px;
    margin-top: 10px;
}

.tab-nav button {
    font-family: 'Space Mono', monospace !important;
    font-size: 0.78rem !important;
    color: #6b6b80 !important;
    background: transparent !important;
    border: none !important;
    border-bottom: 2px solid transparent !important;
    padding: 10px 18px !important;
}
.tab-nav button.selected {
    color: #00e5ff !important;
    border-bottom: 2px solid #00e5ff !important;
}

.block, .panel, .wrap {
    background: #111118 !important;
    border: 1px solid #2a2a38 !important;
    border-radius: 12px !important;
}

button.primary {
    background: #00e5ff !important;
    color: #000 !important;
    border: none !important;
    font-family: 'Space Mono', monospace !important;
    font-weight: 700 !important;
    font-size: 0.82rem !important;
    letter-spacing: 0.06em !important;
    border-radius: 8px !important;
    padding: 10px 28px !important;
    transition: all 0.2s !important;
}
button.primary:hover {
    background: #00b8d9 !important;
    box-shadow: 0 4px 20px rgba(0,229,255,0.3) !important;
    transform: translateY(-1px) !important;
}

label span {
    font-family: 'Syne', sans-serif !important;
    font-size: 0.8rem !important;
    color: #6b6b80 !important;
    text-transform: uppercase;
    letter-spacing: 0.08em;
}

table { font-family: 'Space Mono', monospace !important; font-size: 0.8rem !important; }
th { color: #6b6b80 !important; text-transform: uppercase; letter-spacing: 0.08em; }
td { color: #e8e8f0 !important; }

footer { display: none !important; }
"""

# ─────────────────────────────────────────────
# BUILD GRADIO UI
# ─────────────────────────────────────────────
with gr.Blocks(css=css, title="ANPR System") as demo:

    gr.HTML("""
    <div id="hero">
        <h1>Number Plate<br>Recognition</h1>
        <p>// detect · read · log</p>
        <span class="tag">YOLOv8</span>
        <span class="tag">PaddleOCR</span>
        <span class="tag">Computer Vision</span>
        <span class="tag">Gradio</span>
    </div>
    """)

    with gr.Tabs():

        # ── IMAGE TAB ──────────────────────────────
        with gr.Tab("📷  Image Detection"):
            with gr.Row():
                with gr.Column(scale=1):
                    img_input   = gr.Image(type="pil", label="Upload Image")
                    conf_slider = gr.Slider(0.10, 0.95, value=0.30, step=0.05,
                                            label="Confidence Threshold")
                    run_img_btn = gr.Button("▶  Detect Plates", variant="primary")

                with gr.Column(scale=1):
                    img_output  = gr.Image(type="pil", label="Annotated Result")

            status_img  = gr.Textbox(label="Status", interactive=False)
            plate_table = gr.Dataframe(
                headers=["#", "Plate Text", "Detection Conf.", "OCR Conf."],
                label="Detected Plates",
                interactive=False,
            )

            run_img_btn.click(
                fn=process_image,
                inputs=[img_input, conf_slider],
                outputs=[img_output, plate_table, status_img],
            )

        # ── VIDEO TAB ──────────────────────────────
        with gr.Tab("🎬  Video Detection"):
            with gr.Row():
                with gr.Column(scale=1):
                    vid_input    = gr.Video(label="Upload Video")
                    conf_slider2 = gr.Slider(0.10, 0.95, value=0.30, step=0.05,
                                              label="Confidence Threshold")
                    frame_skip   = gr.Slider(1, 30, value=5, step=1,
                                              label="Process Every N Frames  (higher = faster)")
                    run_vid_btn  = gr.Button("▶  Process Video", variant="primary")

                with gr.Column(scale=1):
                    vid_output = gr.Video(label="Annotated Video")

            status_vid  = gr.Textbox(label="Status", interactive=False)
            vid_table   = gr.Dataframe(
                headers=["#", "Plate Text", "Best OCR Conf."],
                label="Unique Plates Found",
                interactive=False,
            )

            run_vid_btn.click(
                fn=process_video,
                inputs=[vid_input, conf_slider2, frame_skip],
                outputs=[vid_output, vid_table, status_vid],
            )

        # ── HOW IT WORKS TAB ───────────────────────
        with gr.Tab("📖  How It Works"):
            gr.Markdown("""
## Pipeline

This ANPR system runs a **two-stage deep learning pipeline**:

---

### Stage 1 — Plate Detection (YOLOv8)
YOLOv8 scans the full image in a single forward pass and outputs:
- Bounding box coordinates for each detected plate
- A confidence score per detection

### Stage 2 — Text Recognition (PaddleOCR)
Each detected plate region is cropped and passed to PaddleOCR which:
1. Detects text regions inside the crop
2. Classifies orientation (fixes rotated plates)
3. Reads characters using a CRNN-based model

### Video Processing
For videos, every N-th frame is sampled (configurable). Each frame goes through
the same pipeline. Results are deduplicated to surface unique plates.

---

### Models Used

| Model | Role | Source |
|-------|------|--------|
| YOLOv8n | Licence plate detection | Ultralytics |
| PaddleOCR | Text recognition | PaddlePaddle |

---

### Tip
Swap `yolov8n.pt` with a fine-tuned licence-plate weights file
(e.g. from Roboflow Universe) for significantly better plate-specific accuracy.
            """)

    gr.HTML("""
    <div style="text-align:center; padding: 1.5rem 0 0.5rem 0;
                font-family: 'Space Mono', monospace; font-size: 0.7rem; color: #3a3a50;">
        Built with YOLOv8 · PaddleOCR · Gradio &nbsp;·&nbsp; Hosted on Hugging Face Spaces
    </div>
    """)


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