import os
import sys
from pathlib import Path

import numpy as np
import torch

from BaseDetector import baseDet
from models.common import DetectMultiBackend
from utils.datasets import letterbox
from utils.general import check_img_size, non_max_suppression, scale_coords
from utils.torch_utils import select_device

FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]  # YOLOv5 root directory
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))  # add ROOT to PATH
ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative


class Detector(baseDet):
    def __init__(self):
        super(Detector, self).__init__()
        self.init_model()
        self.build_config()

    def init_model(self):
        self.weights = "weights/yolov5s.pt"
        self.device = "0" if torch.cuda.is_available() else "cpu"
        self.imgsz = [640, 640]
        self.half = False

        # Load model
        self.device = select_device(self.device)
        self.model = DetectMultiBackend(
            self.weights, device=self.device, dnn=False, data=ROOT / "data/coco128.yaml"
        )
        stride, names, pt, jit, onnx, engine = (
            self.model.stride,
            self.model.names,
            self.model.pt,
            self.model.jit,
            self.model.onnx,
            self.model.engine,
        )
        self.stride = 32
        imgsz = check_img_size(self.imgsz, s=stride)  # check image size
        print(imgsz)

        # Half
        self.half &= (
            pt or jit or onnx or engine
        ) and self.device.type != "cpu"  # FP16 supported on limited backends with CUDA
        if pt or jit:
            self.model.model.half() if self.half else self.model.model.float()

        self.names = self.model.module.names if hasattr(self.model, "module") else names

    def preprocess(self, img):
        img0 = img.copy()
        img = letterbox(img0, self.imgsz, stride=self.stride, auto=False)[0]
        img = img.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
        img = np.ascontiguousarray(img)
        img = torch.from_numpy(img).to(self.device)
        img = img.half() if self.half else img.float()  # # uint8 to fp16/32
        img /= 255  # 0 - 255 to 0.0 - 1.0
        if len(img.shape) == 3:
            img = img[None]  # expand for batch dim

        return img0, img

    def detect(self, im):
        im0, img = self.preprocess(im)

        pred = self.model(img, augment=False)
        # pred = pred.float()
        pred = non_max_suppression(
            pred,
            self.threshold,
            iou_thres=0.45,
            classes=None,
            agnostic=False,
            multi_label=False,
            labels=(),
            max_det=300,
        )

        pred_boxes = []
        for det in pred:
            if det is not None and len(det):
                det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()

                for *x, conf, cls_id in det:
                    lbl = self.names[int(cls_id)]
                    if lbl not in ["person", "car", "truck"]:
                        continue
                    x1, y1 = int(x[0]), int(x[1])
                    x2, y2 = int(x[2]), int(x[3])
                    pred_boxes.append((x1, y1, x2, y2, lbl, conf))

        return im, pred_boxes