# app.py — FRAQX · ANPR (Reconocimiento de Placas)
# Detector: YOLOv8 (pesos desde HF Hub)
# OCR: EasyOCR
# UI: Gradio (upload o Ctrl+V), Textbox con placa, Dataframe con detecciones
# CONTRATO DE SALIDA (NO CAMBIAR):
#   1) Imagen anotada (PIL)
#   2) Texto placa (str)
#   3) Tabla detecciones (lista de listas con [x1,y1,x2,y2,conf_det,ocr_text])

import os
import io
import re
import cv2
import numpy as np
from PIL import Image
import gradio as gr

from huggingface_hub import hf_hub_download
from ultralytics import YOLO
import easyocr

# ----------------------------- Config -----------------------------
REPO_ID  = "yasirfaizahmed/license-plate-object-detection"  # pesos del detector
FILENAME = "best.pt"
PLATE_REGEX = re.compile(r"[A-Z0-9][A-Z0-9\-]{3,8}")
HEADERS = ["x1", "y1", "x2", "y2", "conf_det", "ocr_text"]

def _clean(text: str) -> str:
    if not text:
        return ""
    text = re.sub(r"[^A-Za-z0-9\-]", "", text.upper())
    text = re.sub(r"\s+", "", text)
    return text.strip()

# ---------------------- Fallback OCR-only -------------------------
def _ocr_only_strategy(bgr):
    res = reader.readtext(bgr)  # [(box, text, conf), ...]
    rows, best = [], ""
    annotated = bgr.copy()

    for (pts, text, conf) in res:
        txt = _clean(text)
        if not PLATE_REGEX.fullmatch(txt):
            continue
        pts = np.array(pts).astype(int)
        x1, y1 = pts[:,0].min(), pts[:,1].min()
        x2, y2 = pts[:,0].max(), pts[:,1].max()

        cv2.rectangle(annotated, (x1,y1), (x2,y2), (0,255,0), 2)
        cv2.putText(annotated, f"{txt} ({conf:.2f})", (x1, max(0,y1-8)),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,255,0), 2, cv2.LINE_AA)

        rows.append([int(x1), int(y1), int(x2), int(y2), round(float(conf),3), txt])
        if len(txt) > len(best):
            best = txt

    return annotated, best, rows

# ------------------------ Carga de modelos ------------------------
def _load_detector():
    try:
        weights_path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
    except Exception as e:
        tried = [FILENAME, "best.pt", "weights/best.pt", "model.pt"]
        weights_path = None
        for fn in tried:
            try:
                weights_path = hf_hub_download(repo_id=REPO_ID, filename=fn)
                break
            except Exception:
                pass
        if weights_path is None:
            raise RuntimeError(
                f"No se pudieron descargar los pesos del detector desde {REPO_ID}. "
                f"Error original: {e}"
            )
    return YOLO(weights_path)

yolo = _load_detector()
reader = easyocr.Reader(["en", "es", "fr"], gpu=False)

# --------------------------- Utilidades ---------------------------
def _ensure_pil(img):
    if isinstance(img, Image.Image):
        return img.convert("RGB")
    if isinstance(img, (bytes, bytearray)):
        return Image.open(io.BytesIO(img)).convert("RGB")
    if isinstance(img, np.ndarray):
        if img.ndim == 2:
            img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
        elif img.shape[2] == 4:
            img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)
        return Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    raise ValueError("Formato de imagen no soportado")

def _clean_plate_text(txt: str) -> str:
    txt = (txt or "").upper().strip()
    txt = re.sub(r"[^A-Z0-9\- ]", "", txt)
    txt = re.sub(r"\s+", "", txt)
    return txt

def _ocr_best_text(img_bgr: np.ndarray) -> str:
    gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)
    gray = cv2.bilateralFilter(gray, 9, 75, 75)
    thr1 = cv2.adaptiveThreshold(
        gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
        cv2.THRESH_BINARY, 11, 2
    )
    res1 = reader.readtext(thr1)
    parts1 = [txt for _, txt, conf in res1 if conf >= 0.35]
    text1 = _clean_plate_text(" ".join(parts1))
    if text1:
        return text1

    thr2 = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)[1]
    res2 = reader.readtext(thr2)
    parts2 = [txt for _, txt, conf in res2 if conf >= 0.35]
    text2 = _clean_plate_text(" ".join(parts2))
    return text2

