#!/usr/bin/env python
# coding: utf-8

# # What's this TensorFlow business?
#
# You've written a lot of code in this assignment to provide a whole host of neural network functionality. Dropout, Batch Norm, and 2D convolutions are some of the workhorses of deep learning in computer vision. You've also worked hard to make your code efficient and vectorized.
#
# For the last part of this assignment, though, we're going to leave behind your beautiful codebase and instead migrate to one of two popular deep learning frameworks: in this instance, TensorFlow (or PyTorch, if you choose to work with that notebook).

# #### What is it?
# TensorFlow is a system for executing computational graphs over Tensor objects, with native support for performing backpropogation for its Variables. In it, we work with Tensors which are n-dimensional arrays analogous to the numpy ndarray.
#
# #### Why?
#
# * Our code will now run on GPUs! Much faster training. Writing your own modules to run on GPUs is beyond the scope of this class, unfortunately.
# * We want you to be ready to use one of these frameworks for your project so you can experiment more efficiently than if you were writing every feature you want to use by hand.
# * We want you to stand on the shoulders of giants! TensorFlow and PyTorch are both excellent frameworks that will make your lives a lot easier, and now that you understand their guts, you are free to use them :)
# * We want you to be exposed to the sort of deep learning code you might run into in academia or industry.

# ## How will I learn TensorFlow?
#
# TensorFlow has many excellent tutorials available, including those from [Google themselves](https://www.tensorflow.org/get_started/get_started).
#
# Otherwise, this notebook will walk you through much of what you need to do to train models in TensorFlow. See the end of the notebook for some links to helpful tutorials if you want to learn more or need further clarification on topics that aren't fully explained here.
#
# **NOTE: This notebook is meant to teach you the latest version of Tensorflow 2.0. Most examples on the web today are still in 1.x, so be careful not to confuse the two when looking up documentation**.
#
# ## Install Tensorflow 2.0
# Tensorflow 2.0 is still not in a fully 100% stable release, but it's still usable and more intuitive than TF 1.x. Please make sure you have it installed before moving on in this notebook! Here are some steps to get started:
#
# 1. Have the latest version of Anaconda installed on your machine.
# 2. Create a new conda environment starting from Python 3.7. In this setup example, we'll call it `tf_20_env`.
# 3. Run the command: `source activate tf_20_env`
# 4. Then pip install TF 2.0 as described here: https://www.tensorflow.org/install/pip
#
# A guide on creating Anaconda enviornments: https://uoa-eresearch.github.io/eresearch-cookbook/recipe/2014/11/20/conda/
#
# This will give you an new enviornemnt to play in TF 2.0. Generally, if you plan to also use TensorFlow in your other projects, you might also want to keep a seperate Conda environment or virtualenv in Python 3.7 that has Tensorflow 1.9, so you can switch back and forth at will.

# # Table of Contents
#
# This notebook has 5 parts. We will walk through TensorFlow at **three different levels of abstraction**, which should help you better understand it and prepare you for working on your project.
#
# 1. Part I, Preparation: load the CIFAR-10 dataset.
# 2. Part II, Barebone TensorFlow: **Abstraction Level 1**, we will work directly with low-level TensorFlow graphs.
# 3. Part III, Keras Model API: **Abstraction Level 2**, we will use `tf.keras.Model` to define arbitrary neural network architecture.
# 4. Part IV, Keras Sequential + Functional API: **Abstraction Level 3**, we will use `tf.keras.Sequential` to define a linear feed-forward network very conveniently, and then explore the functional libraries for building unique and uncommon models that require more flexibility.
# 5. Part V, CIFAR-10 open-ended challenge: please implement your own network to get as high accuracy as possible on CIFAR-10. You can experiment with any layer, optimizer, hyperparameters or other advanced features.
#
# We will discuss Keras in more detail later in the notebook.
#
# Here is a table of comparison:
#
# | API           | Flexibility | Convenience |
# |---------------|-------------|-------------|
# | Barebone      | High        | Low         |
# | `tf.keras.Model`     | High        | Medium      |
# | `tf.keras.Sequential` | Low         | High        |

# # Part I: Preparation
#
# First, we load the CIFAR-10 dataset. This might take a few minutes to download the first time you run it, but after that the files should be cached on disk and loading should be faster.
#
# In previous parts of the assignment we used CS231N-specific code to download and read the CIFAR-10 dataset; however the `tf.keras.datasets` package in TensorFlow provides prebuilt utility functions for loading many common datasets.
#
# For the purposes of this assignment we will still write our own code to preprocess the data and iterate through it in minibatches. The `tf.data` package in TensorFlow provides tools for automating this process, but working with this package adds extra complication and is beyond the scope of this notebook. However using `tf.data` can be much more efficient than the simple approach used in this notebook, so you should consider using it for your project.

# In[1]:

import os
import tensorflow as tf
import numpy as np
import math
import timeit
import matplotlib.pyplot as plt

#config = tf.ConfigProto()
#config.gpu_options.allow_growth = True

#get_ipython().run_line_magic('matplotlib', 'inline')

gpu_devices = tf.config.experimental.list_physical_devices('GPU')
for device in gpu_devices:
    tf.config.experimental.set_memory_growth(device, True)

# In[2]:

def load_cifar10(num_training=49000, num_validation=1000, num_test=10000):
    """
    Fetch the CIFAR-10 dataset from the web and perform preprocessing to prepare
    it for the two-layer neural net classifier. These are the same steps as
    we used for the SVM, but condensed to a single function.
    """
    # Load the raw CIFAR-10 dataset and use appropriate data types and shapes
    cifar10 = tf.keras.datasets.cifar10.load_data()
    (X_train, y_train), (X_test, y_test) = cifar10
    X_train = np.asarray(X_train, dtype=np.float32)
    y_train = np.asarray(y_train, dtype=np.int32).flatten()
    X_test = np.asarray(X_test, dtype=np.float32)
    y_test = np.asarray(y_test, dtype=np.int32).flatten()

    # Subsample the data
    mask = range(num_training, num_training + num_validation)
    X_val = X_train[mask]
    y_val = y_train[mask]
    mask = range(num_training)
    X_train = X_train[mask]
    y_train = y_train[mask]
    mask = range(num_test)
    X_test = X_test[mask]
    y_test = y_test[mask]

    # Normalize the data: subtract the mean pixel and divide by std
    mean_pixel = X_train.mean(axis=(0, 1, 2), keepdims=True)
    std_pixel = X_train.std(axis=(0, 1, 2), keepdims=True)
    X_train = (X_train - mean_pixel) / std_pixel
    X_val = (X_val - mean_pixel) / std_pixel
    X_test = (X_test - mean_pixel) / std_pixel

    return X_train, y_train, X_val, y_val, X_test, y_test

# If there are errors with SSL downloading involving self-signed certificates,
# it may be that your Python version was recently installed on the current machine.
# See: https://github.com/tensorflow/tensorflow/issues/10779
# To fix, run the command: /Applications/Python\ 3.7/Install\ Certificates.command
#   ...replacing paths as necessary.

# Invoke the above function to get our data.
NHW = (0, 1, 2)
X_train, y_train, X_val, y_val, X_test, y_test = load_cifar10()
print('Train data shape: ', X_train.shape)
print('Train labels shape: ', y_train.shape, y_train.dtype)
print('Validation data shape: ', X_val.shape)
print('Validation labels shape: ', y_val.shape)
print('Test data shape: ', X_test.shape)
print('Test labels shape: ', y_test.shape)

# In[3]:

class Dataset(object):
    def __init__(self, X, y, batch_size, shuffle=False):
        """
        Construct a Dataset object to iterate over data X and labels y

        Inputs:
        - X: Numpy array of data, of any shape
        - y: Numpy array of labels, of any shape but with y.shape[0] == X.shape[0]
        - batch_size: Integer giving number of elements per minibatch
        - shuffle: (optional) Boolean, whether to shuffle the data on each epoch
        """
        assert X.shape[0] == y.shape[0], 'Got different numbers of data and labels'
        self.X, self.y = X, y
        self.batch_size, self.shuffle = batch_size, shuffle

    def __iter__(self):
        N, B = self.X.shape[0], self.batch_size
        idxs = np.arange(N)
        if self.shuffle:
            np.random.shuffle(idxs)
        return iter((self.X[i:i+B], self.y[i:i+B]) for i in range(0, N, B))

train_dset = Dataset(X_train, y_train, batch_size=64, shuffle=True)
val_dset = Dataset(X_val, y_val, batch_size=64, shuffle=False)
test_dset = Dataset(X_test, y_test, batch_size=64)

# In[4]:

# We can iterate through a dataset like this:
for t, (x, y) in enumerate(train_dset):
    print(t, x.shape, y.shape)
    if t > 5: break

# You can optionally **use GPU by setting the flag to True below**. It's not neccessary to use a GPU for this assignment; if you are working on Google Cloud then we recommend that you do not use a GPU, as it will be significantly more expensive.

# In[5]:

# Set up some global variables
USE_GPU = True

if USE_GPU:
    device = '/device:GPU:0'
else:
    device = '/cpu:0'

# Constant to control how often we print when training models
print_every = 100

print('Using device: ', device)

