import os
import numpy as np
import tensorflow as tf
import utilities as utils
import utilities_image as utils_img
from random import shuffle

# To be changed to 2
NUM_CLASSES = 2

# To be changed to 400, 400
IMG_WIDTH = 400
IMG_HEIGHT = 400

NUM_EPOCHS = 2

# To be changed to 3
NUM_CHANNELS = 3

# To be changed to 1
BATCH_SIZE = 1

TRAIN_SIZE = 30
TEST_SIZE = 10

REGUL_PARAM = 5e-4

# Set to default
LEARNING_RATE = 0.001

DROPOUT = 0.9

STD_VAR_INIT = 0.1

OUTPUT_PATH = './conv_net_output/'
MODEL_PATH = OUTPUT_PATH + 'conv_net_model/conv_net_model.ckpt'
TRAINING_PATH = '../data/training/'

PREDICTION_PATH = './conv_net_prediction/'




def conv_net_model(x, keep_prob):
    """
     data must be a 4D tensor generated by tf.constant()
      applied on a 4D numpy array with [image index, IMAGE_WIDTH, IMAGE_HEIGHT, CHANNEL]
    """

    # Could use tf.truncated_normal() instead of tf.random_normal(),
    # see what is the best choice
    # note tf.truncated() is used in the FCN implementation
    with tf.variable_scope('weights'):
        weights = {'W_d_conv1': utils.weight_def([3, 3, NUM_CHANNELS, 64], stddev=STD_VAR_INIT, name = 'W_d_conv1'),
                   'W_d_conv2': utils.weight_def([3, 3, 64, 64], stddev=STD_VAR_INIT, name = 'W_d_conv2'),
                   'W_d_conv3': utils.weight_def([3, 3, 64, 128], stddev=STD_VAR_INIT, name = 'W_d_conv3'),
                   'W_d_conv4': utils.weight_def([3, 3, 128, 128], stddev=STD_VAR_INIT, name = 'W_d_conv4'),
                   'W_d_conv5': utils.weight_def([3, 3, 128, 256], stddev=STD_VAR_INIT, name = 'W_d_conv5'),
                   'W_d_conv6': utils.weight_def([3, 3, 256, 256], stddev=STD_VAR_INIT, name = 'W_d_conv6'),

                   'W_deconv1': utils.weight_def([3, 3, 256, 256], stddev=STD_VAR_INIT, name = 'W_deconv1'),
                   'W_u_conv1': utils.weight_def([3, 3, 384, 256], stddev=STD_VAR_INIT, name = 'W_u_conv1'),
                   'W_u_conv2': utils.weight_def([3, 3, 256, 128], stddev=STD_VAR_INIT, name = 'W_u_conv2'),

                   'W_deconv2': utils.weight_def([3, 3, 128, 128], stddev=STD_VAR_INIT, name = 'W_deconv2'),
                   'W_u_conv3': utils.weight_def([3, 3, 192, 128], stddev=STD_VAR_INIT, name = 'W_u_conv3'),
                   'W_u_conv4': utils.weight_def([3, 3, 128, 64], stddev=STD_VAR_INIT, name = 'W_u_conv4'),

                   'W_convout': utils.weight_def([1, 1, 64, NUM_CLASSES], stddev=STD_VAR_INIT, name = 'W_convout')}

    with tf.variable_scope('biases'):
        biases = {'B_d_conv1': utils.bias_def([64], name = 'B_d_conv1'),
                   'B_d_conv2': utils.bias_def([64], name = 'B_d_conv2'),
                   'B_d_conv3': utils.bias_def([128], name = 'B_d_conv3'),
                   'B_d_conv4': utils.bias_def([128], name = 'B_d_conv4'),
                   'B_d_conv5': utils.bias_def([256], name = 'B_d_conv5'),
                   'B_d_conv6': utils.bias_def([256], name = 'B_d_conv6'),
                   'B_deconv1': utils.bias_def([256], name = 'B_deconv1'),
                   'B_u_conv1': utils.bias_def([256], name = 'B_u_conv1'),
                   'B_u_conv2': utils.bias_def([128], name = 'B_u_conv2'),
                   'B_deconv2': utils.bias_def([128], name = 'B_deconv2'),
                   'B_u_conv3': utils.bias_def([128], name = 'B_u_conv3'),
                   'B_u_conv4': utils.bias_def([64], name = 'B_u_conv4'),
                   'B_convout': utils.bias_def([NUM_CLASSES], name = 'B_convout')}

    # IMPORTANT STEP
    # From FCN implementation:
    # shape = tf.shape(data)
    # deconv_shape3 = tf.stack([shape[0], shape[1], shape[2], NUM_OF_CLASSESS])
    #data = tf.reshape(data, shape=[BATCH_SIZE, IMG_WIDTH, IMG_HEIGHT, NUM_CHANNELS])

    # Going down
    # First block
    d_conv_dropout_relu1 = utils.conv2d_dropout_relu(x, weights['W_d_conv1'], biases['B_d_conv1'], keep_prob = keep_prob,\
                                                        name = 'd_conv_1')

    d_conv_dropout_relu2 = utils.conv2d_dropout_relu(d_conv_dropout_relu1, weights['W_d_conv2'], biases['B_d_conv2'], keep_prob = keep_prob,\
                                                        name = 'd_conv_2')

    max_pool1 = utils.maxpool2d(d_conv_dropout_relu2, name = 'max_pool')

    # Second block
    d_conv_dropout_relu3 = utils.conv2d_dropout_relu(max_pool1, weights['W_d_conv3'], biases['B_d_conv3'], keep_prob = keep_prob,\
                                                        name = 'd_conv_3')

    d_conv_dropout_relu4 = utils.conv2d_dropout_relu(d_conv_dropout_relu3, weights['W_d_conv4'], biases['B_d_conv4'], keep_prob = keep_prob,\
                                                        name = 'd_conv_4')

    max_pool2 = utils.maxpool2d(d_conv_dropout_relu4, name = 'max_pool2')


    # Transition
    d_conv_dropout_relu5 = utils.conv2d_dropout_relu(max_pool2, weights['W_d_conv5'], biases['B_d_conv5'], keep_prob = keep_prob,\
                                                        name='d_conv_5')

    d_conv_dropout_relu6 = utils.conv2d_dropout_relu(d_conv_dropout_relu5, weights['W_d_conv6'], biases['B_d_conv6'], keep_prob = keep_prob,\
                                                        name ='d_conv_6')


    # Going up
    # First block
    deconv_relu1 = utils.deconv2d_relu(d_conv_dropout_relu6, weights['W_deconv1'], biases['B_deconv1'], name = 'deconv1')

    concat1 = tf.concat([deconv_relu1, d_conv_dropout_relu4], 3, name = 'concat1')

    u_conv_dropout_relu1 = utils.conv2d_dropout_relu(concat1, weights['W_u_conv1'], biases['B_u_conv1'], keep_prob = keep_prob,\
                                                        name ='u_conv_1')

    u_conv_dropout_relu2 = utils.conv2d_dropout_relu(u_conv_dropout_relu1, weights['W_u_conv2'], biases['B_u_conv2'], keep_prob = keep_prob,\
                                                        name ='u_conv_2')

    # Second block
    deconv_relu2 = utils.deconv2d_relu(u_conv_dropout_relu2, weights['W_deconv2'], biases['B_deconv2'], name = 'deconv2')

    concat2 = tf.concat([deconv_relu2, d_conv_dropout_relu2], 3)

    u_conv_dropout_relu3 = utils.conv2d_dropout_relu(concat2, weights['W_u_conv3'], biases['B_u_conv3'], keep_prob = keep_prob,\
                                                        name ='u_conv_3')

    u_conv_dropout_relu4 = utils.conv2d_dropout_relu(u_conv_dropout_relu3, weights['W_u_conv4'], biases['B_u_conv4'], keep_prob = keep_prob,\
                                                        name ='u_conv_4')

    # Convout
    convout = utils.conv2d_dropout_relu(u_conv_dropout_relu4, weights['W_convout'], biases['B_convout'], keep_prob = tf.constant(1.0),\
                                                        name = 'convout')

    # Storing variables into a variables list
    # weights:
    variables = []
    for _, item in weights.items():
        variables.append(item)

    for _, item in biases.items():
        variables.append(item)


    return convout, variables




