import cv2
import numpy as np

from config import NUM_CLASSES, CRNN_WEIGHTS_PATH, YOLO_WEIGHTS_PATH
from license_plate_recognizer import LicensePlateRecognizer

# ------------------------------------------------------------
# Запуск в режиме реального времени (веб-камера)
# ------------------------------------------------------------
if __name__ == "__main__":
    lpr = LicensePlateRecognizer(
        yolo_model_path=YOLO_WEIGHTS_PATH,
        crnn_model_path=CRNN_WEIGHTS_PATH,
        num_classes=NUM_CLASSES,
        device="cpu"
    )

    cap = cv2.VideoCapture(0) 

    # cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
    # cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)

    while True:
        ret, frame = cap.read()
        if not ret:
            print("Не удалось считать кадр с веб-камеры.")
            break

        detections = lpr.detect_and_recognize_frame(frame, padding=5)

        for det in detections:
            x1, y1, x2, y2 = det["bbox"]
            text = det["text"]

            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)

            cv2.putText(
                frame, text, (x1, y1 - 10), 
                cv2.FONT_HERSHEY_SIMPLEX, 
                0.7, (0, 255, 0), 2
            )

        height, width, _ = frame.shape
        black_bar_width = 300
        black_bar = np.zeros((height, black_bar_width, 3), dtype=np.uint8)

        y_start = 40
        for i, det in enumerate(detections):
            txt = det["text"]
            cv2.putText(
                black_bar,
                f"Plate #{i+1}: {txt}", 
                (10, y_start),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.7, (255, 255, 255), 2
            )
            y_start += 40

        display_frame = np.hstack((frame, black_bar))

        cv2.imshow("License Plate Recognition", display_frame)

        key = cv2.waitKey(1) & 0xFF
        if key == 27 or key == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()