# # Part II: Barebones TensorFlow
# TensorFlow ships with various high-level APIs which make it very convenient to define and train neural networks; we will cover some of these constructs in Part III and Part IV of this notebook. In this section we will start by building a model with basic TensorFlow constructs to help you better understand what's going on under the hood of the higher-level APIs.
#
# **"Barebones Tensorflow" is important to understanding the building blocks of TensorFlow, but much of it involves concepts from TensorFlow 1.x.** We will be working with legacy modules such as `tf.Variable`.
#
# Therefore, please read and understand the differences between legacy (1.x) TF and the new (2.0) TF.
#
# ### Historical background on TensorFlow 1.x
#
# TensorFlow 1.x is primarily a framework for working with **static computational graphs**. Nodes in the computational graph are Tensors which will hold n-dimensional arrays when the graph is run; edges in the graph represent functions that will operate on Tensors when the graph is run to actually perform useful computation.
#
# Before Tensorflow 2.0, we had to configure the graph into two phases. There are plenty of tutorials online that explain this two-step process. The process generally looks like the following for TF 1.x:
# 1. **Build a computational graph that describes the computation that you want to perform**. This stage doesn't actually perform any computation; it just builds up a symbolic representation of your computation. This stage will typically define one or more `placeholder` objects that represent inputs to the computational graph.
# 2. **Run the computational graph many times.** Each time the graph is run (e.g. for one gradient descent step) you will specify which parts of the graph you want to compute, and pass a `feed_dict` dictionary that will give concrete values to any `placeholder`s in the graph.
#
# ### The new paradigm in Tensorflow 2.0
# Now, with Tensorflow 2.0, we can simply adopt a functional form that is more Pythonic and similar in spirit to PyTorch and direct Numpy operation. Instead of the 2-step paradigm with computation graphs, making it (among other things) easier to debug TF code. You can read more details at https://www.tensorflow.org/guide/eager.
#
# The main difference between the TF 1.x and 2.0 approach is that the 2.0 approach doesn't make use of `tf.Session`, `tf.run`, `placeholder`, `feed_dict`. To get more details of what's different between the two version and how to convert between the two, check out the official migration guide: https://www.tensorflow.org/alpha/guide/migration_guide
#
# Later, in the rest of this notebook we'll focus on this new, simpler approach.

# ### TensorFlow warmup: Flatten Function
#
# We can see this in action by defining a simple `flatten` function that will reshape image data for use in a fully-connected network.
#
# In TensorFlow, data for convolutional feature maps is typically stored in a Tensor of shape N x H x W x C where:
#
# - N is the number of datapoints (minibatch size)
# - H is the height of the feature map
# - W is the width of the feature map
# - C is the number of channels in the feature map
#
# This is the right way to represent the data when we are doing something like a 2D convolution, that needs spatial understanding of where the intermediate features are relative to each other. When we use fully connected affine layers to process the image, however, we want each datapoint to be represented by a single vector -- it's no longer useful to segregate the different channels, rows, and columns of the data. So, we use a "flatten" operation to collapse the `H x W x C` values per representation into a single long vector.
#
# Notice the `tf.reshape` call has the target shape as `(N, -1)`, meaning it will reshape/keep the first dimension to be N, and then infer as necessary what the second dimension is in the output, so we can collapse the remaining dimensions from the input properly.
#
# **NOTE**: TensorFlow and PyTorch differ on the default Tensor layout; TensorFlow uses N x H x W x C but PyTorch uses N x C x H x W.

# In[6]:

def flatten(x):
    """
    Input:
    - TensorFlow Tensor of shape (N, D1, ..., DM)

    Output:
    - TensorFlow Tensor of shape (N, D1 * ... * DM)
    """
    N = tf.shape(x)[0]
    return tf.reshape(x, (N, -1))

# In[7]:

def test_flatten():
    # Construct concrete values of the input data x using numpy
    x_np = np.arange(24).reshape((2, 3, 4))
    print('x_np:\n', x_np, '\n')
    # Compute a concrete output value.
    x_flat_np = flatten(x_np)
    print('x_flat_np:\n', x_flat_np, '\n')

test_flatten()

# ### Barebones TensorFlow: Define a Two-Layer Network
# We will now implement our first neural network with TensorFlow: a fully-connected ReLU network with two hidden layers and no biases on the CIFAR10 dataset. For now we will use only low-level TensorFlow operators to define the network; later we will see how to use the higher-level abstractions provided by `tf.keras` to simplify the process.
#
# We will define the forward pass of the network in the function `two_layer_fc`; this will accept TensorFlow Tensors for the inputs and weights of the network, and return a TensorFlow Tensor for the scores.
#
# After defining the network architecture in the `two_layer_fc` function, we will test the implementation by checking the shape of the output.
#
# **It's important that you read and understand this implementation.**

# In[8]:

def two_layer_fc(x, params):
    """
    A fully-connected neural network; the architecture is:
    fully-connected layer -> ReLU -> fully connected layer.
    Note that we only need to define the forward pass here; TensorFlow will take
    care of computing the gradients for us.

    The input to the network will be a minibatch of data, of shape
    (N, d1, ..., dM) where d1 * ... * dM = D. The hidden layer will have H units,
    and the output layer will produce scores for C classes.

    Inputs:
    - x: A TensorFlow Tensor of shape (N, d1, ..., dM) giving a minibatch of
      input data.
    - params: A list [w1, w2] of TensorFlow Tensors giving weights for the
      network, where w1 has shape (D, H) and w2 has shape (H, C).

    Returns:
    - scores: A TensorFlow Tensor of shape (N, C) giving classification scores
      for the input data x.
    """
    w1, w2 = params                   # Unpack the parameters
    x = flatten(x)                    # Flatten the input; now x has shape (N, D)
    h = tf.nn.relu(tf.matmul(x, w1))  # Hidden layer: h has shape (N, H)
    scores = tf.matmul(h, w2)         # Compute scores of shape (N, C)
    return scores

# In[9]:

def two_layer_fc_test():
    hidden_layer_size = 42

    # Scoping our TF operations under a tf.device context manager
    # lets us tell TensorFlow where we want these Tensors to be
    # multiplied and/or operated on, e.g. on a CPU or a GPU.
    with tf.device(device):
        x = tf.zeros((64, 32, 32, 3))
        w1 = tf.zeros((32 * 32 * 3, hidden_layer_size))
        w2 = tf.zeros((hidden_layer_size, 10))

        # Call our two_layer_fc function for the forward pass of the network.
        scores = two_layer_fc(x, [w1, w2])

    print(scores.shape)

two_layer_fc_test()

# ### Barebones TensorFlow: Three-Layer ConvNet
# Here you will complete the implementation of the function `three_layer_convnet` which will perform the forward pass of a three-layer convolutional network. The network should have the following architecture:
#
# 1. A convolutional layer (with bias) with `channel_1` filters, each with shape `KW1 x KH1`, and zero-padding of two
# 2. ReLU nonlinearity
# 3. A convolutional layer (with bias) with `channel_2` filters, each with shape `KW2 x KH2`, and zero-padding of one
# 4. ReLU nonlinearity
# 5. Fully-connected layer with bias, producing scores for `C` classes.
#
# **HINT**: For convolutions: https://www.tensorflow.org/versions/r2.0/api_docs/python/tf/nn/conv2d; be careful with padding!
#
# **HINT**: For biases: https://www.tensorflow.org/performance/xla/broadcasting

# In[10]:

def three_layer_convnet(x, params):
    """
    A three-layer convolutional network with the architecture described above.

    Inputs:
    - x: A TensorFlow Tensor of shape (N, H, W, 3) giving a minibatch of images
    - params: A list of TensorFlow Tensors giving the weights and biases for the
      network; should contain the following:
      - conv_w1: TensorFlow Tensor of shape (KH1, KW1, 3, channel_1) giving
        weights for the first convolutional layer.
      - conv_b1: TensorFlow Tensor of shape (channel_1,) giving biases for the
        first convolutional layer.
      - conv_w2: TensorFlow Tensor of shape (KH2, KW2, channel_1, channel_2)
        giving weights for the second convolutional layer
      - conv_b2: TensorFlow Tensor of shape (channel_2,) giving biases for the
        second convolutional layer.
      - fc_w: TensorFlow Tensor giving weights for the fully-connected layer.
        Can you figure out what the shape should be?
      - fc_b: TensorFlow Tensor giving biases for the fully-connected layer.
        Can you figure out what the shape should be?
    """
    conv_w1, conv_b1, conv_w2, conv_b2, fc_w, fc_b = params
    scores = None
    ############################################################################
    # TODO: Implement the forward pass for the three-layer ConvNet.            #
    ############################################################################
    # *****START OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)*****

    # 1.
    w1x1 = tf.nn.conv2d(input=x, filters=conv_w1, strides = 1, padding=[[0, 0], [2, 2], [2, 2], [0, 0]], data_format='NHWC', name="Prvi")
    conv_layer_1 = tf.nn.bias_add(w1x1, conv_b1)

    # 2.
    reLu1 = tf.nn.relu(conv_layer_1)

    # 3.
    w2x2 = tf.nn.conv2d(input=reLu1, filters=conv_w2, strides = 1, padding=[[0, 0], [1, 1], [1, 1], [0, 0]], data_format='NHWC', name="Drugi")
    conv_layer_2 = tf.nn.bias_add(w2x2, conv_b2)

    # 4.
    relu2 = tf.nn.relu(conv_layer_2)

    # 5.
    relu2flatten = flatten(relu2)
    w3x3 = tf.matmul(relu2flatten, fc_w)
    scores = tf.nn.bias_add(w3x3, fc_b)

    # *****END OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)*****
    ############################################################################
    #                              END OF YOUR CODE                            #
    ############################################################################
    return scores

# After defing the forward pass of the three-layer ConvNet above, run the following cell to test your implementation. Like the two-layer network, we run the graph on a batch of zeros just to make sure the function doesn't crash, and produces outputs of the correct shape.
#
# When you run this function, `scores_np` should have shape `(64, 10)`.

# In[11]:

def three_layer_convnet_test():

    with tf.device(device):
        x = tf.zeros((64, 32, 32, 3))
        conv_w1 = tf.zeros((5, 5, 3, 6))
        conv_b1 = tf.zeros((6,))
        conv_w2 = tf.zeros((3, 3, 6, 9))
        conv_b2 = tf.zeros((9,))
        fc_w = tf.zeros((32 * 32 * 9, 10))
        fc_b = tf.zeros((10,))
        params = [conv_w1, conv_b1, conv_w2, conv_b2, fc_w, fc_b]
        scores = three_layer_convnet(x, params)

    # Inputs to convolutional layers are 4-dimensional arrays with shape
    # [batch_size, height, width, channels]
    print('scores_np has shape: ', scores.shape)

