# Uncomment and run if you need to update your Google SDK
# !sudo apt-get update && sudo apt-get --only-upgrade install google-cloud-sdk

"""
Explanation: Introducing AI Platform Training Service
Learning Objectives:
  - Learn how to make code compatible with AI Platform Training Service
  - Train your model using cloud infrastructure via AI Platform Training Service
  - Deploy your model behind a production grade REST API using AI Platform Training Service
Introduction
In this notebook we'll make the jump from training and predicting locally, to do doing both in the cloud. We'll take advantage of Google Cloud's AI Platform Training Service. 
AI Platform Training Service is a managed service that allows the training and deployment of ML models without having to provision or maintain servers. The infrastructure is handled seamlessly by the managed service for us.
End of explanation
"""


PROJECT = "cloud-training-demos"  # Replace with your PROJECT
BUCKET = "cloud-training-bucket"  # Replace with your BUCKET
REGION = "us-central1"            # Choose an available region for AI Platform Training Service
TFVERSION = "1.14"                # TF version for AI Platform Training Service to use

"""
Explanation: Make code compatible with AI Platform Training Service
In order to make our code compatible with AI Platform Training Service we need to make the following changes:

Upload data to Google Cloud Storage 
Move code into a Python package
Modify code to read data from and write checkpoint files to GCS 

Upload data to Google Cloud Storage (GCS)
Cloud services don't have access to our local files, so we need to upload them to a location the Cloud servers can read from. In this case we'll use GCS.
Specify your project name and bucket name in the cell below.
End of explanation
"""


!gcloud config set project {PROJECT}
!gsutil mb -l {REGION} gs://{BUCKET}
!gsutil -m cp *.csv gs://{BUCKET}/taxifare/smallinput/

"""
Explanation: Jupyter allows the subsitution of python variables into bash commands when using the !&lt;cmd&gt; format.
It is also possible using the %%bash magic but requires an additional parameter.
End of explanation
"""


%%bash
mkdir taxifaremodel
touch taxifaremodel/__init__.py

"""
Explanation: Move code into a python package
When you execute a AI Platform Training Service training job, the service zips up your code and ships it to the Cloud so it can be run on Cloud infrastructure. In order to do this AI Platform Training Service requires your code to be a Python package.
A Python package is simply a collection of one or more .py files along with an __init__.py file to identify the containing directory as a package. The __init__.py sometimes contains initialization code but for our purposes an empty file suffices.
Create Package Directory and __init__.py
The bash command touch creates an empty file in the specified location.
End of explanation
"""


%%writefile taxifaremodel/model.py
# TODO: Your code goes here. Import the necessary libraries (e.g. tensorflow, etc)

CSV_COLUMN_NAMES = # TODO: Your code goes here
CSV_DEFAULTS = # TODO: Your code goes here
FEATURE_NAMES = # TODO: Your code goes here

def parse_row(row):
    # TODO: Your code goes here
    return features, label

def read_dataset(csv_path):
    # TODO: Your code goes here
    return dataset

def train_input_fn(csv_path, batch_size = 128):
    # TODO: Your code goes here
    return dataset

def eval_input_fn(csv_path, batch_size = 128):
    # TODO: Your code goes here
    return dataset
  
def serving_input_receiver_fn():
    # TODO: Your code goes here
    return tf.estimator.export.ServingInputReceiver(features = features, receiver_tensors = receiver_tensors)  
    
def my_rmse(labels, predictions):
    # TODO: Your code goes here
    return {"rmse": tf.metrics.root_mean_squared_error(labels = labels, predictions = pred_values)}

def create_model(model_dir, train_steps):
    # TODO: Your code goes here
    return model

