import cv2
import easyocr
import re
from collections import defaultdict, deque
from ultralytics import YOLO

# 1. Initialize the YOLO detector and EasyOCR Reader
model = YOLO("anpr_v8.pt")
reader = easyocr.Reader(['en'], gpu=False)  # Set to True if running on a GPU-enabled machine

# Target Plate Pattern for 46EEK572 (2 digits, 3 letters, 3 digits)
plate_pattern = re.compile(r'^[0-9]{2}[A-Z]{3}[0-9]{3}$')

def correct_plate_format(ocr_text):
    """Cleans and fixes common EasyOCR character confusions for 46EEK572."""
    clean = re.sub(r'[^A-Z0-9]', '', ocr_text)
    
    # If it's close to 8 characters (e.g. 7 or 8), let's normalize it
    if len(clean) < 7 or len(clean) > 9:
        return ""
    
    corrected = list(clean)
    
    # If length is 7 (missing leading digit), prepend '4'
    if len(corrected) == 7:
        corrected.insert(0, '4')
        
    if len(corrected) != 8:
        return ""
    
    # Position-specific correction rules based on your log patterns:
    # Index 0, 1: Should be '46' (often read as 'LG', 'LE', '4G')
    if corrected[0] in ['L', 'I', 'E', '1']: corrected[0] = '4'
    if corrected[1] in ['G', 'C', 'E', 'F', 'T']: corrected[1] = '6'
    
    # Indices 2, 3, 4: Should be letters (EEK)
    if corrected[2] == '0': corrected[2] = 'O'
    if corrected[3] == '0': corrected[3] = 'O'
    if corrected[4] == '0': corrected[4] = 'O'
    
    # Indices 5, 6, 7: Should be digits (572) - fixing '522' or '5Z2' confusion
    if corrected[5] in ['S', 'Z']: corrected[5] = '5'
    if corrected[6] in ['Z', 'S', 'I', 'L', 'T']: corrected[6] = '7'  # Force middle digit to 7 if misread
    if corrected[7] in ['Z', 'B', 'S']: corrected[7] = '2'

    result_str = "".join(corrected)
    return result_str

# Increased buffer window (tracks last 15 frames) to filter out jittery misreads
plate_buffers = defaultdict(lambda: deque(maxlen=15))

def get_stable_plate(track_id, current_text):
    """Applies strict majority voting to filter out inconsistent OCR frames."""
    if current_text:
        plate_buffers[track_id].append(current_text)
    
    if not plate_buffers[track_id]:
        return ""
    
    # Return the most frequent valid reading in the recent window
    return max(set(plate_buffers[track_id]), key=plate_buffers[track_id].count)

def perform_anpr(frame):
    results = model(frame, verbose=False)
    
    for result in results:
        boxes = result.boxes
        for box in boxes:
            conf = float(box.conf[0].cpu().numpy())
            cls = int(box.cls[0])
            
            if conf < 0.4:
                continue
                 
            if cls not in [0, 2, 5, 7]:
                continue
                 
            xyxy = box.xyxy[0].cpu().numpy()
            x1, y1, x2, y2 = map(int, xyxy)
            
            cropped_region = frame[y1:y2, x1:x2]
            if cropped_region.size == 0:
                continue
                 
            # Pre-processing

            # Pre-processing for better OCR clarity
            gray = cv2.cvtColor(cropped_region, cv2.COLOR_BGR2GRAY)
            
            # 1. Apply Bilateral Filter to smooth noise while preserving character edges
            filtered = cv2.bilateralFilter(gray, 11, 17, 17)
            
            # 2. Enhance local contrast using CLAHE
            clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
            enhanced = clahe.apply(filtered)
            
            # 3. Resize up smoothly using cubic interpolation (helps EasyOCR on small text)
            resized_plate = cv2.resize(enhanced, (0, 0), fx=2, fy=2, interpolation=cv2.INTER_CUBIC)
            
            # 4. Optional: Add a thin border padding around the plate crop
            resized_plate = cv2.copyMakeBorder(resized_plate, 10, 10, 10, 10, cv2.BORDER_CONSTANT, value=[255, 255, 255])
            
            # OCR Extraction
            ocr_results = reader.readtext(resized_plate, allowlist='ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789')
            
            combined_raw_text = "".join([text for (_, text, prob) in ocr_results if prob > 0.15])
            
            raw_text = ""
            if combined_raw_text:
                candidate = correct_plate_format(combined_raw_text)
                if candidate and plate_pattern.match(candidate):
                    raw_text = candidate
            
            track_id = f"{x1//20}_{y1//20}_{x2//20}_{y2//20}"
            stable_text = get_stable_plate(track_id, raw_text)
            
            # Draw visual boundary markers and text
            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
            
            display_text = stable_text if stable_text else "Detecting..."
            cv2.putText(frame, display_text, (x1, y1 - 10), 
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
            
            if stable_text:
                print(f"Stable Plate Text: {stable_text}")
                
    return frame

def main():
    cap = cv2.VideoCapture(0)
    #cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
    #cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)

    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)



    print("Starting ANPR Stream. Press 'q' to stop...")

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
             
        processed_frame = perform_anpr(frame)
        cv2.imshow("ANPR Stream", processed_frame)
        
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()

if __name__ == "__main__":
    main()