three_layer_convnet_test()

# ### Barebones TensorFlow: Training Step
#
# We now define the `training_step` function performs a single training step. This will take three basic steps:
#
# 1. Compute the loss
# 2. Compute the gradient of the loss with respect to all network weights
# 3. Make a weight update step using (stochastic) gradient descent.
#
#
# We need to use a few new TensorFlow functions to do all of this:
# - For computing the cross-entropy loss we'll use `tf.nn.sparse_softmax_cross_entropy_with_logits`: https://www.tensorflow.org/versions/r2.0/api_docs/python/tf/nn/sparse_softmax_cross_entropy_with_logits
#
# - For averaging the loss across a minibatch of data we'll use `tf.reduce_mean`:
# https://www.tensorflow.org/versions/r2.0/api_docs/python/tf/reduce_mean
#
# - For computing gradients of the loss with respect to the weights we'll use `tf.GradientTape` (useful for Eager execution):  https://www.tensorflow.org/versions/r2.0/api_docs/python/tf/GradientTape
#
# - We'll mutate the weight values stored in a TensorFlow Tensor using `tf.assign_sub` ("sub" is for subtraction): https://www.tensorflow.org/api_docs/python/tf/assign_sub
#

# In[12]:

def training_step(model_fn, x, y, params, learning_rate):
    with tf.GradientTape() as tape:
        scores = model_fn(x, params) # Forward pass of the model
        loss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=y, logits=scores)
        total_loss = tf.reduce_mean(loss)
        grad_params = tape.gradient(total_loss, params)

        # Make a vanilla gradient descent step on all of the model parameters
        # Manually update the weights using assign_sub()
        for w, grad_w in zip(params, grad_params):
            w.assign_sub(learning_rate * grad_w)

        return total_loss

# In[13]:

def train_part2(model_fn, init_fn, learning_rate):
    """
    Train a model on CIFAR-10.

    Inputs:
    - model_fn: A Python function that performs the forward pass of the model
      using TensorFlow; it should have the following signature:
      scores = model_fn(x, params) where x is a TensorFlow Tensor giving a
      minibatch of image data, params is a list of TensorFlow Tensors holding
      the model weights, and scores is a TensorFlow Tensor of shape (N, C)
      giving scores for all elements of x.
    - init_fn: A Python function that initializes the parameters of the model.
      It should have the signature params = init_fn() where params is a list
      of TensorFlow Tensors holding the (randomly initialized) weights of the
      model.
    - learning_rate: Python float giving the learning rate to use for SGD.
    """

    params = init_fn()  # Initialize the model parameters

    for t, (x_np, y_np) in enumerate(train_dset):
        # Run the graph on a batch of training data.
        loss = training_step(model_fn, x_np, y_np, params, learning_rate)

        # Periodically print the loss and check accuracy on the val set.
        if t % print_every == 0:
            print('Iteration %d, loss = %.4f' % (t, loss))
            check_accuracy(val_dset, x_np, model_fn, params)

# In[14]:

def check_accuracy(dset, x, model_fn, params):
    """
    Check accuracy on a classification model, e.g. for validation.

    Inputs:
    - dset: A Dataset object against which to check accuracy
    - x: A TensorFlow placeholder Tensor where input images should be fed
    - model_fn: the Model we will be calling to make predictions on x
    - params: parameters for the model_fn to work with

    Returns: Nothing, but prints the accuracy of the model
    """
    num_correct, num_samples = 0, 0
    for x_batch, y_batch in dset:
        scores_np = model_fn(x_batch, params).numpy()
        y_pred = scores_np.argmax(axis=1)
        num_samples += x_batch.shape[0]
        num_correct += (y_pred == y_batch).sum()
    acc = float(num_correct) / num_samples
    print('Got %d / %d correct (%.2f%%)' % (num_correct, num_samples, 100 * acc))

# ### Barebones TensorFlow: Initialization
# We'll use the following utility method to initialize the weight matrices for our models using Kaiming's normalization method.
#
# [1] He et al, *Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification
# *, ICCV 2015, https://arxiv.org/abs/1502.01852

# In[15]:

def create_matrix_with_kaiming_normal(shape):
    if len(shape) == 2:
        fan_in, fan_out = shape[0], shape[1]
    elif len(shape) == 4:
        fan_in, fan_out = np.prod(shape[:3]), shape[3]
    return tf.keras.backend.random_normal(shape) * np.sqrt(2.0 / fan_in)

# ### Barebones TensorFlow: Train a Two-Layer Network
# We are finally ready to use all of the pieces defined above to train a two-layer fully-connected network on CIFAR-10.
#
# We just need to define a function to initialize the weights of the model, and call `train_part2`.
#
# Defining the weights of the network introduces another important piece of TensorFlow API: `tf.Variable`. A TensorFlow Variable is a Tensor whose value is stored in the graph and persists across runs of the computational graph; however unlike constants defined with `tf.zeros` or `tf.random_normal`, the values of a Variable can be mutated as the graph runs; these mutations will persist across graph runs. Learnable parameters of the network are usually stored in Variables.
#
# You don't need to tune any hyperparameters, but you should achieve validation accuracies above 40% after one epoch of training.

# In[16]:

def two_layer_fc_init():
    """
    Initialize the weights of a two-layer network, for use with the
    two_layer_network function defined above.
    You can use the `create_matrix_with_kaiming_normal` helper!

    Inputs: None

    Returns: A list of:
    - w1: TensorFlow tf.Variable giving the weights for the first layer
    - w2: TensorFlow tf.Variable giving the weights for the second layer
    """
    hidden_layer_size = 4000
    w1 = tf.Variable(create_matrix_with_kaiming_normal((3 * 32 * 32, 4000)))
    w2 = tf.Variable(create_matrix_with_kaiming_normal((4000, 10)))
    return [w1, w2]

learning_rate = 1e-2
train_part2(two_layer_fc, two_layer_fc_init, learning_rate)

# ### Barebones TensorFlow: Train a three-layer ConvNet
# We will now use TensorFlow to train a three-layer ConvNet on CIFAR-10.
#
# You need to implement the `three_layer_convnet_init` function. Recall that the architecture of the network is:
#
# 1. Convolutional layer (with bias) with 32 5x5 filters, with zero-padding 2
# 2. ReLU
# 3. Convolutional layer (with bias) with 16 3x3 filters, with zero-padding 1
# 4. ReLU
# 5. Fully-connected layer (with bias) to compute scores for 10 classes
#
# You don't need to do any hyperparameter tuning, but you should see validation accuracies above 43% after one epoch of training.

# In[17]:

def three_layer_convnet_init():
    """
    Initialize the weights of a Three-Layer ConvNet, for use with the
    three_layer_convnet function defined above.
    You can use the `create_matrix_with_kaiming_normal` helper!

    Inputs: None

    Returns a list containing:
    - conv_w1: TensorFlow tf.Variable giving weights for the first conv layer
    - conv_b1: TensorFlow tf.Variable giving biases for the first conv layer
    - conv_w2: TensorFlow tf.Variable giving weights for the second conv layer
    - conv_b2: TensorFlow tf.Variable giving biases for the second conv layer
    - fc_w: TensorFlow tf.Variable giving weights for the fully-connected layer
    - fc_b: TensorFlow tf.Variable giving biases for the fully-connected layer
    """
    params = None
    ############################################################################
    # TODO: Initialize the parameters of the three-layer network.              #
    ############################################################################
    # *****START OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)*****

      #     network; should contain the following:
      # - conv_w1: TensorFlow Tensor of shape (KH1, KW1, 3, channel_1) giving
      #   weights for the first convolutional layer.
      # - conv_b1: TensorFlow Tensor of shape (channel_1,) giving biases for the
      #   first convolutional layer.
      # - conv_w2: TensorFlow Tensor of shape (KH2, KW2, channel_1, channel_2)
      #   giving weights for the second convolutional layer
      # - conv_b2: TensorFlow Tensor of shape (channel_2,) giving biases for the
      #   second convolutional layer.
      # - fc_w: TensorFlow Tensor giving weights for the fully-connected layer.
      #   Can you figure out what the shape should be?
      # - fc_b: TensorFlow Tensor giving biases for the fully-connected layer.
      #   Can you figure out what the shape should be?
      #
    #  conv_w1, conv_b1, conv_w2, conv_b2, fc_w, fc_b = params
    conv_w1 = tf.Variable(create_matrix_with_kaiming_normal((5, 5, 3, 32)))
    conv_b1 = tf.Variable(tf.zeros((32,)))
    conv_w2 = tf.Variable(create_matrix_with_kaiming_normal((3, 3, 32, 16)))
    conv_b2 = tf.Variable(tf.zeros((16,)))
    fc_w = tf.Variable(create_matrix_with_kaiming_normal((32 * 32 * 16, 10)))
    fc_b = tf.Variable(tf.zeros((10,)))
    params = [conv_w1, conv_b1, conv_w2, conv_b2, fc_w, fc_b]
    # *****END OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)*****
    ############################################################################
    #                             END OF YOUR CODE                             #
    ############################################################################
    return params

learning_rate = 3e-3
train_part2(three_layer_convnet, three_layer_convnet_init, learning_rate)

