PROJECT = "cloud-training-demos"  # Replace with your PROJECT
BUCKET = "cloud-training-bucket"  # Replace with your BUCKET
REGION = "us-central1"            # Choose an available region for Cloud MLE
TFVERSION = "1.14"                # TF version for CMLE to use

import os
os.environ["BUCKET"] = BUCKET
os.environ["PROJECT"] = PROJECT
os.environ["REGION"] = REGION
os.environ["TFVERSION"] = TFVERSION

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

"""
Explanation: Create TensorFlow Wide and Deep Model
Learning Objective
- Create a Wide and Deep model using the high-level Estimator API 
- Determine which features to use as wide columns and which to use as deep columns
Introduction
In this notebook, we'll explore modeling our data using a Wide & Deep Neural Network. As before, we can do this uisng the high-level Estimator API in Tensorflow. Have a look at the various other models available through the Estimator API in the documentation here. In particular, have a look at the implementation for Wide & Deep models. 
Start by setting the environment variables related to your project.
End of explanation
"""


%%bash
ls *.csv

"""
Explanation: Let's have a look at the csv files we created in the previous notebooks that we will use for training/eval.
End of explanation
"""


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

CSV_COLUMNS = "weight_pounds,is_male,mother_age,plurality,gestation_weeks".split(',')
LABEL_COLUMN = "weight_pounds"

# Set default values for each CSV column
DEFAULTS = [[0.0], ["null"], [0.0], ["null"], [0.0]]
TRAIN_STEPS = 1000

"""
Explanation: Create TensorFlow model using TensorFlow's Estimator API
First, we'll write an input_fn to read the data and define the csv column names and label column. We'll also set the default csv column values and set the number of training steps.
End of explanation
"""


def read_dataset(filename_pattern, mode, batch_size = 512):
    def _input_fn():
        def decode_csv(value_column):
            columns = tf.decode_csv(records = value_column, record_defaults = DEFAULTS)
            features = dict(zip(CSV_COLUMNS, columns))
            label = features.pop(LABEL_COLUMN)
            return features, label
    
        # Create list of files that match pattern
        file_list = tf.gfile.Glob(filename = filename_pattern)

        # Create dataset from file list
        dataset = (tf.data.TextLineDataset(filenames = file_list)  # Read text file
                     .map(map_func = decode_csv))  # Transform each elem by applying decode_csv fn

        if mode == tf.estimator.ModeKeys.TRAIN:
            num_epochs = None # indefinitely
            dataset = dataset.shuffle(buffer_size = 10 * batch_size)
        else:
            num_epochs = 1 # end-of-input after this

        dataset = dataset.repeat(count = num_epochs).batch(batch_size = batch_size)
        return dataset
    return _input_fn

"""
Explanation: Create the input function
Now we are ready to create an input function using the Dataset API.
End of explanation
"""


def get_wide_deep():
    # Define column types
    fc_is_male,fc_plurality,fc_mother_age,fc_gestation_weeks = [\
        tf.feature_column.categorical_column_with_vocabulary_list(key = "is_male", 
                                                                  vocabulary_list = ["True", "False", "Unknown"]),
        tf.feature_column.categorical_column_with_vocabulary_list(key = "plurality", 
                                                                  vocabulary_list = ["Single(1)", "Twins(2)", "Triplets(3)", "Quadruplets(4)", "Quintuplets(5)", "Multiple(2+)"]),
        tf.feature_column.numeric_column(key = "mother_age"),
        tf.feature_column.numeric_column(key = "gestation_weeks")
    ]

    # Bucketized columns
    fc_age_buckets = tf.feature_column.bucketized_column(source_column = fc_mother_age, boundaries = np.arange(start = 15, stop = 45, step = 1).tolist())
    fc_gestation_buckets = tf.feature_column.bucketized_column(source_column = fc_gestation_weeks, boundaries = np.arange(start = 17, stop = 47, step = 1).tolist())

    # Sparse columns are wide, have a linear relationship with the output
    wide = [fc_is_male,
        fc_plurality,
        fc_age_buckets,
        fc_gestation_buckets]

    # Feature cross all the wide columns and embed into a lower dimension
    crossed = tf.feature_column.crossed_column(keys = wide, hash_bucket_size = 20000)
    fc_embed = tf.feature_column.embedding_column(categorical_column = crossed, dimension = 3)

    # Continuous columns are deep, have a complex relationship with the output
    deep = [fc_mother_age,
        fc_gestation_weeks,
        fc_embed]
    
    return wide, deep

"""
Explanation: Create the feature columns
Next, define the feature columns. For a wide and deep model, we need to determine which features we will use as wide features and which to pass as deep features. The function get_wide_deep below will return a tuple containing the wide feature columns and deep feature columns. Have a look at this blog post on wide and deep models to remind yourself how best to describe the features.
End of explanation
"""


def serving_input_fn():
    feature_placeholders = {
        "is_male": tf.placeholder(dtype = tf.string, shape = [None]),
        "mother_age": tf.placeholder(dtype = tf.float32, shape = [None]),
        "plurality": tf.placeholder(dtype = tf.string, shape = [None]),
        "gestation_weeks": tf.placeholder(dtype = tf.float32, shape = [None])
    }
    
    features = {
        key: tf.expand_dims(input = tensor, axis = -1)
        for key, tensor in feature_placeholders.items()
    }
    
    return tf.estimator.export.ServingInputReceiver(features = features, receiver_tensors = feature_placeholders)

"""
Explanation: Create the Serving Input function
To predict with the TensorFlow model, we also need a serving input function. This will allow us to serve prediction later using the predetermined inputs. We will want all the inputs from our user.
End of explanation
"""


def train_and_evaluate(output_dir):
    wide, deep = get_wide_deep()
    
    EVAL_INTERVAL = 300
    
    run_config = tf.estimator.RunConfig(
        save_checkpoints_secs = EVAL_INTERVAL,
        keep_checkpoint_max = 3)
    
    estimator = tf.estimator.DNNLinearCombinedRegressor(
        model_dir = output_dir,
        linear_feature_columns = wide,
        dnn_feature_columns = deep,
        dnn_hidden_units = [64, 32],
        config = run_config)
    
    train_spec = tf.estimator.TrainSpec(
        input_fn = read_dataset("train.csv", mode = tf.estimator.ModeKeys.TRAIN),
        max_steps = TRAIN_STEPS)
    
    exporter = tf.estimator.LatestExporter(name = "exporter", serving_input_receiver_fn = serving_input_fn)
    
    eval_spec = tf.estimator.EvalSpec(
        input_fn = read_dataset("eval.csv", mode = tf.estimator.ModeKeys.EVAL),
        steps = None,
        start_delay_secs = 60, # start evaluating after N seconds
        throttle_secs = EVAL_INTERVAL,  # evaluate every N seconds
        exporters = exporter)
    
    tf.estimator.train_and_evaluate(estimator, train_spec, eval_spec)

"""
Explanation: Create the model and run training and evaluation
Lastly, we'll create the estimator to train and evaluate. In the cell below, we'll set up a Wide & Deep model (i.e. a DNNLinearCombinedRegressor estimator) and the train and evaluation operations.
End of explanation
"""


# Run the model
shutil.rmtree(path = "babyweight_trained_wd", ignore_errors = True) # start fresh each time
train_and_evaluate("babyweight_trained_wd")

"""
Explanation: Finally, we train the model!
End of explanation
"""


