# 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