from paddleocr import PaddleOCR
from ultralytics import YOLO
import torch
import cv2
from scripts.config.app_configurations import *
from scripts.logging.application_logging import logger

device = 'cuda' if torch.cuda.is_available() else 'cpu'

def load_ocr_model():

    """
    Load paddle OCR model

    """
    ocr = PaddleOCR(use_angle_cls=True, lang='en')
    logger.info("Loaded OCR model")
    return ocr

def load_yolo_model():
    """
    Load license plate detection model

    """
    model = YOLO(MODEL_PATH).to(device)
    logger.info("Loaded Yolo model")
    return model


def predict(img_path, ocr, yolo_model):
    try:

        img = cv2.imread(img_path)

        results = yolo_model.predict(img, iou=0.7) # predict r1, r2 and license plate using yolo model r1 - single letter region in license plate, r2 - four or five letter region in license plate

        if results:
            r1_txt = None
            r2_txt = None

            for r in results:



                boxes = r.boxes
                for box in boxes:
                    b = box.xyxy[0]  # get box coordinates in (top, left, bottom, right) format
                    # print("bb:  ",b.tolist())
                    x1 = int(b[0])
                    y1 = int(b[1])
                    x2 = int(b[2])
                    y2 = int(b[3])
                    w = abs(x2 - x1)
                    h = abs(y2 - y1)



                    lp  = img[y1:y1 + h, x1:x1 + w]

                    c = box.cls


                    if int(c) == 0:
                        #extracting text from r1
                        # cv2.imwrite("LP.png", lp)
                        result = ocr.ocr(lp, cls=True)
                        for idx in range(len(result)):
                            res = result[idx]

                            if res:

                                for ind, line in enumerate(res):


                                    text_tuple = line[1]
                                    if r1_txt is None:
                                        r1_txt = text_tuple[0]
                                        logger.info(f"r1 text : {r1_txt}")
                                    text_conf = text_tuple[1]
                    elif int(c) == 1:
                        #extracting text from r2
                        result = ocr.ocr(lp, cls=True)
                        for idx in range(len(result)):
                            res = result[idx]

                            if res:

                                for ind, line in enumerate(res):


                                    text_tuple = line[1]
                                    if r2_txt is None:
                                        r2_txt = text_tuple[0]
                                        logger.info(f"r2 text : {r2_txt}")
                                    text_conf = text_tuple[1]
                                # cv2.putText(lp, txt, (text_x , text_y + ind*50), cv2.FONT_HERSHEY_SIMPLEX , 1, (255, 0, 0), 2, cv2.LINE_AA)

                    # cv2.imshow('detected LPR', lp)
                    # cv2.waitKey(0)
                    #

            if r1_txt is not None and r2_txt is not None:
                LP = f"{r1_txt}  {r2_txt}"

                logger.info(f"License plate reading : {LP}")
                return LP

        return "Image Not Clear"
    except Exception as e:
        logger.info(f"Error while reading license plate : {e}")
        return f"Error while reading license plate : {e}"





