text: # Copyright 2017-2021 Google Inc. 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.
  
  import argparse
  import logging
  import os, time
  import hypertune
  import numpy as np
  import tensorflow as tf
  
  BUCKET = None
  TF_VERSION = '2-' + tf.__version__[2:3]  # needed to choose container
  
  DEVELOP_MODE = True
  NUM_EXAMPLES = 5000 * 1000  # doesn't need to be precise but get order of magnitude right.
  
  NUM_BUCKETS = 5
  NUM_EMBEDS = 3
  TRAIN_BATCH_SIZE = 64
  DNN_HIDDEN_UNITS = '64,32'
  
  CSV_COLUMNS = (
      'ontime,dep_delay,taxi_out,distance,origin,dest,dep_hour,is_weekday,carrier,' +
      'dep_airport_lat,dep_airport_lon,arr_airport_lat,arr_airport_lon,data_split'
  ).split(',')
  
  CSV_COLUMN_TYPES = [
      1.0, -3.0, 5.0, 1037.493622678299, 'OTH', 'DEN', 21, 1.0, 'OO',
      43.41694444, -124.24694444, 39.86166667, -104.67305556, 'TRAIN'
  ]
  
  
  def features_and_labels(features):
      label = features.pop('ontime')  # this is what we will train for
      return features, label
  
  
  def read_dataset(pattern, batch_size, mode=tf.estimator.ModeKeys.TRAIN, truncate=None):
      dataset = tf.data.experimental.make_csv_dataset(
          pattern, batch_size,
          column_names=CSV_COLUMNS,
          column_defaults=CSV_COLUMN_TYPES,
          sloppy=True,
          num_parallel_reads=2,
          ignore_errors=True,
          num_epochs=1)
      dataset = dataset.map(features_and_labels)
      if mode == tf.estimator.ModeKeys.TRAIN:
          dataset = dataset.shuffle(batch_size * 10)
          dataset = dataset.repeat()
      dataset = dataset.prefetch(1)
      if truncate is not None:
          dataset = dataset.take(truncate)
      return dataset
  
  
  def create_model():
      real = {
          colname: tf.feature_column.numeric_column(colname)
          for colname in
          (
                  'dep_delay,taxi_out,distance,dep_hour,is_weekday,' +
                  'dep_airport_lat,dep_airport_lon,' +
                  'arr_airport_lat,arr_airport_lon'
          ).split(',')
      }
      sparse = {
          'carrier': tf.feature_column.categorical_column_with_vocabulary_list('carrier',
                                                                               vocabulary_list='AS,VX,F9,UA,US,WN,HA,EV,MQ,DL,OO,B6,NK,AA'.split(
                                                                                   ',')),
          'origin': tf.feature_column.categorical_column_with_hash_bucket('origin', hash_bucket_size=1000),
          'dest': tf.feature_column.categorical_column_with_hash_bucket('dest', hash_bucket_size=1000),
      }
  
      inputs = {
          colname: tf.keras.layers.Input(name=colname, shape=(), dtype='float32')
          for colname in real.keys()
      }
      inputs.update({
          colname: tf.keras.layers.Input(name=colname, shape=(), dtype='string')
          for colname in sparse.keys()
      })
  
      latbuckets = np.linspace(20.0, 50.0, NUM_BUCKETS).tolist()  # USA
      lonbuckets = np.linspace(-120.0, -70.0, NUM_BUCKETS).tolist()  # USA
      disc = {}
      disc.update({
          'd_{}'.format(key): tf.feature_column.bucketized_column(real[key], latbuckets)
          for key in ['dep_airport_lat', 'arr_airport_lat']
      })
      disc.update({
          'd_{}'.format(key): tf.feature_column.bucketized_column(real[key], lonbuckets)
          for key in ['dep_airport_lon', 'arr_airport_lon']
      })
  
      # cross columns that make sense in combination
      sparse['dep_loc'] = tf.feature_column.crossed_column(
          [disc['d_dep_airport_lat'], disc['d_dep_airport_lon']], NUM_BUCKETS * NUM_BUCKETS)
      sparse['arr_loc'] = tf.feature_column.crossed_column(
          [disc['d_arr_airport_lat'], disc['d_arr_airport_lon']], NUM_BUCKETS * NUM_BUCKETS)
      sparse['dep_arr'] = tf.feature_column.crossed_column([sparse['dep_loc'], sparse['arr_loc']], NUM_BUCKETS ** 4)
  
      # embed all the sparse columns
      embed = {
          'embed_{}'.format(colname): tf.feature_column.embedding_column(col, NUM_EMBEDS)
          for colname, col in sparse.items()
      }
      real.update(embed)
  
      # one-hot encode the sparse columns
      sparse = {
          colname: tf.feature_column.indicator_column(col)
          for colname, col in sparse.items()
      }
  
      model = wide_and_deep_classifier(
          inputs,
          linear_feature_columns=sparse.values(),
          dnn_feature_columns=real.values(),
          dnn_hidden_units=DNN_HIDDEN_UNITS)
  
      return model
  
  
  def wide_and_deep_classifier(inputs, linear_feature_columns, dnn_feature_columns, dnn_hidden_units):
      deep = tf.keras.layers.DenseFeatures(dnn_feature_columns, name='deep_inputs')(inputs)
      layers = [int(x) for x in dnn_hidden_units.split(',')]
      for layerno, numnodes in enumerate(layers):
          deep = tf.keras.layers.Dense(numnodes, activation='relu', name='dnn_{}'.format(layerno + 1))(deep)
      wide = tf.keras.layers.DenseFeatures(linear_feature_columns, name='wide_inputs')(inputs)
      both = tf.keras.layers.concatenate([deep, wide], name='both')
      output = tf.keras.layers.Dense(1, activation='sigmoid', name='pred')(both)
      model = tf.keras.Model(inputs, output)
      model.compile(optimizer='adam',
                    loss='binary_crossentropy',
                    metrics=['accuracy', rmse, tf.keras.metrics.AUC()])
      return model
  
  
  def rmse(y_true, y_pred):
      return tf.sqrt(tf.reduce_mean(tf.square(y_pred - y_true)))
  
  
  def train_and_evaluate(train_data_pattern, eval_data_pattern, test_data_pattern, export_dir, output_dir):
      train_batch_size = TRAIN_BATCH_SIZE
      if DEVELOP_MODE:
          eval_batch_size = 100
          steps_per_epoch = 3
          epochs = 2
          num_eval_examples = eval_batch_size * 10
      else:
          eval_batch_size = 100
          steps_per_epoch = NUM_EXAMPLES // train_batch_size
          epochs = NUM_EPOCHS
          num_eval_examples = eval_batch_size * 100
  
      train_dataset = read_dataset(train_data_pattern, train_batch_size)
      eval_dataset = read_dataset(eval_data_pattern, eval_batch_size, tf.estimator.ModeKeys.EVAL, num_eval_examples)
  
      # checkpoint
      checkpoint_path = '{}/checkpoints/flights.cpt'.format(output_dir)
      logging.info("Checkpointing to {}".format(checkpoint_path))
      cp_callback = tf.keras.callbacks.ModelCheckpoint(checkpoint_path,
                                                       save_weights_only=True,
                                                       verbose=1)
  
      # call back to write out hyperparameter tuning metric
      METRIC = 'val_rmse'
      hpt = hypertune.HyperTune()
  
      class HpCallback(tf.keras.callbacks.Callback):
          def on_epoch_end(self, epoch, logs=None):
              if logs and METRIC in logs:
                  logging.info("Epoch {}: {} = {}".format(epoch, METRIC, logs[METRIC]))
                  hpt.report_hyperparameter_tuning_metric(hyperparameter_metric_tag=METRIC,
                                                          metric_value=logs[METRIC],
                                                          global_step=epoch)
  
      # train the model
      model = create_model()
      logging.info(f"Training on {train_data_pattern}; eval on {eval_data_pattern}; {epochs} epochs; {steps_per_epoch}")
      history = model.fit(train_dataset,
                          validation_data=eval_dataset,
                          epochs=epochs,
                          steps_per_epoch=steps_per_epoch,
                          callbacks=[cp_callback, HpCallback()])
  
      # export
      logging.info('Exporting to {}'.format(export_dir))
      tf.saved_model.save(model, export_dir)
  
      # write out final metric
      final_rmse = history.history[METRIC][-1]
      logging.info("Validation metric {} on {} samples = {}".format(METRIC, num_eval_examples, final_rmse))
  
      if (not DEVELOP_MODE) and (test_data_pattern is not None) and (not SKIP_FULL_EVAL):
          logging.info("Evaluating over full test dataset")
          test_dataset = read_dataset(test_data_pattern, eval_batch_size, tf.estimator.ModeKeys.EVAL, None)
          final_metrics = model.evaluate(test_dataset)
          logging.info("Final metrics on full test dataset = {}".format(final_metrics))
      else:
          logging.info("Skipping evaluation on full test dataset")
  
  
  if __name__ == '__main__':
      logging.info("Tensorflow version " + tf.__version__)
      parser = argparse.ArgumentParser()
  
      parser.add_argument(
          '--bucket',
          help='Data will be read from gs://BUCKET/ch9/data and output will be in gs://BUCKET/ch9/trained_model',
          required=True
      )
  
      parser.add_argument(
          '--num_examples',
          help='Number of examples per epoch. Get order of magnitude correct.',
          type=int,
          default=5000000
      )
  
      # for hyper-parameter tuning
      parser.add_argument(
          '--train_batch_size',
          help='Number of examples to compute gradient on',
          type=int,
          default=256  # originally 64
      )
      parser.add_argument(
          '--nbuckets',
          help='Number of bins into which to discretize lats and lons',
          type=int,
          default=10  # originally 5
      )
      parser.add_argument(
          '--nembeds',
          help='Embedding dimension for categorical variables',
          type=int,
          default=3
      )
      parser.add_argument(
          '--num_epochs',
          help='Number of epochs (used only if --develop is not set)',
          type=int,
          default=10
      )
      parser.add_argument(
          '--dnn_hidden_units',
          help='Architecture of DNN part of wide-and-deep network',
          default='64,64,64,8'  # originally '64,32'
      )
      parser.add_argument(
          '--develop',
          help='Train on a small subset in development',
          dest='develop',
          action='store_true')
      parser.set_defaults(develop=False)
      parser.add_argument(
          '--skip_full_eval',
          help='Just train. Do not evaluate on test dataset.',
          dest='skip_full_eval',
          action='store_true')
      parser.set_defaults(skip_full_eval=False)
  
      # parse args
      args = parser.parse_args().__dict__
      logging.getLogger().setLevel(logging.INFO)
  
      # The Vertex AI contract. If not running in Vertex AI Training, these will be None
      OUTPUT_MODEL_DIR = os.getenv("AIP_MODEL_DIR")  # or None
      TRAIN_DATA_PATTERN = os.getenv("AIP_TRAINING_DATA_URI")
      EVAL_DATA_PATTERN = os.getenv("AIP_VALIDATION_DATA_URI")
      TEST_DATA_PATTERN = os.getenv("AIP_TEST_DATA_URI")
  
      # set top-level output directory for checkpoints, etc.
      BUCKET = args['bucket']
      OUTPUT_DIR = 'gs://{}/ch9/train_output'.format(BUCKET)
      # During hyperparameter tuning, we need to make sure different trials don't clobber each other
      # https://cloud.google.com/ai-platform/training/docs/distributed-training-details#tf-config-format
      # This doesn't exist in Vertex AI
      # OUTPUT_DIR = os.path.join(
      #     OUTPUT_DIR,
      #     json.loads(
      #         os.environ.get('TF_CONFIG', '{}')
      #     ).get('task', {}).get('trial', '')
      # )
      if OUTPUT_MODEL_DIR:
          # convert gs://ai-analytics-solutions-dsongcp2/aiplatform-custom-job-2021-11-13-22:22:46.175/1/model/
          # to gs://ai-analytics-solutions-dsongcp2/aiplatform-custom-job-2021-11-13-22:22:46.175/1
          OUTPUT_DIR = os.path.join(
              os.path.dirname(OUTPUT_MODEL_DIR if OUTPUT_MODEL_DIR[-1] != '/' else OUTPUT_MODEL_DIR[:-1]),
              'train_output')
      logging.info('Writing checkpoints and other outputs to {}'.format(OUTPUT_DIR))
  
      # Set default values for the contract variables in case we are not running in Vertex AI Training
      if not OUTPUT_MODEL_DIR:
          OUTPUT_MODEL_DIR = os.path.join(OUTPUT_DIR,
                                          'export/flights_{}'.format(time.strftime("%Y%m%d-%H%M%S")))
      if not TRAIN_DATA_PATTERN:
          TRAIN_DATA_PATTERN = 'gs://{}/ch9/data/train*'.format(BUCKET)
          CSV_COLUMNS.pop()  # the data_split column won't exist
          CSV_COLUMN_TYPES.pop()  # the data_split column won't exist
      if not EVAL_DATA_PATTERN:
          EVAL_DATA_PATTERN = 'gs://{}/ch9/data/eval*'.format(BUCKET)
      logging.info('Exporting trained model to {}'.format(OUTPUT_MODEL_DIR))
      logging.info("Reading training data from {}".format(TRAIN_DATA_PATTERN))
      logging.info('Writing trained model to {}'.format(OUTPUT_MODEL_DIR))
  
      # other global parameters
      NUM_BUCKETS = args['nbuckets']
      NUM_EMBEDS = args['nembeds']
      NUM_EXAMPLES = args['num_examples']
      NUM_EPOCHS = args['num_epochs']
      TRAIN_BATCH_SIZE = args['train_batch_size']
      DNN_HIDDEN_UNITS = args['dnn_hidden_units']
      DEVELOP_MODE = args['develop']
      SKIP_FULL_EVAL = args['skip_full_eval']
  
      # run
      train_and_evaluate(TRAIN_DATA_PATTERN, EVAL_DATA_PATTERN, TEST_DATA_PATTERN, OUTPUT_MODEL_DIR, OUTPUT_DIR)
  
      logging.info("Done")
  