# # Part III: Keras Model Subclassing API
#
# Implementing a neural network using the low-level TensorFlow API is a good way to understand how TensorFlow works, but it's a little inconvenient - we had to manually keep track of all Tensors holding learnable parameters. This was fine for a small network, but could quickly become unweildy for a large complex model.
#
# Fortunately TensorFlow 2.0 provides higher-level APIs such as `tf.keras` which make it easy to build models out of modular, object-oriented layers. Further, TensorFlow 2.0 uses eager execution that evaluates operations immediately, without explicitly constructing any computational graphs. This makes it easy to write and debug models, and reduces the boilerplate code.
#
# In this part of the notebook we will define neural network models using the `tf.keras.Model` API. To implement your own model, you need to do the following:
#
# 1. Define a new class which subclasses `tf.keras.Model`. Give your class an intuitive name that describes it, like `TwoLayerFC` or `ThreeLayerConvNet`.
# 2. In the initializer `__init__()` for your new class, define all the layers you need as class attributes. The `tf.keras.layers` package provides many common neural-network layers, like `tf.keras.layers.Dense` for fully-connected layers and `tf.keras.layers.Conv2D` for convolutional layers. Under the hood, these layers will construct `Variable` Tensors for any learnable parameters. **Warning**: Don't forget to call `super(YourModelName, self).__init__()` as the first line in your initializer!
# 3. Implement the `call()` method for your class; this implements the forward pass of your model, and defines the *connectivity* of your network. Layers defined in `__init__()` implement `__call__()` so they can be used as function objects that transform input Tensors into output Tensors. Don't define any new layers in `call()`; any layers you want to use in the forward pass should be defined in `__init__()`.
#
# After you define your `tf.keras.Model` subclass, you can instantiate it and use it like the model functions from Part II.
#
# ### Keras Model Subclassing API: Two-Layer Network
#
# Here is a concrete example of using the `tf.keras.Model` API to define a two-layer network. There are a few new bits of API to be aware of here:
#
# We use an `Initializer` object to set up the initial values of the learnable parameters of the layers; in particular `tf.initializers.VarianceScaling` gives behavior similar to the Kaiming initialization method we used in Part II. You can read more about it here: https://www.tensorflow.org/versions/r2.0/api_docs/python/tf/initializers/VarianceScaling
#
# We construct `tf.keras.layers.Dense` objects to represent the two fully-connected layers of the model. In addition to multiplying their input by a weight matrix and adding a bias vector, these layer can also apply a nonlinearity for you. For the first layer we specify a ReLU activation function by passing `activation='relu'` to the constructor; the second layer uses softmax activation function. Finally, we use `tf.keras.layers.Flatten` to flatten the output from the previous fully-connected layer.

# In[18]:

class TwoLayerFC(tf.keras.Model):
    def __init__(self, hidden_size, num_classes):
        super(TwoLayerFC, self).__init__()
        initializer = tf.initializers.VarianceScaling(scale=2.0)
        self.fc1 = tf.keras.layers.Dense(hidden_size, activation='relu',
                                   kernel_initializer=initializer)
        self.fc2 = tf.keras.layers.Dense(num_classes, activation='softmax',
                                   kernel_initializer=initializer)
        self.flatten = tf.keras.layers.Flatten()

    def call(self, x, training=False):
        x = self.flatten(x)
        x = self.fc1(x)
        x = self.fc2(x)
        return x

def test_TwoLayerFC():
    """ A small unit test to exercise the TwoLayerFC model above. """
    input_size, hidden_size, num_classes = 50, 42, 10
    x = tf.zeros((64, input_size))
    model = TwoLayerFC(hidden_size, num_classes)
    with tf.device(device):
        scores = model(x)
        print(scores.shape)

test_TwoLayerFC()

# ### Keras Model Subclassing API: Three-Layer ConvNet
# Now it's your turn to implement a three-layer ConvNet using the `tf.keras.Model` API. Your model should have the same architecture used in Part II:
#
# 1. Convolutional layer with 5 x 5 kernels, with zero-padding of 2
# 2. ReLU nonlinearity
# 3. Convolutional layer with 3 x 3 kernels, with zero-padding of 1
# 4. ReLU nonlinearity
# 5. Fully-connected layer to give class scores
# 6. Softmax nonlinearity
#
# You should initialize the weights of your network using the same initialization method as was used in the two-layer network above.
#
# **Hint**: Refer to the documentation for `tf.keras.layers.Conv2D` and `tf.keras.layers.Dense`:
#
# +
# https://www.tensorflow.org/versions/r2.0/api_docs/python/tf/keras/layers/Dense

# In[19]:

class ThreeLayerConvNet(tf.keras.Model):
    def __init__(self, channel_1, channel_2, num_classes):
        super(ThreeLayerConvNet, self).__init__()
        ########################################################################
        # TODO: Implement the __init__ method for a three-layer ConvNet. You   #
        # should instantiate layer objects to be used in the forward pass.     #
        ########################################################################
        # *****START OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)*****

        initializer = tf.initializers.VarianceScaling(scale=2.0)

        self.conv1 = tf.keras.layers.Conv2D(filters=channel_1, kernel_size=(5,5), strides=1, padding='same', activation='relu', use_bias=True, kernel_initializer=initializer, bias_initializer='zeros')
        self.conv2 = tf.keras.layers.Conv2D(filters=channel_2, kernel_size=(3,3), strides=1, padding='same', activation='relu', use_bias=True, kernel_initializer=initializer, bias_initializer='zeros')
        self.flatten = tf.keras.layers.Flatten()
        self.fc1 = tf.keras.layers.Dense(num_classes, activation='softmax',
                                   kernel_initializer=initializer)

        # *****END OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)*****
        ########################################################################
        #                           END OF YOUR CODE                           #
        ########################################################################

    def call(self, x, training=False):
        scores = None
        ########################################################################
        # TODO: Implement the forward pass for a three-layer ConvNet. You      #
        # should use the layer objects defined in the __init__ method.         #
        ########################################################################
        # *****START OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)*****

        c1 = self.conv1(x)
        c2 = self.conv2(c1)
        c2 = self.flatten(c2)
        scores =  self.fc1(c2)

        # *****END OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)*****
        ########################################################################
        #                           END OF YOUR CODE                           #
        ########################################################################
        return scores

# Once you complete the implementation of the `ThreeLayerConvNet` above you can run the following to ensure that your implementation does not crash and produces outputs of the expected shape.

# In[20]:

def test_ThreeLayerConvNet():
    channel_1, channel_2, num_classes = 12, 8, 10
    model = ThreeLayerConvNet(channel_1, channel_2, num_classes)
    with tf.device(device):
        x = tf.zeros((64, 3, 32, 32))
        scores = model(x)
        print(scores.shape)

test_ThreeLayerConvNet()

# ### Keras Model Subclassing API: Eager Training
#
# While keras models have a builtin training loop (using the `model.fit`), sometimes you need more customization. Here's an example, of a training loop implemented with eager execution.
#
# In particular, notice `tf.GradientTape`. Automatic differentiation is used in the backend for implementing backpropagation in frameworks like TensorFlow. During eager execution, `tf.GradientTape` is used to trace operations for computing gradients later. A particular `tf.GradientTape` can only compute one gradient; subsequent calls to tape will throw a runtime error.
#
# TensorFlow 2.0 ships with easy-to-use built-in metrics under `tf.keras.metrics` module. Each metric is an object, and we can use `update_state()` to add observations and `reset_state()` to clear all observations. We can get the current result of a metric by calling `result()` on the metric object.

# In[21]:

def train_part34(model_init_fn, optimizer_init_fn, num_epochs=1, is_training=False):
    """
    Simple training loop for use with models defined using tf.keras. It trains
    a model for one epoch on the CIFAR-10 training set and periodically checks
    accuracy on the CIFAR-10 validation set.

    Inputs:
    - model_init_fn: A function that takes no parameters; when called it
      constructs the model we want to train: model = model_init_fn()
    - optimizer_init_fn: A function which takes no parameters; when called it
      constructs the Optimizer object we will use to optimize the model:
      optimizer = optimizer_init_fn()
    - num_epochs: The number of epochs to train for

    Returns: Nothing, but prints progress during trainingn
    """
    with tf.device(device):

        # Compute the loss like we did in Part II
        loss_fn = tf.keras.losses.SparseCategoricalCrossentropy()

        model = model_init_fn()
        optimizer = optimizer_init_fn()

        train_loss = tf.keras.metrics.Mean(name='train_loss')
        train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='train_accuracy')

        val_loss = tf.keras.metrics.Mean(name='val_loss')
        val_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='val_accuracy')

        t = 0
        for epoch in range(num_epochs):

            # Reset the metrics - https://www.tensorflow.org/alpha/guide/migration_guide#new-style_metrics
            train_loss.reset_states()
            train_accuracy.reset_states()

            for x_np, y_np in train_dset:
                with tf.GradientTape() as tape:

                    # Use the model function to build the forward pass.
                    scores = model(x_np, training=is_training)
                    loss = loss_fn(y_np, scores)

                    gradients = tape.gradient(loss, model.trainable_variables)
                    optimizer.apply_gradients(zip(gradients, model.trainable_variables))

                    # Update the metrics
                    train_loss.update_state(loss)
                    train_accuracy.update_state(y_np, scores)

                    if t % print_every == 0:
                        val_loss.reset_states()
                        val_accuracy.reset_states()
                        for test_x, test_y in val_dset:
                            # During validation at end of epoch, training set to False
                            prediction = model(test_x, training=False)
                            t_loss = loss_fn(test_y, prediction)

                            val_loss.update_state(t_loss)
                            val_accuracy.update_state(test_y, prediction)

                        template = 'Iteration {}, Epoch {}, Loss: {}, Accuracy: {}, Val Loss: {}, Val Accuracy: {}'
                        print (template.format(t, epoch+1,
                                             train_loss.result(),
                                             train_accuracy.result()*100,
                                             val_loss.result(),
                                             val_accuracy.result()*100))
                    t += 1

# ### Keras Model Subclassing API: Train a Two-Layer Network
# We can now use the tools defined above to train a two-layer network on CIFAR-10. We define the `model_init_fn` and `optimizer_init_fn` that construct the model and optimizer respectively when called. Here we want to train the model using stochastic gradient descent with no momentum, so we construct a `tf.keras.optimizers.SGD` function; you can [read about it here](https://www.tensorflow.org/versions/r2.0/api_docs/python/tf/optimizers/SGD).
#
# You don't need to tune any hyperparameters here, but you should achieve validation accuracies above 40% after one epoch of training.

