import cv2
import pandas as pd
import numpy as np
import pytesseract
import os
from ultralytics import YOLO
from datetime import datetime

# Configure Tesseract executable path
pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe'

# Load YOLO model
model = YOLO('best.pt')

# Define polygon area for detection (adjust as needed)
area = [(27, 250), (16, 310), (1015, 310), (992, 250)]

# Track processed numbers to avoid duplicates
processed_numbers = set()

# Initialize video capture
cap = cv2.VideoCapture('mycarplate.mp4')  # Change this to your video file path

# Load class list from file
with open("coco1.txt", "r") as file:
    class_list = file.read().splitlines()

# Define CSV file and write headers if not already present
csv_file = "car_plate_data_stored.csv"

# Ensure the file has headers if it doesn't already exist
try:
    pd.read_csv(csv_file)
except FileNotFoundError:
    with open(csv_file, "w") as file:
        file.write("ImageFile,Date,Time,Confidence\n")

frame_count = 0

# Function to correct commonly misrecognized characters
def correct_characters(text):
    replacements = {
        '0': 'O',  # Replace '0' with 'O' if detected, to avoid confusion with 'D'
        'O': 'D',  # Replace 'O' with 'D' in typical license plate context
        'I': '1',  # Replace 'I' with '1' if detected
        'Q': '0'   # Replace 'Q' with '0' if detected
    }
    corrected_text = ''.join(replacements.get(c, c) for c in text)
    return corrected_text

# Function to format license plate text
def format_plate_text(text):
    text = ''.join(filter(str.isalnum, text))
    if len(text) == 10:  # Assuming full-length plates are correctly detected
        return f"{text[:2]} {text[2:4]} {text[4:6]} {text[6:]}"
    return text  # If not 10 characters, return as-is

# Deskew function to correct any skewed license plates
def deskew_image(image):
    # Check if the image is already grayscale (single channel)
    if len(image.shape) == 3:
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    else:
        gray = image  # It's already grayscale, no need to convert

    # Apply Canny edge detection
    edges = cv2.Canny(gray, 50, 150, apertureSize=3)

    # Use Hough Transform to detect lines
    lines = cv2.HoughLines(edges, 1, np.pi / 180, 100)
    
    if lines is not None:
        for line in lines:
            rho, theta = line[0]
            a = np.cos(theta)
            b = np.sin(theta)
            x0 = a * rho
            y0 = b * rho
            x1 = int(x0 + 1000 * (-b))
            y1 = int(y0 + 1000 * (a))
            x2 = int(x0 - 1000 * (-b))
            y2 = int(y0 - 1000 * (a))
            cv2.line(image, (x1, y1), (x2, y2), (0, 0, 255), 2)
            
    return image

# Processing video frames
while True:
    ret, frame = cap.read()
    frame_count += 1

    # Skip frames to process every third frame for efficiency
    if frame_count % 3 != 0:
        continue
    if not ret:
        break

    # Resize frame for consistent processing
    frame = cv2.resize(frame, (1020, 500))

    # Perform YOLO object detection
    results = model.predict(frame)
    detected_boxes = results[0].boxes.data
    detections = pd.DataFrame(detected_boxes).astype("float")

    for _, row in detections.iterrows():
        x1, y1, x2, y2 = int(row[0]), int(row[1]), int(row[2]), int(row[3])
        class_id = int(row[5])
        class_name = class_list[class_id]

        # Calculate center point of the detected box
        cx, cy = (x1 + x2) // 2, (y1 + y2) // 2

        # Check if center is within the defined polygon area
        if cv2.pointPolygonTest(np.array(area, np.int32), (cx, cy), False) >= 0:
            # Crop detected area
            crop = frame[y1:y2, x1:x2]
            gray_crop = cv2.cvtColor(crop, cv2.COLOR_BGR2GRAY)

            # Apply bilateral filter for smoothing, and threshold for clarity
            gray_crop = cv2.bilateralFilter(gray_crop, 10, 20, 20)
            _, threshold_crop = cv2.threshold(gray_crop, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)

            # Deskew the image if needed
            threshold_crop = deskew_image(threshold_crop)

            # Generate unique identifier for the detection
            current_date = datetime.now().strftime("%d-%m-%Y")
            current_time = datetime.now().strftime("%H:%M:%S")
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
            
            # Save the cropped license plate image
            detection_folder = "detected_plates"
            if not os.path.exists(detection_folder):
                os.makedirs(detection_folder)
            
            image_filename = f"plate_{timestamp}.jpg"
            image_path = os.path.join(detection_folder, image_filename)
            cv2.imwrite(image_path, crop)  # Save original crop
            
            # Save detection information to CSV
            if image_filename not in processed_numbers:  # Use filename as unique identifier
                processed_numbers.add(image_filename)
                with open(csv_file, "a") as file:
                    confidence = float(row[4])  # Get detection confidence
                    file.write(f"{image_filename},{current_date},{current_time},{confidence:.2f}\n")

                # Draw bounding box and display cropped image
                cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 1)
                cv2.imshow('Detected Plate', threshold_crop)

    # Draw the detection area on the main frame
    cv2.polylines(frame, [np.array(area, np.int32)], True, (255, 0, 0), 2)
    cv2.imshow("RGB", frame)

    # Break loop on 'ESC' key press
    if cv2.waitKey(1) & 0xFF == 27:
        break

# Release resources
cap.release()
cv2.destroyAllWindows()