import tensorflow as tf
import pylab
import numpy as np
import os

%matplotlib inline
%config InlineBackend.figure_format = 'retina'

tf.logging.set_verbosity(tf.logging.INFO)

"""
Explanation: Train the Model
End of explanation
"""


tf.reset_default_graph()

sess = tf.Session()
print(sess)

"""
Explanation: Reset TensorFlow Graph
Useful in Jupyter Notebooks
Create TensorFlow Session
End of explanation
"""


num_samples = 100000

x_train = np.random.rand(num_samples).astype(np.float32)
print(x_train)

noise = np.random.normal(scale=0.01, size=len(x_train))

y_train = x_train * 0.1 + 0.3 + noise
print(y_train)

pylab.plot(x_train, y_train, '.')

x_test = np.random.rand(len(x_train)).astype(np.float32)
print(x_test)

noise = np.random.normal(scale=.01, size=len(x_train))

y_test = x_test * 0.1 + 0.3 + noise
print(y_test)

pylab.plot(x_test, y_test, '.')

with tf.device("/cpu:0"):
    W = tf.get_variable(shape=[], name='weights')
    print(W)

    b = tf.get_variable(shape=[], name='bias')
    print(b)

    x_observed = tf.placeholder(shape=[None], 
                                dtype=tf.float32, 
                                name='x_observed')
    print(x_observed)

    y_pred = W * x_observed + b
    print(y_pred)

learning_rate = 0.025

with tf.device("/cpu:0"):
    y_observed = tf.placeholder(shape=[None], dtype=tf.float32, name='y_observed')
    print(y_observed)

    loss_op = tf.reduce_mean(tf.square(y_pred - y_observed))
    optimizer_op = tf.train.GradientDescentOptimizer(learning_rate)
    
    train_op = optimizer_op.minimize(loss_op)  

    print("Loss Scalar: ", loss_op)
    print("Optimizer Op: ", optimizer_op)
    print("Train Op: ", train_op)

"""
Explanation: Load Model Training and Test/Validation Data
End of explanation
"""


with tf.device("/cpu:0"):
    init_op = tf.global_variables_initializer()
    print(init_op)

sess.run(init_op)
print("Initial random W: %f" % sess.run(W))
print("Initial random b: %f" % sess.run(b))

"""
Explanation: Randomly Initialize Variables (Weights and Bias)
The goal is to learn more accurate Weights and Bias during training.
End of explanation
"""


def test(x, y):
    return sess.run(loss_op, feed_dict={x_observed: x, y_observed: y})

test(x_train, y_train)

"""
Explanation: View Accuracy of Pre-Training, Initial Random Variables
We want this to be close to 0, but it's relatively far away.  This is why we train!
End of explanation
"""


loss_summary_scalar_op = tf.summary.scalar('loss', loss_op)
loss_summary_merge_all_op = tf.summary.merge_all()

from datetime import datetime 
version = int(datetime.now().strftime("%s"))

train_summary_writer = tf.summary.FileWriter('/root/tensorboard/linear/cpu/%s/train' % version, 
                                            graph=tf.get_default_graph())

test_summary_writer = tf.summary.FileWriter('/root/tensorboard/linear/cpu/%s/test' % version,
                                            graph=tf.get_default_graph())

"""
Explanation: Setup Loss Summary Operations for Tensorboard
End of explanation
"""


%%time

with tf.device("/cpu:0"):
    run_metadata = tf.RunMetadata()
    max_steps = 401
    for step in range(max_steps):
        if (step < max_steps - 1):
            test_summary_log, _ = sess.run([loss_summary_merge_all_op, loss_op], feed_dict={x_observed: x_test, y_observed: y_test})
            train_summary_log, _ = sess.run([loss_summary_merge_all_op, train_op], feed_dict={x_observed: x_train, y_observed: y_train})
        else:  
            test_summary_log, _ = sess.run([loss_summary_merge_all_op, loss_op], feed_dict={x_observed: x_test, y_observed: y_test})
            train_summary_log, _ = sess.run([loss_summary_merge_all_op, train_op], feed_dict={x_observed: x_train, y_observed: y_train}, 
                                            options=tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE), 
                                            run_metadata=run_metadata)
        if step % 10 == 0:
            print(step, sess.run([W, b]))
            train_summary_writer.add_summary(train_summary_log, step)
            train_summary_writer.flush()
            test_summary_writer.add_summary(test_summary_log, step)
            test_summary_writer.flush()

pylab.plot(x_train, y_train, '.', label="target")
pylab.plot(x_train, sess.run(y_pred, 
                             feed_dict={x_observed: x_train, 
                                        y_observed: y_train}), 
           ".", 
           label="predicted")
pylab.legend()
pylab.ylim(0, 1.0)

"""
Explanation: Train Model
End of explanation
"""


from tensorflow.python.saved_model import utils
from tensorflow.python.saved_model import signature_constants
from tensorflow.python.saved_model import signature_def_utils

graph = tf.get_default_graph()

x_observed = graph.get_tensor_by_name('x_observed:0')
y_pred = graph.get_tensor_by_name('add:0')

tensor_info_x_observed = utils.build_tensor_info(x_observed)
print(tensor_info_x_observed)

tensor_info_y_pred = utils.build_tensor_info(y_pred)
print(tensor_info_y_pred)

prediction_signature = signature_def_utils.build_signature_def(inputs = 
                {'x_observed': tensor_info_x_observed}, 
                outputs = {'y_pred': tensor_info_y_pred}, 
                method_name = signature_constants.PREDICT_METHOD_NAME)

%%bash 

rm -rf ./export
rm -rf ./variables
rm saved_model.pb

from tensorflow.python.saved_model import builder as saved_model_builder
from tensorflow.python.saved_model import tag_constants

fully_optimized_saved_model_path = './export'
print(fully_optimized_saved_model_path)

builder = saved_model_builder.SavedModelBuilder(fully_optimized_saved_model_path)
builder.add_meta_graph_and_variables(sess, 
                                     [tag_constants.SERVING],
                                     signature_def_map={'predict':prediction_signature,                                     
signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:prediction_signature}, 
                                     clear_devices=True,
)

builder.save(as_text=False)

%%bash

ls -l ./export

%%bash

mv ./export/* .
rm -rf ./export/

%%bash

ls -l ./

class PioBundle(object):

    def __init__(self):
        pass

    def transform_request(self,
                          request):
        import tensorflow as tf

        import json
        import numpy as np

        request_str = request.decode('utf-8')
        request_json = json.loads(request_str)
        request_np = np.asarray([request_json['x_observed']])

        return request_np
    
    
    def transform_response(self,
                           response):
        import json
        return json.dumps({"y_pred": response.tolist()[0]})

import dill as pickle

pio_bundle = PioBundle()

pio_bundle_pkl_path = 'pio_bundle.pkl'

with open(pio_bundle_pkl_path, 'wb') as fh:
    pickle.dump(pio_bundle, fh)

%%bash

ls -l pio_bundle.pkl 
ls -l saved_model.pb 
ls -l variables/

sess.close()

"""
Explanation: View loss in Tensorboard
Navigate to the Scalars and Graphs tab at this URL:
http://[ip-address]:6006
End of explanation
"""


