#@title 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
#
# https://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.

"""
Explanation: Copyright 2018 The AdaNet Authors.
End of explanation
"""


#@test {"skip": true}
# If you're running this in Colab, first install the adanet package:
!pip install adanet

import functools
import os
import shutil

import adanet
from adanet.examples import simple_dnn
import matplotlib.pyplot as plt
import tensorflow.compat.v1 as tf


# The random seed to use.
RANDOM_SEED = 42

LOG_DIR = '/tmp/models'

"""
Explanation: Customizing AdaNet
<table class="tfo-notebook-buttons" align="left">
  <td>
    <a target="_blank" href="https://colab.research.google.com/github/tensorflow/adanet/blob/master/adanet/examples/tutorials/customizing_adanet.ipynb"><img src="https://www.tensorflow.org/images/colab_logo_32px.png" />Run in Google Colab</a>
  </td>
  <td>
    <a target="_blank" href="https://github.com/tensorflow/adanet/blob/master/adanet/examples/tutorials/customizing_adanet.ipynb"><img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />View source on GitHub</a>
  </td>
</table>

Often times, as a researcher or machine learning practitioner, you will have
some prior knowledge about a dataset. Ideally you should be able to encode that
knowledge into your machine learning algorithm. With adanet, you can do so by
defining the neural architecture search space that the AdaNet algorithm should
explore.
In this tutorial, we will explore the flexibility of the adanet framework, and
create a custom search space for an image-classification dataset using high-level
TensorFlow libraries like the
tf.keras.layers
functional API.
End of explanation
"""


(x_train, y_train), (x_test, y_test) = (
    tf.keras.datasets.fashion_mnist.load_data())

"""
Explanation: Fashion MNIST dataset
In this example, we will use the Fashion MNIST dataset
[Xiao et al., 2017] for classifying fashion
apparel images into one of ten categories:

T-shirt/top
Trouser
Pullover
Dress
Coat
Sandal
Shirt
Sneaker
Bag
Ankle boot


Download the data
Conveniently, the data is available via Keras:
End of explanation
"""


FEATURES_KEY = "images"


def generator(images, labels):
  """Returns a generator that returns image-label pairs."""

  def _gen():
    for image, label in zip(images, labels):
      yield image, label

  return _gen


def preprocess_image(image, label):
  """Preprocesses an image for an `Estimator`."""
  # First let's scale the pixel values to be between 0 and 1.
  image = image / 255.
  # Next we reshape the image so that we can apply a 2D convolution to it.
  image = tf.reshape(image, [28, 28, 1])
  # Finally the features need to be supplied as a dictionary.
  features = {FEATURES_KEY: image}
  return features, label


def input_fn(partition, training, batch_size):
  """Generate an input_fn for the Estimator."""

  def _input_fn():
    if partition == "train":
      dataset = tf.data.Dataset.from_generator(
          generator(x_train, y_train), (tf.float32, tf.int32), ((28, 28), ()))
    elif partition == "predict":
      dataset = tf.data.Dataset.from_generator(
          generator(x_test[:10], y_test[:10]), (tf.float32, tf.int32), ((28,28), ()))
    else:
      dataset = tf.data.Dataset.from_generator(
          generator(x_test, y_test), (tf.float32, tf.int32), ((28, 28), ()))

    # We call repeat after shuffling, rather than before, to prevent separate
    # epochs from blending together.
    if training:
      dataset = dataset.shuffle(10 * batch_size, seed=RANDOM_SEED).repeat()

    dataset = dataset.map(preprocess_image).batch(batch_size)
    iterator = dataset.make_one_shot_iterator()
    features, labels = iterator.get_next()
    return features, labels

  return _input_fn

"""
Explanation: Supply the data in TensorFlow
Our first task is to supply the data in TensorFlow. Using the
tf.estimator.Estimator covention, we will define a function that returns an
input_fn which returns feature and label Tensors.
We will also use the tf.data.Dataset API to feed the data into our models.
End of explanation
"""


#@test {"skip": true}

get_ipython().system_raw(
    'tensorboard --logdir {} --host 0.0.0.0 --port 6006 &'
    .format(LOG_DIR)
)

