# Copyright 2018 The TensorFlow Hub Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://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.
# ==============================================================================

"""
Explanation: <a href="https://colab.research.google.com/github/tensorflow/tpu/blob/master/tools/colab/shakespeare_with_tpu_and_keras.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>
End of explanation
"""


import tensorflow as tf

SEQ_LEN = 100
BATCH_SIZE = 128

"""
Explanation: Predict Shakespeare with Cloud TPUs and Keras
Overview
This example uses tf.keras to build a language model and train it on a Cloud TPU. This language model predicts the next character of text given the text so far. The trained model can generate new snippets of text that read in a similar style to the text training data.
The model trains for 10 epochs and completes in approximately 5 minutes.
This notebook is hosted on GitHub. To view it in its original repository, after opening the notebook in Colab, select File > View on GitHub.
Learning objectives
In this notebook, you will learn how to:
*   Build a two-layer, forward-LSTM model.
*   Convert a tf.keras model to an equivalent TPU version and then use the standard Keras methods to train: fit, predict, and evaluate.
*   Use the trained model to make predictions and generate your own Shakespeare-esque play.
Instruction to Train on Deep Learning VM


Start a VM attached to a TPU:
```
    INSTANCE_NAME=mytpu   # CHANGE THIS
    GCP_LOGIN_NAME=google-cloud-customer@gmail.com  # CHANGE THIS
    TPU_NAME=$INSTANCE_NAME
gcloud compute instances create $INSTANCE_NAME \
--machine-type n1-standard-8 \
--image-project deeplearning-platform-release \
--image-family tf-1-12-cpu \
--scopes cloud-platform \
--metadata proxy-user-mail="${GCP_LOGIN_NAME}",\
startup-script="echo export TPU_NAME=$TPU_NAME > /etc/profile.d/tpu-env.sh"
gcloud compute tpus create $TPU_NAME \
 --network default \
 --range 10.240.1.0 \
 --version 1.12
* Open JupyterLab from the GCP web console and clone the repo that this notebook is hosted in:
   git clone https://www.github.com/GoogleCloudPlatform/training-data-analyst
   ```
   and navigate to courses/fast-and-lean-data-science/09_lstm_keras_tpu.ipynb


Alternate method: Train on Colab

On the main menu, click Runtime and select Change runtime type. Set "TPU" as the hardware accelerator.
Click Runtime again and select Runtime > Run All. You can also run the cells manually with Shift-ENTER. 

Note that the Colab method is suitable for educational use, but you will not be able to use it for long-term work.
End of explanation
"""


!gsutil cat gs://cloud-training-demos/tpudemo/shakespeare.txt | head -1000 | tail -10

"""
Explanation: Data, model, and training
In this example, you train the model on the combined works of William Shakespeare, then use the model to compose a play in the style of The Great Bard:
<blockquote>
Loves that led me no dumbs lack her Berjoy's face with her to-day.  
The spirits roar'd; which shames which within his powers  
    Which tied up remedies lending with occasion,  
A loud and Lancaster, stabb'd in me  
    Upon my sword for ever: 'Agripo'er, his days let me free.  
    Stop it of that word, be so: at Lear,  
    When I did profess the hour-stranger for my life,  
    When I did sink to be cried how for aught;  
    Some beds which seeks chaste senses prove burning;  
But he perforces seen in her eyes so fast;  
And _  
</blockquote>

Data
We have downloaded The Complete Works of William Shakespeare as a single text file from Project Gutenberg and stored in a GCS bucket. Because TPUs are located in Google Cloud, for optimal performance, they read data directly from Google Cloud Storage (GCS).
We will use snippets from this file as the training data for the model. The target snippet is offset by one character.
End of explanation
"""


import numpy as np
import six
import tensorflow as tf
import time
import os

SHAKESPEARE_TXT = 'gs://cloud-training-demos/tpudemo/shakespeare.txt'

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

def transform(txt, pad_to=None):
  # drop any non-ascii characters
  output = np.asarray([ord(c) for c in txt if ord(c) < 255], dtype=np.int32)
  if pad_to is not None:
    output = output[:pad_to]
    output = np.concatenate([
        np.zeros([pad_to - len(txt)], dtype=np.int32),
        output,
    ])
  return output

def training_generator(seq_len=SEQ_LEN, batch_size=BATCH_SIZE):
  """A generator yields (source, target) arrays for training."""
  with tf.gfile.GFile(SHAKESPEARE_TXT, 'r') as f:
    txt = f.read()

  tf.logging.info('Input text [%d] %s', len(txt), txt[:50])
  source = transform(txt)
  while True:
    offsets = np.random.randint(0, len(source) - seq_len, batch_size)

    # Our model uses sparse crossentropy loss, but Keras requires labels
    # to have the same rank as the input logits.  We add an empty final
    # dimension to account for this.
    yield (
        np.stack([source[idx:idx + seq_len] for idx in offsets]),
        np.expand_dims(
            np.stack([source[idx + 1:idx + seq_len + 1] for idx in offsets]),
            -1),
    )

