import torch
from ultralytics import YOLO
import easyocr
import cv2
import numpy as np
import gradio as gr
import os
import logging
import re

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Ensure directories exist
os.makedirs(os.getenv('EASYOCR_MODULE_PATH', '/app/.EasyOCR'), exist_ok=True)
os.makedirs(os.getenv('YOLO_CONFIG_DIR', '/app/.config/Ultralytics'), exist_ok=True)

# Download pretrained model
ANPR_WEIGHTS = "anpr_yolov8.pt"
if not os.path.exists(ANPR_WEIGHTS):
    logger.info(f"Downloading model weights to {ANPR_WEIGHTS}")
    os.system(f"wget -O {ANPR_WEIGHTS} https://github.com/ultralytics/assets/releases/download/v8.2.0/yolov8n.pt")

# Load YOLO model
try:
    model = YOLO(ANPR_WEIGHTS)
    logger.info(f"Successfully loaded YOLO model from {ANPR_WEIGHTS}")
except Exception as e:
    logger.error(f"Error loading YOLO model from {ANPR_WEIGHTS}: {str(e)}")
    raise

# Load OCR reader
try:
    reader = easyocr.Reader(['en'], model_storage_directory=os.getenv('EASYOCR_MODULE_PATH', '/app/.EasyOCR'))
    logger.info("Successfully initialized EasyOCR reader")
except Exception as e:
    logger.error(f"Error initializing EasyOCR reader: {str(e)}")
    raise

def detect_and_read_plate(image):
    logger.info("Starting image processing for license plate detection")
    try:
        detected_texts = []
        results = model(image, conf=0.1)  # Low confidence for broader detection
        logger.info(f"YOLO model returned {len(results)} results")
        
        for result in results:
            boxes = result.boxes.xyxy.cpu().numpy()
            confidences = result.boxes.conf.cpu().numpy()
            logger.info(f"Detected {len(boxes)} bounding boxes")
            
            for box, conf in zip(boxes, confidences):
                x1, y1, x2, y2 = map(int, box)
                # Minimal size filter
                if (x2 - x1) < 20 or (y2 - y1) < 10:
                    logger.warning(f"Skipping small box: ({x1}, {y1}, {x2}, {y2})")
                    continue
                
                # Crop the detected license plate
                plate_img = image[y1:y2, x1:x2]
                if plate_img.size == 0:
                    logger.warning("Empty cropped image, skipping")
                    continue
                
                # Run OCR
                logger.info("Running EasyOCR on cropped plate")
                ocr_result = reader.readtext(plate_img)
                if ocr_result:
                    for res in ocr_result:
                        text = res[1]
                        confidence = res[2]
                        # Light filtering: prefer alphanumeric, avoid overly long text
                        if len(text) <= 12 and confidence > 0.2 and re.match(r'^[A-Z0-9\s\-]+$', text):
                            detected_texts.append(f"{text} (conf: {conf:.2f})")
                            logger.info(f"Detected Plate: {text} (YOLO conf: {conf:.2f}, OCR conf: {confidence:.2f})")
                        else:
                            logger.info(f"Rejected text: {text} (OCR conf: {confidence:.2f})")
                
                # Draw bounding box
                cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
        
        output_text = "\n".join(detected_texts) if detected_texts else "No license plate detected"
        logger.info(f"Output text: {output_text}")
        return image, output_text
    except Exception as e:
        logger.error(f"Error during detection: {str(e)}")
        return image, f"Error processing image: {str(e)}"

# Create Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Automatic Number Plate Recognition (ANPR)")
    gr.Markdown("Upload an image of a car to detect and read its license plate. Results may take a few seconds.")
    
    with gr.Row():
        image_input = gr.Image(type="numpy", label="Upload an image of a car")
    with gr.Row():
        image_output = gr.Image(type="numpy", label="Detected License Plate Image")
        text_output = gr.Textbox(label="Detected License Plate Number")
    
    submit_button = gr.Button("Detect License Plate")
    submit_button.click(
        fn=detect_and_read_plate,
        inputs=image_input,
        outputs=[image_output, text_output],
        show_progress=True
    )

demo.launch(server_name="0.0.0.0", server_port=7860)