# Install ngrok binary.
! wget https://bin.equinox.io/c/4VmDzA7iaHb/ngrok-stable-linux-amd64.zip
! unzip ngrok-stable-linux-amd64.zip

# Delete old logs dir.
shutil.rmtree(LOG_DIR, ignore_errors=True)

print("Follow this link to open TensorBoard in a new tab.")
get_ipython().system_raw('./ngrok http 6006 &')
! curl -s http://localhost:4040/api/tunnels | python3 -c \
    "import sys, json; print(json.load(sys.stdin)['tunnels'][0]['public_url'])"



"""
Explanation: Launch TensorBoard
Let's run TensorBoard to visualize model training over time. We'll use ngrok to tunnel traffic to localhost.
The instructions for setting up Tensorboard were obtained from https://www.dlology.com/blog/quick-guide-to-run-tensorboard-in-google-colab/
Run the next cells and follow the link to see the TensorBoard in a new tab.
End of explanation
"""


# The number of classes.
NUM_CLASSES = 10

# A `Head` instance defines the loss function and metrics for `Estimators`.
head = tf.estimator.MultiClassHead(NUM_CLASSES)

# Some `Estimators` use feature columns for understanding their input features.
feature_columns = [
    tf.feature_column.numeric_column(FEATURES_KEY, shape=[28, 28, 1])
]

def make_config(experiment_name):
  # Estimator configuration.
  return tf.estimator.RunConfig(
    save_checkpoints_steps=1000,
    save_summary_steps=1000,
    tf_random_seed=RANDOM_SEED,
    model_dir=os.path.join(LOG_DIR, experiment_name))

"""
Explanation: Establish baselines
The next task should be to get somes baselines to see how our model performs on
this dataset.
Let's define some information to share with all our tf.estimator.Estimators:
End of explanation
"""


#@test {"skip": true}
#@title Parameters
LEARNING_RATE = 0.001  #@param {type:"number"}
TRAIN_STEPS = 5000  #@param {type:"integer"}
BATCH_SIZE = 64  #@param {type:"integer"}

estimator = tf.estimator.LinearClassifier(
    feature_columns=feature_columns,
    n_classes=NUM_CLASSES,
    optimizer=tf.train.RMSPropOptimizer(learning_rate=LEARNING_RATE),
    config=make_config("linear"))

tf.estimator.train_and_evaluate(
    estimator,
    train_spec=tf.estimator.TrainSpec(
        input_fn=input_fn("train", training=True, batch_size=BATCH_SIZE),
        max_steps=TRAIN_STEPS),
    eval_spec=tf.estimator.EvalSpec(
        input_fn=input_fn("test", training=False, batch_size=BATCH_SIZE),
        steps=None,
        start_delay_secs=1,
        throttle_secs=1,  
    ))
results = estimator.evaluate(
    input_fn("test", training=False, batch_size=BATCH_SIZE),
    steps=None)
print("Accuracy:", results["accuracy"])
print("Loss:", results["average_loss"])

"""
Explanation: Let's start simple, and train a linear model:
End of explanation
"""


#@test {"skip": true}
#@title Parameters
LEARNING_RATE = 0.003  #@param {type:"number"}
TRAIN_STEPS = 5000  #@param {type:"integer"}
BATCH_SIZE = 64  #@param {type:"integer"}
ADANET_ITERATIONS = 2  #@param {type:"integer"}

estimator = adanet.Estimator(
    head=head,
    subnetwork_generator=simple_dnn.Generator(
        feature_columns=feature_columns,
        optimizer=tf.train.RMSPropOptimizer(learning_rate=LEARNING_RATE),
        seed=RANDOM_SEED),
    max_iteration_steps=TRAIN_STEPS // ADANET_ITERATIONS,
    evaluator=adanet.Evaluator(
        input_fn=input_fn("train", training=False, batch_size=BATCH_SIZE),
        steps=None),
    config=make_config("simple_dnn"))

tf.estimator.train_and_evaluate(
    estimator,
    train_spec=tf.estimator.TrainSpec(
        input_fn=input_fn("train", training=True, batch_size=BATCH_SIZE),
        max_steps=TRAIN_STEPS),
    eval_spec=tf.estimator.EvalSpec(
        input_fn=input_fn("test", training=False, batch_size=BATCH_SIZE),
        steps=None,
        start_delay_secs=1,
        throttle_secs=1,  
    ))