def train_and_evaluate(params):
    OUTDIR = params["output_dir"]
    TRAIN_DATA_PATH = params["train_data_path"]
    EVAL_DATA_PATH = params["eval_data_path"]
    TRAIN_STEPS = params["train_steps"]

    model = # TODO: Your code goes here. 

    train_spec = # TODO: Your code goes here
    
    exporter = # TODO: Your code goes here

    eval_spec = # TODO: Your code goes here

    tf.logging.set_verbosity(tf.logging.INFO) 
    shutil.rmtree(path = OUTDIR, ignore_errors = True)

    tf.estimator.train_and_evaluate(estimator = model, train_spec = train_spec, eval_spec = eval_spec)

"""
Explanation: Paste existing code into model.py
A Python package requires our code to be in a .py file, as opposed to notebook cells. So we simply copy and paste our existing code for the previous notebook into a single file.
The %%writefile magic writes the contents of its cell to disk with the specified name.
Exercise 1
In the cell below, write the content of the model.py to the file taxifaremodel/model.py. This will allow us to package the model we 
developed in the previous labs so that we can deploy it to AI Platform Training Service. You'll also need to reuse the input functions and the EvalSpec, TrainSpec, RunConfig, etc. that we implemented in the previous labs.
Complete all the TODOs in the cell below by copy/pasting the code we developed in the previous labs. This will write all the necessary components we developed in our notebook to a single model.py file. 
Once we have the code running well locally, we will execute the next cells to train and deploy your packaged model to AI Platform Training Service.
End of explanation
"""


%%writefile taxifaremodel/task.py
import argparse
import json
import os

from . import model

if __name__ == "__main__":
    
    parser = argparse.ArgumentParser()
    
    parser.add_argument(
        "--train_data_path",
        help = "GCS or local path to training data",
        required = True
    )
    parser.add_argument(
        "--train_steps",
        help = "Steps to run the training job for (default: 1000)",
        type = int,
        default = 1000
    )
    parser.add_argument(
        "--eval_data_path",
        help = "GCS or local path to evaluation data",
        required = True
    )
    parser.add_argument(
        # TODO: Your code goes here
    )
    parser.add_argument(
        # TODO: Your code goes here
    )
    args = parser.parse_args().__dict__

    model.train_and_evaluate(args)

"""
Explanation: Modify code to read data from and write checkpoint files to GCS
If you look closely above, you'll notice two changes to the code

The input function now supports reading a list of files matching a file name pattern instead of just a single CSV
This is useful because large datasets tend to exist in shards.
The train and evaluate portion is wrapped in a function that takes a parameter dictionary as an argument.
This is useful because the output directory, data paths and number of train steps will be different depending on whether we're training locally or in the cloud. Parametrizing allows us to use the same code for both.

We specify these parameters at run time via the command line. Which means we need to add code to parse command line parameters and invoke train_and_evaluate() with those params. This is the job of the task.py file. 
Exposing parameters to the command line also allows us to use AI Platform Training Service's automatic hyperparameter tuning feature which we'll cover in a future lesson.
Exercise 2
Add two additional command line parameter parsers to the list we've started below. You should add code to parse command line parameters for the output_dir and the job-dir. Look at the examples below to make sure you have the correct format, including a help description and required specification.
End of explanation
"""


%%time
!gcloud ai-platform local train \
    --package-path=taxifaremodel \
    --module-name=taxifaremodel.task \
    -- \
    --train_data_path=taxi-train.csv \
    --eval_data_path=taxi-valid.csv  \
    --train_steps=1 \
    --output_dir=taxi_trained 

"""
Explanation: Train using AI Platform Training Service (local)
AI Platform Training Service comes with a local test tool (gcloud ai-platform local train) to ensure we've packaged our code directly. It's best to first run that for a few steps before trying a Cloud job. 
The arguments before -- \ are for AI Platform Training Service
- package-path: speficies the location of the Python package
- module-name: specifies which .py file should be run within the package. task.py is our entry point so we specify that
The arguments after -- \ are sent to our task.py.
End of explanation
"""


OUTDIR = "gs://{}/taxifare/trained_small".format(BUCKET)

