import tensorflow as tf
import time
import numpy as np


class MLPerceptron:
    """
    A trainable version VGG19.
    """

    def __init__(self, mlp_npy_path=None, trainable=True, learning_rate=0.05, dropout=0.5, size_layer_fc=1024):
        if mlp_npy_path is not None:
            self.data_dict = np.load(mlp_npy_path, encoding='latin1').item()
            print("npy file loaded")
        else:
            self.data_dict = None
            print("random weight")

        self.var_dict = {}
        self.trainable = trainable
        self.learning_rate = learning_rate
        self.dropout = dropout
        self.size_layer = size_layer_fc

    def build(self, input_batch, target, train_mode=None):

        start_time = time.time()

        self.fc7 = self.fc_layer(input_batch, 2048, self.size_layer, "fc7")
        self.relu7 = tf.nn.relu(self.fc7)

        # DROPOUT
        if train_mode is not None:
            # train_mode: True[train] -> dropout activate | False[test] -> dropout deactivate
            self.relu7 = tf.cond(train_mode, lambda: tf.nn.dropout(self.relu7, self.dropout), lambda: self.relu7)
        elif self.trainable:
            # train_mode: None -> Default [test or classification]
            self.relu7 = tf.nn.dropout(self.relu7, self.dropout)

        self.fc8 = self.fc_layer(self.relu7, self.size_layer, 2, "fc8")
        self.prob = tf.nn.softmax(self.fc8, name="prob")

        # COST - TRAINING
        # if train_mode is True:
        self.cost = tf.reduce_mean((self.prob - target) ** 2)
        self.train = tf.train.AdamOptimizer(self.learning_rate).minimize(self.cost)

        self.data_dict = None
        print(("build model finished: %ds" % (time.time() - start_time)))

    def fc_layer(self, bottom, in_size, out_size, name):
        with tf.variable_scope(name):
            weights, biases = self.get_fc_var(in_size, out_size, name)

            x = tf.reshape(bottom, [-1, in_size])
            fc = tf.nn.bias_add(tf.matmul(x, weights), biases)
            return fc

    def get_fc_var(self, in_size, out_size, name):
        initial_value = tf.truncated_normal([in_size, out_size], 0.0, 0.001)
        weights = self.get_var_fc(initial_value, name, 0, name + "_weights")

        initial_value = tf.truncated_normal([out_size], .0, .001)
        biases = self.get_var_fc(initial_value, name, 1, name + "_biases")

        return weights, biases

    def get_var_fc(self, initial_value, name, idx, var_name):
        if self.data_dict is not None and name in self.data_dict:
            value = self.data_dict[name][idx]
        else:
            value = initial_value

        if self.trainable:
            var = tf.Variable(value, name=var_name)
        else:
            var = tf.constant(value, dtype=tf.float32, name=var_name)

        self.var_dict[(name, idx)] = var
        assert var.get_shape() == initial_value.get_shape()
        return var

    def save_npy(self, sess, npy_path="./mlp-save.npy"):
        assert isinstance(sess, tf.Session)

        data_dict = {}

        for (name, idx), var in list(self.var_dict.items()):
            var_out = sess.run(var)
            if name not in data_dict:
                data_dict[name] = {}
            data_dict[name][idx] = var_out

        np.save(npy_path, data_dict)
        print("File saved", npy_path)
        return npy_path
