# --- 1. Import all the necessary tools ---
import gradio as gr
from ultralytics import YOLO
from huggingface_hub import hf_hub_download
import numpy as np
import cv2
import roboflow
from collections import Counter
import re

# --- 2. Load BOTH of your AI models ---
print("Downloading and loading models...")

# --- Model 1: The Character Detector (from Hugging Face) ---
character_model_path = hf_hub_download(
    repo_id="MKgoud/License-Plate-Character-Detector", 
    filename="Charcter-LP.pt"
)
character_model = YOLO(character_model_path)
print("✅ Character Detector loaded.")

# --- Model 2: The Plate Detector (from Roboflow) ---
ROBOFLOW_API_KEY = "YfKCsreNkoXYFD1CfMBY"
DETECTOR_WORKSPACE_ID = "mylprproject"
DETECTOR_PROJECT_ID = "license-plate-yuw1z-kirke"
DETECTOR_VERSION_NUMBER = 1

rf = roboflow.Roboflow(api_key=ROBOFLOW_API_KEY)
project_detector = rf.workspace(DETECTOR_WORKSPACE_ID).project(DETECTOR_PROJECT_ID)
plate_model = project_detector.version(DETECTOR_VERSION_NUMBER).model
print("✅ Plate Detector loaded.")

# --- 3. Enhanced preprocessing functions ---
def enhance_plate_image(plate_crop):
    """
    Apply multiple enhancement techniques to improve character visibility
    """
    enhanced_crops = []
    
    # Original image
    enhanced_crops.append(plate_crop)
    
    # Convert to grayscale and back to RGB for consistent processing
    gray = cv2.cvtColor(plate_crop, cv2.COLOR_RGB2GRAY)
    
    # Enhancement 1: Adaptive histogram equalization
    clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8,8))
    enhanced_gray = clahe.apply(gray)
    enhanced_crops.append(cv2.cvtColor(enhanced_gray, cv2.COLOR_GRAY2RGB))
    
    # Enhancement 2: Gaussian blur + unsharp mask
    blurred = cv2.GaussianBlur(gray, (3, 3), 0)
    unsharp = cv2.addWeighted(gray, 1.5, blurred, -0.5, 0)
    unsharp = np.clip(unsharp, 0, 255).astype(np.uint8)
    enhanced_crops.append(cv2.cvtColor(unsharp, cv2.COLOR_GRAY2RGB))
    
    # Enhancement 3: Morphological operations
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (2, 2))
    morph = cv2.morphologyEx(gray, cv2.MORPH_CLOSE, kernel)
    enhanced_crops.append(cv2.cvtColor(morph, cv2.COLOR_GRAY2RGB))
    
    # Enhancement 4: Bilateral filter
    bilateral = cv2.bilateralFilter(gray, 9, 75, 75)
    enhanced_crops.append(cv2.cvtColor(bilateral, cv2.COLOR_GRAY2RGB))
    
    return enhanced_crops

def post_process_text(raw_text):
    """
    Apply license plate specific formatting and corrections
    """
    if not raw_text:
        return raw_text
    
    # Remove any spaces first
    text = raw_text.replace(" ", "")
    
    # Common character corrections for license plates
    corrections = {
        '0': 'O',  # In letter context
        'O': '0',  # In number context  
        'I': '1',
        '1': 'I',
        'S': '5',
        '5': 'S',
        'Z': '2',
        'B': '8',
        '8': 'B',
        'G': '6',
        '6': 'G'
    }
    
    # For Philippine plates, common format is 3 letters + 3 numbers (like NOV706)
    if len(text) >= 6:
        corrected_chars = list(text)
        
        # First 3 should typically be letters
        for i in range(min(3, len(corrected_chars))):
            char = corrected_chars[i]
            if char.isdigit():
                # Convert common digit misreads to letters
                if char in ['0', '1', '5', '8']:
                    letter_map = {'0': 'O', '1': 'I', '5': 'S', '8': 'B'}
                    corrected_chars[i] = letter_map.get(char, char)
        
        # Last 3 should typically be numbers  
        for i in range(3, min(6, len(corrected_chars))):
            char = corrected_chars[i]
            if char.isalpha():
                # Convert common letter misreads to numbers
                if char in ['O', 'I', 'S', 'B', 'G', 'Z']:
                    number_map = {'O': '0', 'I': '1', 'S': '5', 'B': '8', 'G': '6', 'Z': '2'}
                    corrected_chars[i] = number_map.get(char, char)
        
        text = ''.join(corrected_chars)
    
    return text

