import cv2
import time
import math
import os
import numpy as np
import tensorflow as tf
import zipfile
import locality_aware_nms as nms_locality
import lanms
from matplotlib import pyplot as plt
import imutils

# fixme: 1. 按比例预留：
tf_config = tf.ConfigProto()
tf_config.gpu_options.per_process_gpu_memory_fraction = 0.30 # 分配50%
# session = tf.Session(config=tf_config)
# fixme: 2. 或者干脆自适应然后自动增长：
# tf_config = tf.ConfigProto()
# tf_config.gpu_options.allow_growth = True # 自适应
# session = tf.Session(config=tf_config)


trainID_list = ['M0101','M0201','M0202','M0204','M0206','M0207',
                'M0210','M0301','M0401','M0402','M0501','M0603',
                'M0604','M0605','M0702','M0703','M0704','M0901',
                'M0902','M1002','M1003','M1005','M1006','M1008',
                'M1102','M1201','M1202','M1304','M1305','M1306']

# testID_list = [ 'M0203','M0205','M0208','M0209','M0403','M0601',
#                 'M0602','M0606','M0701','M0801','M0802','M1001',
#                 'M1004','M1007','M1009','M1101','M1301','M1302',
#                 'M1303','M1401']

testID_list = ["Test_images"]

tf.app.flags.DEFINE_string('test_data_path', './UAV-benchmark-M', '')
tf.app.flags.DEFINE_string('gpu_list', '1', '')
tf.app.flags.DEFINE_integer('batch_size_per_gpu', 1, '')
tf.app.flags.DEFINE_string('checkpoint_path', 'GAnet_Vgg16_B10', '')
tf.app.flags.DEFINE_string('Img_output_dir', './GAnet_Vgg16_B10_B10_Res/img_B10_50001', '')
tf.app.flags.DEFINE_string('Txt_output_dir', './GAnet_Vgg16_B10_Res/txt_B10_50001', '')
tf.app.flags.DEFINE_string('ImgHeatMap_output_dir', './GAnet_Vgg16_B10_B10_Res/imgHeatMap_B10_50001', '')
tf.app.flags.DEFINE_bool('no_write_images', False, 'do not write images')

import GAnet_model_UAV_Vgg16 as model
from GAnet_data import restore_rectangle


FLAGS = tf.app.flags.FLAGS


def get_images(data_path):
    '''
    find image files in test data path
    :return: list of files found
    '''
    files = []
    exts = ['jpg', 'png', 'jpeg', 'JPG']
    for parent, dirnames, filenames in os.walk(data_path):
        for filename in filenames:
            for ext in exts:
                if filename.endswith(ext):
                    files.append(os.path.join(parent, filename))
                    break
    print('Find {} images'.format(len(files)))
    return files


