import tensorflow as tf
import numpy as np

from tflearn.layers.conv import global_avg_pool
from tensorflow.contrib import slim


tf.app.flags.DEFINE_integer('text_scale', 512, '')

from nets_tensorflow import vgg, vgg_P14
from nets import resnet_v1

FLAGS = tf.app.flags.FLAGS


def unpool(inputs):
    return tf.image.resize_bilinear(inputs, size=[tf.shape(inputs)[1]*2,  tf.shape(inputs)[2]*2])


def mean_image_subtraction(images, means=[123.68, 116.78, 103.94]):
    '''
    image normalization
    :param images:
    :param means:
    :return:
    '''
    num_channels = images.get_shape().as_list()[-1]
    if len(means) != num_channels:
      raise ValueError('len(means) must match the number of channels')
    channels = tf.split(axis=3, num_or_size_splits=num_channels, value=images)
    for i in range(num_channels):
        channels[i] -= means[i]
    return tf.concat(axis=3, values=channels)




def model(images, weight_decay=1e-5, is_training=True):
    '''
    define the model, we use slim's implemention of resnet
    '''
    images = mean_image_subtraction(images)
    batch_size = FLAGS.batch_size_per_gpu

    # FIXME: CAI note, Feature extractor stem in the EAST paper.
    with slim.arg_scope(vgg_P14.vgg_arg_scope(weight_decay=weight_decay)):
        logits, end_points= vgg_P14.vgg_16(images, is_training=is_training, scope='vgg_16')

    # with slim.arg_scope(resnet_v1.resnet_arg_scope(weight_decay=weight_decay)):
    #     logits, end_points = resnet_v1.resnet_v1_50(images, is_training=is_training, scope='resnet_v1_50')



    with tf.variable_scope('feature_fusion', values=[end_points.values]):
        batch_norm_params = {
        'decay': 0.997,
        'epsilon': 1e-5,
        'scale': True,
        'is_training': is_training
        }
        with slim.arg_scope([slim.conv2d],
                            activation_fn=tf.nn.relu,
                            normalizer_fn=slim.batch_norm,
                            normalizer_params=batch_norm_params,
                            weights_regularizer=slim.l2_regularizer(weight_decay)):

            # First-level scale parsing
            # 512 512 256 128
            f = [end_points['pool4'], end_points['pool3'],
                 end_points['pool2'], end_points['pool1']]
            for i in range(4):
                print('Shape of f_{} {}'.format(i, f[i].shape))

            g = [None, None, None, None]
            h = [None, None, None, None]
            gap = [None, None, None]

            num_outputs = [512, 256, 128, 64]
            for i in range(4):
                f[i] = slim.conv2d(f[i], num_outputs[i], 1)
                if i == 0:
                    h[i] = f[i]
                else:
                    # c1_1 = slim.conv2d(tf.concat([g[i-1], f[i]], axis=-1), num_outputs[i], 1)

                    global_pool = global_avg_pool(h[i-1])
                    gap[i-1] = global_pool

                    global_pool_weight = tf.reshape(global_pool,(batch_size, 1, 1, num_outputs[i-1]))

                    reduce_dim = slim.conv2d(global_pool_weight, num_outputs[i]/4, 1,activation_fn=tf.nn.relu )
                    increase_dim = slim.conv2d(reduce_dim, num_outputs[i], 1, activation_fn=tf.nn.sigmoid )
                    f_w = tf.multiply(f[i], increase_dim )

                    c1_1 = slim.conv2d(tf.concat([g[i - 1], f_w ], axis=-1), num_outputs[i]*2, 1)
                    # c1_1 = tf.concat([g[i - 1], f_w], axis=-1)

                    h[i] = slim.conv2d(c1_1, num_outputs[i], 3)
                if i <= 2:
                    g[i] = unpool(h[i])
                    g[i] = slim.conv2d(g[i],num_outputs[i],3)   # Author should focus on here.
                else:
                    g[i] = slim.conv2d(h[i], num_outputs[i], 3)
                print('Shape of h_{} {}, g_{} {}'.format(i, h[i].shape, i, g[i].shape))


            # here we use a slightly different way for regression part,
            # we first use a sigmoid to limit the regression range, and also
            # this is do with the angle map

            F_score = slim.conv2d(g[3], 1, 1, activation_fn=tf.nn.sigmoid, normalizer_fn=None)
            # 4 channel of axis aligned bbox and 1 channel rotation angle
            geo_map = slim.conv2d(g[3], 4, 1, activation_fn=tf.nn.sigmoid, normalizer_fn=None) * FLAGS.text_scale
            angle_map = (slim.conv2d(g[3], 1, 1, activation_fn=tf.nn.sigmoid, normalizer_fn=None) - 0.5) * np.pi/2 # angle is between [-45, 45]
            F_geometry = tf.concat([geo_map, angle_map], axis=-1)
            # Key region
            F_score_KR = slim.conv2d(g[3], 4, 1, activation_fn=tf.nn.sigmoid, normalizer_fn=None)

            # Tegother Classify
            gaps = tf.concat([gap[0],gap[1],gap[2]],axis=-1)
            logits = slim.fully_connected(gaps,2,activation_fn=None, scope='gaps2cls')

    return F_score, F_geometry, F_score_KR, logits


