#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""Trains, evaluates and saves the model network using a queue."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import sys
import os
import numpy as np
import scipy as scp
import random
sys.path.insert(1, 'incl')
# sys.path.append('/home/yu/projects/KittiSeg/submodules/evaluation/kitti_devkit')
from incl.seg_utils import seg_utils as seg
# import seg_utils as seg

import tensorflow as tf
import time
import cv2
import incl.tensorvision
import incl.tensorvision.utils as utils
sys.path.append('/home/yu/projects/FCN_GoogLeNet')
# import post_crf

def eval_image(hypes, gt_image, cnn_image):
    """."""
    thresh = np.array(range(0, 256))/255.0

    road_color = np.array(hypes['data']['road_color'])
    background_color = np.array(hypes['data']['background_color'])
    gt_road = np.all(gt_image == road_color, axis=2)
    gt_bg = np.all(gt_image == background_color, axis=2)
    valid_gt = gt_road + gt_bg

    FN, FP, posNum, negNum = seg.evalExp(gt_road, cnn_image,
                                         thresh, validMap=None,
                                         validArea=valid_gt)

    return FN, FP, posNum, negNum


def resize_label_image(image, gt_image, image_height, image_width):
    image = scp.misc.imresize(image, size=(image_height, image_width),
                              interp='cubic')
    shape = gt_image.shape
    gt_image = scp.misc.imresize(gt_image, size=(image_height, image_width),
                                 interp='nearest')

    return image, gt_image


def evaluate(hypes, sess, image_pl, inf_out):

    softmax = inf_out['softmax']
    data_dir = hypes['dirs']['data_dir']

    eval_dict = {}
    for phase in ['train', 'val']:
        data_file = hypes['data']['{}_file'.format(phase)]
        data_file = os.path.join(data_dir, data_file)
        image_dir = os.path.dirname(data_file)

        thresh = np.array(range(0, 256))/255.0
        total_fp = np.zeros(thresh.shape)
        total_fn = np.zeros(thresh.shape)
        total_posnum = 0
        total_negnum = 0

        image_list = []

        with open(data_file) as file:
            for i, datum in enumerate(file):
                    datum = datum.rstrip()
                    image_file, gt_file = datum.split(" ")
                    image_file = os.path.join(image_dir, image_file)
                    gt_file = os.path.join(image_dir, gt_file)

                    image = scp.misc.imread(image_file, mode='RGB')
                    gt_image = scp.misc.imread(gt_file, mode='RGB')

                    if hypes['jitter']['fix_shape']:
                        shape = image.shape
                        image_height = hypes['jitter']['image_height']
                        image_width = hypes['jitter']['image_width']
                        assert(image_height >= shape[0])
                        assert(image_width >= shape[1])

                        offset_x = (image_height - shape[0])//2
                        offset_y = (image_width - shape[1])//2
                        new_image = np.zeros([image_height, image_width, 3])
                        new_image[offset_x:offset_x+shape[0],
                                  offset_y:offset_y+shape[1]] = image
                        input_image = new_image
                    elif hypes['jitter']['reseize_image']:
                        image_height = hypes['jitter']['image_height']
                        image_width = hypes['jitter']['image_width']
                        gt_image_old = gt_image
                        image, gt_image = resize_label_image(image, gt_image,
                                                             image_height,
                                                             image_width)
                        input_image = image
                    else:
                        input_image = image

                    shape = input_image.shape

                    feed_dict = {image_pl: input_image}

                    output = sess.run([softmax], feed_dict=feed_dict)
                    output_im = output[0][:, 1].reshape(shape[0], shape[1])

                    if hypes['jitter']['fix_shape']:
                        gt_shape = gt_image.shape
                        output_im = output_im[offset_x:offset_x+gt_shape[0],
                                              offset_y:offset_y+gt_shape[1]]

                    if phase == 'val':
                        # Saving RB Plot
                        ov_image = seg.make_overlay(image, output_im)
                        name = os.path.basename(image_file)
                        image_list.append((name, ov_image))

                        name2 = name.split('.')[0] + '_green.png'

                        hard = output_im > 0.5
                        green_image = utils.fast_overlay(image, hard)
                        image_list.append((name2, green_image))

                    FN, FP, posNum, negNum = eval_image(hypes,
                                                        gt_image, output_im)

                    total_fp += FP
                    total_fn += FN
                    total_posnum += posNum
                    total_negnum += negNum

        eval_dict[phase] = seg.pxEval_maximizeFMeasure(
            total_posnum, total_negnum, total_fn, total_fp, thresh=thresh)

        if phase == 'val':
            start_time = time.time()
            for i in xrange(10):
                sess.run([softmax], feed_dict=feed_dict)
            dt = (time.time() - start_time)/10
            print ('val for 10 times , the average is '+ str(dt))
    eval_list = []

    for phase in ['train', 'val']:
        eval_list.append(('[{}] MaxF1'.format(phase),
                          100*eval_dict[phase]['MaxF']))
        eval_list.append(('[{}] BestThresh'.format(phase),
                          100*eval_dict[phase]['BestThresh']))
        eval_list.append(('[{}] Average Precision'.format(phase),
                          100*eval_dict[phase]['AvgPrec']))
    eval_list.append(('Speed (msec)', 1000*dt))
    eval_list.append(('Speed (fps)', 1/dt))

    return eval_list, image_list



