import gradio as gr
import cv2
import numpy as np
from PIL import Image
from typing import get_args

from fast_alpr import ALPR
from fast_alpr.default_detector import PlateDetectorModel
from fast_alpr.default_ocr import OcrModel

# Default models
DETECTOR_MODELS = list(get_args(PlateDetectorModel))
OCR_MODELS = list(get_args(OcrModel))
# Put global OCR first
OCR_MODELS.remove("cct-s-v1-global-model")
OCR_MODELS.insert(0, "cct-s-v1-global-model")

def process_image(image, detector_model, ocr_model):
    """
    Process an image with ALPR system
    
    Args:
        image: PIL Image or numpy array
        detector_model: Selected detector model
        ocr_model: Selected OCR model
    
    Returns:
        tuple: (annotated_image, results_text)
    """
    if image is None:
        return None, "Please upload an image to continue."
    
    try:
        # Convert image to numpy array if it's a PIL Image
        if isinstance(image, Image.Image):
            img_array = np.array(image.convert("RGB"))
        else:
            img_array = image
        
        # Initialize ALPR with selected models
        alpr = ALPR(detector_model=detector_model, ocr_model=ocr_model)
        
        # Run ALPR on the image
        results = alpr.predict(img_array)
        
        # Draw predictions on the image
        annotated_img_array = alpr.draw_predictions(img_array)
        
        # Convert back to PIL Image for Gradio
        annotated_img = Image.fromarray(annotated_img_array)
        
        # Format results text
        if results:
            results_text = "**OCR Results:**\n"
            for i, result in enumerate(results):
                plate_text = result.ocr.text if result.ocr else "N/A"
                plate_confidence = result.ocr.confidence if result.ocr else 0.0
                results_text += f"{i+1}. Detected Plate: `{plate_text}` with confidence `{plate_confidence:.2f}`\n"
        else:
            results_text = "No license plate detected."
        
        return annotated_img, results_text
        
    except Exception as e:
        return None, f"Error processing image: {str(e)}"

# Create Gradio interface
with gr.Blocks(
    title="Automatic License Plate Recognition (ALPR)",
    theme=gr.themes.Soft(
        primary_hue="green",
        secondary_hue="blue",
        neutral_hue="slate"
    )
) as demo:
    
    gr.Markdown("""
    # Automatic License Plate Recognition (ALPR)
    
    An automatic license plate recognition (ALPR) system with customizable detector and OCR models.
    
    This system uses the FastALPR library to detect and recognize license plates in images.
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### Model Selection")
            detector_dropdown = gr.Dropdown(
                choices=DETECTOR_MODELS,
                value=DETECTOR_MODELS[0],
                label="Choose Detector Model",
                info="Select the model for license plate detection"
            )
            
            ocr_dropdown = gr.Dropdown(
                choices=OCR_MODELS,
                value=OCR_MODELS[0],
                label="Choose OCR Model",
                info="Select the model for text recognition"
            )
            
            gr.Markdown("### Upload Image")
            image_input = gr.Image(
                label="Upload an image of a vehicle with a license plate",
                type="pil",
                height=300
            )
            
            process_btn = gr.Button(
                "Process Image",
                variant="primary",
                size="lg"
            )
        
        with gr.Column(scale=1):
            gr.Markdown("### Results")
            image_output = gr.Image(
                label="Annotated Image with OCR Results",
                height=300
            )
            
            text_output = gr.Markdown(
                label="OCR Results",
                value="Upload an image and click 'Process Image' to see results."
            )
    
    # Add some examples
    gr.Markdown("### Example Images")
    gr.Examples(
        examples=[],
        inputs=image_input,
        label="Try with these examples (if available)"
    )
    
    # Connect the interface
    process_btn.click(
        fn=process_image,
        inputs=[image_input, detector_dropdown, ocr_dropdown],
        outputs=[image_output, text_output]
    )
    
    gr.Markdown("""
    ---
    ### How to Use
    
    1. **Select Models**: Choose your preferred detector and OCR models from the dropdowns
    2. **Upload Image**: Upload an image containing a vehicle with a license plate
    3. **Process**: Click the "Process Image" button to run ALPR
    4. **View Results**: See the annotated image and OCR text results
    
    ###  Supported Models
    
    - **Detector Models**: Various YOLO-based detection models
    - **OCR Models**: CCT OCR models for text recognition including global and specialized models
    
    ###  Features
    
    - Real-time license plate detection
    - Customizable model selection
    - Visual annotations with bounding boxes
    - Confidence scores for OCR results
    - Support for multiple image formats
    """)

# Launch the app
if __name__ == "__main__":
    demo.launch(share=False)