# change these to try this notebook out
# In "production", these will be replaced by the parameters passed to papermill
BUCKET = 'cloud-training-demos-ml'
PROJECT = 'cloud-training-demos'
REGION = 'us-central1'
DEVELOP_MODE = True
EAGER_MODE = False
NBUCKETS = 5

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

%%bash
gcloud config set project $PROJECT
gcloud config set compute/region $REGION

"""
Explanation: Developing, Training, and Deploying a TensorFlow model on Google Cloud Platform (completely within Jupyter)
In Chapter 9 of Data Science on the Google Cloud Platform, I trained a TensorFlow Estimator model to predict flight delays.
In this notebook, we'll modernize the workflow:
* Use eager mode for TensorFlow development
* Use tf.data to write the input pipeline
* Run the notebook as-is on Cloud using Deep Learning VM or Kubeflow pipelines
* Deploy the trained model to Cloud ML Engine as a web service
The combination of eager mode, tf.data and DLVM/KFP makes this workflow a lot easier.
We don't need to deal with Python packages or Docker containers.
End of explanation
"""


DATA_BUCKET = "gs://cloud-training-demos/flights/chapter8/output/"
TRAIN_DATA_PATTERN = DATA_BUCKET + "train*"
VALID_DATA_PATTERN = DATA_BUCKET + "test*"

!gsutil ls $DATA_BUCKET

"""
Explanation: Creating the input data pipeline
End of explanation
"""


import os, json, math
import numpy as np
import tensorflow as tf
print("Tensorflow version " + tf.__version__)

if DEVELOP_MODE and EAGER_MODE:
    print("Enabling Eager mode for development only")
    tf.enable_eager_execution()
else:
    EAGER_MODE = False

CSV_COLUMNS  = ('ontime,dep_delay,taxiout,distance,avg_dep_delay,avg_arr_delay' + \
                ',carrier,dep_lat,dep_lon,arr_lat,arr_lon,origin,dest').split(',')
LABEL_COLUMN = 'ontime'
DEFAULTS     = [[0.0],[0.0],[0.0],[0.0],[0.0],[0.0],\
                ['na'],[0.0],[0.0],[0.0],[0.0],['na'],['na']]

def decode_csv(line):
  column_values = tf.decode_csv(line, DEFAULTS)
  column_names = CSV_COLUMNS
  decoded_line = dict(zip(column_names, column_values)) # create a dictionary {'column_name': value, ...} for each line  
  return decoded_line

def load_dataset(pattern):
  filenames = tf.data.Dataset.list_files(pattern)
  dataset = filenames.interleave(tf.data.TextLineDataset, cycle_length=16)
  dataset = dataset.map(decode_csv)
  return dataset

if EAGER_MODE:
    dataset = load_dataset(TRAIN_DATA_PATTERN)
    for n, data in enumerate(dataset):
      numpy_data = {k: v.numpy() for k, v in data.items()} # .numpy() works only in eager mode
      print(numpy_data)
      if n>3: break

%%writefile example_input.json
{"dep_delay": 14.0, "taxiout": 13.0, "distance": 319.0, "avg_dep_delay": 25.863039, "avg_arr_delay": 27.0, "carrier": "WN", "dep_lat": 32.84722, "dep_lon": -96.85167, "arr_lat": 31.9425, "arr_lon": -102.20194, "origin": "DAL", "dest": "MAF"}
{"dep_delay": -9.0, "taxiout": 21.0, "distance": 301.0, "avg_dep_delay": 41.050808, "avg_arr_delay": -7.0, "carrier": "EV", "dep_lat": 29.984444, "dep_lon": -95.34139, "arr_lat": 27.544167, "arr_lon": -99.46167, "origin": "IAH", "dest": "LRD"}

def features_and_labels(features):
  label = features.pop('ontime') # this is what we will train for
  return features, label
  
def prepare_dataset(pattern, batch_size, truncate=None, mode=tf.estimator.ModeKeys.TRAIN):
  dataset = load_dataset(pattern)
  dataset = dataset.map(features_and_labels)
  dataset = dataset.cache()
  if mode == tf.estimator.ModeKeys.TRAIN:
    dataset = dataset.shuffle(1000)
    dataset = dataset.repeat()
  dataset = dataset.batch(batch_size)
  dataset = dataset.prefetch(10)
  if truncate is not None:
    dataset = dataset.take(truncate)
  return dataset

if EAGER_MODE:
    print("Calling prepare")
    one_item = prepare_dataset(TRAIN_DATA_PATTERN, batch_size=2, truncate=1)
    print(list(one_item)) # should print one batch of 2 items

"""
Explanation: Use tf.data to read the CSV files
Note the use of Eager mode to develop the code.  I turn off eager mode once I get to the next section
End of explanation
"""


import tensorflow.feature_column as fc

real = {
    colname : fc.numeric_column(colname) \
          for colname in \
            ('dep_delay,taxiout,distance,avg_dep_delay,avg_arr_delay' + 
             ',dep_lat,dep_lon,arr_lat,arr_lon').split(',')
}
sparse = {
      'carrier': fc.categorical_column_with_vocabulary_list('carrier',
                  vocabulary_list='AS,VX,F9,UA,US,WN,HA,EV,MQ,DL,OO,B6,NK,AA'.split(',')),
      'origin' : fc.categorical_column_with_hash_bucket('origin', hash_bucket_size=1000),
      'dest'   : fc.categorical_column_with_hash_bucket('dest', hash_bucket_size=1000)
}