# def evaluate_with_crf(hypes,sess,image_pl,inf_out):
#
#     softmax = inf_out['softmax']
#     data_dir = hypes['dirs']['data_dir']
#
#     eval_dict = {}
#     for phase in ['train', 'val']:
#         data_file = hypes['data']['{}_file'.format(phase)]
#         data_file = os.path.join(data_dir, data_file)
#         image_dir = os.path.dirname(data_file)
#
#         thresh = np.array(range(0, 256))/255.0
#         total_fp = np.zeros(thresh.shape)
#         total_fn = np.zeros(thresh.shape)
#         total_posnum = 0
#         total_negnum = 0
#
#         image_list = []
#
#         with open(data_file) as file:
#             for i, datum in enumerate(file):
#                     datum = datum.rstrip()
#                     image_file, gt_file = datum.split(" ")
#                     image_file = os.path.join(image_dir, image_file)
#                     gt_file = os.path.join(image_dir, gt_file)
#
#                     image = scp.misc.imread(image_file, mode='RGB')
#                     gt_image = scp.misc.imread(gt_file, mode='RGB')
#
#                     if hypes['jitter']['fix_shape']:
#                         shape = image.shape
#                         image_height = hypes['jitter']['image_height']
#                         image_width = hypes['jitter']['image_width']
#                         assert(image_height >= shape[0])
#                         assert(image_width >= shape[1])
#
#                         offset_x = (image_height - shape[0])//2
#                         offset_y = (image_width - shape[1])//2
#                         new_image = np.zeros([image_height, image_width, 3])
#                         new_image[offset_x:offset_x+shape[0],
#                                   offset_y:offset_y+shape[1]] = image
#                         input_image = new_image
#                     elif hypes['jitter']['reseize_image']:
#                         image_height = hypes['jitter']['image_height']
#                         image_width = hypes['jitter']['image_width']
#                         gt_image_old = gt_image
#                         image, gt_image = resize_label_image(image, gt_image,
#                                                              image_height,
#                                                              image_width)
#                         input_image = image
#                     else:
#                         input_image = image
#
#                     shape = input_image.shape
#                     image_to_input=np.expand_dims(input_image,axis=0)
#
#                     feed_dict = {image_pl: image_to_input}
#
#                     output = sess.run([softmax], feed_dict=feed_dict)
#                     # output_im = output[0][:, 1].reshape(shape[0], shape[1])
#                     output_im = output[0][:, :].reshape(shape[0], shape[1],2)
#                     output_im = post_crf.post_process_crf(input_image, output_im, 2)
#                     if phase == 'val':
#                         _,_,_,_,num=image_file.split("/")
#                         dataset_tag,class_tag=num.split('_')
#                         output_im_=output_im*255.0
#                         output_im_=output_im_.reshape(shape[0],shape[1],-1)
#                         output_path=image_dir+'/gt_image_3/'+dataset_tag+'_'+'road'+'_'+class_tag
#                         # np.save(image_dir+'/gt_image_1/'+num,output_im_)
#                         cv2.imwrite(output_path,output_im_)
#
#                     if hypes['jitter']['fix_shape']:
#                         gt_shape = gt_image.shape
#                         output_im = output_im[offset_x:offset_x+gt_shape[0],
#                                               offset_y:offset_y+gt_shape[1]]
#
#                     if phase == 'val':
#                         # Saving RB Plot
#                         ov_image = seg.make_overlay(image, output_im)
#                         name = os.path.basename(image_file)
#                         image_list.append((name, ov_image))
#
#                         name2 = name.split('.')[0] + '_green.png'
#
#                         hard = output_im > 0.5
#                         green_image = utils.fast_overlay(image, hard)
#                         image_list.append((name2, green_image))
#
#                     FN, FP, posNum, negNum = eval_image(hypes,
#                                                         gt_image, output_im)
#
#                     total_fp += FP
#                     total_fn += FN
#                     total_posnum += posNum
#                     total_negnum += negNum
#
#         eval_dict[phase] = seg.pxEval_maximizeFMeasure(
#             total_posnum, total_negnum, total_fn, total_fp, thresh=thresh)
#
#         if phase == 'val':
#             start_time = time.time()
#             for i in xrange(10):
#                 sess.run([softmax], feed_dict=feed_dict)
#             dt = (time.time() - start_time)/10
#             print ('val for 10 times , the average is '+ str(dt))
#     eval_list = []
#
#     for phase in ['train', 'val']:
#         eval_list.append(('[{}] MaxF1'.format(phase),
#                           100*eval_dict[phase]['MaxF']))
#         eval_list.append(('[{}] BestThresh'.format(phase),
#                           100*eval_dict[phase]['BestThresh']))
#         eval_list.append(('[{}] Average Precision'.format(phase),
#                           100*eval_dict[phase]['AvgPrec']))
#     eval_list.append(('Speed (msec)', 1000*dt))
#     eval_list.append(('Speed (fps)', 1/dt))
#
#     return eval_list, image_list


