# -*- coding: utf-8 -*-
import tensorflow as tf
import numpy as np
import PEM_load_data


def abs_smooth(x):
    """Smoothed absolute function. Useful to compute an L1 smooth error.
    """
    absx = tf.abs(x)
    minx = tf.minimum(absx, 1)
    r = 0.5 * ((absx - 1) * minx + absx)
    return r


def PEM_loss(anchors_iou, match_iou, config):
    # iou regressor
    u_hmask = tf.cast(match_iou > 0.6, dtype=tf.float32)
    u_mmask = tf.cast(tf.logical_and(match_iou <= 0.6, match_iou > 0.2), dtype=tf.float32)
    u_lmask = tf.cast(match_iou < 0.2, dtype=tf.float32)

    num_h = tf.reduce_sum(u_hmask)
    num_m = tf.reduce_sum(u_mmask)
    num_l = tf.reduce_sum(u_lmask)

    r_m = config.u_ratio_m * num_h / (num_m)
    r_m = tf.minimum(r_m, 1)
    u_smmask = tf.random_uniform([tf.shape(u_hmask)[0]], dtype=tf.float32)
    u_smmask = u_smmask * u_mmask
    u_smmask = tf.cast(u_smmask > (1. - r_m), dtype=tf.float32)

    r_l = config.u_ratio_l * num_h / (num_l)
    r_l = tf.minimum(r_l, 1)
    u_slmask = tf.random_uniform([tf.shape(u_hmask)[0]], dtype=tf.float32)
    u_slmask = u_slmask * u_lmask
    u_slmask = tf.cast(u_slmask > (1. - r_l), dtype=tf.float32)

    iou_weights = u_hmask + u_smmask + u_slmask
    iou_loss = abs_smooth(match_iou - anchors_iou)
    print(match_iou.get_shape(), anchors_iou.get_shape())
    iou_loss = tf.losses.compute_weighted_loss(iou_loss, iou_weights)

    num_iou = [tf.reduce_sum(u_hmask), tf.reduce_sum(u_smmask), tf.reduce_sum(u_slmask)]
    loss = {'iou_loss': iou_loss, 'num_iou': num_iou}
    return loss


def PEM_Train(X, Y_iou, LR, config):
    net = 0.1 * tf.matmul(X, config.W["iou_0"]) + config.biases["iou_0"]
    net = tf.nn.relu(net)
    net = 0.1 * tf.matmul(net, config.W["iou_1"]) + config.biases["iou_1"]
    net = tf.nn.sigmoid(net)

    anchors_iou = net
    anchors_iou = tf.reshape(anchors_iou, [-1])
    loss = PEM_loss(anchors_iou, Y_iou, config)
    PEM_trainable_variables = tf.trainable_variables()

    l2 = 0.0001 * sum(tf.nn.l2_loss(tf_var) for tf_var in PEM_trainable_variables)
    cost = 10 * loss["iou_loss"] + l2
    optimizer = tf.train.AdamOptimizer(learning_rate=LR).minimize(cost, var_list=PEM_trainable_variables)
    loss["l2"] = l2
    return optimizer, loss


class Config(object):
    """
    define a class to store parameters,
    the input should be feature mat of training and testing
    """

    def __init__(self):
        self.training_epochs = 61

        self.input_steps = 256
        self.learning_rates = [0.001] * 10 + [0.0001] * 10
        self.training_epochs = len(self.learning_rates)
        self.num_random = 10
        self.batch_size = 16
        self.u_ratio_m = 1
        self.u_ratio_l = 2

        with tf.variable_scope("latent_net"):
            self.W = {
                'iou_0': tf.Variable(tf.truncated_normal([32, 256])),
                'iou_1': tf.Variable(tf.truncated_normal([256, 1]))}
            self.biases = {
                'iou_0': tf.Variable(tf.truncated_normal([256])),
                'iou_1': tf.Variable(tf.truncated_normal([1]))}


if __name__ == "__main__":
    config = Config()

    X_feature = tf.placeholder(tf.float32, [None, 32])
    Y_iou = tf.placeholder(tf.float32, [None])
    LR = tf.placeholder(tf.float32)

    optimizer, loss = PEM_Train(X_feature, Y_iou, LR, config)

    model_saver = tf.train.Saver(var_list=tf.trainable_variables(), max_to_keep=80)
    tf_config = tf.ConfigProto()
    tf_config.gpu_options.allow_growth = True
    tf_config.log_device_placement = True
    sess = tf.InteractiveSession(config=tf_config)
    tf.global_variables_initializer().run()

    train_dict, val_dict, test_dict = PEM_load_data.getDatasetDict()

    train_data = PEM_load_data.getTrainData(config.batch_size, "train")
    val_data = PEM_load_data.getTrainData(config.batch_size, "validation")

    train_info = {"iou_loss": [], "l2": []}
    val_info = {"iou_loss": [], "l2": []}

    best_val_cost = 1000000

    for epoch in range(0, config.training_epochs):
        ## TRAIN ##

        mini_info = {"iou_loss": [], "l2": []}
        for idx in range(len(train_data)):
            prop_dict = train_data[idx]
            batch_feature, batch_iou_list, batch_ioa_list = PEM_load_data.prop_dict_data(prop_dict)
            _, out_loss = sess.run([optimizer, loss], feed_dict={X_feature: batch_feature,
                                                                 Y_iou: batch_iou_list,
                                                                 LR: config.learning_rates[epoch]})
            mini_info["iou_loss"].append(out_loss["iou_loss"])
            mini_info["l2"].append(out_loss["l2"])

        train_info["iou_loss"].append(np.mean(mini_info["iou_loss"]))
        train_info["l2"].append(np.mean(mini_info["l2"]))

        mini_info = {"iou_loss": [], "l2": []}
        for idx in range(len(val_data)):
            prop_dict = val_data[idx]
            batch_feature, batch_iou_list, batch_ioa_list = PEM_load_data.prop_dict_data(prop_dict)
            out_loss = sess.run(loss, feed_dict={X_feature: batch_feature,
                                                 Y_iou: batch_iou_list,
                                                 LR: config.learning_rates[epoch]})
            mini_info["iou_loss"].append(out_loss["iou_loss"])
            mini_info["l2"].append(out_loss["l2"])

        val_info["iou_loss"].append(np.mean(mini_info["iou_loss"]))
        val_info["l2"].append(np.mean(mini_info["l2"]))

        print("Epoch-%d Train Loss:  %.04f" % (epoch, train_info["iou_loss"][-1]))
        print("Epoch-%d Val   Loss:  %.04f" % (epoch, val_info["iou_loss"][-1]))

        model_saver.save(sess, "models/PEM/pem_model_checkpoint")
        if val_info["iou_loss"][-1] < best_val_cost:
            best_val_cost = val_info["iou_loss"][-1]
            model_saver.save(sess, "models/PEM/pem_model_best")
