# blahmports
import tensorflow as tf
import numpy as np

# define hyperparameters
learning_rate = 0.01
training_epochs = 1000
batch_size = 100
display_step = 10

# Step 1:
# Calculating the layer's output.
# (now dynamic)
def get_layer_op(input_layer, weight_shape, bias_shape):

    # first initialize constants
    # (since we are not doing that anywhere)
    # weight_stddev = (2.0/weight_shape[0])**0.5
    weight_init = tf.random_uniform_initializer(minval=-1, maxval=1)
    bias_init = tf.constant_initializer(value=0)

    # initialize weights
    # remember scoping & get_variable?
    W = tf.get_variable("W", weight_shape, initializer=weight_init)

    # initialize biases
    b = tf.get_variable("b", bias_shape, initializer=bias_init)

    # return the output
    return tf.nn.relu(tf.matmul(input_layer, W) + b)

# Step 2:
# Computing the losses
# by comparing layer output and actual output.
def compute_losses(op, y):

    xentropy = tf.nn.softmax_cross_entropy_with_logits(logits=op, labels=y)
    #~

    loss = tf.reduce_mean(xentropy)
    #~

    return loss

# Step 3:
# Training This Shit
# Over Step 1 & 2.
def training(cost,  global_step):

    tf.summary.scalar("cost", cost)

    optimizer = tf.train.GradientDescentOptimizer(learning_rate)
    #~  optimizer?

    train_op = optimizer.minimize(cost, global_step=global_step)
    #~  minimize opt?

    return train_op

# Step 4:
# Evaluate the samples, return accuracy
def evaluate(op, y):

    predicted_op = tf.argmax(op, 1)
    actual_op = tf.argmax(y, 1)
    #~  argmaxes?

    correct_pred = tf.equal(predicted_op, actual_op)
    #~  check correctness of predictions

    accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32))
    #~  reduce mean?

    return accuracy

# an aggreagation of layers here
def actual_network(x):

    #~~# what shapes to use?
    with tf.variable_scope('layer_1'):
        layer_1_output = get_layer_op(x, [batch_size, 3, 256], [256])

    with tf.variable_scope('layer_2'):
        layer_2_output = get_layer_op(layer_1_output, [batch_size, 256, 256], [256])

    with tf.variable_scope('layer_3'):
        layer_3_output = get_layer_op(layer_2_output, [batch_size, 256, 256], [256])

    with tf.variable_scope('layer_out'):
        layer_4_output = get_layer_op(layer_3_output, [batch_size, 256, 1], [1])

    return layer_4_output

# reutrns batches
def get_batch(batch_size, x_train, y_train, i):
    #
    x_train_new = {}
    for each_ in x_train:
        try:
            x_train_new['path'] += [each_['path']]
            x_train_new['belief'] += [each_['belief']]
            x_train_new['problem'] += [each_['problem']]
        except:
            x_train_new['path'] = [each_['path']]
            x_train_new['belief'] = [each_['belief']]
            x_train_new['problem'] = [each_['problem']]
    return x_train_new, y_train[i:batch_size+i]

# create input sequences
def create_input_sets():

    problems = list(np.random.randint(0, vocab_size, train_test_size))
    traversal_paths = []
    for i in range(train_test_size):
            path_traversal = []
            for i in range(10):
                path_traversal += [list(np.random.randint(0, train_test_size, 3))]
            traversal_paths += [list(path_traversal)]
    beliefs = []
    for i in range(train_test_size):
        beliefs += [list(np.random.randint(0, train_test_size, 3))]

    full_set = []
    for i in range(train_test_size):
        full_set += [{'problem' : problems[i],
                'path' : traversal_paths[i],
                'belief' : beliefs[i]}]

    return full_set[:split], full_set[split:]

# create wanted sequences
def create_output_sets():

    traversal_paths = []
    for i in range(train_test_size):
            path_traversal = []
            for i in range(10):
                path_traversal += [list(np.random.randint(0, train_test_size, 3))]
            traversal_paths += [list(path_traversal)]

    return traversal_paths[:split], traversal_paths[split:]

if __name__ == '__main__':
    #
    x_train, x_test = create_input_sets()
    #
    y_train, y_test = create_output_sets()
    #
    # create the graph
    with tf.Graph().as_default():
#       #
        # launch into our scope
        with tf.variable_scope('creativity_approximator'):
            #
            # input placeholders
            # ###################
            #~~# are shapes okay for these?
            #
            # x = list of lists (all probable traversal paths for a solution 'a')
            # [
            #   [
            #       [node, edge_property, node], [n, p, n], ...
            #   ],
            #   [
            #       [node, edge_property, node], [n, p, n], ...
            #   ],
            # ]
            x = tf.placeholder(tf.float32, shape=[None, None, None, None], name='probable_traversals')
            #
            # y = a singular traversal path
            # [
            #   [node, edge_property, node], [n, p, n], ...
            # ]
            y = tf.placeholder(tf.float32, shape=[None], name='core_belief')
            #
            # z = a word
            # 'this_problem'
            z = tf.placeholder(tf.float32, shape=None, name='problem')
            #
            # output placeholder
            # ###################
            #~~# are shapes okay for these?
            # a = a singular traversal path
            # [
            #   [node, edge_property, node], [n, p, n], ...
            # ]
            a = tf.placeholder(tf.float32, shape=[None], name='perfect_solution')
            #
            # merge all input placeholders
            # since those are inputs to
            # the neural network
            w = tf.stack([x, y, z])
            #~~# this obviously doesnt work
            # because SHAPES!

            # get network output
            output = actual_network(w)
            #
            # get loss
            cost = compute_losses(output, y)
            #
            # define global step to keep
            # track of training process
            global_step = tf.Variable(0, name='global_step', trainable=False)
            #
            # change weights
            train_op = training(cost, global_step)
            #
            # evaluate changed weights
            eval_op = evaluate(output, y)

            # initiate saver and summarizer
            saver = tf.train.Saver()
            summary_op = tf.summary.merge_all()
            #
            # initiate and run session
            sess = tf.Session()
            summary_writer = tf.summary.FileWriter('model/', graph_def=sess.graph_def)
            init_op = tf.global_variables_initializer()
            sess.run(init_op)

            i = 1

            # training cycle
            for each_epoch in range(training_epochs):

                avg_cost = 0.
                total_batches = int(len(x_train)/batch_size)

                # loop over all batches now
                for i in range(total_batches):

                    batch_x, batch_y = get_batch(batch_size, x_train, y_train, i)
                    feed_dict = {x : batch_x['path'],
                            y : batch_x['belief'],
                            z: batch_x['problem'],
                            a: batch_y['solution']}
                    sess.run(train_op, feed_dict=feed_dict)

                    batch_cost = sess.run(cost, feed_dict=feed_dict)
                    avg_cost += batch_cost/total_batches

                    i += (i * batch_size)

                # do updates
                if each_epoch % display_step == 0:

                    # show stats
                    validate_feed = {x: x_test, y: y_test}
                    accuracy = sess.run(eval_op, feed_dict=validate_feed)
                    print('Model Error Rate:', (1 - accuracy) * 100, '%')

                    # save summary and model
                    summary_str = sess.run(summary_op, feed_dict=feed_dict)
                    summary_writer.add_summary(summary_str, sess.run(global_step))
                    saver.save(sess, "model/model-checkpoint", global_step=global_step)

            print('Iterations Complete.')

            test_feed = {x: x_test, y: y_test}
            test_accuracy = sess.run(eval_op, feed_dict=test_feed)
            print(accuracy * 100, '% Accuracy.')

            # you can see tensorboard by -
            # tensorboard --logdir=model