def evaluate_without_crf(hypes,sess,image_pl,inf_out):

    softmax = inf_out['softmax']
    data_dir = hypes['dirs']['data_dir']

    eval_dict = {}
    for phase in ['train', 'val']:
        data_file = hypes['data']['{}_file'.format(phase)]
        data_file = os.path.join(data_dir, data_file)
        image_dir = os.path.dirname(data_file)

        thresh = np.array(range(0, 256))/255.0
        total_fp = np.zeros(thresh.shape)
        total_fn = np.zeros(thresh.shape)
        total_posnum = 0
        total_negnum = 0

        image_list = []

        with open(data_file) as file:
            for i, datum in enumerate(file):
                    datum = datum.rstrip()
                    image_file, gt_file = datum.split(" ")
                    image_file = os.path.join(image_dir, image_file)
                    gt_file = os.path.join(image_dir, gt_file)

                    image = scp.misc.imread(image_file, mode='RGB')
                    gt_image = scp.misc.imread(gt_file, mode='RGB')

                    if hypes['jitter']['fix_shape']:
                        shape = image.shape
                        image_height = hypes['jitter']['image_height']
                        image_width = hypes['jitter']['image_width']
                        assert(image_height >= shape[0])
                        assert(image_width >= shape[1])

                        offset_x = (image_height - shape[0])//2
                        offset_y = (image_width - shape[1])//2
                        new_image = np.zeros([image_height, image_width, 3])
                        new_image[offset_x:offset_x+shape[0],
                                  offset_y:offset_y+shape[1]] = image
                        input_image = new_image
                    elif hypes['jitter']['reseize_image']:
                        image_height = hypes['jitter']['image_height']
                        image_width = hypes['jitter']['image_width']
                        gt_image_old = gt_image
                        image, gt_image = resize_label_image(image, gt_image,
                                                             image_height,
                                                             image_width)
                        input_image = image
                    else:
                        input_image = image

                    shape = input_image.shape
                    image_to_input=np.expand_dims(input_image,axis=0)

                    feed_dict = {image_pl: image_to_input}

                    output = sess.run([softmax], feed_dict=feed_dict)
                    output_im = output[0][:, 1].reshape(shape[0], shape[1])
                    # output_im = output[0][:, :].reshape(shape[0], shape[1],2)
                    # output_im = post_crf.post_process_crf(input_image, output_im, 2)
                    if phase == 'val':
                        _,_,_,_,num=image_file.split("/")
                        dataset_tag,class_tag=num.split('_')
                        output_im_=output_im*255.0
                        output_im_=output_im_.reshape(shape[0],shape[1],-1)
                        output_path=image_dir+'/gt_image_1/'+dataset_tag+'_'+'road'+'_'+class_tag
                        # np.save(image_dir+'/gt_image_1/'+num,output_im_)
                        cv2.imwrite(output_path,output_im_)



                    if hypes['jitter']['fix_shape']:
                        gt_shape = gt_image.shape
                        output_im = output_im[offset_x:offset_x+gt_shape[0],
                                              offset_y:offset_y+gt_shape[1]]

                    if phase == 'val':
                        # Saving RB Plot
                        ov_image = seg.make_overlay(image, output_im)
                        name = os.path.basename(image_file)
                        image_list.append((name, ov_image))

                        name2 = name.split('.')[0] + '_green.png'

                        hard = output_im > 0.5
                        green_image = utils.fast_overlay(image, hard)
                        image_list.append((name2, green_image))

                    FN, FP, posNum, negNum = eval_image(hypes,
                                                        gt_image, output_im)

                    total_fp += FP
                    total_fn += FN
                    total_posnum += posNum
                    total_negnum += negNum

        eval_dict[phase] = seg.pxEval_maximizeFMeasure(
            total_posnum, total_negnum, total_fn, total_fp, thresh=thresh)

        if phase == 'val':
            start_time = time.time()
            for i in xrange(10):
                sess.run([softmax], feed_dict=feed_dict)
            dt = (time.time() - start_time)/10
            print ('val for 10 times , the average is '+ str(dt))
    eval_list = []

    for phase in ['train', 'val']:
        eval_list.append(('[{}] MaxF1'.format(phase),
                          100*eval_dict[phase]['MaxF']))
        eval_list.append(('[{}] BestThresh'.format(phase),
                          100*eval_dict[phase]['BestThresh']))
        eval_list.append(('[{}] Average Precision'.format(phase),
                          100*eval_dict[phase]['AvgPrec']))
    eval_list.append(('Speed (msec)', 1000*dt))
    eval_list.append(('Speed (fps)', 1/dt))

    return eval_list, image_list

