import os.path
from collections import defaultdict

import cv2
import numpy as np

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 capture(image, rectangle):
    x1, y1, x2, y2 = rectangle
    height, width = image.shape[:2]
    if x1 < 0:
        x1 = 0
    if y1 < 0:
        y1 = 0
    if x2 > width:
        x2 = width
    if y2 > height:
        y2 = height
    return image[int(y1):int(y2), int(x1):int(x2)]


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(capture(image, book_area['box']))
    return result