# Ensure the right version of Tensorflow is installed.
!pip freeze | grep tensorflow==2.1

# change these to try this notebook out
BUCKET = 'cloud-training-demos-ml'
PROJECT = 'cloud-training-demos'
REGION = 'us-east1' #'us-central1'

import os
os.environ['BUCKET'] = BUCKET
os.environ['PROJECT'] = PROJECT
os.environ['REGION'] = REGION

%%bash
if ! gsutil ls | grep -q gs://${BUCKET}/; then
  gsutil mb -l ${REGION} gs://${BUCKET}
fi

%%bash
ls *.csv

"""
Explanation: <h1> Create Keras DNN model </h1>

<h2>Learning Objectives</h2>
<ol>
<li>Use the tf.data API to create our ML datasets</li>
<li>Use the Keras functional API to create our DNN architecture</li>
<li>Deploy a version of the model on Cloud AI Platform</li>
</ol>

Note: This notebook requires TensorFlow 2.1 as we are creating a model using Keras. 
TODO: Complete the lab notebook #TODO sections. You can refer to the ../solutions/3_keras_dnn.ipynb notebook for reference.
End of explanation
"""


import shutil
import numpy as np
import tensorflow as tf
print(tf.__version__)

# Determine CSV, label, and key columns
CSV_COLUMNS = 'weight_pounds,is_male,mother_age,plurality,gestation_weeks,key'.split(',')
LABEL_COLUMN = 'weight_pounds'
KEY_COLUMN = 'key'

# Set default values for each CSV column. Treat is_male and plurality as strings.
DEFAULTS = [[0.0], ['null'], [0.0], ['null'], [0.0], ['nokey']]

# load the training data
# use the tf.data API to create and load a dataset. Be sure to map features and labels.
# https://www.tensorflow.org/api_docs/python/tf/data/experimental
def load_dataset(pattern, batch_size=1, mode=tf.estimator.ModeKeys.EVAL):
  dataset = (
              # TODO create the dataset
              )
 
  if mode == tf.estimator.ModeKeys.TRAIN:
        # TODO add shuffling to the dataset if it's in training mode:
        dataset = dataset
  dataset = dataset.prefetch(1) # take advantage of multi-threading; 1=AUTOTUNE
  return dataset

"""
Explanation: Create Keras model
<p>
First, write an input_fn to read the data.
End of explanation
"""


## Build a simple Keras DNN using its Functional API
def rmse(y_true, y_pred):
    return tf.sqrt(tf.reduce_mean(tf.square(y_pred - y_true))) 

# Helper function to handle categorical columns
def categorical_fc(name, values):
  return tf.feature_column.indicator_column(
    tf.feature_column.categorical_column_with_vocabulary_list(name, values))

def build_dnn_model():
    # input layer
    inputs = {
        colname : tf.keras.layers.Input(name=colname, shape=(), dtype='float32')
           for colname in [] # TODO complete array of numeric input columns
    }
    inputs.update({
        colname : tf.keras.layers.Input(name=colname, shape=(), dtype='string')
            for colname in [] # TODO complete array of string input columns   
    })

    # feature columns from inputs
    feature_columns = {
        colname : tf.feature_column.numeric_column(colname)
            for colname in [] # TODO complete array of numeric feature columns
    }
    if False:
        # Until TF-serving supports 2.0, so as to get servable model
        feature_columns['is_male'] = categorical_fc('is_male', ['True', 'False', 'Unknown'])
        feature_columns['plurality'] = categorical_fc('plurality',
                          ['Single(1)', 'Twins(2)', 'Triplets(3)',
                           'Quadruplets(4)', 'Quintuplets(5)','Multiple(2+)'])

    # TODO: Use the feature columns and inputs above to create a DNN of [64, 32]

    

    model = tf.keras.models.Model(inputs, output)
    model.compile(optimizer='adam', loss='mse', metrics=[rmse, 'mse'])    
    return model

print("Here is our DNN architecture so far:\n")

# note how to use strategy to do distributed training
strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
    model = build_dnn_model()
print(model.summary())

"""
Explanation: Next, define the feature columns. mother_age and gestation_weeks should be numeric.
The others (is_male, plurality) should be categorical.
End of explanation
"""


tf.keras.utils.plot_model(model, 'dnn_model.png', show_shapes=False, rankdir='LR')

"""
Explanation: We can visualize the DNN using the Keras plot_model utility.
End of explanation
"""


TRAIN_BATCH_SIZE = 32
NUM_TRAIN_EXAMPLES = 10000 * 5 # training dataset repeats, so it will wrap around
NUM_EVALS = 5  # how many times to evaluate
NUM_EVAL_EXAMPLES = 10000 # enough to get a reasonable sample, but not so much that it slows down

