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

device = 'cuda' if torch.cuda.is_available() else 'cpu'
# print(f'Using device: {device}')
def load_ocr_model():
    ocr = PaddleOCR(use_angle_cls=True, lang='en')
    return ocr

def load_yolo_model():
    model = YOLO(MODEL_PATH).to(device)
    return model


def predict(img_path, ocr, yolo_model):
    img = cv2.imread(img_path)
    results = yolo_model.predict(img, iou=0.7)
    if results:
        r1_txt = None
        r2_txt = None

        for r in results:

            # annotator = Annotator(img)

            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)

                cv2.rectangle(img, (x1, y1), (x2, y2), (255, 0, 0), 1)

                # lp = img[max(y1 - 30, 0):min(y1+h + 30, img.shape[0]), max(x1 - 30, 0):min(x1+w + 30, img.shape[1])]

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

                c = box.cls

                # print("class:   ", c)
                # annotator.box_label(b, model.names[int(c)])

                # img = annotator.result()
                # cv2.imshow('YOLO V8 Detection', cv2.resize(img, (900, 600)))
                if int(c) == 0:
                    result = ocr.ocr(lp, cls=True)
                    for idx in range(len(result)):
                        res = result[idx]
                        text_x = 10
                        text_y = 30
                        if res:

                            for ind, line in enumerate(res):
                                # print("line")
                                # print(line)

                                text_tuple = line[1]
                                if r1_txt is None:
                                    r1_txt = text_tuple[0]
                                    print("R1  TEXT:   ", r1_txt)
                                text_conf = text_tuple[1]
                elif int(c) == 1:

                    result = ocr.ocr(lp, cls=True)
                    for idx in range(len(result)):
                        res = result[idx]
                        text_x = 10
                        text_y = 30
                        if res:

                            for ind, line in enumerate(res):
                                # print("line")
                                # print(line)

                                text_tuple = line[1]
                                if r2_txt is None:
                                    r2_txt = text_tuple[0]
                                    print("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}"

            # print("LP:  ", LP)
            return LP

    return "Image Not Clear"
    # result = ocr.ocr(img_path, cls=False,det=True,rec=True)[0]
    # txts = [line[1][0] for line in result]




