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 puzzle_detect.model_260 import tiny_yolo_body
from puzzle_detect.post_process import yolo_eval, letterbox_image
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/puzzle_yolo_loss_4.15.h5"
anchors = get_anchors(package_dir + "/yolo_data/my_anchors_puzzle.txt")
classes = get_classes(package_dir + "/yolo_data/my_classes_puzzle.txt")
colors = get_colors(len(classes))
image_shape = (160, 256, 3)


def detect(image_np, model=None, sess=None, draw=False):
    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]
    best_h, best_w = image_shape[:2]
    image_resize = pil_resize(image_np, best_h, best_w)
    # image_pil = np2pil(image_resize)
    image_resize = cv2.cvtColor(image_resize, cv2.COLOR_BGR2GRAY)
    image_resize = 255. - image_resize
    image_resize = np.uint8(image_resize)
    image_resize = image_resize / 255.
    image_resize = np.expand_dims(image_resize, 0)
    image_resize = np.expand_dims(image_resize, -1)

    # 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("out_boxes", out_boxes)
    print("out_scores", out_scores)
    out_boxes = out_boxes.astype(np.int32)
    out_classes = out_classes.astype(np.int32)

    # 还原
    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
        # image_np = cv2.cvtColor(np.array(image_pil), cv2.COLOR_RGB2BGR)
        class_names = get_classes("yolo_data/my_classes_puzzle.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]
    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, 1))
    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_12.jpg"
    detect(cv2.imread(image_path))