# ------------------------ Inferencia principal --------------------
def detect_and_read(image, conf_thres=0.25, iou_thres=0.45, debug=False):
    """
    Respeta los UMBRALES seleccionados por el usuario.
    Devuelve:
      - Imagen anotada (PIL)
      - Placa principal (str)
      - Registros [[x1,y1,x2,y2,conf_det,ocr_text], ...]
    """
    if image is None:
        return None, "", []

    pil = _ensure_pil(image)
    bgr = cv2.cvtColor(np.array(pil, dtype=np.uint8), cv2.COLOR_RGB2BGR)

    # YOLO → si falla o no hay pesos, usar OCR-only
    if yolo is None:
        annotated, best, rows = _ocr_only_strategy(bgr)
        out = Image.fromarray(cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB))
        return out, best, rows

    # Detección con umbrales del usuario
    results = yolo.predict(source=bgr, conf=conf_thres, iou=iou_thres, verbose=False)
    bboxes = []
    if len(results):
        r0 = results[0]
        for b in r0.boxes:
            x1, y1, x2, y2 = map(int, b.xyxy[0].tolist())
            score = float(b.conf[0].item()) if hasattr(b, "conf") else None
            bboxes.append((x1, y1, x2, y2, score))

    # Fallback si no hay cajas
    if not bboxes:
        annotated, best, rows = _ocr_only_strategy(bgr)
        out = Image.fromarray(cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB))
        return out, best, rows

    # OCR por recorte (con padding para no cortar caracteres)
    annotated = bgr.copy()
    records = []
    for (x1, y1, x2, y2, score) in bboxes:
        pad = 10
        x1 = max(0, x1 - pad); y1 = max(0, y1 - pad)
        x2 = min(bgr.shape[1]-1, x2 + pad); y2 = min(bgr.shape[0]-1, y2 + pad)

        crop = bgr[y1:y2, x1:x2]

        gray = cv2.cvtColor(crop, cv2.COLOR_BGR2GRAY)
        gray = cv2.bilateralFilter(gray, 9, 75, 75)
        kernel = np.array([[0,-1,0],[-1,5,-1],[0,-1,0]])
        gray = cv2.filter2D(gray, -1, kernel)
        thr1 = cv2.adaptiveThreshold(gray,255,cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                                     cv2.THRESH_BINARY,11,2)
        res1 = reader.readtext(thr1)
        parts1 = [t for _,t,c in res1 if c>=0.35]
        text = _clean(" ".join(parts1))
        if not text:
            thr2 = cv2.threshold(gray,0,255,cv2.THRESH_BINARY+cv2.THRESH_OTSU)[1]
            res2 = reader.readtext(thr2)
            parts2 = [t for _,t,c in res2 if c>=0.35]
            text = _clean(" ".join(parts2))

        cv2.rectangle(annotated, (x1,y1), (x2,y2), (0,255,0), 2)
        label = text if text else "PLATE"
        if score is not None:
            label = f"{label} ({score:.2f})"
        cv2.putText(annotated, label, (x1, max(0, y1-8)),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,255,0), 2, cv2.LINE_AA)

        records.append([x1, y1, x2, y2, round(score,3) if score is not None else None, text])

    best_plate = max([r[-1] for r in records if r[-1]], key=len, default="")
    out = Image.fromarray(cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB))
    return out, best_plate, records

# ------------------------------ UI -------------------------------
EXAMPLE_URLS = []  # puedes añadir URLs públicas a .jpg/.png para pruebas

with gr.Blocks(title="FRAQX · ANPR — Reconocimiento de Placas") as demo:
    gr.Markdown(
        """
# FRAQX · Reconocimiento de Placas (ANPR)
Sube una imagen o **pega (Ctrl+V)**. El sistema detecta la placa y lee el texto (OCR).  
Detector: **YOLOv8** · OCR: **EasyOCR** — *Demo*
        """
    )

    with gr.Row():
        with gr.Column(scale=6):
            img_in = gr.Image(
                sources=["upload", "clipboard"],
                type="pil",
                label="Imagen (upload o Ctrl+V)",
                interactive=True
            )
            conf = gr.Slider(0.05, 0.9, value=0.25, step=0.05, label="Confianza (detección)")
            iou  = gr.Slider(0.1,  0.9, value=0.45, step=0.05, label="IoU (NMS)")
            debug = gr.Checkbox(value=False, label="Mostrar logs en consola (debug)")
            btn = gr.Button("Detectar y reconocer", variant="primary")
        with gr.Column(scale=6):
            img_out   = gr.Image(type="pil", label="Resultado anotado")
            plate_box = gr.Textbox(label="Placa reconocida (OCR)", show_copy_button=True, interactive=False, placeholder="—")
            table     = gr.Dataframe(
                headers=HEADERS,
                datatype=["number","number","number","number","number","str"],
                label="Placas detectadas",
                wrap=True
            )

    btn.click(
        detect_and_read,
        inputs=[img_in, conf, iou, debug],
        outputs=[img_out, plate_box, table]
    )

    if EXAMPLE_URLS:
        gr.Examples(label="Ejemplos", examples=[[u] for u in EXAMPLE_URLS], inputs=[img_in])

    gr.Markdown(
        """
---
<div style="text-align:center;opacity:.8;font-size:14px;">
Powered by <strong>FRAQX AI Vision</strong> ·
<span style="opacity:.7">Demo ilustrativa. Verifica normativa aplicable (privacidad, videovigilancia, datos personales).</span>
</div>
        """
    )

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