import os
import sys
import cv2
import torch
import numpy as np


from utils.general import Profile, non_max_suppression, scale_boxes, check_img_size
from utils.augmentations import letterbox
from models.common import DetectMultiBackend
from utils.torch_utils import select_device, smart_inference_mode
from ultralytics.utils.plotting import Annotator, colors, save_one_box

from pathlib import Path
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 yolov5(object):
    def __init__(self, weights, conf_thresh=0.45, iou_thresh = 0.45, classes=None, imgsz=(640,640), img_size=640, stride=32, autoSize = True, device=0) -> None:
        # Load model
        device = select_device(device)
        data = ROOT / 'data/coco128.yaml'
        self.img_size = img_size
        self.stride = stride
        self.auto = autoSize
        self.model = DetectMultiBackend(weights, device=device, dnn=False, data=data, fp16=False)
        stride, self.names, pt = self.model.stride, self.model.names, self.model.pt
        imgsz = check_img_size(imgsz, s=stride)  # check image size
        self.seen, self.windows, self.dt = 0, [], (Profile(), Profile(), Profile())
        self.augment = False
        self.visualize = False,  # visualize features
        self.conf_thres = conf_thresh
        self.iou_thres = iou_thresh
        self.classes = classes
        self.hide_conf = False
        self.hide_labels = False
        self.transforms = None
        self.imgs = [None] * 1
        


    def run(self, img):
        s = ''
        self.imgs[0] = img.copy()
        im0 = img.copy()
        annotator = Annotator(self.imgs[0], line_width=3, example=str(self.names))
        # if self.transforms:
        #     im = np.stack([self.transforms(x) for x in im0])  # transforms
        # else:
        im = np.stack([letterbox(x, self.img_size, stride=self.stride, auto=self.auto)[0] for x in self.imgs])  # resize
        im = im[..., ::-1].transpose((0, 3, 1, 2))  # BGR to RGB, BHWC to BCHW
        im = np.ascontiguousarray(im)  # contiguous

        with self.dt[0]:
            im = torch.from_numpy(im).to(self.model.device)
            im = im.half() if self.model.fp16 else im.float()  # uint8 to fp16/32
            im /= 255  # 0 - 255 to 0.0 - 1.0
            if len(im.shape) == 3:
                im = im[None]  # expand for batch dim

        # Inference
        with self.dt[1]:
            # visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False
            pred = self.model(im, augment=self.augment, visualize=False)

        # NMS
        with self.dt[2]:
            pred = non_max_suppression(pred, self.conf_thres, self.iou_thres, self.classes, agnostic=False, max_det=100)

        # Process predictions
        dets = []
        for i, det in enumerate(pred):  # per image
            if len(det):
                # Rescale boxes from img_size to im0 size
                det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()

                # Print results
                for c in det[:, 5].unique():
                    n = (det[:, 5] == c).sum()  # detections per class
                    s += f"{n} {self.names[int(c)]}{'s' * (n > 1)}, "  # add to string

                # Write results
                for *xyxy, conf, cls in reversed(det):
                    c = int(cls)  # integer class
                    label = self.names[c] if self.hide_conf else f'{self.names[c]}'
                    confidence = float(conf)
                    # box = xyxy.tolist()
                    dets.append([int(xyxy[0].cpu().numpy()), 
                                 int(xyxy[1].cpu().numpy()), 
                                 int(xyxy[2].cpu().numpy()), 
                                 int(xyxy[3].cpu().numpy()), 
                                 confidence,
                                 c])
                    
                    # confidence_str = f'{confidence:.2f}'
                    # label = None if self.hide_labels else (self.names[c] if self.hide_conf else f'{self.names[c]} {conf:.2f}')
                    # annotator.box_label(xyxy, label, color=colors(c, True))

            # im0 = annotator.result()
            # cv2.imshow("yolov5", im0)
            # cv2.waitKey(1)  # 1 millisecond

            return np.array(dets)