a = six.next(training_generator(seq_len=10, batch_size=1))
print(a)
#print(tf.convert_to_tensor(a[1]))

def create_dataset():
    return tf.data.Dataset.from_generator(training_generator, 
                                          (tf.int32, tf.int32),
                                          (tf.TensorShape([BATCH_SIZE, SEQ_LEN]), 
                                           tf.TensorShape([BATCH_SIZE, SEQ_LEN, 1]))
                                         )
                                          

"""
Explanation: Build the tf.data.Dataset
End of explanation
"""


EMBEDDING_DIM = 512

def lstm_model(seq_len, batch_size, stateful):
  """Language model: predict the next word given the current word."""
  source = tf.keras.Input(
      name='seed', shape=(seq_len,), batch_size=batch_size, dtype=tf.int32)

  embedding = tf.keras.layers.Embedding(input_dim=256, output_dim=EMBEDDING_DIM)(source)
  lstm_1 = tf.keras.layers.LSTM(EMBEDDING_DIM, stateful=stateful, return_sequences=True)(embedding)
  lstm_2 = tf.keras.layers.LSTM(EMBEDDING_DIM, stateful=stateful, return_sequences=True)(lstm_1)
  predicted_char = tf.keras.layers.TimeDistributed(tf.keras.layers.Dense(256, activation='softmax'))(lstm_2)
  model = tf.keras.Model(inputs=[source], outputs=[predicted_char])
  model.compile(
      optimizer=tf.train.RMSPropOptimizer(learning_rate=0.01),
      loss='sparse_categorical_crossentropy',
      metrics=['sparse_categorical_accuracy'])
  return model

"""
Explanation: Build the model
The model is defined as a two-layer, forward-LSTM—with two changes from the tf.keras standard LSTM definition:

Define the input shape of the model to comply with the XLA compiler's static shape requirement.
Use tf.train.Optimizer instead of a standard Keras optimizer (Keras optimizer support is still experimental).
End of explanation
"""


tf.keras.backend.clear_session()

training_model = lstm_model(seq_len=SEQ_LEN, batch_size=BATCH_SIZE, stateful=False)

# Use TPU if it exists, else fall back to GPU
try: # TPU detection
  tpu = tf.contrib.cluster_resolver.TPUClusterResolver()
  training_model = tf.contrib.tpu.keras_to_tpu_model(
        training_model,
        strategy=tf.contrib.tpu.TPUDistributionStrategy(tpu))
  training_input = create_dataset  # Function that returns a dataset
  print('Running on TPU ', tpu.cluster_spec().as_dict()['worker'])
except ValueError:
  tpu = None
  training_input = create_dataset()  # The dataset itself
  print("Running on GPU or CPU")

# Run fit()
training_model.fit(
    training_input,
    steps_per_epoch=100,
    epochs=10,
)
tpu_model.save_weights('/tmp/bard.h5', overwrite=True)

"""
Explanation: Train the model
The tf.contrib.tpu.keras_to_tpu_model function converts a tf.keras model to an equivalent TPU version. You then use the standard Keras methods to train: fit, predict, and evaluate.
End of explanation
"""


BATCH_SIZE = 5
PREDICT_LEN = 250

# Keras requires the batch size be specified ahead of time for stateful models.
# We use a sequence length of 1, as we will be feeding in one character at a 
# time and predicting the next character.
prediction_model = lstm_model(seq_len=1, batch_size=BATCH_SIZE, stateful=True)
prediction_model.load_weights('/tmp/bard.h5')

# We seed the model with our initial string, copied BATCH_SIZE times

seed_txt = 'Looks it not like the king?  Verily, we must go! '
seed = transform(seed_txt)
seed = np.repeat(np.expand_dims(seed, 0), BATCH_SIZE, axis=0)

# First, run the seed forward to prime the state of the model.
prediction_model.reset_states()
for i in range(len(seed_txt) - 1):
  prediction_model.predict(seed[:, i:i + 1])

# Now we can accumulate predictions!
predictions = [seed[:, -1:]]
for i in range(PREDICT_LEN):
  last_word = predictions[-1]
  next_probits = prediction_model.predict(last_word)[:, 0, :]
  
  # sample from our output distribution
  next_idx = [
      np.random.choice(256, p=next_probits[i])
      for i in range(BATCH_SIZE)
  ]
  predictions.append(np.asarray(next_idx, dtype=np.int32))
  

for i in range(BATCH_SIZE):
  print('PREDICTION %d\n\n' % i)
  p = [predictions[j][i] for j in range(PREDICT_LEN)]
  generated = ''.join([chr(c) for c in p])
  print(generated)
  print()
  assert len(generated) == PREDICT_LEN, 'Generated text too short'

"""
Explanation: Make predictions with the model
Use the trained model to make predictions and generate your own Shakespeare-esque play.
Start the model off with a seed sentence, then generate 250 characters from it. The model makes five predictions from the initial seed.
End of explanation
"""