def improved_filtering(boxes, character_results, plate_crop_shape, min_confidence=0.3):
    """
    Enhanced filtering focusing on main license plate number only
    """
    if len(boxes) == 0:
        return []
    
    detections = []
    
    # Extract all detection info
    for box in boxes:
        confidence = float(box.conf[0])
        if confidence < min_confidence:
            continue
            
        class_id = int(box.cls[0])
        character = character_results[0].names[class_id]
        x1, y1, x2, y2 = box.xyxy[0]
        
        detections.append({
            'char': character,
            'conf': confidence,
            'x1': float(x1), 'y1': float(y1), 'x2': float(x2), 'y2': float(y2),
            'width': float(x2 - x1),
            'height': float(y2 - y1),
            'center_x': float((x1 + x2) / 2),
            'center_y': float((y1 + y2) / 2)
        })
    
    if len(detections) == 0:
        return []
    
    # MAIN IMPROVEMENT: Focus on the upper portion of the plate
    # Most license plates have the main number in the top 60% of the plate
    plate_height = plate_crop_shape[0]
    upper_threshold = plate_height * 0.70  # Only consider top 65% of plate
    
    # Filter out detections in lower portion (subsidiary text area)
    upper_detections = [d for d in detections if d['center_y'] <= upper_threshold]
    
    if len(upper_detections) == 0:
        # Fallback: if no detections in upper area, use all but be more selective
        upper_detections = detections
        print("Warning: No detections in upper area, using all detections")
    
    print(f"Filtered to upper area: {len(upper_detections)}/{len(detections)} detections")
    
    # Calculate statistics for filtering (now only on upper detections)
    heights = [d['height'] for d in upper_detections]
    widths = [d['width'] for d in upper_detections]
    y_centers = [d['center_y'] for d in upper_detections]
    
    if len(heights) == 0:
        return []
    
    median_height = np.median(heights)
    median_width = np.median(widths)
    median_y_center = np.median(y_centers)
    
    # More aggressive filtering for main plate numbers
    filtered_detections = []
    for detection in upper_detections:
        # Size filtering (tighter for main numbers)
        height_ratio = detection['height'] / median_height
        width_ratio = detection['width'] / median_width
        
        # Alignment filtering (tighter)
        y_deviation = abs(detection['center_y'] - median_y_center)
        max_y_deviation = median_height * 0.4  # Reduced from 0.6 to 0.4
        
        # Height-based filtering: main numbers are usually larger
        min_height_threshold = plate_height * 0.15  # At least 15% of plate height
        
        # Keep detection if it passes all criteria
        if (0.5 <= height_ratio <= 2.0 and      # Tighter height range
            0.4 <= width_ratio <= 2.5 and       # Tighter width range  
            y_deviation <= max_y_deviation and   # Better alignment
            detection['height'] >= min_height_threshold):  # Minimum size
            filtered_detections.append(detection)
    
    return filtered_detections

