# Imports
import numpy as np
import tensorflow as tf
import os
import input_data # custom input for my load_dataset() Function

tf.logging.set_verbosity (tf.logging.DEBUG) # Sets verbosity for TF

def main(argv):
    (train_data, train_labels) = input_data.load_dataset()
    train_data = np.asarray(train_data, dtype=np.float32)
    train_labels = np.asarray(train_labels, dtype=np.float32)

    classifier = tf.estimator.Estimator(
        model_fn=ffnn_model_fn, model_dir="/tmp/ffnn2")
    print("Classifier set up")

    train_input_fn = tf.estimator.inputs.numpy_input_fn(
        x={"x":train_data},
        y=train_labels,
        batch_size=100,
        num_epochs=None,
        shuffle=True)
    print("Training set up")

    classifier.train(
        input_fn=train_input_fn,
        steps=500)
    print("Training done")


    eval_input_fn = tf.estimator.inputs.numpy_input_fn(
        x={"x":train_data},
        y=train_labels,
        num_epochs=1,
        shuffle=False)

    eval_results = classifier.evaluate(input_fn=eval_input_fn)
    print("Eval done")
    print(eval_results)

def ffnn_model_fn(features, labels, mode):
    input_layer = features["x"]

    hidden_layer1 = tf.layers.dense(
        inputs=input_layer,
        units=3,
        activation=tf.nn.relu)

    hidden_layer2 = tf.layers.dense(
        inputs=hidden_layer1,
        units=2,
        activation=tf.nn.relu,
        name="hidden_layer2")

    with tf.variable_scope("hidden_layer2", reuse=True):
        w = tf.get_variable("kernel")

    ### Not Working
    sess = tf.InteractiveSession();
    a = tf.Print(w,[w]);
    print(a.eval());
    ###

    out1 = tf.layers.dense(inputs=hidden_layer2, units=1)

    if mode == tf.estimator.ModeKeys.PREDICT:
        print("Not implemented yet")

    labels = tf.reshape(labels, [-1,1])
    loss = tf.losses.mean_squared_error(labels=labels, predictions=out1)

    if mode == tf.estimator.ModeKeys.TRAIN:
        optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.001)
        train_op = optimizer.minimize(
            loss=loss,
            global_step=tf.train.get_global_step())
        return tf.estimator.EstimatorSpec(mode=mode, loss=loss, train_op=train_op)

    eval_metric_ops = {
        "accuracy" : tf.metrics.accuracy(labels=labels, predictions=out1)
    }
    return tf.estimator.EstimatorSpec(mode=mode, loss=loss, eval_metric_ops=eval_metric_ops)

if __name__ == "__main__":
    tf.app.run()