#@title Licensed under the Apache License, Version 2.0 (the "License"); { display-mode: "form" }
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

! pip install tf-nightly

from __future__ import division, print_function, absolute_import

import tensorflow as tf
from tensorflow.contrib import autograph

import matplotlib.pyplot as plt

tf.enable_eager_execution()

def g(x):
  if x > 0:
    x = x * x
  else:
    x = 0.0
  return x

print(autograph.to_code(g))

print('Eager results: %2.2f, %2.2f' % (g(tf.constant(9.0)), g(tf.constant(-9.0))))

tf_g = autograph.to_graph(g)

with tf.Graph().as_default():  
  # The result works like a regular op: takes tensors in, returns tensors.
  # You can inspect the graph using tf.get_default_graph().as_graph_def()
  g_out1 = tf_g(tf.constant( 9.0))
  g_out2 = tf_g(tf.constant(-9.0))
  with tf.Session() as sess:
    print('Graph results: %2.2f, %2.2f\n' % (sess.run(g_out1), sess.run(g_out2)))

# Continue in a loop
def f(l):
  s = 0
  for c in l:
    if c % 2 > 0:
      continue
    s += c
  return s

print('Eager result: %d' % f(tf.constant([10,12,15,20])))

tf_f = autograph.to_graph(f)

with tf.Graph().as_default():  
  with tf.Session():
    print('Graph result: %d\n\n' % tf_f(tf.constant([10,12,15,20])).eval())

print(autograph.to_code(f))

@autograph.convert()
def fizzbuzz(num):
  if num % 3 == 0 and num % 5 == 0:
      print('FizzBuzz')
  elif num % 3 == 0:
      print('Fizz')
  elif num % 5 == 0:
      print('Buzz')
  else:
      print(num)
  return num


with tf.Graph().as_default():
  # The result works like a regular op: takes tensors in, returns tensors.
  # You can inspect the graph using tf.get_default_graph().as_graph_def()
  num = tf.placeholder(tf.int32)
  result = fizzbuzz(num)
  with tf.Session() as sess:
    for n in range(10,16):
      sess.run(result, feed_dict={num:n})

@autograph.convert()
def f(x):
  assert x != 0, 'Do not pass zero!'
  return x * x

with tf.Graph().as_default():  
  with tf.Session():
    try:
      print(f(tf.constant(0)).eval())
    except tf.errors.InvalidArgumentError as e:
      print('Got error message:\n    %s' % e.message)

@autograph.convert()
def f(n):
  if n >= 0:
    while n < 5:
      n += 1
      print(n)
  return n
    
with tf.Graph().as_default():
  with tf.Session():
    f(tf.constant(0)).eval()

@autograph.convert()
def f(n):
  z = []
  # We ask you to tell us the element dtype of the list
  autograph.utils.set_element_type(z, tf.int32)
  
  for i in range(n):
    z.append(i)
  # when you're done with the list, stack it
  # (this is just like np.stack)
  return autograph.stack(z) 

#tf_f = autograph.to_graph(f)

with tf.Graph().as_default():  
  with tf.Session():
    print(f(tf.constant(3)).eval())

@autograph.convert()
def nearest_odd_square(x):
  if x > 0:
    x = x * x
    if x % 2 == 0:
      x = x + 1
  return x

with tf.Graph().as_default():  
  with tf.Session() as sess:
    print(sess.run(nearest_odd_square(tf.constant(4))))
    print(sess.run(nearest_odd_square(tf.constant(5))))
    print(sess.run(nearest_odd_square(tf.constant(6))))

@autograph.convert()
def square_until_stop(x, y):
  while x < y:
    x = x * x
  return x
    
with tf.Graph().as_default():  
  with tf.Session() as sess:
    print(sess.run(square_until_stop(tf.constant(4), tf.constant(100))))

@autograph.convert()
def argwhere_cumsum(x, threshold):
  current_sum = 0.0
  idx = 0
  for i in range(len(x)):
    idx = i
    if current_sum >= threshold:
      break
    current_sum += x[i]
  return idx

N = 10
with tf.Graph().as_default():  
  with tf.Session() as sess:
    idx = argwhere_cumsum(tf.ones(N), tf.constant(float(N/2)))
    print(sess.run(idx))

