def parse(line):
  """Parses a line from the colors dataset."""
  items = tf.string_split([line], ",").values
  rgb = tf.string_to_number(items[1:], out_type=tf.float32) / 255.0
  color_name = items[0]
  chars = tf.one_hot(tf.decode_raw(color_name, tf.uint8), depth=256)
  length = tf.cast(tf.shape(chars)[0], dtype=tf.int64)
  return rgb, chars, length


def set_static_batch_shape(batch_size):
  def apply(rgb, chars, length):
    rgb.set_shape((batch_size, None))
    chars.set_shape((batch_size, None, 256))
    length.set_shape((batch_size,))
    return rgb, chars, length
  return apply


def load_dataset(data_dir, url, batch_size, training=True):
  """Loads the colors data at path into a tf.PaddedDataset."""
  path = tf.keras.utils.get_file(os.path.basename(url), url, cache_dir=data_dir)
  dataset = tf.data.TextLineDataset(path)
  dataset = dataset.skip(1)
  dataset = dataset.map(parse)
  dataset = dataset.cache()
  dataset = dataset.repeat()
  if training:
    dataset = dataset.shuffle(buffer_size=3000)
  dataset = dataset.padded_batch(
      batch_size, padded_shapes=((None,), (None, 256), ()))
  # To simplify the model code, we statically set as many of the shapes that we
  # know.
  dataset = dataset.map(set_static_batch_shape(batch_size))
  return dataset

"""
Explanation: Case study: training a custom RNN, using Keras and Estimators
In this section, we show how you can use AutoGraph to build RNNColorbot, an RNN that takes as input names of colors and predicts their corresponding RGB tuples. The model will be trained by a custom Estimator.
To get started, set up the dataset. The following cells defines methods that download and format the data needed for RNNColorbot; the details aren't important (read them in the privacy of your own home if you so wish), but make sure to run the cells before proceeding.
End of explanation
"""


@autograph.convert()
class RnnColorbot(tf.keras.Model):
  """RNN Colorbot model."""

  def __init__(self):
    super(RnnColorbot, self).__init__()
    self.lower_cell = tf.contrib.rnn.LSTMBlockCell(256, dtype=tf.float32)
    self.upper_cell = tf.contrib.rnn.LSTMBlockCell(128, dtype=tf.float32)
    self.relu_layer = tf.layers.Dense(3, activation=tf.nn.relu)

  def _rnn_layer(self, chars, cell, batch_size, training):
    """A single RNN layer.

    Args:
      chars: A Tensor of shape (max_sequence_length, batch_size, input_size)
      cell: An object of type tf.contrib.rnn.LSTMBlockCell
      batch_size: Int, the batch size to use
      training: Boolean, whether the layer is used for training

    Returns:
      A Tensor of shape (max_sequence_length, batch_size, output_size).
    """
    hidden_outputs = tf.TensorArray(tf.float32, 0, True)
    state, output = cell.zero_state(batch_size, tf.float32)
    for ch in chars:
      cell_output, (state, output) = cell.call(ch, (state, output))
      hidden_outputs.append(cell_output)
    hidden_outputs = autograph.stack(hidden_outputs)
    if training:
      hidden_outputs = tf.nn.dropout(hidden_outputs, 0.5)
    return hidden_outputs

  def build(self, _):
    """Creates the model variables. See keras.Model.build()."""
    self.lower_cell.build(tf.TensorShape((None, 256)))
    self.upper_cell.build(tf.TensorShape((None, 256)))
    self.relu_layer.build(tf.TensorShape((None, 128)))    
    self.built = True


  def call(self, inputs, training=False):
    """The RNN model code. Uses Eager.

    The model consists of two RNN layers (made by lower_cell and upper_cell),
    followed by a fully connected layer with ReLU activation.

    Args:
      inputs: A tuple (chars, length)
      training: Boolean, whether the layer is used for training

    Returns:
      A Tensor of shape (batch_size, 3) - the model predictions.
    """
    chars, length = inputs
    batch_size = chars.shape[0]
    seq = tf.transpose(chars, (1, 0, 2))

    seq = self._rnn_layer(seq, self.lower_cell, batch_size, training)
    seq = self._rnn_layer(seq, self.upper_cell, batch_size, training)

    # Grab just the end-of-sequence from each output.
    indices = (length - 1, list(range(batch_size)))
    indices = tf.stack(indices, 1)
    sequence_ends = tf.gather_nd(seq, indices)
    return self.relu_layer(sequence_ends)

