import os
import sys
import tensorflow as tf
import cv2
from keras.layers import Lambda, Input
from keras.models import Model
import numpy as np
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/../")
from chinese_detect.model_260 import tiny_yolo_body
from chinese_detect.post_process import yolo_eval
from utils import get_classes, get_colors, draw_boxes, get_anchors, pil_resize, np2pil


package_dir = os.path.abspath(os.path.dirname(__file__))
model_path = package_dir + "/models/char_yolo_loss_39.90.h5"
anchors = get_anchors(package_dir + "/yolo_data/my_anchors.txt")
classes = get_classes(package_dir + "/yolo_data/my_classes.txt")
colors = get_colors(len(classes))
image_shape = (160, 256, 3)
tips_shape = (40, 160, 3)


def detect(image_np, model=None, sess=None, draw=False, is_tips=0):
    if sess is None:
        sess = tf.compat.v1.Session(graph=tf.Graph())

    if model is None:
        with sess.as_default():
            with sess.graph.as_default():
                model = get_tiny_inference_model(anchors, len(classes), weights_path=model_path)

    # create image input
    h, w = image_np.shape[:2]
    # print("h, w", h, w)
    best_h, best_w = image_shape[:2]

    # 如是提示图片需paste
    if is_tips:
        image_tips = pil_resize(image_np, tips_shape[0], tips_shape[1])
        image_resize = np.zeros(image_shape, dtype=np.uint8)
        image_resize[:tips_shape[0], :tips_shape[1], :] = image_tips[:, :, :]
    else:
        image_resize = pil_resize(image_np, best_h, best_w)
    # image_pil = np2pil(image_resize)
    image_resize = np.array(image_resize, dtype='float32')
    image_resize = image_resize.astype('float32') / 255.
    image_resize = np.expand_dims(image_resize, 0)

    # create image shape input
    need_shape = np.array([best_h, best_w])
    need_shape = np.expand_dims(need_shape, 0)

    # inference data
    with sess.as_default():
        with sess.graph.as_default():
            out_boxes, out_scores, out_classes = model.predict([image_resize, need_shape], steps=1)
    # print("image_size", need_shape)
    print("out_boxes", out_boxes)
    print("out_scores", out_scores)
    # print("out_classes", out_classes)
    out_boxes = out_boxes.astype(np.int32)
    out_classes = out_classes.astype(np.int32)

    # 还原
    if is_tips:
        out_boxes[:, 0] = h * out_boxes[:, 0] / tips_shape[0]
        out_boxes[:, 2] = h * out_boxes[:, 2] / tips_shape[0]
        out_boxes[:, 1] = w * out_boxes[:, 1] / tips_shape[1]
        out_boxes[:, 3] = w * out_boxes[:, 3] / tips_shape[1]
    else:
        out_boxes[:, 0] = h * out_boxes[:, 0] / best_h
        out_boxes[:, 2] = h * out_boxes[:, 2] / best_h
        out_boxes[:, 1] = w * out_boxes[:, 1] / best_w
        out_boxes[:, 3] = w * out_boxes[:, 3] / best_w

    image_pil = np2pil(image_np)
    if draw:
        # draw
        class_names = get_classes("yolo_data/my_classes.txt")
        colors = get_colors(len(class_names))
        image_resize, out_boxes = draw_boxes(image_pil, out_boxes, out_classes, out_scores, class_names, colors)
        image_np_result = cv2.cvtColor(np.array(image_resize), cv2.COLOR_RGB2BGR)
        cv2.imshow("result", image_np_result)
        cv2.waitKey(0)
    else:
        temp_boxes = []
        for i, c in reversed(list(enumerate(out_classes))):
            top, left, bottom, right = out_boxes[i]
            top = max(0, np.floor(top + 0.5).astype('int32'))
            left = max(0, np.floor(left + 0.5).astype('int32'))
            bottom = min(image_pil.size[1], np.floor(bottom + 0.5).astype('int32'))
            right = min(image_pil.size[0], np.floor(right + 0.5).astype('int32'))
            temp_boxes.append([(left, top), (right, bottom)])
        out_boxes = temp_boxes
    out_boxes = [[int(x[0][0]), int(x[0][1]), int(x[1][0]), int(x[1][1])] for x in out_boxes]
    out_boxes.sort(key=lambda x: (x[0], x[1], x[2], x[3]))
    return image_np, out_boxes, out_classes


def get_tiny_inference_model(anchors, num_classes, weights_path='models/tiny_yolo_weights.h5'):
    """create the inference model, for Tiny YOLOv3"""
    image_input = Input(shape=(None, None, 3))
    need_shape = Input(shape=(2,), dtype='int64', name='image_shape')
    num_anchors = len(anchors)

    model_body = tiny_yolo_body(image_input, num_anchors//2, num_classes)
    print('Create Tiny YOLOv3 model with {} anchors and {} classes.'.format(num_anchors, num_classes))

    model_body.load_weights(weights_path)
    print('Load weights {}.'.format(weights_path))

    boxes, scores, classes = Lambda(yolo_eval,
                                    name='yolo_eval',
                                    arguments={'anchors': anchors,
                                               'num_classes': num_classes}
                                    )([model_body.output, need_shape])
    model = Model([model_body.input, need_shape], [boxes, scores, classes])
    # model.summary(120)
    return model


if __name__ == '__main__':
    image_path = "D:/Project/captcha/data/test/yolo_3.jpg"
    _img = cv2.imread(image_path)
    cv2.imshow("origin_image", _img)
    _, boxes, _ = detect(_img, is_tips=1, draw=False)
    for box in boxes:
        cv2.imshow("sub", _img[box[1]:box[3], box[0]:box[2], :])
        cv2.waitKey(0)