def evaluate_test(hypes,sess,image_pl,inf_out):

    softmax = inf_out['softmax']
    data_dir = hypes['dirs']['data_dir']

    eval_dict = {}
    for phase in ['test']:
        data_file = hypes['data']['{}_file'.format(phase)]
        data_file = os.path.join(data_dir, data_file)
        image_dir = os.path.dirname(data_file)

        thresh = np.array(range(0, 256))/255.0
        total_fp = np.zeros(thresh.shape)
        total_fn = np.zeros(thresh.shape)
        total_posnum = 0
        total_negnum = 0

        image_list = []

        with open(data_file) as file:
            for i, datum in enumerate(file):
                    datum = datum.rstrip()
                    image_file = datum
                    # image_file, gt_file = datum.split(" ")
                    image_file = os.path.join(image_dir, image_file)
                    # gt_file = os.path.join(image_dir, gt_file)

                    image = scp.misc.imread(image_file, mode='RGB')
                    # gt_image = scp.misc.imread(gt_file, mode='RGB')

                    if hypes['jitter']['fix_shape']:
                        shape = image.shape
                        image_height = hypes['jitter']['image_height']
                        image_width = hypes['jitter']['image_width']
                        assert(image_height >= shape[0])
                        assert(image_width >= shape[1])

                        offset_x = (image_height - shape[0])//2
                        offset_y = (image_width - shape[1])//2
                        new_image = np.zeros([image_height, image_width, 3])
                        new_image[offset_x:offset_x+shape[0],
                                  offset_y:offset_y+shape[1]] = image
                        input_image = new_image
                    elif hypes['jitter']['reseize_image']:
                        image_height = hypes['jitter']['image_height']
                        image_width = hypes['jitter']['image_width']
                        gt_image_old = gt_image
                        image, gt_image = resize_label_image(image, gt_image,
                                                             image_height,
                                                             image_width)
                        input_image = image
                    else:
                        input_image = image

                    shape = input_image.shape
                    image_to_input=np.expand_dims(input_image,axis=0)

                    feed_dict = {image_pl: image_to_input}

                    output = sess.run([softmax], feed_dict=feed_dict)
                    output_im = output[0][:, 1].reshape(shape[0], shape[1])
                    # output_im = output[0][:, :].reshape(shape[0], shape[1],2)
                    # output_im = post_crf.post_process_crf(input_image, output_im, 2)
                    if phase == 'test':
                        _,_,_,_,num=image_file.split("/")
                        dataset_tag,class_tag=num.split('_')
                        output_im_=output_im*255.0
                        output_im_=output_im_.reshape(shape[0],shape[1],-1)
                        output_path=image_dir+'/gt_image_3/'+dataset_tag+'_'+'road'+'_'+class_tag
                        # np.save(image_dir+'/gt_image_1/'+num,output_im_)
                        cv2.imwrite(output_path,output_im_)



        #             if hypes['jitter']['fix_shape']:
        #                 gt_shape = gt_image.shape
        #                 output_im = output_im[offset_x:offset_x+gt_shape[0],
        #                                       offset_y:offset_y+gt_shape[1]]
        #
        #             if phase == 'val':
        #                 # Saving RB Plot
        #                 ov_image = seg.make_overlay(image, output_im)
        #                 name = os.path.basename(image_file)
        #                 image_list.append((name, ov_image))
        #
        #                 name2 = name.split('.')[0] + '_green.png'
        #
        #                 hard = output_im > 0.5
        #                 green_image = utils.fast_overlay(image, hard)
        #                 image_list.append((name2, green_image))
        #
        #             FN, FP, posNum, negNum = eval_image(hypes,
        #                                                 gt_image, output_im)
        #
        #             total_fp += FP
        #             total_fn += FN
        #             total_posnum += posNum
        #             total_negnum += negNum
        #
        # eval_dict[phase] = seg.pxEval_maximizeFMeasure(
        #     total_posnum, total_negnum, total_fn, total_fp, thresh=thresh)
        #
        # if phase == 'val':
        #     start_time = time.time()
        #     for i in xrange(10):
        #         sess.run([softmax], feed_dict=feed_dict)
        #     dt = (time.time() - start_time)/10
        #     print ('val for 10 times , the average is '+ str(dt))
    eval_list = []


    return eval_list, image_list