trainds = load_dataset('train*', TRAIN_BATCH_SIZE, tf.estimator.ModeKeys.TRAIN)
evalds = load_dataset('eval*', 1000, tf.estimator.ModeKeys.EVAL).take(NUM_EVAL_EXAMPLES//1000)

steps_per_epoch = NUM_TRAIN_EXAMPLES // (TRAIN_BATCH_SIZE * NUM_EVALS)

history = model.fit(trainds, 
                    validation_data=evalds,
                    epochs=NUM_EVALS, 
                    steps_per_epoch=steps_per_epoch)

"""
Explanation: Train and evaluate
End of explanation
"""


# plot
import matplotlib.pyplot as plt
nrows = 1
ncols = 2
fig = plt.figure(figsize=(10, 5))

for idx, key in enumerate(['loss', 'rmse']):
    ax = fig.add_subplot(nrows, ncols, idx+1)
    plt.plot(history.history[key])
    plt.plot(history.history['val_{}'.format(key)])
    plt.title('model {}'.format(key))
    plt.ylabel(key)
    plt.xlabel('epoch')
    plt.legend(['train', 'validation'], loc='upper left');

"""
Explanation: Visualize loss curve
End of explanation
"""


# Serving function that passes through keys
@tf.function(input_signature=[{
      'is_male': tf.TensorSpec([None,], dtype=tf.string, name='is_male'),
      'mother_age': tf.TensorSpec([None,], dtype=tf.float32, name='mother_age'),
      'plurality': tf.TensorSpec([None,], dtype=tf.string, name='plurality'),
      'gestation_weeks': tf.TensorSpec([None,], dtype=tf.float32, name='gestation_weeks'),
      'key': tf.TensorSpec([None,], dtype=tf.string, name='key')
}])
def my_serve(inputs):
    feats = inputs.copy()
    key = feats.pop('key')
    output = model(feats)
    return {'key': key, 'babyweight': output}

import shutil, os, datetime
OUTPUT_DIR = './export/babyweight'
shutil.rmtree(OUTPUT_DIR, ignore_errors=True)
EXPORT_PATH = os.path.join(OUTPUT_DIR, datetime.datetime.now().strftime('%Y%m%d%H%M%S'))
tf.saved_model.save(model, EXPORT_PATH, signatures={'serving_default': my_serve})
print("Exported trained model to {}".format(EXPORT_PATH))
os.environ['EXPORT_PATH'] = EXPORT_PATH

!find $EXPORT_PATH

"""
Explanation: Save the model
Let's wrap the model so that we can supply keyed predictions, and get the key back in our output
End of explanation
"""


!saved_model_cli show --tag_set serve --signature_def serving_default --dir {EXPORT_PATH}

%%bash
MODEL_NAME="babyweight"
VERSION_NAME="dnn"
MODEL_LOCATION=$EXPORT_PATH
echo "Deleting and deploying $MODEL_NAME $MODEL_VERSION from $MODEL_LOCATION ... this will take a few minutes"

if [[ $(gcloud ai-platform models list --format='value(name)' | grep $MODEL_NAME) ]]; then
    echo "The model named $MODEL_NAME already exists."
else
    # create model
    echo "Creating $MODEL_NAME model now."
    gcloud ai-platform models create --regions=$REGION $MODEL_NAME
fi

if [[ $(gcloud ai-platform versions list --model $MODEL_NAME --format='value(name)' | grep $VERSION_NAME) ]]; then
    echo "Deleting already the existing model $MODEL_NAME:$VERSION_NAME ... "
    gcloud ai-platform versions delete --model=$MODEL_NAME $VERSION_NAME
    echo "Please run this cell again if you don't see a Creating message ... "
    sleep 2
fi

# TODO create model on Cloud AI Platform. Use python-version 3.5 and runtime-version 1.14
# https://cloud.google.com/sdk/gcloud/reference/ai-platform/versions/create
echo "Creating $MODEL_NAME:$VERSION_NAME"
gcloud ai-platform versions create # TODO complete the statement

"""
Explanation: Deploy trained model to Cloud AI Platform
End of explanation
"""


%%writefile input.json
{"key": "b1", "is_male": "True", "mother_age": 26.0, "plurality": "Single(1)", "gestation_weeks": 39}
{"key": "b2", "is_male": "True", "mother_age": 33.0, "plurality": "Single(1)", "gestation_weeks": 41}
{"key": "g1", "is_male": "False", "mother_age": 26.0, "plurality": "Single(1)", "gestation_weeks": 39}
{"key": "g2", "is_male": "False", "mother_age": 33.0, "plurality": "Single(1)", "gestation_weeks": 41}

!gcloud ai-platform predict --model babyweight --json-instances input.json --version dnn

"""
Explanation: Monitor the model creation at GCP Console > AI Platform and once the model version dnn is created, proceed to the next cell.
End of explanation
"""


from oauth2client.client import GoogleCredentials
from googleapiclient import discovery

credentials = GoogleCredentials.get_application_default()
api = discovery.build('ml', 'v1', credentials=credentials)
project = PROJECT
model_name = 'babyweight'
version_name = 'dnn'

input_data = {
  'instances': [
    {
      'key': 'b1',
      'is_male': 'True',
      'mother_age': 26.0,
      'plurality': 'Single(1)',
      'gestation_weeks': 39
    },
    {
      'key': 'g1',
      'is_male': 'False',
      'mother_age': 29.0,
      'plurality': 'Single(1)',
      'gestation_weeks': 38
    },
    {
      'key': 'b2',
      'is_male': 'True',
      'mother_age': 26.0,
      'plurality': 'Triplets(3)',
      'gestation_weeks': 39
    },
    {
      'key': 'u1',
      'is_male': 'Unknown',
      'mother_age': 29.0,
      'plurality': 'Multiple(2+)',
      'gestation_weeks': 38
    },
  ]
}

parent = 'projects/%s/models/%s/versions/%s' % (project, model_name, version_name)
prediction = api.projects().predict(body=input_data, name=parent).execute()
print(prediction)
print(prediction['predictions'][0]['babyweight'][0])

"""
Explanation: main.py
This is the code that exists in serving/application/main.py, i.e. the code in the web application that accesses the ML API.
End of explanation
"""


