import threading
from typing import List
from numpy import ndarray
from core.detectors.base import LayoutDetectorBase
from core.detectors.paddle_yolo.infer import Detector, print_arguments
from core.predictor import LayoutBox


class FLAGS:
    batch_size = 1
    camera_id = -1
    collect_trt_shape_info = False
    combine_method = "nms"
    cpu_threads = 1
    device = "cpu"
    enable_mkldnn = False
    enable_mkldnn_bfloat16 = False
    image_dir = None
    image_file = None
    match_metric = "ios"
    match_threshold = "0.6"
    model_dir = "models/yolov7p6_e6_300e_coco"
    output_dir = "output"
    overlap_ratio = [0.25, 0.25]
    run_benchmark = False
    run_mode = "paddle"
    save_images = False
    save_results = False
    slice_infer = False
    slice_size = [640, 640]
    threshold = 0.5
    trt_calib_mode = False
    trt_max_shape = 1280
    trt_min_shape = 1
    trt_opt_shape = 640
    tuned_trt_shape_file = "shape_range_info.pbtxt"
    use_coco_category = False
    use_gpu = False
    video_file = None


class PaddleYoloDetector(LayoutDetectorBase):
    print("======加载 PaddleYOLO 模型======")
    _detector = Detector(
        model_dir=FLAGS.model_dir,
        device=FLAGS.device.upper(),
        run_mode=FLAGS.run_mode,
        batch_size=FLAGS.batch_size,
        trt_min_shape=FLAGS.trt_min_shape,
        trt_max_shape=FLAGS.trt_max_shape,
        trt_opt_shape=FLAGS.trt_opt_shape,
        trt_calib_mode=FLAGS.trt_calib_mode,
        cpu_threads=FLAGS.cpu_threads,
        enable_mkldnn=FLAGS.enable_mkldnn,
        enable_mkldnn_bfloat16=FLAGS.enable_mkldnn_bfloat16,
        threshold=FLAGS.threshold,
        output_dir=FLAGS.output_dir,
    )
    print_arguments(_detector)
    print("========>模型加载成功")

    lock = threading.Lock()

    @classmethod
    def predict(cls, img: ndarray, **kwargs) -> List[LayoutBox]:
        try:
            cls.lock.acquire()
            predicts = cls._detector.predict_image(
                image_list=[img], repeats=1, save_results=False
            )
            results = []
            boxes = predicts["boxes"].tolist()  # type: ignore
            for box in boxes:
                clazz = int(box[0])
                conf = box[1]
                bbox = box[2:6]
                results.append(
                    LayoutBox(
                        clazz=clazz,
                        clazz_name=None,
                        bbox=bbox,
                        conf=conf,
                    )
                )
        finally:
            cls.lock.release()
        return results