(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.mnist.load_data()

def mlp_model(input_shape):
  model = tf.keras.Sequential((
      tf.keras.layers.Flatten(),
      tf.keras.layers.Dense(100, activation='relu', input_shape=input_shape),
      tf.keras.layers.Dense(100, activation='relu'),
      tf.keras.layers.Dense(10, activation='softmax')))
  model.build()
  return model


def predict(m, x, y):
  y_p = m(x)
  losses = tf.keras.losses.categorical_crossentropy(y, y_p)
  l = tf.reduce_mean(losses)
  accuracies = tf.keras.metrics.categorical_accuracy(y, y_p)
  accuracy = tf.reduce_mean(accuracies)
  return l, accuracy


def fit(m, x, y, opt):
  l, accuracy = predict(m, x, y)
  # Autograph automatically adds the necessary `tf.control_dependencies` here.
  # (Without them nothing depends on `opt.minimize`, so it doesn't run.)
  # This makes it much more like eager-code.
  opt.minimize(l)
  return l, accuracy


def setup_mnist_data(is_training, batch_size):
  if is_training:
    ds = tf.data.Dataset.from_tensor_slices((train_images, train_labels))
    ds = ds.shuffle(batch_size * 10)
  else:
    ds = tf.data.Dataset.from_tensor_slices((test_images, test_labels))

  ds = ds.repeat()
  ds = ds.batch(batch_size)
  return ds


def get_next_batch(ds):
  itr = ds.make_one_shot_iterator()
  image, label = itr.get_next()
  x = tf.to_float(image)/255.0
  y = tf.one_hot(tf.squeeze(label), 10)
  return x, y 

# Use `recursive = True` to recursively convert functions called by this one.
@autograph.convert(recursive=True)
def train(train_ds, test_ds, hp):
  m = mlp_model((28 * 28,))
  opt = tf.train.MomentumOptimizer(hp.learning_rate, 0.9)
  
  # We'd like to save our losses to a list. In order for AutoGraph
  # to convert these lists into their graph equivalent,
  # we need to specify the element type of the lists.
  train_losses = []
  autograph.utils.set_element_type(train_losses, tf.float32)
  test_losses = []
  autograph.utils.set_element_type(test_losses, tf.float32)
  train_accuracies = []
  autograph.utils.set_element_type(train_accuracies, tf.float32)
  test_accuracies = []
  autograph.utils.set_element_type(test_accuracies, tf.float32)
  
  # This entire training loop will be run in-graph.
  i = tf.constant(0)
  while i < hp.max_steps:
    train_x, train_y = get_next_batch(train_ds)
    test_x, test_y = get_next_batch(test_ds)

    step_train_loss, step_train_accuracy = fit(m, train_x, train_y, opt)
    step_test_loss, step_test_accuracy = predict(m, test_x, test_y)
    if i % (hp.max_steps // 10) == 0:
      print('Step', i, 'train loss:', step_train_loss, 'test loss:',
            step_test_loss, 'train accuracy:', step_train_accuracy,
            'test accuracy:', step_test_accuracy)
    train_losses.append(step_train_loss)
    test_losses.append(step_test_loss)
    train_accuracies.append(step_train_accuracy)
    test_accuracies.append(step_test_accuracy)
    i += 1
  
  # We've recorded our loss values and accuracies 
  # to a list in a graph with AutoGraph's help.
  # In order to return the values as a Tensor, 
  # we need to stack them before returning them.
  return (autograph.stack(train_losses), autograph.stack(test_losses),  
          autograph.stack(train_accuracies), autograph.stack(test_accuracies))

with tf.Graph().as_default() as g:
  hp = tf.contrib.training.HParams(
      learning_rate=0.05,
      max_steps=500,
  )
  train_ds = setup_mnist_data(True, 50)
  test_ds = setup_mnist_data(False, 1000)
  (train_losses, test_losses, train_accuracies,
   test_accuracies) = train(train_ds, test_ds, hp)

  init = tf.global_variables_initializer()
  
with tf.Session(graph=g) as sess:
  sess.run(init)
  (train_losses, test_losses, train_accuracies,
   test_accuracies) = sess.run([train_losses, test_losses, train_accuracies,
                                test_accuracies])
  
plt.title('MNIST train/test losses')
plt.plot(train_losses, label='train loss')
plt.plot(test_losses, label='test loss')
plt.legend()
plt.xlabel('Training step')
plt.ylabel('Loss')
plt.show()
plt.title('MNIST train/test accuracies')
plt.plot(train_accuracies, label='train accuracy')
plt.plot(test_accuracies, label='test accuracy')
plt.legend(loc='lower right')
plt.xlabel('Training step')
plt.ylabel('Accuracy')
plt.show()