results = estimator.evaluate(
    input_fn("test", training=False, batch_size=BATCH_SIZE),
    steps=None)
print("Accuracy:", results["accuracy"])
print("Loss:", results["average_loss"])

"""
Explanation: The linear model with default parameters achieves about 84.14% accuracy.
Let's see if we can do better with the simple_dnn AdaNet:
End of explanation
"""


class SimpleCNNBuilder(adanet.subnetwork.Builder):
  """Builds a CNN subnetwork for AdaNet."""

  def __init__(self, learning_rate, max_iteration_steps, seed):
    """Initializes a `SimpleCNNBuilder`.

    Args:
      learning_rate: The float learning rate to use.
      max_iteration_steps: The number of steps per iteration.
      seed: The random seed.

    Returns:
      An instance of `SimpleCNNBuilder`.
    """
    self._learning_rate = learning_rate
    self._max_iteration_steps = max_iteration_steps
    self._seed = seed

  def build_subnetwork(self,
                       features,
                       logits_dimension,
                       training,
                       iteration_step,
                       summary,
                       previous_ensemble=None):
    """See `adanet.subnetwork.Builder`."""
    images = list(features.values())[0]
    
    # Visualize some of the input images in TensorBoard.
    summary.image("images", images)
    
    kernel_initializer = tf.keras.initializers.he_normal(seed=self._seed)
    x = tf.keras.layers.Conv2D(
        filters=16,
        kernel_size=3,
        padding="same",
        activation="relu",
        kernel_initializer=kernel_initializer)(
            images)
    x = tf.keras.layers.MaxPool2D(pool_size=2, strides=2)(x)
    x = tf.keras.layers.Flatten()(x)
    x = tf.keras.layers.Dense(
        units=64, activation="relu", kernel_initializer=kernel_initializer)(
            x)

    # The `Head` passed to adanet.Estimator will apply the softmax activation.
    logits = tf.keras.layers.Dense(
        units=10, activation=None, kernel_initializer=kernel_initializer)(
            x)

    # Use a constant complexity measure, since all subnetworks have the same
    # architecture and hyperparameters.
    complexity = tf.constant(1)

    return adanet.Subnetwork(
        last_layer=x,
        logits=logits,
        complexity=complexity,
        persisted_tensors={})

  def build_subnetwork_train_op(self,
                                subnetwork,
                                loss,
                                var_list,
                                labels,
                                iteration_step,
                                summary,
                                previous_ensemble=None):
    """See `adanet.subnetwork.Builder`."""

    # Momentum optimizer with cosine learning rate decay works well with CNNs.
    learning_rate = tf.train.cosine_decay(
        learning_rate=self._learning_rate,
        global_step=iteration_step,
        decay_steps=self._max_iteration_steps)
    optimizer = tf.train.MomentumOptimizer(learning_rate, .9)
    # NOTE: The `adanet.Estimator` increments the global step.
    return optimizer.minimize(loss=loss, var_list=var_list)

  def build_mixture_weights_train_op(self, loss, var_list, logits, labels,
                                     iteration_step, summary):
    """See `adanet.subnetwork.Builder`."""
    return tf.no_op("mixture_weights_train_op")

  @property
  def name(self):
    """See `adanet.subnetwork.Builder`."""
    return "simple_cnn"

"""
Explanation: The simple_dnn AdaNet model with default parameters achieves about 85.64%
accuracy.
This improvement can be attributed to simple_dnn searching over
fully-connected neural networks which have more expressive power than the linear
model due to their non-linear activations.
Fully-connected layers are permutation invariant to their inputs, meaning that
if we consistently swapped two pixels before training, the final model would
perform identically. However, there is spatial and locality information in
images that we should try to capture. Applying a few convolutions to our inputs
will allow us to do so, and that will require defining a custom
adanet.subnetwork.Builder and adanet.subnetwork.Generator.
Define a convolutional AdaNet model
Creating a new search space for AdaNet to explore is straightforward. There are
two abstract classes you need to extend:

adanet.subnetwork.Builder
adanet.subnetwork.Generator

Similar to the tf.estimator.Estimator model_fn, adanet.subnetwork.Builder
allows you to define your own TensorFlow graph for creating a neural network,
and specify the training operations.
Below we define one that applies a 2D convolution, max-pooling, and then a
fully-connected layer to the images:
End of explanation
"""