# In[22]:

hidden_size, num_classes = 4000, 10
learning_rate = 1e-2

def model_init_fn():
    return TwoLayerFC(hidden_size, num_classes)

def optimizer_init_fn():
    return tf.keras.optimizers.SGD(learning_rate=learning_rate)

train_part34(model_init_fn, optimizer_init_fn)

# ### Keras Model Subclassing  API: Train a Three-Layer ConvNet
# Here you should use the tools we've defined above to train a three-layer ConvNet on CIFAR-10. Your ConvNet should use 32 filters in the first convolutional layer and 16 filters in the second layer.
#
# To train the model you should use gradient descent with Nesterov momentum 0.9.
#
# **HINT**: https://www.tensorflow.org/versions/r2.0/api_docs/python/tf/optimizers/SGD
#
# You don't need to perform any hyperparameter tuning, but you should achieve validation accuracies above 50% after training for one epoch.

# In[23]:

learning_rate = 3e-3
channel_1, channel_2, num_classes = 32, 16, 10

def model_init_fn():
    model = None
    ############################################################################
    # TODO: Complete the implementation of model_fn.                           #
    ############################################################################
    # *****START OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)*****

    model = ThreeLayerConvNet(channel_1, channel_2, num_classes)

    # *****END OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)*****
    ############################################################################
    #                           END OF YOUR CODE                               #
    ############################################################################
    return model

def optimizer_init_fn():
    optimizer = None
    ############################################################################
    # TODO: Complete the implementation of model_fn.                           #
    ############################################################################
    # *****START OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)*****

    optimizer = tf.keras.optimizers.SGD(learning_rate=learning_rate, momentum=0.9, nesterov=True)

    # *****END OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)*****
    ############################################################################
    #                           END OF YOUR CODE                               #
    ############################################################################
    return optimizer

train_part34(model_init_fn, optimizer_init_fn)

# # Part IV: Keras Sequential API
# In Part III we introduced the `tf.keras.Model` API, which allows you to define models with any number of learnable layers and with arbitrary connectivity between layers.
#
# However for many models you don't need such flexibility - a lot of models can be expressed as a sequential stack of layers, with the output of each layer fed to the next layer as input. If your model fits this pattern, then there is an even easier way to define your model: using `tf.keras.Sequential`. You don't need to write any custom classes; you simply call the `tf.keras.Sequential` constructor with a list containing a sequence of layer objects.
#
# One complication with `tf.keras.Sequential` is that you must define the shape of the input to the model by passing a value to the `input_shape` of the first layer in your model.
#
# ### Keras Sequential API: Two-Layer Network
# In this subsection, we will rewrite the two-layer fully-connected network using `tf.keras.Sequential`, and train it using the training loop defined above.
#
# You don't need to perform any hyperparameter tuning here, but you should see validation accuracies above 40% after training for one epoch.

# In[24]:

learning_rate = 1e-2

def model_init_fn():
    input_shape = (32, 32, 3)
    hidden_layer_size, num_classes = 4000, 10
    initializer = tf.initializers.VarianceScaling(scale=2.0)
    layers = [
        tf.keras.layers.Flatten(input_shape=input_shape),
        tf.keras.layers.Dense(hidden_layer_size, activation='relu',
                              kernel_initializer=initializer),
        tf.keras.layers.Dense(num_classes, activation='softmax',
                              kernel_initializer=initializer),
    ]
    model = tf.keras.Sequential(layers)
    return model

def optimizer_init_fn():
    return tf.keras.optimizers.SGD(learning_rate=learning_rate)

train_part34(model_init_fn, optimizer_init_fn)

# ### Abstracting Away the Training Loop
# In the previous examples, we used a customised training loop to train models (e.g. `train_part34`). Writing your own training loop is only required if you need more flexibility and control during training your model. Alternately, you can also use  built-in APIs like `tf.keras.Model.fit()` and `tf.keras.Model.evaluate` to train and evaluate a model. Also remember to configure your model for training by calling `tf.keras.Model.compile.
#
# You don't need to perform any hyperparameter tuning here, but you should see validation and test accuracies above 42% after training for one epoch.

# In[25]:

model = model_init_fn()
model.compile(optimizer=tf.keras.optimizers.SGD(learning_rate=learning_rate),
              loss='sparse_categorical_crossentropy',
              metrics=[tf.keras.metrics.sparse_categorical_accuracy])
model.fit(X_train, y_train, batch_size=64, epochs=1, validation_data=(X_val, y_val))
model.evaluate(X_test, y_test)

# ### Keras Sequential API: Three-Layer ConvNet
# Here you should use `tf.keras.Sequential` to reimplement the same three-layer ConvNet architecture used in Part II and Part III. As a reminder, your model should have the following architecture:
#
# 1. Convolutional layer with 32 5x5 kernels, using zero padding of 2
# 2. ReLU nonlinearity
# 3. Convolutional layer with 16 3x3 kernels, using zero padding of 1
# 4. ReLU nonlinearity
# 5. Fully-connected layer giving class scores
# 6. Softmax nonlinearity
#
# You should initialize the weights of the model using a `tf.initializers.VarianceScaling` as above.
#
# You should train the model using Nesterov momentum 0.9.
#
# You don't need to perform any hyperparameter search, but you should achieve accuracy above 45% after training for one epoch.

# In[26]:

def model_init_fn():
    model = None
    ############################################################################
    # TODO: Construct a three-layer ConvNet using tf.keras.Sequential.         #
    ############################################################################
    # *****START OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)*****
    channel_1, channel_2, num_classes = 32, 16, 10
    initializer = tf.initializers.VarianceScaling(scale=2.0)
    conv1 = tf.keras.layers.Conv2D(filters=channel_1, kernel_size=(5,5), strides=1, padding='same', activation='relu', use_bias=True, kernel_initializer=initializer, bias_initializer='zeros')
    conv2 = tf.keras.layers.Conv2D(filters=channel_2, kernel_size=(3,3), strides=1, padding='same', activation='relu', use_bias=True, kernel_initializer=initializer, bias_initializer='zeros')
    flatten = tf.keras.layers.Flatten()
    fc1 = tf.keras.layers.Dense(num_classes, activation='softmax', kernel_initializer=initializer)
    layers = [conv1, conv2, flatten, fc1]
    model = tf.keras.Sequential(layers)
    # *****END OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)*****
    ############################################################################
    #                            END OF YOUR CODE                              #
    ############################################################################
    return model

learning_rate = 5e-4
def optimizer_init_fn():
    optimizer = None
    ############################################################################
    # TODO: Complete the implementation of model_fn.                           #
    ############################################################################
    # *****START OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)*****

    optimizer = tf.keras.optimizers.SGD(learning_rate=learning_rate, momentum=0.9, nesterov=True)

    # *****END OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)*****
    ############################################################################
    #                           END OF YOUR CODE                               #
    ############################################################################
    return optimizer

train_part34(model_init_fn, optimizer_init_fn)

# We will also train this model with the built-in training loop APIs provided by TensorFlow.

# In[27]:

model = model_init_fn()
model.compile(optimizer='sgd',
              loss='sparse_categorical_crossentropy',
              metrics=[tf.keras.metrics.sparse_categorical_accuracy])
model.fit(X_train, y_train, batch_size=64, epochs=1, validation_data=(X_val, y_val))
model.evaluate(X_test, y_test)

# ##  Part IV: Functional API
# ### Demonstration with a Two-Layer Network
#
# In the previous section, we saw how we can use `tf.keras.Sequential` to stack layers to quickly build simple models. But this comes at the cost of losing flexibility.
#
# Often we will have to write complex models that have non-sequential data flows: a layer can have **multiple inputs and/or outputs**, such as stacking the output of 2 previous layers together to feed as input to a third! (Some examples are residual connections and dense blocks.)
#
# In such cases, we can use Keras functional API to write models with complex topologies such as:
#
#  1. Multi-input models
#  2. Multi-output models
#  3. Models with shared layers (the same layer called several times)
#  4. Models with non-sequential data flows (e.g. residual connections)
#
# Writing a model with Functional API requires us to create a `tf.keras.Model` instance and explicitly write input tensors and output tensors for this model.

# In[28]:

def two_layer_fc_functional(input_shape, hidden_size, num_classes):
    initializer = tf.initializers.VarianceScaling(scale=2.0)
    inputs = tf.keras.Input(shape=input_shape)
    flattened_inputs = tf.keras.layers.Flatten()(inputs)
    fc1_output = tf.keras.layers.Dense(hidden_size, activation='relu',
                                 kernel_initializer=initializer)(flattened_inputs)
    scores = tf.keras.layers.Dense(num_classes, activation='softmax',
                             kernel_initializer=initializer)(fc1_output)

    # Instantiate the model given inputs and outputs.
    model = tf.keras.Model(inputs=inputs, outputs=scores)
    return model

def test_two_layer_fc_functional():
    """ A small unit test to exercise the TwoLayerFC model above. """
    input_size, hidden_size, num_classes = 50, 42, 10
    input_shape = (50,)

    x = tf.zeros((64, input_size))
    model = two_layer_fc_functional(input_shape, hidden_size, num_classes)

    with tf.device(device):
        scores = model(x)
        print(scores.shape)

test_two_layer_fc_functional()

# ### Keras Functional API: Train a Two-Layer Network
# You can now train this two-layer network constructed using the functional API.
#
# You don't need to perform any hyperparameter tuning here, but you should see validation accuracies above 40% after training for one epoch.

# In[29]:

input_shape = (32, 32, 3)
hidden_size, num_classes = 4000, 10
learning_rate = 1e-2

def model_init_fn():
    return two_layer_fc_functional(input_shape, hidden_size, num_classes)

def optimizer_init_fn():
    return tf.keras.optimizers.SGD(learning_rate=learning_rate)

train_part34(model_init_fn, optimizer_init_fn)

