import os

import cv2
from PIL import Image
from keras.layers import Lambda, Input
from keras.models import Model
import numpy as np
from click_captcha.model import tiny_yolo_body
from click_captcha.post_process import yolo_eval, letterbox_image
from click_captcha.utils import get_classes, get_colors, draw_boxes, get_anchors, pil_resize, np2pil


def detect_char(image_path, model):
    image_np = cv2.imread(image_path)

    # create image input
    h, w = image_np.shape[:2]
    best_h, best_w = 160, 256
    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
    image_shape = np.array([best_h, best_w])
    image_shape = np.expand_dims(image_shape, 0)

    # inference data
    out_boxes, out_scores, out_classes = model.predict([image_resize, image_shape])
    print("image_size", image_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)

    # draw
    # image_np = cv2.cvtColor(np.array(image_pil), cv2.COLOR_RGB2BGR)
    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)
    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))
    image_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, image_shape])
    # boxes, scores, classes = yolo_eval([model_body.output, image_shape], anchors, num_classes)
    model = Model([model_body.input, image_shape], [boxes, scores, classes])
    model.summary(120)
    return model


if __name__ == '__main__':
    _dir = os.path.abspath(os.path.dirname(os.path.abspath(__file__)))
    model_path = _dir + "/models/e14-loss17.08-char.h5"
    anchors = get_anchors(_dir + "/yolo_data/my_anchors.txt")
    class_names = get_classes(_dir + "/yolo_data/my_classes.txt")
    colors = get_colors(len(class_names))

    yolo_model = get_tiny_inference_model(anchors, len(class_names), weights_path=model_path)
    yolo_model.load_weights(model_path)

    image_path = "D:/Project/captcha/data/test/yolo_20.jpg"
    # image_path = "../data/detect/1.jpg"
    detect_char(image_path, yolo_model)