class SimpleCNNGenerator(adanet.subnetwork.Generator):
  """Generates a `SimpleCNN` at each iteration.
  """

  def __init__(self, learning_rate, max_iteration_steps, seed=None):
    """Initializes a `Generator` that builds `SimpleCNNs`.

    Args:
      learning_rate: The float learning rate to use.
      max_iteration_steps: The number of steps per iteration.
      seed: The random seed.

    Returns:
      An instance of `Generator`.
    """
    self._seed = seed
    self._dnn_builder_fn = functools.partial(
        SimpleCNNBuilder,
        learning_rate=learning_rate,
        max_iteration_steps=max_iteration_steps)

  def generate_candidates(self, previous_ensemble, iteration_number,
                          previous_ensemble_reports, all_reports):
    """See `adanet.subnetwork.Generator`."""
    seed = self._seed
    # Change the seed according to the iteration so that each subnetwork
    # learns something different.
    if seed is not None:
      seed += iteration_number
    return [self._dnn_builder_fn(seed=seed)]

"""
Explanation: Next, we extend a adanet.subnetwork.Generator, which defines the search
space of candidate SimpleCNNBuilders to consider including the final network.
It can create one or more at each iteration with different parameters, and the
AdaNet algorithm will select the candidate that best improves the overall neural
network's adanet_loss on the training set.
The one below is very simple: it always creates the same architecture, but gives
it a different random seed at each iteration:
End of explanation
"""


#@title Parameters
LEARNING_RATE = 0.05  #@param {type:"number"}
TRAIN_STEPS = 5000  #@param {type:"integer"}
BATCH_SIZE = 64  #@param {type:"integer"}
ADANET_ITERATIONS = 2  #@param {type:"integer"}

max_iteration_steps = TRAIN_STEPS // ADANET_ITERATIONS
estimator = adanet.Estimator(
    head=head,
    subnetwork_generator=SimpleCNNGenerator(
        learning_rate=LEARNING_RATE,
        max_iteration_steps=max_iteration_steps,
        seed=RANDOM_SEED),
    max_iteration_steps=max_iteration_steps,
    evaluator=adanet.Evaluator(
        input_fn=input_fn("train", training=False, batch_size=BATCH_SIZE),
        steps=None),
    adanet_loss_decay=.99,
    config=make_config("simple_cnn"))

results, _ = tf.estimator.train_and_evaluate(
    estimator,
    train_spec=tf.estimator.TrainSpec(
        input_fn=input_fn("train", training=True, batch_size=BATCH_SIZE),
        max_steps=TRAIN_STEPS),
    eval_spec=tf.estimator.EvalSpec(
        input_fn=input_fn("test", training=False, batch_size=BATCH_SIZE),
        steps=None,
        start_delay_secs=1,
        throttle_secs=1,  
    ))
print("Accuracy:", results["accuracy"])
print("Loss:", results["average_loss"])

"""
Explanation: With these defined, we pass them into a new adanet.Estimator:
End of explanation
"""


predictions = estimator.predict(input_fn=input_fn("predict", training=False, batch_size=1))

for i, val in enumerate(predictions):
    predicted_class = val['class_ids'][0]
    prediction_confidence = val['probabilities'][predicted_class] * 100
    
    # Display the image
    plt.imshow(x_test[i])
    plt.rcParams['figure.figsize'] = (1,1)
    plt.show()
    
    print('Predicted class: %s, confidence: %s%%' % (predicted_class, round(prediction_confidence, 3)))
    print('Actual class: %s \n\n' % y_test[i])

"""
Explanation: Our SimpleCNNGenerator code achieves 90.21% accuracy.
Generating predictions on our trained model
Now that we've got a trained model, we can use it to generate predictions on new input. To keep things simple, here we'll generate predictions on our estimator using the first 10 examples from the test set.
End of explanation
"""


