"""
A pure TensorFlow implementation of a neural network. This can be
used as a drop-in replacement for a Keras model.
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

from cleverhans.picklable_model import MLP, Conv2D, ReLU, Flatten, Linear
from cleverhans.picklable_model import Softmax


def make_basic_picklable_cnn(
    nb_filters=64, nb_classes=10, input_shape=(None, 28, 28, 1)
):
    """The model for the picklable models tutorial."""
    layers = [
        Conv2D(nb_filters, (8, 8), (2, 2), "SAME"),
        ReLU(),
        Conv2D(nb_filters * 2, (6, 6), (2, 2), "VALID"),
        ReLU(),
        Conv2D(nb_filters * 2, (5, 5), (1, 1), "VALID"),
        ReLU(),
        Flatten(),
        Linear(nb_classes),
        Softmax(),
    ]
    model = MLP(layers, input_shape)
    return model