class ConvNet(object):

    def __init__(self):
        self.x = tf.placeholder(tf.float32, shape = [None, IMG_WIDTH, IMG_HEIGHT, NUM_CHANNELS])
        self.y = tf.placeholder(tf.float32, shape = [None, IMG_WIDTH, IMG_HEIGHT, NUM_CLASSES])
        self.keep_prob = tf.placeholder(tf.float32)

        logits, self.variables = conv_net_model(self.x, self.keep_prob)

        self.cost = self._get_cost(logits)

        self.gradients_node = tf.gradients(self.cost, self.variables)

        # Only computed for the sake of summary in TensorBoard
        self.cross_entropy = tf.reduce_mean(utils.cross_entropy(tf.reshape(self.y, [-1, NUM_CLASSES]),
                                tf.reshape(utils.pixel_wise_softmax_2(logits), [-1, NUM_CLASSES])))

        self.predicter = utils.pixel_wise_softmax_2(logits)

        self.correct_pred = tf.equal(tf.argmax(self.predicter, 3), tf.argmax(self.y, 3))

        self.accuracy = tf.reduce_mean(tf.cast(self.correct_pred, tf.float32))

    def _get_cost(self, logits):
        flat_logits = tf.reshape(logits, [-1, NUM_CLASSES])
        flat_labels = tf.reshape(self.y, [-1, NUM_CLASSES])

        loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=flat_logits,
                                                                        labels=flat_labels))
        # add regularization
        regularizers = sum([tf.nn.l2_loss(variable) for variable in self.variables])
        loss += REGUL_PARAM * regularizers

        return loss

    def predict(self, x_test):
        init = tf.global_variables_initializer()

        with tf.Session as sess:
            sess.run(init)

            # Restoring the model from previous train
            # We take the model from the path specified in MODEL_PATH
            self.restore(sess, MODEL_PATH)

            y_useless = np.empty((BATCH_SIZE, IMG_WIDTH, IMG_HEIGHT, NUM_CLASSES))
            prediction = sess.run(self.predicter, feed_dict={self.x : x_test, self.y : y_useless, self.keep_prob : 1.})

        return prediction

    def save(self, sess, model_save_path):
        saver = tf.train.Saver()
        save_path = saver.save(sess, model_save_path)

        return save_path

    def restore(self, sess):
        saver = tf.train.Saver()
        saver.restore(sess, MODEL_PATH)