# --- 4. Enhanced main prediction function ---
def detect_license_plate(input_image):
    """
    Enhanced version with multi-enhancement processing and ensemble voting
    """
    print("New image received. Starting enhanced 2-stage pipeline...")
    output_image = input_image.copy()
    
    # --- STAGE 1: Find the license plate ---
    plate_predictions = plate_model.predict(input_image, confidence=40, overlap=30).json()['predictions']

    if not plate_predictions:
        return output_image, "No license plate found."

    # Get the highest confidence plate detection
    plate_box = max(plate_predictions, key=lambda x: x['confidence'])
    x1, y1, x2, y2 = [int(p) for p in [plate_box['x'] - plate_box['width'] / 2, 
                                       plate_box['y'] - plate_box['height'] / 2,
                                       plate_box['x'] + plate_box['width'] / 2, 
                                       plate_box['y'] + plate_box['height'] / 2]]
    
    # Add some padding to the plate crop, but reduce vertical padding to avoid extra text
    h_padding = 8  # Horizontal padding
    v_padding = 3  # Minimal vertical padding to avoid bottom text
    y1 = max(0, y1 - v_padding)
    x1 = max(0, x1 - h_padding)
    y2 = min(input_image.shape[0], y2 + v_padding)
    x2 = min(input_image.shape[1], x2 + h_padding)
    
    plate_crop = input_image[y1:y2, x1:x2]
    
    # Crop to focus on upper portion where main numbers are located
    plate_height = plate_crop.shape[0]
    # Keep top 70% of the plate to exclude bottom text area
    main_number_crop = plate_crop[:int(plate_height * 0.7), :]
    
    # --- STAGE 2: Multi-enhancement character detection ---
    enhanced_crops = enhance_plate_image(main_number_crop)
    
    all_detections = []
    character_votes = {}
    
    # Process each enhanced version
    for i, enhanced_crop in enumerate(enhanced_crops):
        try:
            character_results = character_model(enhanced_crop, conf=0.3, iou=0.4)
            
            if character_results and hasattr(character_results[0], 'boxes') and len(character_results[0].boxes) > 0:
                boxes = character_results[0].boxes.cpu().numpy()
                filtered_detections = improved_filtering(boxes, character_results, 
                                                       main_number_crop.shape, min_confidence=0.3)
                
                print(f"Enhancement {i}: {len(boxes)} raw -> {len(filtered_detections)} filtered detections")
                
                for detection in filtered_detections:
                    # Add enhancement method info
                    detection['enhancement'] = i
                    all_detections.append(detection)
                    
                    # Collect votes for ensemble
                    x_pos = int(detection['center_x'] / 8) * 8  # Tighter grouping
                    key = f"x{x_pos}"
                    if key not in character_votes:
                        character_votes[key] = []
                    character_votes[key].append((detection['char'], detection['conf']))
                    
        except Exception as e:
            print(f"Error processing enhancement {i}: {e}")
            continue
    
    # --- STAGE 3: Ensemble voting and final selection ---
    final_detections = []
    
    if character_votes:
        for x_key in sorted(character_votes.keys()):
            votes = character_votes[x_key]
            
            # Weight votes by confidence and count
            char_scores = {}
            for char, conf in votes:
                if char not in char_scores:
                    char_scores[char] = []
                char_scores[char].append(conf)
            
            # Calculate weighted scores
            best_char = None
            best_score = 0
            
            for char, confs in char_scores.items():
                # Score = average confidence * count weight
                avg_conf = np.mean(confs)
                count_weight = min(len(confs) / len(enhanced_crops), 1.0)
                score = avg_conf * (0.7 + 0.3 * count_weight)
                
                if score > best_score:
                    best_score = score
                    best_char = char
            
            if best_char and best_score > 0.3:
                # Find representative detection for drawing
                x_pos = int(x_key[1:])
                representative = min([d for d in all_detections if abs(d['center_x'] - x_pos) < 15],
                                   key=lambda x: abs(x['center_x'] - x_pos), default=None)
                
                if representative:
                    representative['final_char'] = best_char
                    representative['final_conf'] = best_score
                    final_detections.append(representative)
    
    # --- STAGE 4: Draw results and generate text ---
    # Draw the main plate box
    cv2.rectangle(output_image, (x1, y1), (x2, y2), (0, 0, 255), 2)
    cv2.putText(output_image, f"Plate Conf: {plate_box['confidence']:.2f}", 
                (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)
    
    # Draw character boxes (adjust coordinates back to main number crop area)
    for detection in final_detections:
        abs_x1 = x1 + int(detection['x1'])
        abs_y1 = y1 + int(detection['y1'])
        abs_x2 = x1 + int(detection['x2'])
        abs_y2 = y1 + int(detection['y2'])
        
        cv2.rectangle(output_image, (abs_x1, abs_y1), (abs_x2, abs_y2), (0, 255, 0), 2)
        cv2.putText(output_image, f"{detection['final_char']} {detection['final_conf']:.2f}", 
                   (abs_x1, abs_y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
    
    # Draw a line to show the detection area boundary
    main_area_y = y1 + int(plate_height * 0.7)
    cv2.line(output_image, (x1, main_area_y), (x2, main_area_y), (255, 255, 0), 2)
    
    # Sort by x position and create final text
    final_detections.sort(key=lambda x: x['center_x'])
    raw_text = "".join([d['final_char'] for d in final_detections])
    
    # Apply post-processing
    final_text = post_process_text(raw_text)
    
    result_text = f"Raw: {raw_text}\nProcessed: {final_text}" if raw_text != final_text else final_text
    
    print(f"Prediction complete. Final result: {result_text}")
    print(f"Used {len(final_detections)} characters from {len(all_detections)} total detections")
    
    return output_image, result_text

# --- 5. Create the Gradio Web Interface ---
with gr.Blocks() as demo:
    gr.Markdown("# Enhanced High-Accuracy License Plate Detector")
    gr.Markdown("""
    This system uses an advanced 2-stage AI pipeline with:
    - Multiple image enhancement techniques
    - Ensemble voting across different processed versions
    - Smart filtering and post-processing
    - Common license plate character corrections
    """)
    
    with gr.Row():
        image_input = gr.Image(type="numpy", label="Upload License Plate Image")
        image_output = gr.Image(type="numpy", label="Detection Results")
    
    text_output = gr.Textbox(label="Detected Characters", lines=3)
    predict_button = gr.Button(value="Detect Characters", variant="primary")
    
    predict_button.click(
        fn=detect_license_plate, 
        inputs=image_input, 
        outputs=[image_output, text_output]
    )

# --- 6. Launch the application ---
demo.launch()