def resize_image(im, max_side_len=1024):
    '''
    resize image to a size multiple of 32 which is required by the network
    :param im: the resized image
    :param max_side_len: limit of max image size to avoid out of memory in gpu
    :return: the resized image and the resize ratio
    '''
    h, w, _ = im.shape

    resize_w = w
    resize_h = h

    # limit the max side
    if max(resize_h, resize_w) > max_side_len:
        ratio = float(max_side_len) / resize_h if resize_h > resize_w else float(max_side_len) / resize_w
    else:
        ratio = 1.
    resize_h = int(resize_h * ratio)
    resize_w = int(resize_w * ratio)

    resize_h = resize_h if resize_h % 32 == 0 else (resize_h // 32 - 1) * 32
    resize_w = resize_w if resize_w % 32 == 0 else (resize_w // 32 - 1) * 32
    resize_h = max(32, resize_h)
    resize_w = max(32, resize_w)
    im = cv2.resize(im, (int(resize_w), int(resize_h)))

    ratio_h = resize_h / float(h)
    ratio_w = resize_w / float(w)

    return im, (ratio_h, ratio_w)


def detect(score_map, geo_map, timer, scale = 4, score_map_thresh=0.8, box_thresh=0.1, nms_thres=0.2):
    '''
    restore text boxes from score map and geo map
    :param score_map:
    :param geo_map:
    :param timer:
    :param score_map_thresh: threshhold for score map
    :param box_thresh: threshhold for boxes
    :param nms_thres: threshold for nms
    :return:
    '''
    if len(score_map.shape) == 4:
        score_map = score_map[0, :, :, 0]
        geo_map = geo_map[0, :, :, ]
    # filter the score map
    xy_text = np.argwhere(score_map > score_map_thresh)
    # sort the text boxes via the y axis
    xy_text = xy_text[np.argsort(xy_text[:, 0])]
    # restore
    start = time.time()
    text_box_restored = restore_rectangle(xy_text[:, ::-1]*4, geo_map[xy_text[:, 0], xy_text[:, 1], :]) # N*4*2
    print('{} text boxes before nms'.format(text_box_restored.shape[0]))
    boxes = np.zeros((text_box_restored.shape[0], 9), dtype=np.float32)
    boxes[:, :8] = text_box_restored.reshape((-1, 8))
    boxes[:, 8] = score_map[xy_text[:, 0], xy_text[:, 1]]
    timer['restore'] = time.time() - start
    # nms part
    start = time.time()
    #boxes = nms_locality.nms_locality(boxes.astype(np.float64), nms_thres)
    boxes = lanms.merge_quadrangle_n9(boxes.astype('float32'), nms_thres)
    timer['nms'] = time.time() - start

    if boxes.shape[0] == 0:
        return np.array([]), timer

    boxes_10 = np.zeros((boxes.shape[0], 10), dtype=np.float32)
    boxes_10[:, :9] = boxes
    boxes_10[:, 9] = boxes[:, 8]
    for i, box in enumerate(boxes_10):
        mask = np.zeros_like(score_map, dtype=np.uint8)
        cv2.fillPoly(mask, box[:8].reshape((-1, 4, 2)).astype(np.int32) // scale, 1)
        boxes_10[i, 8] = cv2.mean(score_map, mask)[0]
    boxes_10 = boxes_10[boxes_10[:, 8] > box_thresh]
    boxes_10[:, 8] = boxes_10[:, 9]
    boxes = boxes_10[:, :9]

    return boxes, timer


def detect_HeatMap(score_map, geo_map, timer, scale = 2, score_map_thresh=0.8, box_thresh=0.1, nms_thres=0.2):
    '''
    restore text boxes from score map and geo map
    :param score_map:  four scale 128,64,32,16
    :param geo_map:    four scale 128,64,32,16
    :param feature:    four scale pool2, pool3, pool4, pool5
    :param timer:
    :param scale: based on the ratio of original image,
                  for examle: 4 denotes the size of score_map needs to be magnified 4 times to
                  reach the original size.    # 0, 1, 2, 3  scale = 32, 16, 8, 4
    :param score_map_thresh: threshhold for score map
    :param box_thresh: threshhold for boxes
    :param nms_thres: threshold for nms
    :return:
    '''
    # heat_scoreMap = [] # scale = 0, 1, 2, 3 with 32, 16, 8, 4
    # heat_feature = []  # scale = 0, 1, 2, 3 with 32, 16, 8, 4
    if len(score_map.shape) == 4:
        score_map = score_map[0, :, :, 0]
        geo_map = geo_map[0, :, :, ]

    # first,copy the score_map; second, resize to range 0-255; finally, show and return.
    # In fact, the map can be shown via plt.imshow() with float parameter [0-1.0].
    heat_map = np.copy(score_map)

    # filter the score map
    xy_text = np.argwhere(score_map > score_map_thresh)
    # sort the text boxes via the y axis
    xy_text = xy_text[np.argsort(xy_text[:, 0])]
    # restore
    start = time.time()
    text_box_restored = restore_rectangle(xy_text[:, ::-1]*scale, geo_map[xy_text[:, 0], xy_text[:, 1], :]) # N*4*2
    print('{} text boxes before nms'.format(text_box_restored.shape[0]))
    boxes = np.zeros((text_box_restored.shape[0], 9), dtype=np.float32)
    boxes[:, :8] = text_box_restored.reshape((-1, 8))
    boxes[:, 8] = score_map[xy_text[:, 0], xy_text[:, 1]]
    timer['restore'] = time.time() - start
    # nms part
    start = time.time()

    # FIXME: CAI note, 这里有两个版本的 NMS，第一个使用python写的标准版本，第二个是c++写的加速版本的。
    # boxes = nms_locality.nms_locality(boxes.astype(np.float64), nms_thres)              # Pure python
    boxes = lanms.merge_quadrangle_n9(boxes.astype('float32'), nms_thres)             # Python + C++
    boxes = boxes.astype(np.float64)

    timer['nms'] = time.time() - start

    if boxes.shape[0] == 0:
        return np.array([]), timer, heat_map


    for i, box in enumerate(boxes):
        mask = np.zeros_like(score_map, dtype=np.uint8)
        cv2.fillPoly(mask, box[:8].reshape((-1, 4, 2)).astype(np.int32) // scale, 1)
        boxes[i, 8] = cv2.mean(score_map, mask)[0]
    # boxes = boxes[boxes[:, 8] > box_thresh]

    return boxes, timer, heat_map



def sort_poly(p):
    min_axis = np.argmin(np.sum(p, axis=1))
    p = p[[min_axis, (min_axis+1)%4, (min_axis+2)%4, (min_axis+3)%4]]
    if abs(p[0, 0] - p[1, 0]) > abs(p[0, 1] - p[1, 1]):
        return p
    else:
        return p[[0, 3, 2, 1]]


def main(argv=None):
    import os
    os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu_list


    if not os.path.exists(FLAGS.Img_output_dir):
        os.makedirs(FLAGS.Img_output_dir)
    if not os.path.exists(FLAGS.Txt_output_dir):
        os.makedirs(FLAGS.Txt_output_dir)


    with tf.get_default_graph().as_default():
        input_images = tf.placeholder(tf.float32, shape=[None, None, None, 3], name='input_images')
        global_step = tf.get_variable('global_step', [], initializer=tf.constant_initializer(0), trainable=False)

        f_score, f_geometry = model.model(input_images, is_training=False)

        variable_averages = tf.train.ExponentialMovingAverage(0.997, global_step)
        saver = tf.train.Saver(variable_averages.variables_to_restore())

        #with tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) as sess:
        with tf.Session(config=tf_config) as sess:
            ckpt_state = tf.train.get_checkpoint_state(FLAGS.checkpoint_path)
            model_path = os.path.join(FLAGS.checkpoint_path, os.path.basename(ckpt_state.model_checkpoint_path))
            print('Restore from {}'.format(model_path))
            saver.restore(sess, model_path)

            im_fn_list = get_images()
            sum_time = 0.0
            for im_fn in im_fn_list:
                im = cv2.imread(im_fn)[:, :, ::-1]
                start_time = time.time()
                im_resized, (ratio_h, ratio_w) = resize_image(im)

                timer = {'net': 0, 'restore': 0, 'nms': 0}
                start = time.time()
                score, geometry = sess.run([f_score, f_geometry], feed_dict={input_images: [im_resized]})
                timer['net'] = time.time() - start

                boxes, timer = detect(score_map=score, geo_map=geometry, timer=timer)
                print('{} : net {:.0f}ms, restore {:.0f}ms, nms {:.0f}ms'.format(
                    im_fn, timer['net']*1000, timer['restore']*1000, timer['nms']*1000))

                # if boxes is not None:
                if len(boxes) != 0:
                    scores = boxes[:, 8]
                    boxes = boxes[:, :8].reshape((-1, 4, 2))
                    boxes[:, :, 0] /= ratio_w
                    boxes[:, :, 1] /= ratio_h


                duration = time.time() - start_time
                print('[timing] {}'.format(duration))
                sum_time += duration

                # save to file
                # if boxes is not None:
                if True:
                    res_file = os.path.join(
                        FLAGS.Txt_output_dir,
                        '{}.txt'.format(
                            os.path.basename(im_fn).split('.')[0]))

                    with open(res_file, 'w') as f:
                        for idx, box in enumerate(boxes):
                            # to avoid submitting errors
                            box = sort_poly(box.astype(np.int32))
                            if np.linalg.norm(box[0] - box[1]) < 5 or np.linalg.norm(box[3]-box[0]) < 5:
                                continue
                            # Fixme: remove score
                            f.write('{},{},{},{},{},{},{},{},{}\r\n'.format(
                                box[0, 0], box[0, 1], box[1, 0], box[1, 1], box[2, 0], box[2, 1], box[3, 0], box[3, 1],scores[idx]
                            )),
                            cv2.polylines(im[:, :, ::-1], [box.astype(np.int32).reshape((-1, 1, 2))], True, color=(0, 255, 0), thickness=3)
                if not FLAGS.no_write_images:
                    img_path = os.path.join(FLAGS.Img_output_dir, os.path.basename(im_fn))
                    cv2.imwrite(img_path, im[:, :, ::-1])

            print('Average second per frame is {}'.format(sum_time / len(im_fn_list)))
            print('Average frame per second is {}'.format(1.0 / (sum_time / len(im_fn_list))))
            zip_ForHRN(FLAGS.Txt_output_dir, FLAGS.Txt_output_dir + '.zip')


def main_BatchEva(argv=None):
    import os
    os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu_list

    print("hello world")

    with tf.get_default_graph().as_default():
        input_images = tf.placeholder(tf.float32, shape=[None, None, None, 3], name='input_images')
        global_step = tf.get_variable('global_step', [], initializer=tf.constant_initializer(0), trainable=False)

        f_score, f_geometry, F_score_KR, F_logits = model.model(input_images, is_training=False)

        variable_averages = tf.train.ExponentialMovingAverage(0.997, global_step)
        saver = tf.train.Saver(variable_averages.variables_to_restore())

        with tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) as sess:
            # For batch evaluation
            pattern_str_original = "model_checkpoint_path:'model.ckpt-step'"
            point_list = [ str(e) for e in range(50001, 30000, -5000)]
            for step in point_list:
                pattern_str = pattern_str_original.replace('step', step)
                with open(os.path.join(FLAGS.checkpoint_path, 'checkpoint'), 'w') as fp:
                    fp.write(pattern_str)
                idx = FLAGS.Txt_output_dir.rfind("_")
                Local_txt_output_dir = FLAGS.Txt_output_dir.replace(FLAGS.Txt_output_dir[idx + 1:], step)
                Local_img_output_dir_tmp = FLAGS.Img_output_dir.replace(FLAGS.Img_output_dir[idx + 1:], step)
                Local_ImgHeatMap_output_dir_tmp = FLAGS.ImgHeatMap_output_dir.replace(FLAGS.ImgHeatMap_output_dir[idx+1:], step )

                if not os.path.exists(Local_txt_output_dir):
                    os.makedirs(Local_txt_output_dir)
                if not os.path.exists(Local_img_output_dir_tmp):
                    os.makedirs(Local_img_output_dir_tmp)
                if not os.path.exists(Local_ImgHeatMap_output_dir_tmp):
                    os.makedirs(Local_ImgHeatMap_output_dir_tmp)

                ckpt_state = tf.train.get_checkpoint_state(FLAGS.checkpoint_path)
                model_path = os.path.join(FLAGS.checkpoint_path, os.path.basename(ckpt_state.model_checkpoint_path))
                print('Restore from {}'.format(model_path))
                saver.restore(sess, model_path)
                sum_time = 0.0
                im_num = 0

                for video_id in testID_list:
                    video_frame_dir = os.path.join( FLAGS.test_data_path, video_id)

                    fw = open(os.path.join(Local_txt_output_dir, video_id+".txt"), 'w')
                    Local_img_output_dir = os.path.join( Local_img_output_dir_tmp, video_id)
                    if not os.path.exists(Local_img_output_dir):
                        os.makedirs(Local_img_output_dir)
                    Local_ImgHeatMap_output_dir = os.path.join(Local_ImgHeatMap_output_dir_tmp, video_id)
                    if not os.path.exists(Local_ImgHeatMap_output_dir):
                        os.makedirs(Local_ImgHeatMap_output_dir)

                    im_fn_list = get_images(video_frame_dir)
                    im_num += len(im_fn_list)
                    for im_fn in im_fn_list:
                        im_name = os.path.split(im_fn)[1]
                        im_name_fg = os.path.splitext(im_name)[0].replace('img','')  # img000123.jpg==>img000123==>000123
                        im = cv2.imread(im_fn)[:, :, ::-1]
                        height, width, _ = im.shape
                        start_time = time.time()
                        im_resized, (ratio_h, ratio_w) = resize_image(im)

                        timer = {'net': 0, 'restore': 0, 'nms': 0}
                        start = time.time()

                        score, geometry, score_KR, logits = \
                            sess.run([f_score, f_geometry, F_score_KR, F_logits],
                                               feed_dict={input_images: [im_resized]})

                        timer['net'] = time.time() - start

                        # boxes, timer = detect(score_map=score, geo_map=geometry, timer=timer)
                        boxes, timer, heat_map = detect_HeatMap(score_map=score, geo_map=geometry, timer=timer)
                        print('{} : net {:.0f}ms, restore {:.0f}ms, nms {:.0f}ms'.format(
                            im_fn, timer['net']*1000, timer['restore']*1000, timer['nms']*1000))

                        # if boxes is not None:
                        if len(boxes) != 0:
                            scores = boxes[:, 8]
                            boxes = boxes[:, :8].reshape((-1, 4, 2))
                            boxes[:, :, 0] /= ratio_w
                            boxes[:, :, 1] /= ratio_h

                        duration = time.time() - start_time
                        print('[timing] {}'.format(duration))
                        sum_time += duration

                        if len(boxes) == 0:
                            continue
                        for idx, box in enumerate(boxes):
                            # to avoid submitting errors
                            box = sort_poly(box.astype(np.int32))
                            # box = check_and_validate(box)
                            if np.linalg.norm(box[0] - box[1]) < 5 or np.linalg.norm(box[3]-box[0]) < 5:
                                continue

                            box[0, 0] = max(0, min(box[0, 0], width))
                            box[1, 0] = max(0, min(box[1, 0], width))
                            box[2, 0] = max(0, min(box[2, 0], width))
                            box[3, 0] = max(0, min(box[3, 0], width))
                            box[0, 1] = max(0, min(box[0, 1], height))
                            box[1, 1] = max(0, min(box[1, 1], height))
                            box[2, 1] = max(0, min(box[2, 1], height))
                            box[3, 1] = max(0, min(box[3, 1], height))
                            x_min = min(box[:,0]) + 0.01
                            x_max = max(box[:,0]) + 0.01
                            y_min = min(box[:,1]) + 0.01
                            y_max = max(box[:, 1]) + 0.01
                            box_width = x_max - x_min
                            box_height = y_max - y_min

                            # 24,-1, 195.0,156.0,54.0,56.0, 0.021061,1,-1
                            fw.write('{},{},{},{},{},{},{},{},{}\r\n'.format(
                                int(im_name_fg), -1, x_min, y_min, box_width, box_height,scores[idx],1,-1
                            ))
                            cv2.polylines(im[:, :, ::-1], [box.astype(np.int32).reshape((-1, 1, 2))], True, color=(0, 255, 0), thickness=2)
                        # Fixme: For each image
                        if not FLAGS.no_write_images:
                            img_path = os.path.join(Local_img_output_dir, os.path.basename(im_fn))
                            cv2.imwrite(img_path, im[:, :, ::-1])

                            # Fixme: heat_map
                            # heat_map_path = os.path.join(Local_ImgHeatMap_output_dir,
                            #                              os.path.splitext(os.path.basename(im_fn))[0] + "_heatMap.png")
                            # fig, ax = plt.subplots()
                            # plt.subplots_adjust(top=1, bottom=0, right=1, left=0, hspace=0, wspace=0)
                            # plt.margins(0, 0)
                            # ax.set_xticks([])
                            # ax.set_yticks([])
                            # # plt.imshow(imutils.opencv2matplotlib(source_image.copy()[:, :, ::-1]))
                            # plt.imshow(imutils.opencv2matplotlib(im[:, :, ::-1].copy()[:, :, ::-1]))
                            # plt.imshow(cv2.resize(heat_map, (width, height), interpolation=cv2.INTER_LINEAR),
                            #            cmap=plt.get_cmap('jet'), interpolation='bilinear', alpha=0.5)
                            # # plt.show()
                            # fig.savefig(heat_map_path)
                            # fig.clear()
                    fw.close()

                        # save to file
                        # if boxes is not None:
                        # if True:
                        #     res_file = os.path.join(
                        #         Local_txt_output_dir,
                        #         '{}.txt'.format(
                        #             os.path.basename(im_fn).split('.')[0]))
                        #
                        #     with open(res_file, 'w') as f:
                        #         # if boxes is None:
                        #         if len(boxes) == 0:
                        #             continue
                        #         for idx, box in enumerate(boxes):
                        #             # to avoid submitting errors
                        #             box = sort_poly(box.astype(np.int32))
                        #             # box = check_and_validate(box)
                        #             if np.linalg.norm(box[0] - box[1]) < 5 or np.linalg.norm(box[3]-box[0]) < 5:
                        #                 continue
                        #
                        #             box[0, 0] = max(0, min(box[0, 0], width))
                        #             box[1, 0] = max(0, min(box[1, 0], width))
                        #             box[2, 0] = max(0, min(box[2, 0], width))
                        #             box[3, 0] = max(0, min(box[3, 0], width))
                        #             box[0, 1] = max(0, min(box[0, 1], height))
                        #             box[1, 1] = max(0, min(box[1, 1], height))
                        #             box[2, 1] = max(0, min(box[2, 1], height))
                        #             box[3, 1] = max(0, min(box[3, 1], height))
                        #             x_min = min(box[:,0])
                        #             x_max = max(box[:,0])
                        #             y_min = min(box[:,1])
                        #             y_max = max(box[:, 1])
                        #             box_width = x_max - x_min
                        #             box_height = y_max - y_min
                        #
                        #             # 24,-1, 195.0,156.0,54.0,56.0, 0.021061,1,-1
                        #             f.write('{},{},{},{},{},{},{},{},{}\r\n'.format(
                        #                 int(im_name_fg), -1, x_min, y_min, box_width, box_height,scores[idx],1,-1
                        #             ))
                        #             cv2.polylines(im[:, :, ::-1], [box.astype(np.int32).reshape((-1, 1, 2))], True, color=(0, 255, 0), thickness=2)
                        # if not FLAGS.no_write_images:
                        #     img_path = os.path.join(Local_img_output_dir, os.path.basename(im_fn))
                        #     cv2.imwrite(img_path, im[:, :, ::-1])
                        #
                        #     # Fixme: heat_map
                        #     # heat_map_path = os.path.join(Local_ImgHeatMap_output_dir,
                        #     #                              os.path.splitext(os.path.basename(im_fn))[0] + "_heatMap.png")
                        #     # fig, ax = plt.subplots()
                        #     # plt.subplots_adjust(top=1, bottom=0, right=1, left=0, hspace=0, wspace=0)
                        #     # plt.margins(0, 0)
                        #     # ax.set_xticks([])
                        #     # ax.set_yticks([])
                        #     # # plt.imshow(imutils.opencv2matplotlib(source_image.copy()[:, :, ::-1]))
                        #     # plt.imshow(imutils.opencv2matplotlib(im[:, :, ::-1].copy()[:, :, ::-1]))
                        #     # plt.imshow(cv2.resize(heat_map, (width, height), interpolation=cv2.INTER_LINEAR),
                        #     #            cmap=plt.get_cmap('jet'), interpolation='bilinear', alpha=0.5)
                        #     # # plt.show()
                        #     # fig.savefig(heat_map_path)
                        #     # fig.clear()

                print('Average second per frame is {}'.format(sum_time/im_num))
                print('Average frame per second is {}'.format(1.0 / (sum_time / im_num)))
                # zip_ForHRN(Local_txt_output_dir, Local_txt_output_dir + '.zip')


def zip_ForHRN(startdir, file_names):
    '''
    compress the folder of results for evaluating
    :param startdir:
    :param file_names:
    :return:
    '''
    z = zipfile.ZipFile(file_names, 'w', zipfile.ZIP_DEFLATED)
    for dirpath, dirnames, filenames in os.walk(startdir):
        fpath = dirpath.replace(startdir,'')
        fpath = fpath and fpath + os.sep or ''
        for filename in filenames:
            z.write(os.path.join(dirpath, filename), fpath+filename)
    print('compression', startdir)
    z.close()


if __name__ == '__main__':
    # tf.app.run()
    main_BatchEva()