# # Part V: CIFAR-10 open-ended challenge
#
# In this section you can experiment with whatever ConvNet architecture you'd like on CIFAR-10.
#
# You should experiment with architectures, hyperparameters, loss functions, regularization, or anything else you can think of to train a model that achieves **at least 70%** accuracy on the **validation** set within 10 epochs. You can use the built-in train function, the `train_part34` function from above, or implement your own training loop.
#
# Describe what you did at the end of the notebook.
#
# ### Some things you can try:
# - **Filter size**: Above we used 5x5 and 3x3; is this optimal?
# - **Number of filters**: Above we used 16 and 32 filters. Would more or fewer do better?
# - **Pooling**: We didn't use any pooling above. Would this improve the model?
# - **Normalization**: Would your model be improved with batch normalization, layer normalization, group normalization, or some other normalization strategy?
# - **Network architecture**: The ConvNet above has only three layers of trainable parameters. Would a deeper model do better?
# - **Global average pooling**: Instead of flattening after the final convolutional layer, would global average pooling do better? This strategy is used for example in Google's Inception network and in Residual Networks.
# - **Regularization**: Would some kind of regularization improve performance? Maybe weight decay or dropout?
#
# ### NOTE: Batch Normalization / Dropout
# If you are using Batch Normalization and Dropout, remember to pass `is_training=True` if you use the `train_part34()` function. BatchNorm and Dropout layers have different behaviors at training and inference time. `training` is a specific keyword argument reserved for this purpose in any `tf.keras.Model`'s `call()` function. Read more about this here : https://www.tensorflow.org/versions/r2.0/api_docs/python/tf/keras/layers/BatchNormalization#methods
# https://www.tensorflow.org/versions/r2.0/api_docs/python/tf/keras/layers/Dropout#methods
#
# ### Tips for training
# For each network architecture that you try, you should tune the learning rate and other hyperparameters. When doing this there are a couple important things to keep in mind:
#
# - If the parameters are working well, you should see improvement within a few hundred iterations
# - Remember the coarse-to-fine approach for hyperparameter tuning: start by testing a large range of hyperparameters for just a few training iterations to find the combinations of parameters that are working at all.
# - Once you have found some sets of parameters that seem to work, search more finely around these parameters. You may need to train for more epochs.
# - You should use the validation set for hyperparameter search, and save your test set for evaluating your architecture on the best parameters as selected by the validation set.
#
# ### Going above and beyond
# If you are feeling adventurous there are many other features you can implement to try and improve your performance. You are **not required** to implement any of these, but don't miss the fun if you have time!
#
# - Alternative optimizers: you can try Adam, Adagrad, RMSprop, etc.
# - Alternative activation functions such as leaky ReLU, parametric ReLU, ELU, or MaxOut.
# - Model ensembles
# - Data augmentation
# - New Architectures
#   - [ResNets](https://arxiv.org/abs/1512.03385) where the input from the previous layer is added to the output.
#   - [DenseNets](https://arxiv.org/abs/1608.06993) where inputs into previous layers are concatenated together.
#   - [This blog has an in-depth overview](https://chatbotslife.com/resnets-highwaynets-and-densenets-oh-my-9bb15918ee32)
#
# ### Have fun and happy training!

# In[30]:

class CustomConvNet(tf.keras.Model):
    def __init__(self):
        super(CustomConvNet, self).__init__()
        ############################################################################
        # TODO: Construct a model that performs well on CIFAR-10                   #
        ############################################################################
        # *****START OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)*****

        initializer = tf.initializers.VarianceScaling(scale=2.0)

        self.conv1 = tf.keras.layers.Conv2D(filters=32, kernel_size=(5,5), strides=1, padding='same', activation='linear', use_bias=True, kernel_initializer=initializer, bias_initializer='zeros')
        self.batch1 = tf.keras.layers.BatchNormalization()
        self.relu1 = tf.keras.layers.ReLU()
        self.conv1 = tf.keras.layers.Conv2D(filters=32, kernel_size=(5,5), strides=1, padding='same', activation='linear', use_bias=True, kernel_initializer=initializer, bias_initializer='zeros')
        self.batch1 = tf.keras.layers.BatchNormalization()
        self.relu1 = tf.keras.layers.ReLU()
        self.drop1 = tf.keras.layers.Dropout(0.1)
        self.pool = tf.keras.layers.MaxPool2D()

        self.conv2 = tf.keras.layers.Conv2D(filters=16, kernel_size=(3,3), strides=1, padding='same', activation='linear', use_bias=True, kernel_initializer=initializer, bias_initializer='zeros')
        self.batch2 = tf.keras.layers.BatchNormalization()
        self.relu2 = tf.keras.layers.ReLU()
        self.drop2 = tf.keras.layers.Dropout(0.1)
        self.flatten = tf.keras.layers.Flatten()

        # self.fc1 = tf.keras.layers.Dense(num_classes, activation='linear',
        #                            kernel_initializer=initializer)
        # self.fcbatch1 = tf.keras.layers.BatchNormalization()
        # self.fcrelu1 = tf.keras.layers.ReLU()
        # self.fcdrop1 = tf.keras.layers.Dropout(0.1)
        self.fc1 = tf.keras.layers.Dense(num_classes, activation='softmax',
                                   kernel_initializer=initializer)

        # *****END OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)*****
        ############################################################################
        #                            END OF YOUR CODE                              #
        ############################################################################

    def call(self, input_tensor, training=False):
        ############################################################################
        # TODO: Construct a model that performs well on CIFAR-10                   #
        ############################################################################
        # *****START OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)*****

        x = self.conv1(input_tensor)
        x = self.batch1(x, training=training)
        x = self.relu1(x)
        x = self.drop1(x, training=training)
        x = self.conv2(x)
        x = self.batch2(x, training=training)
        x = self.relu2(x)
        x = self.drop2(x, training=training)
        x = self.flatten(x)
        x = self.fc1(x)
        # x = self.fcbatch1(x)
        # x = self.fcrelu1(x)
        # x = self.fcdrop1(x)
        # x = self.fc2(x)

        # *****END OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)*****
        ############################################################################
        #                            END OF YOUR CODE                              #
        ############################################################################

        return x

# device = '/device:GPU:0'   # Change this to a CPU/GPU as you wish!
device = '/cpu:0'        # Change this to a CPU/GPU as you wish!
print_every = 700
num_epochs = 10

model = CustomConvNet()

def model_init_fn():
    return CustomConvNet()

def optimizer_init_fn():
    learning_rate = 1e-3
    return tf.keras.optimizers.Adam(learning_rate)

train_part34(model_init_fn, optimizer_init_fn, num_epochs=num_epochs, is_training=True)

# ## Describe what you did
#
# In the cell below you should write an explanation of what you did, any additional features that you implemented, and/or any graphs that you made in the process of training and evaluating your network.

# THIS OVERFITS, I'll come back later when I have a GPU.


r"""
Random testing

Some Sage modules do random testing in their doctests; that is, they
construct test cases using a random number generator.  To get the
broadest possible test coverage, we want everybody who runs the
doctests to use a different random seed; but we also want to be able
to reproduce the problems when debugging.  This module provides a
decorator to help write random testers that meet these goals.
"""

from functools import wraps

def random_testing(fn):
    r"""
    This decorator helps create random testers.  These can be run as
    part of the standard Sage test suite; everybody who runs the test
    will use a different random number seed, so many different random
    tests will eventually be run.

    INPUT:

        - ``fn`` - The function that we are wrapping for random testing.

    The resulting function will take two additional arguments, *seed*
    (default ``None``) and *print_seed* (default ``False``).  The
    result will set the random number seed to the given seed value (or
    to a truly random value, if *seed* is not specified), then call
    the original function.  If *print_seed* is true, then the seed will
    be printed before calling the original function.  If the original
    function raises an exception, then the random seed that was used
    will be displayed, along with a message entreating the user to
    submit a bug report.  All other arguments will be passed through
    to the original function.

    Here is a set of recommendations for using this wrapper.

    The function to be tested should take arguments specifying the
    difficulty of the test (size of the test cases, number of
    iterations, etc.), as well as an argument *verbose* (defaulting to
    false).  With *verbose* true, it should print the values being
    tested.  Suppose ``test_foo()`` takes an argument for number of
    iterations.  Then the doctests could be::

        test_foo(2, verbose=True, seed=0)
        test_foo(10)
        test_foo(100) # long time

    The first doctest, with the specified seed and ``verbose=True``, simply
    verifies that the tests really are reproducible (that ``test_foo``
    is correctly using the :mod:`randstate` framework).  The next two tests
    use truly random seeds, and will print out the seed used if the test
    fails (raises an exception).

    If you want a very long-running test using this setup, you should do
    something like::

        for _ in xrange(10^10): test_foo(100)

    instead of::

        test_foo(10^12)

    If the test fails after several hours, the latter snippet would
    make you rerun the test for several hours while reproducing and
    debugging the problem.  With the former snippet, you only need to
    rerun ``test_foo(100)`` with a known-failing random seed.

    See :func:`sage.misc.random_testing.test_add_commutes` for a
    simple example using this decorator, and :mod:`sage.rings.tests`
    for realistic uses.

    Setting *print_seed* to true is useless in doctests, because the
    random seed printed will never match the expected doctest result
    (and using ``# random`` means the doctest framework will never
    report an error even if one happens).  However, it is useful if
    you have a random test that sometimes segfaults.  The normal
    print-the-random-seed-on-exceptions won't work then, so you can
    run::

        while True: test_foo(print_seed=True)

    and look at the last seed that was printed before it crashed.

    TESTS::

        sage: from sage.misc.random_testing import random_testing
        sage: def foo(verbose=False):
        ...       'oh look, a docstring'
        ...       n = ZZ.random_element(2^50)
        ...       if verbose:
        ...           print "Random value: %s" % n
        ...       assert(n == 49681376900427)
        sage: foo = random_testing(foo)
        sage: foo(seed=0, verbose=True)
        Random value: 49681376900427
        sage: foo(seed=15, verbose=True)
        Random value: 1049538412064764
        Random testing has revealed a problem in foo
        Please report this bug!  You may be the first
        person in the world to have seen this problem.
        Please include this random seed in your bug report:
        Random seed: 15
        AssertionError()
        sage: foo() # random
        Random testing has revealed a problem in foo
        Please report this bug!  You may be the first
        person in the world to have seen this problem.
        Please include this random seed in your bug report:
        Random seed: 272500700755151445506092479579811710040
        AssertionError()
        sage: foo.__doc__
        'oh look, a docstring'
        sage: foo.__name__
        'foo'
        sage: def bar(): pass
        sage: bar = random_testing(bar)
        sage: bar(print_seed=True) # random
        Random seed: 262841091890156346923539765543814146051
    """
    from sage.misc.randstate import seed, initial_seed
    from sys import stdout
    @wraps(fn)
    def wrapped_fun(*args, **kwargs):
        arg_seed = None
        if 'seed' in kwargs:
            arg_seed = kwargs['seed']
            del kwargs['seed']
        with seed(arg_seed):
            used_seed = initial_seed()
            if 'print_seed' in kwargs:
                if kwargs['print_seed']:
                    print("Random seed: {}".format(used_seed))
                    del kwargs['print_seed']
                # I don't know if this line is necessary, but it can't
                # hurt; and it would be a real pity to lose the
                # information you need to reproduce a segfault because
                # it was missing...
                stdout.flush()
            try:
                fn(*args, **kwargs)
            except Exception as e:
                # We treat any sort of Exception as a doctest
                # failure.  (We have to eat the exception, because if
                # doctesting sees an exception, it doesn't display
                # whatever was printed before the exception happened
                # -- so the text we print here would be lost.)  Note
                # that KeyboardInterrupt is not an Exception, so
                # pressing Control-C doesn't print this message.
                print("Random testing has revealed a problem in " + fn.__name__)
                print("Please report this bug!  You may be the first")
                print("person in the world to have seen this problem.")
                print("Please include this random seed in your bug report:")
                print("Random seed: {}".format(used_seed))
                print(repr(e))
    return wrapped_fun

