import itertools
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,
        conf_threshold: float = 0.0,
        overlaps_iou_threshold: float = 0.9,
        overlaps_max_count: int = 5,
        **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:
                print(box)
                clazz = int(box[0])
                conf = box[1]
                bbox = box[2:6]
                results.append(
                    LayoutBox(
                        clazz=clazz, clazz_name=None, bbox=bbox, conf=conf
                    )
                )
            results = _filter_by_conf(results, conf_threshold)
            results = _filter_by_overlaps(
                results, overlaps_iou_threshold, overlaps_max_count
            )
        finally:
            cls.lock.release()
        return results


def _filter_by_conf(
    boxes: List[LayoutBox],
    conf_threshold: float,
) -> List[LayoutBox]:
    # 按置信度过滤LayoutBox
    boxes = list(filter(lambda e: e.conf >= conf_threshold, boxes))
    return boxes


def _filter_by_overlaps(
    boxes: List[LayoutBox],
    overlaps_iou_threshold: float,
    overlaps_max_count: int,
) -> List[LayoutBox]:
    """
    对多个 iou 大于 `overlaps_iou_threshold` 的区域，仅保留 `overlaps_max_count` 个置信度最高的。
    """
    # 按置信度进行排序
    boxes = sorted(boxes, key=lambda e: e.conf, reverse=True)
    # 每一个桶中都是重叠区域较大的LayoutBox
    buckets: List[List[LayoutBox]] = []

    # 将目标于每一个桶中的每一个LayoutBox进行比较，找到目标应该存在于哪一个桶
    def get_bucket(box: LayoutBox):
        for bucket in buckets:
            for e in bucket:
                if box.iou(e) >= overlaps_iou_threshold:
                    return bucket
        return None

    for layout in boxes:
        bucket = get_bucket(layout)
        # 若当前不存在于目标layout重叠的内容，则新建一个桶
        if not bucket:
            buckets.append([layout])
        # 若找到目标应该位于的桶，则只收取置信度较高的overlaps_max_count个框选区域
        elif len(bucket) < overlaps_max_count:
            bucket.append(layout)
    # 将所用桶中的数据合为一个列表
    new_layouts = list(itertools.chain.from_iterable(buckets))
    return new_layouts