@autograph.convert()
def loss_fn(labels, predictions):
  return tf.reduce_mean((predictions - labels) ** 2)

"""
Explanation: To show the use of control flow, we write the RNN loop by hand, rather than using a pre-built RNN model.
Note how we write the model code in Eager style, with regular if and while statements. Then, we annotate the functions with @autograph.convert to have them automatically compiled to run in graph mode.
We use Keras to define the model, and we will train it using Estimators.
End of explanation
"""


def model_fn(features, labels, mode, params):
  """Estimator model function."""
  chars = features['chars']
  sequence_length = features['sequence_length']
  inputs = (chars, sequence_length)

  # Create the model. Simply using the AutoGraph-ed class just works!
  colorbot = RnnColorbot()
  colorbot.build(None)

  if mode == tf.estimator.ModeKeys.TRAIN:
    predictions = colorbot(inputs, training=True)
    loss = loss_fn(labels, predictions)

    learning_rate = params['learning_rate']
    optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)
    global_step = tf.train.get_global_step()
    train_op = optimizer.minimize(loss, global_step=global_step)
    return tf.estimator.EstimatorSpec(mode, loss=loss, train_op=train_op)

  elif mode == tf.estimator.ModeKeys.EVAL:
    predictions = colorbot(inputs)
    loss = loss_fn(labels, predictions)

    return tf.estimator.EstimatorSpec(mode, loss=loss)

  elif mode == tf.estimator.ModeKeys.PREDICT:
    predictions = colorbot(inputs)

    predictions = tf.minimum(predictions, 1.0)
    return tf.estimator.EstimatorSpec(mode, predictions=predictions)

"""
Explanation: We will now create the model function for the custom Estimator.
In the model function, we simply use the model class we defined above - that's it!
End of explanation
"""


def input_fn(data_dir, data_url, params, training=True):
  """An input function for training"""
  batch_size = params['batch_size']
  
  # load_dataset defined above
  dataset = load_dataset(data_dir, data_url, batch_size, training=training)

  # Package the pipeline end in a format suitable for the estimator.
  labels, chars, sequence_length = dataset.make_one_shot_iterator().get_next()
  features = {
      'chars': chars,
      'sequence_length': sequence_length
  }

  return features, labels

"""
Explanation: We'll create an input function that will feed our training and eval data.
End of explanation
"""


params = {
    'batch_size': 64,
    'learning_rate': 0.01,
}

train_url = "https://raw.githubusercontent.com/random-forests/tensorflow-workshop/master/archive/extras/colorbot/data/train.csv"
test_url = "https://raw.githubusercontent.com/random-forests/tensorflow-workshop/master/archive/extras/colorbot/data/test.csv"
data_dir = "tmp/rnn/data"

regressor = tf.estimator.Estimator(
    model_fn=model_fn,
    params=params)

regressor.train(
    input_fn=lambda: input_fn(data_dir, train_url, params),
    steps=100)
eval_results = regressor.evaluate(
    input_fn=lambda: input_fn(data_dir, test_url, params, training=False),
    steps=2
)

print('Eval loss at step %d: %s' % (eval_results['global_step'], eval_results['loss']))

"""
Explanation: We now have everything in place to build our custom estimator and use it for training and eval!
End of explanation
"""


def predict_input_fn(color_name):
  """An input function for prediction."""
  _, chars, sequence_length = parse(color_name)

  # We create a batch of a single element.
  features = {
      'chars': tf.expand_dims(chars, 0),
      'sequence_length': tf.expand_dims(sequence_length, 0)
  }
  return features, None


def draw_prediction(color_name, pred):
  pred = pred * 255
  pred = pred.astype(np.uint8)
  plt.axis('off')
  plt.imshow(pred)
  plt.title(color_name)
  plt.show()


def predict_with_estimator(color_name, regressor):
  predictions = regressor.predict(
      input_fn=lambda:predict_input_fn(color_name))
  pred = next(predictions)
  predictions.close()
  pred = np.minimum(pred, 1.0)
  pred = np.expand_dims(np.expand_dims(pred, 0), 0)

  draw_prediction(color_name, pred)

tb = widgets.TabBar(["RNN Colorbot"])
while True:
  with tb.output_to(0):
    try:
      color_name = six.moves.input("Give me a color name (or press 'enter' to exit): ")
    except (EOFError, KeyboardInterrupt):
      break
  if not color_name:
    break
  with tb.output_to(0):
    tb.clear_tab()
    predict_with_estimator(color_name, regressor)
  

"""
Explanation: And here's the same estimator used for inference.
End of explanation
"""


