import os.path
from collections import defaultdict

import cv2
import numpy as np

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


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):
    infer_cfg = os.path.join(MODEL_DIR, 'infer_cfg.yml')
    # load infer config
    infer_config = PredictConfig(infer_cfg)

    return predict_image(infer_config, PADDLE_DET, img_path)


def get_book_areas(img_path):
    detect_result = detect_image(img_path)
    book_areas = detect_result[73]
    result = []
    image = cv2.imread(img_path)
    for book_area in book_areas:
        result.append(image_util.capture(image, book_area['box']))
    return result