@random_testing
def test_add_commutes(trials, verbose=False):
    r"""
    This is a simple demonstration of the :func:`random_testing` decorator and
    its recommended usage.

    We test that addition is commutative over rationals.

    EXAMPLES::

        sage: from sage.misc.random_testing import test_add_commutes
        sage: test_add_commutes(2, verbose=True, seed=0)
        a == -4, b == 0 ...
        Passes!
        a == -1/2, b == -1/95 ...
        Passes!
        sage: test_add_commutes(10)
        sage: test_add_commutes(1000) # long time
    """
    from sage.rings.all import QQ
    for _ in xrange(trials):
        a = QQ.random_element()
        b = QQ.random_element()
        if verbose:
            print("a == {}, b == {} ...".format(a, b))
        assert(a+b == b+a)
        if verbose:
            print("Passes!")

@random_testing
def test_add_is_mul(trials, verbose=False):
    r"""
    This example demonstrates a failing :func:`random_testing` test,
    and shows how to reproduce the error.

    DO NOT USE THIS AS AN EXAMPLE OF HOW TO USE
    :func:`random_testing`!  Instead, look at
    :func:`sage.misc.random_testing.test_add_commutes`.

    We test that ``a+b == a*b``, for *a*, *b* rational.  This is of
    course false, so the test will almost always fail.

    EXAMPLES::

        sage: from sage.misc.random_testing import test_add_is_mul

    We start by testing that we get reproducible results when setting
    *seed* to 0.

    ::

        sage: test_add_is_mul(2, verbose=True, seed=0)
        a == -4, b == 0 ...
        Random testing has revealed a problem in test_add_is_mul
        Please report this bug!  You may be the first
        person in the world to have seen this problem.
        Please include this random seed in your bug report:
        Random seed: 0
        AssertionError()

    Normally in a ``@random_testing`` doctest, we would leave off the
    ``verbose=True`` and the ``# random``.  We put it in here so that we can
    verify that we are seeing the exact same error when we reproduce
    the error below.

    ::

        sage: test_add_is_mul(10, verbose=True) # random
        a == -2/7, b == 1 ...
        Random testing has revealed a problem in test_add_is_mul
        Please report this bug!  You may be the first
        person in the world to have seen this problem.
        Please include this random seed in your bug report:
        Random seed: 216390410596009428782506007128692114173
        AssertionError()

    OK, now assume that some user has reported a
    :func:`test_add_is_mul` failure.  We can specify the same
    *random_seed* that was found in the bug report, and we will get the
    exact same failure so that we can debug the "problem".

    ::

        sage: test_add_is_mul(10, verbose=True, seed=216390410596009428782506007128692114173)
        a == -2/7, b == 1 ...
        Random testing has revealed a problem in test_add_is_mul
        Please report this bug!  You may be the first
        person in the world to have seen this problem.
        Please include this random seed in your bug report:
        Random seed: 216390410596009428782506007128692114173
        AssertionError()
    """
    from sage.rings.all import QQ
    for _ in xrange(trials):
        a = QQ.random_element()
        b = QQ.random_element()
        if verbose:
            print("a == {}, b == {} ...".format(a, b))
        assert(a+b == a*b)
        if verbose:
            print("Passes!")


import time
import pandas as pd
import numpy as np

CITY_DATA = { 'chicago': 'chicago.csv',
              'new york city': 'new_york_city.csv',
              'washington': 'washington.csv' }

month_list = ['January', 'February', 'March', 'April', 'May', 'June']
day_list = ['Monday', 'Tuesday', 'Wednesday', 'Thursday', 'Friday', 'Saturday', 'Sunday']

def get_filters():
    """
    Asks user to specify a city, month, and day to analyze.

    Returns:
        (str) city - name of the city to analyze
        (str) month - name of the month to filter by, or "all" to apply no month filter
        (str) day - name of the day of week to filter by, or "all" to apply no day filter
    """
    print('Hello! Let\'s explore some US bikeshare data!')
    city_list = list(CITY_DATA.keys())
    # This block gets the user's desired city, formats it, and checks it against the city_list variable.
    city = input('Would you like to investigate Chicago, Washington, or New York City? ').lower()
    while city not in city_list:
        print('That\'s not a valid city for this program.')
        city = input('Would you like to investigate Chicago, Washington, or New York City? ').lower()
    # This block asks if the month needs to be filtered.  If yes, it asks for the month name and checks it against the month_list variable.
    # If no, it sets month = 'all'.
    month_filter = input('Would you like to filter data by month? Enter yes or no: ').lower()
    while month_filter != 'yes' and month_filter != 'no':
        print('Not a valid input. Please specify yes or no')
        month_filter = input('Would you like to filter data by month? Enter yes or no: ').lower()
    if month_filter == 'yes':
        month = input('Filter by which month? This program has data for January through June: ').title()
        while month not in month_list:
            print('Not a valid month for this program.')
            month = input('Filter by which month? This program has data for January through June: ').title()
    else:
        month = 'all'
    # This block asks if the day needs to be filtered. If yes, it checks the user's input of a day name against the day_list variable.
    # If no, it sets the day variable = 'all'.
    day_filter = input('Would you like to filter data by day? Enter yes or no: ').lower()
    while day_filter != 'yes' and day_filter != 'no':
        print('Please specify yes or no.')
        day_filter = input('Would you like to filter data by day? Enter yes or no: ').lower()
    if day_filter == 'yes':
        day = input('Filter by which day? ').title()
        while day not in day_list:
            print('Not a valid day.')
            day = input('Filter by which day? ').title()
    else:
        day = 'all'

    print('-'*40)
    return city, month, day

def load_data(city, month, day):
    """
    Loads data for the specified city and filters by month and day if applicable.

    Args:
        (str) city - name of the city to analyze
        (str) month - name of the month to filter by, or "all" to apply no month filter
        (str) day - name of the day of week to filter by, or "all" to apply no day filter
    Returns:
        df - Pandas DataFrame containing city data filtered by month and day
    """
    # This block loads the user's desired city's table into a DataFrame and uses Pandas datetime functionality to append additional columns
    # needed later for time based statistics.  It then filters the DataFrame by month and day if specified by the user.
    df = pd.read_csv(CITY_DATA[city])
    df['Start Time'] = pd.to_datetime(df['Start Time'])
    df['month'] = df['Start Time'].dt.month
    df['day_of_week'] = df['Start Time'].dt.weekday_name
    df['hour'] = df['Start Time'].dt.hour
    if month != 'all':
        month = month_list.index(month)+1
        df = df[df['month'] == month]

    if day != 'all':
        df = df[df['day_of_week'] == day]

    return df

def condisplay_message(city, month, day):
    """
    Prints a line with the user's selected filters in a string message.

    Args:
        (str) city - name of the city to analyze
        (str) month - name of the month to filter by, or "all" to apply no month filter
        (str) day - name of the day of week to filter by, or "all" to apply no day filter

    """
    # This block prints a string with the user's selected filters. It's not completely necessary,
    # but is useful as an easily read reminder of the filters used, located at the top of the program's output.
    if month == 'all' and day == 'all':
        print('For all months and all days in {}:\n'.format(city.title()))
    elif month == 'all' and day != 'all':
        print('For all months and all {}s in {}:\n'.format(day, city.title()))
    elif month != 'all' and day == 'all':
        print('For all days in the month of {} in {}:\n'.format(month, city.title()))
    else:
        print('For all {}s in the month of {} in {}:\n'.format(day, month, city.title()))