class Trainer(object):

    def __init__(self, conv_net):
        self.conv_net = conv_net

    def _get_optimizer(self, global_step):
        self.learning_rate_node = tf.Variable(LEARNING_RATE)
        optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate_node)\
                                            .minimize(self.conv_net.cost, global_step=global_step)
        return optimizer

    def _initialize(self):
        global_step = tf.Variable(0)

        self.norm_gradients_node = tf.Variable(tf.constant(0.0, shape=[len(self.conv_net.gradients_node)]))

        # Definition of the summaries
        tf.summary.histogram('norm_grads', self.norm_gradients_node)
        tf.summary.scalar('loss', self.conv_net.cost)
        tf.summary.scalar('cross_entropy', self.conv_net.cross_entropy)
        tf.summary.scalar('accuracy', self.conv_net.accuracy)

        self.optimizer = self._get_optimizer(global_step)
        tf.summary.scalar('learning_rate', self.learning_rate_node)

        # Merging all summaries for tensor flow
        self.summary_op = tf.summary.merge_all()
        init = tf.global_variables_initializer()

        return init


    def train(self, restore=False):

        if not os.path.exists(OUTPUT_PATH + 'train_predictions/'):
            os.makedirs(OUTPUT_PATH + 'train_predictions/')

        print('-> Loading data:')
        data = utils_img.load_images(TRAINING_PATH + "images/", TRAIN_SIZE)
        labels = utils_img.load_groundtruths(TRAINING_PATH + "groundtruth/", TRAIN_SIZE)

        print("train data shape: ", data.shape)
        print("labels shape: ", labels.shape)


        init = self._initialize()

        with tf.Session() as sess:
            # writing the graph
            tf.train.write_graph(sess.graph_def, OUTPUT_PATH, 'graph.pb', False)

            sess.run(init)

            if restore:
                ckpt = tf.train.get_checkpoint_state(OUTPUT_PATH)
                if ckpt and ckpt.model_checkpoint_path:
                    self.conv_net.restore(sess, ckpt.model_checkpoint_path)

            summary_writer = tf.summary.FileWriter(OUTPUT_PATH + 'summary/', graph=sess.graph)


            hm_epochs = NUM_EPOCHS
            image_indices = np.array(range(data.shape[0]))
            for epoch in range(hm_epochs):
                shuffle(image_indices)
                count_batch = 0
                epoch_loss = 0
                for step in range(int(TRAIN_SIZE/BATCH_SIZE)):
                    # feeding epoch_x and epoch_y for training current batch (to replace with our own )
                    epoch_x = data[[image_indices[count_batch:count_batch+BATCH_SIZE]]]
                    epoch_y = labels[[image_indices[count_batch:count_batch + BATCH_SIZE]]]
                    count_batch += BATCH_SIZE
                    img_idx = [image_indices[count_batch:count_batch+BATCH_SIZE]][0]

                    # sees.run() evaluate a tensor
                    # the first argument of sess.run() is an array corresponding to every operation needed
                    # Feeding sess.run() with y is necessary because cost_op needs it
                    _, loss, lr, gradients = sess.run((self.optimizer, self.conv_net.cost, self.learning_rate_node, self.conv_net.gradients_node),\
                                                        feed_dict={self.conv_net.x: epoch_x, self.conv_net.y: epoch_y,
                                                                    self.conv_net.keep_prob: DROPOUT})
                    epoch_loss += loss
                    print('Step', step+1, 'completed out of', int(TRAIN_SIZE/BATCH_SIZE), '/ step loss:', loss)

                print('epoch_x shape:', epoch_x.shape)
                print('epoch_y shape:', epoch_y.shape)


                print('-> Epoch', epoch+1, 'completed out of', hm_epochs, '/ epoch loss:', epoch_loss)
                print('Current batch size:', BATCH_SIZE)

                # In the future convert it in test set
                for i in range(TRAIN_SIZE):
                    img = data[[i]]
                    groundtruth = labels[[i]]
                    accuracy = self.output_stats(sess, summary_writer, step, img, groundtruth)
                    print('Image' + str(i+1) + ' accuracy:', accuracy, '%%')
                    self.store_prediction(sess, img, groundtruth,\
                                            save_path = OUTPUT_PATH + "train_predictions/image_{a}_epoch_{b}.png".format(a=(i+1), b=epoch))

                model_path = self.conv_net.save(sess, MODEL_PATH)
            return model_path


    def output_stats(self, sess, summary_writer, step, batch_x, batch_y):
        # Calculate batch loss and accuracy
        summary_str, loss, acc, predictions = sess.run([self.summary_op, self.conv_net.cost, self.conv_net.accuracy, self.conv_net.predicter],
                                                        feed_dict={self.conv_net.x: batch_x, self.conv_net.y: batch_y, self.conv_net.keep_prob: 1.})
        summary_writer.add_summary(summary_str, step)
        summary_writer.flush()
        return utils.accuracy(predictions, batch_y)



    def store_prediction(self, sess, test_x, test_y, save_path):
        predictions = sess.run(self.conv_net.predicter, feed_dict={self.conv_net.x: test_x, self.conv_net.y: test_y, self.conv_net.keep_prob: 1.})

        utils_img.save_image_from_proba_pred(predictions, save_path)




def main():
    conv_net = ConvNet()
    trainer = Trainer(conv_net)
    save_model_path = trainer.train()
    print('Model saved in:', save_model_path)


if __name__ == '__main__':
    main()
