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

# --- Inisialisasi Model ---
print("Mulai inisialisasi model...")

# Inisialisasi YOLO
try:
    print("Memuat model YOLOv8 ('best2.pt')...")
    yolo_model = YOLO('best2.pt')  # Ganti path jika modelmu di folder lain
    print("Model YOLOv8 berhasil dimuat.")
except Exception as e:
    print(f"Error saat memuat model kustom 'best2.pt': {e}")
    print("Mencoba memuat model YOLOv8 standar ('yolov8n-obb.pt')...")
    yolo_model = YOLO('yolov8n-obb.pt')
    print("Model YOLOv8 standar berhasil dimuat.")

# Inisialisasi PaddleOCR
ocr_engine = None
try:
    print("Memuat model PaddleOCR...")
    ocr_engine = PaddleOCR(
        use_angle_cls=True,
        lang='en',
        det_db_box_thresh=0.3,
        show_log=False
    )
    print("Model PaddleOCR berhasil dimuat.")
except Exception as e:
    print("!!! Gagal memuat PaddleOCR:", e)

def find_plate_in_text(text):
    cleaned = re.sub(r'[^A-Z0-9]', '', text.upper())
    match = re.match(r'^([A-Z]{1,2})(\d{1,4})([A-Z]{1,3})', cleaned)
    if match:
        return f"{match.group(1)} {match.group(2)} {match.group(3)}"
    return f"[Tidak Valid]: {cleaned}"

def detect_crop_and_ocr(image, yolo, ocr, margin=10):
    print("-> Mulai deteksi...")
    image_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) if isinstance(image, Image.Image) else image

    if image_cv is None or image_cv.size == 0:
        return None, None, "Gambar tidak valid."

    results = yolo(image_cv, task='obb')[0]
    img_annotated_rgb = cv2.cvtColor(image_cv, cv2.COLOR_BGR2RGB)
    crop_rgb = None
    ocr_text = "Plat tidak terdeteksi."

    if len(results.obb.xyxyxyxy) > 0:
        obb = results.obb.xyxyxyxy[0].cpu().numpy().reshape((4, 2))
        xs, ys = obb[:, 0], obb[:, 1]
        x1, x2 = int(xs.min()), int(xs.max())
        y1, y2 = int(ys.min()), int(ys.max())
        h, w = image_cv.shape[:2]
        x1m, y1m = max(0, x1 - margin), max(0, y1 - margin)
        x2m, y2m = min(w, x2 + margin), min(h, y2 + margin)
        cv2.rectangle(img_annotated_rgb, (x1m, y1m), (x2m, y2m), (0, 255, 0), 3)

        crop_cv = image_cv[y1m:y2m, x1m:x2m]
        if crop_cv.size > 0 and ocr is not None:
            h, w = crop_cv.shape[:2]
            crop_resized = cv2.resize(crop_cv, (w * 2, h * 2), interpolation=cv2.INTER_CUBIC)
            blurred = cv2.GaussianBlur(crop_resized, (0, 0), 3)
            sharpened = cv2.addWeighted(crop_resized, 1.5, blurred, -0.5, 0)
            crop_rgb = cv2.cvtColor(sharpened, cv2.COLOR_BGR2RGB)

            try:
                ocr_results = ocr.ocr(crop_rgb, cls=True)
                if ocr_results and ocr_results[0]:
                    texts = [line[1][0] for line in ocr_results[0] if line and line[1]]
                    raw_text = "".join(texts)
                    ocr_text = find_plate_in_text(raw_text)
                else:
                    ocr_text = "OCR tidak dapat membaca teks."
            except Exception as e:
                ocr_text = f"Error OCR: {e}"
    else:
        print("-> Tidak ada plat yang terdeteksi.")

    status_text = "Plat Terdeteksi" if crop_rgb is not None else "Plat Tidak Terdeteksi"
    cv2.putText(img_annotated_rgb, status_text, (15, 40), cv2.FONT_HERSHEY_SIMPLEX, 1.5, (255, 0, 0), 3)
    return img_annotated_rgb, crop_rgb, ocr_text

def process_image_for_gradio(uploaded_image):
    print("\n==============================")
    print("Request baru diterima...")

    if ocr_engine is None:
        placeholder_img = np.zeros((300, 500, 3), dtype=np.uint8)
        cv2.putText(placeholder_img, "OCR Engine Gagal", (50, 150), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 2)
        return placeholder_img, None, "ERROR: OCR Engine gagal dimuat."

    if uploaded_image is None:
        return None, None, "Mohon unggah gambar terlebih dahulu."

    try:
        full_img, crop_img, text = detect_crop_and_ocr(uploaded_image, yolo_model, ocr_engine)
        if crop_img is None:
            crop_img = np.zeros((100, 300, 3), dtype=np.uint8)
            cv2.putText(crop_img, "Tidak ada plat", (40, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 2)
        return full_img, crop_img, text
    except Exception as e:
        error_msg = f"Terjadi kesalahan: {e}"
        placeholder_img = np.zeros((300, 500, 3), dtype=np.uint8)
        return placeholder_img, None, error_msg

# === INTERFACE GRADIO ===
example_img_path = "AB2638XU.jpg"
example_img = None
if os.path.exists(example_img_path):
    example_img = Image.open(example_img_path)

with gr.Blocks(theme=gr.themes.Soft()) as iface:
    gr.Markdown("""
    # 🚗 Deteksi & OCR Plat Nomor Kendaraan Indonesia
    Unggah gambar kendaraan, deteksi plat dengan YOLOv8 OBB, dan baca teks dengan PaddleOCR.
    """)
    with gr.Row():
        with gr.Column(scale=2):
            image_input = gr.Image(type="pil", label="Unggah Gambar", value=example_img)
            submit_button = gr.Button("Proses", variant="primary")
        with gr.Column(scale=3):
            image_output = gr.Image(type="numpy", label="Hasil Deteksi")
    with gr.Row():
        crop_output = gr.Image(type="numpy", label="Plat Terpotong")
        text_output = gr.Textbox(label="Teks Plat Nomor")

    submit_button.click(
        fn=process_image_for_gradio,
        inputs=image_input,
        outputs=[image_output, crop_output, text_output]
    )

if __name__ == '__main__':
    iface.launch(debug=True)