def time_stats(df, month, day):
    """
    Displays statistics on the most frequent times of travel.

    Args:
        (DataFrame) df - pre-filtered Pandas DataFrame
        (str) month - name of the month to filter by, or "all" to apply no month filter
        (str) day - name of the day of week to filter by, or "all" to apply no day filter
    """

    print('\nCalculating The Most Frequent Times of Travel...\n')
    start_time = time.time()

    # This block finds and prints the most popular month statistics if a month filter was not selected by the user.
    if month == 'all':
        month_value_counts = df['month'].value_counts()
        pop_month = month_list[month_value_counts.keys()[0]-1]
        pop_month_count = month_value_counts.iloc[0]
        print('The most popular month was {}, with {} bikerides.'.format(pop_month, pop_month_count))

    # This block finds and prints the most popular day statistics if a day filter was not selected by the user.
    if day == 'all':
        day_value_counts = df['day_of_week'].value_counts()
        pop_day = day_value_counts.keys()[0]
        pop_day_count = day_value_counts.iloc[0]
        print('The most popular day was {}, with {} bikerides.'.format(pop_day, pop_day_count))

    # This block finds and prints the most popular hour statistics for any and all filters selected by the user.
    hour_value_counts = df['hour'].value_counts()
    pop_hour = hour_value_counts.keys()[0]
    pop_hour_count = hour_value_counts.iloc[0]
    print('The most popular hour was {} o\'clock with {} bikerides.'.format(pop_hour, pop_hour_count))

    print("\nThis took %s seconds." % (time.time() - start_time))
    print('-'*40)

def station_stats(df):
    """
    Displays statistics on the most popular stations and trip.

    Args:
    (DataFrame) df - pre-filtered Pandas DataFrame
    """

    print('\nCalculating The Most Popular Stations and Trip...\n')
    start_time = time.time()
    # Adds a column to the DataFrame to help find popular start to end trip locations.
    df['Trip'] = df['Start Station'].str.cat(df['End Station'], sep=' to ')

    # This block finds and prints the name of the most popular starting place, as well as total number of trips from that station.
    start_station_counts = df['Start Station'].value_counts()
    pop_start_station = start_station_counts.keys()[0]
    pop_start_station_count = start_station_counts.iloc[0]
    print('The most popular start station was {} with {} trips starting there.'.format(pop_start_station, pop_start_station_count))

    # This block finds and prints the name of the most popular ending place, as well as total number of trips that end at that station.
    end_station_counts = df['End Station'].value_counts()
    pop_end_station = end_station_counts.keys()[0]
    pop_end_station_count = end_station_counts.iloc[0]
    print('The most popular ending station was {} with {} trips ending there.'.format(pop_end_station, pop_end_station_count))

    # This block finds and prints the most popular trip, start to end, based on the new column created.
    trip_counts = df['Trip'].value_counts()
    pop_trip = trip_counts.keys()[0]
    pop_trip_count = trip_counts.iloc[0]
    print('The most popular trip was {} with {} trips.'.format(pop_trip, pop_trip_count))

    print("\nThis took %s seconds." % (time.time() - start_time))
    print('-'*40)

def trip_time_units(trip_time_type, trip_time):
    """
    This program helps the trip_duration_stats() function decide what time units to use. The function checks if the time amount is
    less than or equal to 2 times the next biggest unit. If it is, it checks the next size unit, if it's not, it prints the time with
    the current unit. For example if its 1000 seconds, it checks if its less than or equal to 2 of the next biggest unit (minutes).
    It is not, so it checks if its less than or equal to 2 of the next biggest unit (hours). It is so it stops at minutes and prints
    the desired response.
    It also only prints to 2 decimal places, for more easily read numbers.

    Args:
        (str) trip_time_type - indicates the type of travel time being calculated for the printed output. Either 'total', or 'mean'
        (int) trip_time - trip time in seconds
    """
    if trip_time <= 120:
        print('The {} travel duration was {:.2f} seconds.'.format(trip_time_type, trip_time))
    elif trip_time <= 120*60:
        print('The {} travel duration was {:.2f} minutes'.format(trip_time_type, trip_time/60))
    elif trip_time <= 48*60*60:
        print('The {} travel duration was {:.2f} hours'.format(trip_time_type, trip_time/(60*60)))
    elif trip_time <= 60*24*60*60:
        print('The {} travel duration was {:.2f} days'.format(trip_time_type, trip_time/(60*60*24)))
    elif trip_time <= 48*30*24*60*60:
        print('The {} travel duration was {:.2f} months'.format(trip_time_type, trip_time/(60*60*24*30)))
    else:
        print('The {} travel duration was {:.2f} years'.format(trip_time_type, trip_time/(60*60*24*365)))

def trip_duration_stats(df):
    """
    Displays statistics on the total and average trip duration.

    Args:
    (DataFrame) df - pre-filtered Pandas DataFrame

    """

    print('\nCalculating Trip Duration...\n')
    start_time = time.time()

    # This block finds the total trip duration of all trips taken for the user's selected filters. It then displays that time
    # in seconds, minutes, hours, days, months, or years, rounded to 2 decimal places.
    total_trip_time = df['Trip Duration'].sum()
    total_trip_type = 'total'
    trip_time_units(total_trip_type, total_trip_time)

    # This block finds the mean trip duration for all trips taken for the user's selected filters. It then displays that time
    # in seconds, minutes, hours, days, months, or years, rounded to 2 decimal places.
    mean_trip_time = df['Trip Duration'].mean()
    mean_trip_type = 'mean'
    trip_time_units(mean_trip_type, mean_trip_time)

    print("\nThis took %s seconds." % (time.time() - start_time))
    print('-'*40)

def user_stats(df, city):
    """Displays statistics on bikeshare users.

    Args:
        (DataFrame) df - pre-filtered Pandas DataFrame
        (str) city - name of the city to analyze

    """

    print('\nCalculating User Stats...\n')
    start_time = time.time()

    # Prints a Pandas Series of all user types and their counts, as well as filling all NaNs with N/A (not available).
    print(df['User Type'].fillna('N/A').value_counts())
    print()

    # Because the Washington data table doesn't have a gender column, a try/except function is used to handle any KeyErrors.
    # If there is no KeyError exception thrown up, this block will print a Pandas series with a count of gender types (male and female)
    # as well as a column for N/A, filled with any NaN values.
    try:
        print(df['Gender'].fillna('N/A').value_counts())
        print()
    except KeyError:
        print('{} has no gender data available.'.format(city.title()))
        print()

    # The Washington data table also has no birth year column, so another try/except block is used. If no KeyError comes up, this block will find and
    # print values for the oldest birth year, most recent birth year, and the most common birth year, for the user's selected filters.
    try:
        oldest_birth_year = int(df['Birth Year'].min())
        youngest_birth_year = int(df['Birth Year'].max())
        common_birth_year = int(df['Birth Year'].mode()[0])
        print('The earliest user birth year is {}. The most recent user birth year is {}. The most common user birth year is {}.'.format(
            oldest_birth_year, youngest_birth_year, common_birth_year))
    except KeyError:
        print('{} has no birth year data available.'.format(city.title()))

    print("\nThis took %s seconds." % (time.time() - start_time))
    print('-'*40)

def display_data(df):
    """
    This function gives the user the option to see 5 rows of the raw data from their filtered DataFrame
    and then see 5 more rows until they decide to stop.

    Args:
        (DataFrame) df - the Pandas DataFrame loaded based on the user's filters.
    """

    # This code gets users yes or no input on seeing 5 rows of data, then asks again and continuously shows
    # the next 5 rows until the user inputs no which stops the function.
    display_input = input('Would you like to see the raw data (5 rows)? Enter yes or no: ').lower()
    while display_input != 'yes' and display_input != 'no':
        display_input = input('Not a valid input. Please enter yes or no: ').lower()
    if display_input == 'yes':
        i=0
        while True:
            print(df.iloc[i:i+5,:])
            display_input = input('Would you like to see 5 more rows? Enter yes or no: ').lower()
            while display_input != 'yes' and display_input != 'no':
                display_input = input('Not a valid input. Please enter yes or no: ').lower()
            if display_input == 'yes':
                i += 5
            else:
                break

def main():
    while True:
        city, month, day = get_filters()
        df = load_data(city, month, day)
        condisplay_message(city, month, day)
        time_stats(df, month, day)
        station_stats(df)
        trip_duration_stats(df)
        user_stats(df, city)
        display_data(df)

        restart = input('\nWould you like to restart? Enter yes or no.\n')
        if restart.lower() != 'yes':
            break

if __name__ == "__main__":
	main()


#!/usr/bin/env python3

import unittest
from unittest.mock import patch
import numpy as np
import pandas as pd

from tmc import points

from tmc.utils import load, get_out, patch_helper

module_name="src.special_missing_values"
special_missing_values = load(module_name, "special_missing_values")
main = load(module_name, "main")
ph = patch_helper(module_name)

@points('p04-14.1')
class SpecialMissingValues(unittest.TestCase):

    def test_shape(self):
        df = special_missing_values()
        self.assertEqual(df.shape, (17,7), msg="Incorrect shape!")

    def test_columns(self):
        df = special_missing_values()
        np.testing.assert_array_equal(df.columns,
                                      ["Pos", "LW", "Title", "Artist", "Publisher", "Peak Pos", "WoC"],
                                      err_msg="Incorrect column names!")

    def test_called(self):
        with patch(ph("special_missing_values"), wraps=special_missing_values) as psmv,\
             patch(ph("pd.read_csv"), wraps=pd.read_csv) as prc:
            main()
            psmv.assert_called()
            prc.assert_called()

    def test_content(self):
        df = special_missing_values()
        np.testing.assert_array_equal(df["Pos"], [3,4,6,9,10,12,15,16,21,22,24,30,31,34,35,38,39],
                                   err_msg="The values in position column were incorrect!")

if __name__ == '__main__':
    unittest.main()


#
# File: __init__.py
#

from project0_py.helpers import *

## top-level submission file

### imports here
import numpy as np
###

def f(a, b):
	'''
	Function for adding two numbers
	Args:
		a (float): first number
		b (float): second number
	Returns:
		c (float): result
	'''

	## Your code goes here

	##

	return c