import torch
import numpy as np
import cv2
import os

BASE_DIR = os.path.abspath(os.getcwd())

model_plates = torch.hub.load('ultralytics/yolov5', 'custom',
                              path=os.path.join(BASE_DIR, 'detector', 'static', 'plates.pt'))
# model.conf = 0.60  # NMS confidence threshold
# model.iou = 0.60  # NMS IoU threshold
# model.agnostic = False  # NMS class-agnostic
# model.multi_label = False  # NMS multiple labels per box
# model.classes = None  # (optional list) filter by class, i.e. = [0, 15, 16] for COCO persons, cats and dogs
# model.max_det = 1  # maximum number of detections per image
# model.amp = False  # Automatic Mixed Precision (AMP) inference

model_chars = torch.hub.load('ultralytics/yolov5', 'custom',
                             path=os.path.join(BASE_DIR, 'detector', 'static', 'chars.pt'))


def pad_img_to_fit_bbox(img, x1, x2, y1, y2):
    img = np.pad(img, ((np.abs(np.minimum(0, y1)), np.maximum(y2 - img.shape[0], 0)),
                       (np.abs(np.minimum(0, x1)), np.maximum(x2 - img.shape[1], 0)), (0, 0)), mode="constant")
    y1 += np.abs(np.minimum(0, y1))
    y2 += np.abs(np.minimum(0, y1))
    x1 += np.abs(np.minimum(0, x1))
    x2 += np.abs(np.minimum(0, x1))
    return img, x1, x2, y1, y2


def imcrop(img, bbox):
    x1, y1, x2, y2 = bbox
    if x1 < 0 or y1 < 0 or x2 > img.shape[1] or y2 > img.shape[0]:
        img, x1, x2, y1, y2 = pad_img_to_fit_bbox(img, x1, x2, y1, y2)
    return img[y1:y2, x1:x2, :]


def detect_plates(img):
    detect = model_plates(img)
    records = detect.pandas().xyxy[0].to_dict(orient='records')
    plates = []
    if records:
        for plate in records:
            xi, yi, xf, yf = int(plate['xmin']), int(plate['ymin']), int(plate['xmax']), int(plate['ymax'])
            crop = imcrop(img, (xi, yi, xf, yf))
            plates.append(((xi, yi), (xf, yf), crop))
    return plates


def detect_chars(img):
    img = cv2.resize(img, (640, 480))
    detect = model_chars(img)
    records = detect.pandas().xyxy[0].to_dict(orient='records')
    yolo = np.squeeze(detect.render())
    text = ''
    if records:
        records = sorted(records, key=lambda d: d['xmin'])
        text = ''.join([i.get('name') for i in records])
    return text, img


def save_plates(img):
    detect = model_plates(img)
    detect.crop(save=True)


# def yolo_detections_to_norfair_detections(yolo_detections, track_points="centroid"):
#     """convert detections_as_xywh to norfair detections"""
#     norfair_detections = []
#
#     if track_points == "centroid":
#         detections_as_xywh = yolo_detections.xywh[0]
#         for detection_as_xywh in detections_as_xywh:
#             centroid = np.array(
#                 [detection_as_xywh[0].item(), detection_as_xywh[1].item()]
#             )
#             scores = np.array([detection_as_xywh[4].item()])
#             norfair_detections.append(
#                 Detection(
#                     points=centroid,
#                     scores=scores,
#                     label=int(detection_as_xywh[-1].item()),
#                 )
#             )
#     elif track_points == "bbox":
#         detections_as_xyxy = yolo_detections.xyxy[0]
#         for detection_as_xyxy in detections_as_xyxy:
#             bbox = np.array(
#                 [
#                     [detection_as_xyxy[0].item(), detection_as_xyxy[1].item()],
#                     [detection_as_xyxy[2].item(), detection_as_xyxy[3].item()],
#                 ]
#             )
#             scores = np.array(
#                 [detection_as_xyxy[4].item(), detection_as_xyxy[4].item()]
#             )
#             norfair_detections.append(
#                 Detection(
#                     points=bbox, scores=scores, label=int(detection_as_xyxy[-1].item())
#                 )
#             )
#
#     return norfair_detections