"""
Explanation: Create TensorFlow wide-and-deep model
We'll create feature columns, and do some discretization and feature engineering.
See the book for details.
End of explanation
"""


latbuckets = np.linspace(20.0, 50.0, NBUCKETS).tolist()  # USA
lonbuckets = np.linspace(-120.0, -70.0, NBUCKETS).tolist() # USA
disc = {}
disc.update({
       'd_{}'.format(key) : fc.bucketized_column(real[key], latbuckets) \
          for key in ['dep_lat', 'arr_lat']
})
disc.update({
       'd_{}'.format(key) : fc.bucketized_column(real[key], lonbuckets) \
          for key in ['dep_lon', 'arr_lon']
})

# cross columns that make sense in combination
sparse['dep_loc'] = fc.crossed_column([disc['d_dep_lat'], disc['d_dep_lon']], NBUCKETS*NBUCKETS)
sparse['arr_loc'] = fc.crossed_column([disc['d_arr_lat'], disc['d_arr_lon']], NBUCKETS*NBUCKETS)
sparse['dep_arr'] = fc.crossed_column([sparse['dep_loc'], sparse['arr_loc']], NBUCKETS ** 4)
sparse['ori_dest'] = fc.crossed_column(['origin', 'dest'], hash_bucket_size=1000)

# embed all the sparse columns
embed = {
       colname : fc.embedding_column(col, 10) \
          for colname, col in sparse.items()
}
real.update(embed)

if DEVELOP_MODE:
    print(sparse.keys())
    print(real.keys())

"""
Explanation: Feature engineering
End of explanation
"""


def serving_input_fn():
    feature_placeholders = {
        # All the real-valued columns
        column: tf.placeholder(tf.float32, [None]) \
             for column in ('dep_delay,taxiout,distance,avg_dep_delay,avg_arr_delay' + 
                            ',dep_lat,dep_lon,arr_lat,arr_lon').split(',')
    }
    feature_placeholders.update({
        column: tf.placeholder(tf.string, [None]) for column in ['carrier', 'origin', 'dest']
    })
    features = feature_placeholders # no transformations
    return tf.estimator.export.ServingInputReceiver(features, feature_placeholders)

"""
Explanation: Serving
This serving input function is how the model will be deployed for prediction. We require these fields for prediction
End of explanation
"""


model_dir='gs://{}/flights/trained_model'.format(BUCKET)
os.environ['OUTDIR'] = model_dir  # needed for deployment
print('Writing trained model to {}'.format(model_dir))

!gsutil -m rm -rf $OUTDIR

estimator = tf.estimator.DNNLinearCombinedClassifier(
        model_dir = model_dir,
        linear_feature_columns = sparse.values(),
        dnn_feature_columns = real.values(),
        dnn_hidden_units = [64, 32]
)

train_batch_size = 64
train_input_fn = lambda: prepare_dataset(TRAIN_DATA_PATTERN, train_batch_size).make_one_shot_iterator().get_next()
eval_batch_size = 100 if DEVELOP_MODE else 10000
eval_input_fn = lambda: prepare_dataset(VALID_DATA_PATTERN, eval_batch_size, eval_batch_size*10, tf.estimator.ModeKeys.EVAL).make_one_shot_iterator().get_next()
num_steps = 10 if DEVELOP_MODE else (1000000 // train_batch_size)

train_spec = tf.estimator.TrainSpec(train_input_fn, max_steps = num_steps)
exporter = tf.estimator.LatestExporter('exporter', serving_input_fn)
eval_spec = tf.estimator.EvalSpec(eval_input_fn, steps=10, exporters=exporter)
tf.estimator.train_and_evaluate(estimator, train_spec, eval_spec)

"""
Explanation: Train the model and evaluate once in a while
Also checkpoint
End of explanation
"""


%%bash
model_dir=$(gsutil ls ${OUTDIR}/export/exporter | tail -1)
echo $model_dir
saved_model_cli show --dir ${model_dir} --all

%%bash

MODEL_NAME="flights"
MODEL_VERSION="kfp"
TFVERSION="1.10"
MODEL_LOCATION=$(gsutil ls ${OUTDIR}/export/exporter | tail -1)
echo "Run these commands one-by-one (the very first time, you'll create a model and then create a version)"
yes | gcloud ml-engine versions delete ${MODEL_VERSION} --model ${MODEL_NAME}
#gcloud ml-engine models delete ${MODEL_NAME}
#gcloud ml-engine models create ${MODEL_NAME} --regions $REGION
gcloud ml-engine versions create ${MODEL_VERSION} --model ${MODEL_NAME} --origin ${MODEL_LOCATION} --runtime-version $TFVERSION

!gcloud ml-engine predict --model=flights --version=kfp --json-instances=example_input.json

"""
Explanation: Deploy the trained model
End of explanation
"""


