__author__ = 'rogerjiang'

'''
This file performs the inference based on the learned mode
parameters (checkpoint). Each image is cut into 4 quarters due
to the limit of gpu memory. Vertical and horizonta reflections,
and [0, 90, 180, 270] degrees rotations are performed for each
quarter, and the outputs of sigmoid predictions are later
arithmetically averaged.
'''

import tensorflow as tf
import simplejson
from utils import data_utils, train_utils
import os
import numpy as np
import cv2
import train
import pandas as pd
from shapely import wkt
import time
import sys



def test_input(img, img_size, H):
    '''
    This function cuts each image into 4 quarters:
    ([upper, lower] * [left, right]), performs vertical and horizontal
    reflections, and [0, 90, 180, 270] degrees rotations for each quarter.
    It yields (4 * 2 * 4 =)32 images.
    :param img:
    :param img_size:
    :return:
    '''
    [img_width, img_height] = img_size
    [crop_width, crop_height] = H['crop_size']
    pad = H['pad']
    x_width = H['x_width']
    x_height = H['x_height']

    for [x_start, x_end, y_start, y_end] in [
        [0, crop_width, 0, crop_height],
        [0, crop_width, crop_height, img_height],
        [crop_width, img_width, 0, crop_height],
        [crop_width, img_width, crop_height, img_height],
    ]:
        feature = img[x_start: x_end, y_start:y_end, :]
        for feat_trans in [feature, np.rollaxis(feature, 1, 0)]:
            for [x_step, y_step] in [[1, 1], [-1, 1], [1, -1], [-1, -1]]:
                feature_w_padding = cv2.copyMakeBorder(
                    feat_trans[::x_step, ::y_step, :],
                    pad, x_width - pad - feat_trans.shape[0],
                    pad, x_height - pad - feat_trans.shape[1],
                    cv2.BORDER_REFLECT_101)
                yield feat_trans.shape, feature_w_padding



def pred_for_each_quarter(sess, img_in, pred, img_data, H):
    '''

    :param sess:
    :param img_in:
    :param pred:
    :param img_data:
    :param H:
    :return:
    '''
    num_channel = H['num_channel']
    x_width = H['x_width']
    x_height = H['x_height']
    batch_size = H['batch_size']

    mask_stack, shape_stack = [], []
    for feat_shape, img in test_input(
            img_data.train_feature, img_data.image_size, H):
        predictions, = sess.run(
            [pred],
            feed_dict={img_in: np.reshape(img,
                                          [batch_size,
                                           x_width,
                                           x_height,
                                           num_channel])})
        mask_stack.append(predictions)
        shape_stack.append(feat_shape)
    return mask_stack, shape_stack

def stitch_mask(img_stack, img_size, feat_shape, H):
    '''
    img_stack is the stack of pixel-wise inference of input images from
    test_input. This function reverts the reflection and rotations and
    stitches the 4 quarters together.
    :param img_stack:
    :param img_size:
    :param feat_shape:
    :return:
    '''
    mask = np.zeros([8, img_size[0], img_size[1]])
    [img_width, img_height] = img_size
    [crop_width, crop_height] = H['crop_size']
    pad = H['pad']

    idx = 0
    for [x_start, x_end, y_start, y_end] in [
        [0, crop_width, 0, crop_height],
        [0, crop_width, crop_height, img_height],
        [crop_width, img_width, 0, crop_height],
        [crop_width, img_width, crop_height, img_height],
    ]:
        quarter = 0
        for feat_trans in range(2):
            for [x_step, y_step] in [[1, 1], [-1, 1], [1, -1], [-1, -1]]:

                img_stack[idx] = img_stack[idx] \
                    [pad: pad + feat_shape[idx][0],