import tempfile
from collections import defaultdict

import cv2
import numpy as np
from onnxruntime import InferenceSession

from paddle_detection.deploy.third_engine.onnx.infer import PredictConfig
from paddle_detection.deploy.third_engine.onnx.preprocess import Compose
from util import image_util, util


def predict_image(infer_config, predictor, img_path):
    # load preprocess transforms
    transforms = Compose(infer_config.preprocess_infos)
    # predict image
    inputs = transforms(img_path)
    inputs["image"] = np.array(inputs["image"]).astype('float32')
    inputs_name = [var.name for var in predictor.get_inputs()]
    inputs = {k: inputs[k][None,] for k in inputs_name}

    outputs = predictor.run(output_names=None, input_feed=inputs)

    bboxes = np.array(outputs[0])
    result = defaultdict(list)
    for bbox in bboxes:
        if bbox[0] > -1 and bbox[1] > infer_config.draw_threshold:
            result[bbox[0]].append({"score": bbox[1], "box": bbox[2:]})
    return result


def detect_image(img_path):
    onnx_file = "model/object_det_model/ppyoloe_plus_crn_x_80e_coco_w_nms.onnx"
    infer_cfg = "model/object_det_model/infer_cfg.yml"
    # load predictor
    predictor = InferenceSession(onnx_file)
    # load infer config
    infer_config = PredictConfig(infer_cfg)

    return predict_image(infer_config, predictor, img_path)


def get_book_areas(image):
    with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file:
        cv2.imwrite(temp_file.name, image)
    detect_result = detect_image(temp_file.name)
    util.delete_temp_file(temp_file.name)
    book_areas = detect_result[73]
    result = []
    for book_area in book_areas:
        result.append(image_util.capture(image, book_area["box"]))
    return result