import os
import tensorflow as tf
from tensorflow.core.protobuf import saver_pb2
import driving_data
import model

os.environ["CUDA_VISIBLE_DEVICES"]="1"

LOGDIR = './save'

sess = tf.InteractiveSession()

L2NormConst = 0.001

train_vars = tf.trainable_variables()

loss1 = tf.reduce_mean(tf.square(tf.subtract(model.y1_, model.y_1))) 
loss2 = tf.reduce_mean(tf.square(tf.subtract(model.y2_, model.y_2)))
loss3 = tf.reduce_mean(tf.square(tf.subtract(model.y3_, model.y_3)))
loss = tf.divide((loss1 + loss2 + loss3), 3)

# L2 - Regularization 
L2_decay = tf.add_n([tf.nn.l2_loss(v) for v in train_vars]) * L2NormConst
# Total Loss
loss = loss + L2_decay

train_step = tf.train.AdamOptimizer(1e-4).minimize(loss)
sess.run(tf.initialize_all_variables())

# create a summary to monitor cost tensor
tf.summary.scalar("loss", loss)
# merge all summaries into a single op
merged_summary_op = tf.summary.merge_all()

saver = tf.train.Saver(write_version = saver_pb2.SaverDef.V1)

# op to write logs to Tensorboard
logs_path = './logs'
summary_writer = tf.summary.FileWriter(logs_path, graph=tf.get_default_graph())

epochs = 30
batch_size = 100

# train over the dataset about 30 times
for epoch in range(epochs):
  for i in range(int(driving_data.num_images/batch_size)):
    xs1, xs2, ys1, ys2, ys3 = driving_data.LoadTrainBatch(batch_size)
    train_step.run(feed_dict={model.x1:xs1, model.x2:xs2 , model.y1_: ys1, model.y2_: ys2 , model.y3_: ys3, model.keep_prob: 0.8})
    if i % 10 == 0:
      xs1, xs2, ys1, ys2, ys3 = driving_data.LoadValBatch(batch_size)
      loss_value = loss.eval(feed_dict={model.x1:xs1, model.x2:xs2 , model.y1_: ys1, model.y2_: ys2 , model.y3_: ys3, model.keep_prob: 1.0})
      print("Epoch: %d, Step: %d, Loss: %g" % (epoch, epoch * batch_size + i, loss_value))

    # write logs at every iteration
    summary = merged_summary_op.eval(feed_dict={model.x1:xs1, model.x2:xs2 , model.y1_: ys1, model.y2_: ys2 , model.y3_: ys3, model.keep_prob: 1.0})
    summary_writer.add_summary(summary, epoch * driving_data.num_images/batch_size + i)

    if i % batch_size == 0:
      if not os.path.exists(LOGDIR):
        os.makedirs(LOGDIR)
      checkpoint_path = os.path.join(LOGDIR, "model.ckpt")
      filename = saver.save(sess, checkpoint_path)
  print("Model saved in file: %s" % filename)

print("Run the command line:\n" \
          "--> tensorboard --logdir=./logs " \
          "\nThen open http://0.0.0.0:6006/ into your web browser")
