# YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
import argparse
import csv
import os
import platform
import sys
from pathlib import Path
import glob
import numpy as np
import pickle
import torch
from tools.config import cfg

sys.path.append('./ytracking')
from models.common import DetectMultiBackend
from utils.dataloaders import LoadImages
from utils.general import (LOGGER, Profile, check_img_size, check_requirements, colorstr, cv2,
                           increment_path, non_max_suppression, scale_boxes, strip_optimizer)
from utils.torch_utils import select_device, smart_inference_mode

'''集成跟踪模块，输出跟踪结果文件 .npy'''
# from ultralytics.engine.results import Boxes    # Results
# from ultralytics.utils import IterableSimpleNamespace, yaml_load
from tracking.utils.plotting import Annotator, colors
from tracking.utils import Boxes, IterableSimpleNamespace, yaml_load
from tracking.trackers import BOTSORT, BYTETracker
from tracking.utils.showtrack import drawtracks
import time


def init_trackers(tracker_yaml=None, bs=1):
    """
    Initialize trackers for object tracking during prediction.
    """
    # tracker_yaml = r"./tracking/trackers/cfg/botsort.yaml"

    TRACKER_MAP = {'bytetrack': BYTETracker, 'botsort': BOTSORT}
    cfg = IterableSimpleNamespace(**yaml_load(tracker_yaml))
    trackers = []
    for _ in range(bs):
        tracker = TRACKER_MAP[cfg.tracker_type](args=cfg, frame_rate=30)
        trackers.append(tracker)

    return trackers


@smart_inference_mode()
def run(
        # weights=cfg.tracking_model,  # model path or triton URL
        Model,  # model path or triton URL
        source=None,  # file/dir/URL/glob/screen/0(webcam)

        project=r'./runs/detect',  # save results to project/name

        tracker_yaml=cfg.botsort,
        imgsz=(640, 640),  # inference size (height, width)
        conf_thres=0.25,  # confidence threshold
        iou_thres=0.45,  # NMS IOU threshold
        max_det=1000,  # maximum detections per image
        device='',  # cuda device, i.e. 0 or 0,1,2,3 or cpu
        bs=1,  # batch_size
        save_img=True,  # do not save images/videos
        classes=None,  # filter by class: --class 0, or --class 0 2 3
        agnostic_nms=False,  # class-agnostic NMS
        augment=False,  # augmented inference
        visualize=False,  # visualize features
        line_thickness=3,  # bounding box thickness (pixels)
        half=False,  # use FP16 half-precision inference
        dnn=False,  # use OpenCV DNN for ONNX inference
        vid_stride=1,  # video frame-rate stride
):
    if source is None:
        raise ValueError("Have to provide --source argument")

    # Load model
    # device = select_device(device)
    # model = DetectMultiBackend(weights, device=device, dnn=dnn, fp16=half)
    if Model is None:
        raise ValueError("Have to provide --model argument")
    model = Model.yoloModel
    print(model.stride, model.names, model.pt)
    stride, names, pt = model.stride, model.names, model.pt
    imgsz = check_img_size(imgsz, s=stride)  # check image size

    # Run inference
    model.warmup(imgsz=(1 if pt or model.triton else bs, 3, *imgsz))  # warmup

    ##=============================生成文件夹 save_dir，存储检测跟踪图像 
    source = str(source)
    save_dir = Path(project) / Path(source).stem

    # Dataloader
    seen, dt = 0, (Profile(), Profile(), Profile())
    dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)

    ## ================================================= 生成跟踪器对象
    tracker = init_trackers(tracker_yaml, bs)[0]
    track_boxes = np.empty((0, 9), dtype=np.float32)
    features_dict = {}
    frameid_img = {}

    for path, im, im0s, vid_cap, s in dataset:
        # img preprocess
        with dt[0]:
            im = torch.from_numpy(im).to(model.device)
            im = im.half() if 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 dt[1]:
            visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False
            pred = model(im, augment=augment, visualize=visualize)

        # NMS
        with dt[2]:
            pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)

        # Process predictions
        for i, det in enumerate(pred):  # per image
            seen += 1
            frameid_img[seen] = im0s.copy()

            p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0)

            s += '%gx%g ' % im.shape[2:]  # print string

            if len(det):
                # Rescale boxes from img_size to im0 size
                det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()

                # boxes_and_imgs.append((det.cpu().numpy(), im0, frame))
                ## ================================================================ writed by WQG
                det_tracking = Boxes(det, im0.shape).cpu().numpy()
                tracks = tracker.update(det_tracking, im0)

                if len(tracks):
                    tracks[:, 7] = seen
                    track_boxes = np.concatenate([track_boxes, tracks], axis=0)
                    feat_dict = {int(x.idx): x.curr_feat for x in tracker.tracked_stracks if x.is_activated}
                    frame_id = tracks[0, 7]
                    features_dict.update({int(frame_id): feat_dict})

    return track_boxes, features_dict, frameid_img


def main():
    ROOT = Path(Path.cwd())
    check_requirements(ROOT / 'requirements.txt', exclude=('tensorboard', 'thop'))

    optdict = {'weights': r"D:/Project/ieemoo-ai/tools/ckpts/best_158734_cls11_noaug10.pt",
               'source': r"D:/Project/ieemoo-ai/testdata/88.mp4",
               }
    run(**optdict)


if __name__ == '__main__':
    main()