def dice_coefficient(y_true_cls, y_pred_cls,
                     training_mask):
    '''
    dice loss
    :param y_true_cls:
    :param y_pred_cls:
    :param training_mask:
    :return:
    '''
    eps = 1e-5
    intersection = tf.reduce_sum(y_true_cls * y_pred_cls * training_mask)
    union = tf.reduce_sum(y_true_cls * training_mask) + tf.reduce_sum(y_pred_cls * training_mask) + eps
    loss = 1. - (2 * intersection / union)
    # tf.summary.scalar('classification_dice_loss', loss)
    return loss


def loss(y_true_cls, y_pred_cls,
         y_true_geo, y_pred_geo,
         training_mask):
    '''
    define the loss used for training, contraning two part,
    the first part we use dice loss instead of weighted logloss,
    the second part is the iou loss defined in the paper
    :param y_true_cls: ground truth of text
    :param y_pred_cls: prediction os text
    :param y_true_geo: ground truth of geometry
    :param y_pred_geo: prediction of geometry
    :param training_mask: mask used in training, to ignore some text annotated by ###
    :return:
    '''
    classification_loss = dice_coefficient(y_true_cls, y_pred_cls, training_mask)
    # scale classification loss to match the iou loss part
    classification_loss *= 0.01

    # d1 -> top, d2->right, d3->bottom, d4->left
    d1_gt, d2_gt, d3_gt, d4_gt, theta_gt = tf.split(value=y_true_geo, num_or_size_splits=5, axis=3)
    d1_pred, d2_pred, d3_pred, d4_pred, theta_pred = tf.split(value=y_pred_geo, num_or_size_splits=5, axis=3)
    area_gt = (d1_gt + d3_gt) * (d2_gt + d4_gt)
    area_pred = (d1_pred + d3_pred) * (d2_pred + d4_pred)
    w_union = tf.minimum(d2_gt, d2_pred) + tf.minimum(d4_gt, d4_pred)
    h_union = tf.minimum(d1_gt, d1_pred) + tf.minimum(d3_gt, d3_pred)
    area_intersect = w_union * h_union
    area_union = area_gt + area_pred - area_intersect
    L_AABB = -tf.log((area_intersect + 1.0)/(area_union + 1.0))
    L_theta = 1 - tf.cos(theta_pred - theta_gt)
    # tf.summary.scalar('geometry_AABB', tf.reduce_mean(L_AABB * y_true_cls * training_mask))
    # tf.summary.scalar('geometry_theta', tf.reduce_mean(L_theta * y_true_cls * training_mask))
    L_g = L_AABB + 20 * L_theta

    return tf.reduce_mean(L_g * y_true_cls * training_mask) + classification_loss




def loss_KR(y_true_cls, y_pred_cls,
         y_true_geo, y_pred_geo,
         y_true_ps_map,ps_pred_cls,
         training_mask,ps_maps):
    '''
    define the loss used for training, contraning two part,
    the first part we use dice loss instead of weighted logloss,
    the second part is the iou loss defined in the paper
    :param y_true_cls: ground truth of text
    :param y_pred_cls: prediction os text
    :param y_true_geo: ground truth of geometry
    :param y_pred_geo: prediction of geometry
    :param training_mask: mask used in training, to ignore some text annotated by ###
    :return:
    '''

    classification_loss = dice_coefficient(y_true_cls, y_pred_cls, training_mask)
    # scale classification loss to match the iou loss part
    classification_loss *= 0.01

    # d1 -> top, d2->right, d3->bottom, d4->left
    d1_gt, d2_gt, d3_gt, d4_gt, theta_gt = tf.split(value=y_true_geo, num_or_size_splits=5, axis=3)
    d1_pred, d2_pred, d3_pred, d4_pred, theta_pred = tf.split(value=y_pred_geo, num_or_size_splits=5, axis=3)
    area_gt = (d1_gt + d3_gt) * (d2_gt + d4_gt)
    area_pred = (d1_pred + d3_pred) * (d2_pred + d4_pred)
    w_union = tf.minimum(d2_gt, d2_pred) + tf.minimum(d4_gt, d4_pred)
    h_union = tf.minimum(d1_gt, d1_pred) + tf.minimum(d3_gt, d3_pred)
    area_intersect = w_union * h_union
    area_union = area_gt + area_pred - area_intersect
    L_AABB = -tf.log((area_intersect + 1.0)/(area_union + 1.0))
    L_theta = 1 - tf.cos(theta_pred - theta_gt)
    # tf.summary.scalar('geometry_AABB', tf.reduce_mean(L_AABB * y_true_cls * training_mask))
    # tf.summary.scalar('geometry_theta', tf.reduce_mean(L_theta * y_true_cls * training_mask))
    L_g = L_AABB + 20 * L_theta

    loss_ps = 0
    for j in range(4):
        eps = 1e-5

        intersection1 = tf.reduce_sum(ps_maps[:, :, :, j, np.newaxis] * y_true_ps_map * ps_pred_cls[:, :, :, j, np.newaxis] * training_mask)

        union1 = (tf.reduce_sum(ps_maps[:, :, :, j, np.newaxis] * y_true_ps_map * training_mask) + tf.reduce_sum(
            ps_pred_cls[:, :, :, j, np.newaxis] * training_mask)) + eps

        loss_ps += 1. - (2 * (intersection1) / (union1))
    loss_ps *= 0.0025
    # tf.summary.scalar('ps_dice_loss', loss_ps)

    return tf.reduce_mean(L_g * y_true_cls * training_mask) + classification_loss + loss_ps



