text: # Copyright 2019 Google LLC
  #
  # 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.
  
  # AutoEncoder - Pooling with Dense Layers/Hidden Units
  
  import tensorflow as tf
  from tensorflow.keras import Model, Input
  from tensorflow.keras.layers import Dense, ReLU, BatchNormalization, Reshape, Flatten
  
  import sys
  sys.path.append('../')
  from models_c import Composable
  
  class AutoEncoder(Composable):
      ''' Construct an AutoEncoder '''
      # metaparameter: number of filters per layer
      layers = [ {'n_nodes': 256 }, { 'n_nodes': 128 }, { 'n_nodes': 64 } ]
  
      def __init__(self, layers=None, input_shape=(32, 32, 3),
                   init_weights='he_normal', reg=None, relu=None, bias=True):
          ''' Construct an AutoEncoder
              input_shape : input shape to the autoencoder
              layers      : the number of filters per layer
              init_weights: kernel initializer
              reg         : kernel regularizer
              relu        : clip value for ReLU
              bias        : whether to use bias
          '''
          # Configure base (super) class
          super().__init__(init_weights=init_weights, reg=reg, relu=relu, bias=bias)
  
          if layers is None:
             layers = self.layers
  
          # remember the layers
          self.layers = layers
  
          # remember the input shape
          self.input_shape = input_shape
  
          inputs = Input(input_shape)
          encoder = self.encoder(inputs, layers=layers)
          outputs = self.decoder(encoder, layers=layers)
          self._model = Model(inputs, outputs)
  
      def encoder(self, x, **metaparameters):
          ''' Construct the Encoder 
              x     : input to the encoder
              layers: number of nodes per layer
          '''
          layers = metaparameters['layers']
  
          # Flatten the input image
          x = Flatten()(x)
  
          # Progressive Unit Pooling
          for layer in layers:
              n_nodes = layer['n_nodes']
              x = self.Dense(x, n_nodes)
              x = self.BatchNormalization(x)
              x = self.ReLU(x)
  
          # The Encoding
          return x
  
      def decoder(self, x, **metaparameters):
          ''' Construct the Decoder
              x     : input to the decoder
              layers: number of nodes per layer
          '''
          layers = metaparameters['layers']
  
          # Progressive Unit Unpooling
          for _ in range(len(layers)-1, 0, -1):
              n_nodes = layers[_]['n_nodes']
              x = self.Dense(x, n_nodes)
              x = self.BatchNormalization(x)
              x = self.ReLU(x)
  
          # Last unpooling and match shape to input
          units = self.input_shape[0] * self.input_shape[1] * self.input_shape[2]
          print("INPUT", self.input_shape, "UNITS", units)
          x = self.Dense(x, units, activation='sigmoid')
  
          # Reshape back into an image
          x = Reshape(self.input_shape)(x)
  
          # The decoded image
          return x
  
      def compile(self, optimizer='adam'):
          ''' Compile the model ''' 
          self._model.compile(loss='binary_crossentropy', optimizer=optimizer, metrics=['accuracy'])
  
      def extract(self):
          ''' Extract the pretrained encoder
          '''
          # Get the trained weights from the autoencoder
          weights = self._model.get_weights()
  
          # Extract out the weights for just the encoder  (6 sets per layer)
          encoder_weights = weights[0 : int((6 * len(self.layers)))]
    
          # Construct a copy the encoder
          inputs = Input(self.input_shape)
          outputs = self.encoder(inputs, layers=self.layers)
          encoder = Model(inputs, outputs)
  
          # Initialize the encoder with the pretrained weights
          encoder.set_weights(encoder_weights)
  
          return encoder
  
  # Example autoencoder
  # autoencoder = AutoEncoder()
  
  # Train the model, and extract pretrained encoder
  # e = autoencoder.extract()
  
  def example():
      ''' Example for constructing/training an AutoEncoder model on MNIST
      '''
      # Example of constructing an AutoEncoder
      ae = AutoEncoder(input_shape=(28, 28, 1))
      ae.model.summary()
  
      from tensorflow.keras.datasets import mnist
      import numpy as np
      (x_train, y_train), (x_test, y_test) = mnist.load_data()
      x_train = (x_train / 255.0).astype(np.float32)
      x_test  = (x_test  / 255.0).astype(np.float32)
      x_train = np.expand_dims(x_train, axis=-1)
      x_test  = np.expand_dims(x_test, axis=-1)
  
      ae.compile()
      ae.model.fit(x_train, x_train, epochs=10, batch_size=32, validation_split=0.1, verbose=1)
      ae.model.evaluate(x_test, x_test)
  
  
  example()
  