!gsutil -m rm -rf {OUTDIR} # start fresh each time
!gcloud ai-platform jobs submit training taxifare_$(date -u +%y%m%d_%H%M%S) \
    --package-path=taxifaremodel \
    --module-name=taxifaremodel.task \
    --job-dir=gs://{BUCKET}/taxifare \
    --python-version=3.5 \
    --runtime-version={TFVERSION} \
    --region={REGION} \
    -- \
    --train_data_path=gs://{BUCKET}/taxifare/smallinput/taxi-train.csv \
    --eval_data_path=gs://{BUCKET}/taxifare/smallinput/taxi-valid.csv  \
    --train_steps=1000 \
    --output_dir={OUTDIR}

"""
Explanation: Train using AI Platform Training Service (Cloud)
To submit to the Cloud we use gcloud ai-platform jobs submit training [jobname] and simply specify some additional parameters for AI Platform Training Service:
- jobname: A unique identifier for the Cloud job. We usually append system time to ensure uniqueness
- job-dir: A GCS location to upload the Python package to
- runtime-version: Version of TF to use. Defaults to 1.0 if not specified
- python-version: Version of Python to use. Defaults to 2.7 if not specified
- region: Cloud region to train in. See here for supported AI Platform Training Service regions
Below the -- \ note how we've changed our task.py args to be GCS locations
End of explanation
"""


!gsutil ls gs://{BUCKET}/taxifare/trained_small/export/exporter

"""
Explanation: You can track your job and view logs using cloud console. It will take 5-10 minutes to complete. Wait until the job finishes before moving on.
Deploy model
Now let's take our exported SavedModel and deploy it behind a REST API. To do so we'll use AI Platform Training Service's managed TF Serving feature which auto-scales based on load.
End of explanation
"""


VERSION='v1'
!gcloud ai-platform models create taxifare --regions us-central1
!gcloud ai-platform versions delete {VERSION} --model taxifare --quiet
!gcloud ai-platform versions create {VERSION} --model taxifare \
    --origin $(gsutil ls gs://{BUCKET}/taxifare/trained_small/export/exporter | tail -1) \
    --python-version=3.5 \
    --runtime-version {TFVERSION}

"""
Explanation: AI Platform Training Service uses a model versioning system. First you create a model folder, and within the folder you create versions of the model. 
Note: You will see an error below if the model folder already exists, it is safe to ignore
End of explanation
"""


%%writefile ./test.json
{"dayofweek": 1, "hourofday": 0, "pickuplon": -73.885262, "pickuplat": 40.773008, "dropofflon": -73.987232, "dropofflat": 40.732403}

"""
Explanation: Online prediction
Now that we have deployed our model behind a production grade REST API, we can invoke it remotely. 
We could invoke it directly calling the REST API with an HTTP POST request reference docs, however AI Platform Training Service provides an easy way to invoke it via command line.
Invoke prediction REST API via command line
First we write our prediction requests to file in json format
End of explanation
"""


!gcloud ai-platform predict --model=taxifare --json-instances=./test.json

"""
Explanation: Then we use gcloud ai-platform predict and specify the model name and location of the json file. Since we don't explicitly specify --version, the default model version will be used. 
Since we only have one version it is already the default, but if we had multiple model versions we can designate the default using gcloud ai-platform versions set-default or using cloud console
End of explanation
"""


from googleapiclient import discovery
from oauth2client.client import GoogleCredentials
import json

credentials = # TODO: Your code goes here
api = # TODO: Your code goes here

request_data = {"instances":
  [
      {
        # TODO: Your code goes here
      }
  ]
}

parent = # TODO: Your code goes here

response = # TODO: Your code goes here

print("response = {0}".format(response))

"""
Explanation: Invoke prediction REST API via python
Exercise 3
In the cell below, use the Google Python client library to query the model you just deployed on AI Platform Training Service. Find the estimated taxi fare for a ride with the following properties
- ride occurs on Monday
- at 8:00 am
- pick up at (40.773, -73.885)
- drop off at (40.732, -73.987)
Have a look at this post and examples on "Using the Python Client Library" and "Getting Online Predictions" from Google Cloud.
End of explanation
"""


