text: from __future__ import absolute_import, division, print_function
  
  import collections
  import numpy as np
  from six.moves import range
  import tensorflow as tf
  import datetime
  
  from tensorflow_federated import python as tff
  from tensorflow.python.keras.optimizer_v2 import gradient_descent
  from tensorflow.keras import layers
  
  
  tf.compat.v1.enable_v2_behavior()
  
  EXP_CODE = 'rB32C4'
  NUM_EXAMPLES_PER_USER = 2000
  BATCH_SIZE = 32
  USERS = 5
  NUM_EPOCHS = 1
  CLASSES = 10
  
  WIDTH = 32
  HEIGHT = 32
  CHANNELS = 3
  
  def mane():
      """ Run program """
      cifar_train, cifar_test = tf.keras.datasets.cifar10.load_data()
      federated_train_data = [get_distributed(cifar_train, u, 'r') for u in range(USERS)]
      (X_test, y_test) = get_non_distributed(cifar_test)
      sample_batch = federated_train_data[1][-2]
      non_federated_model = create_compiled_keras_model()
  
      def model_fn():
      	keras_model = create_compiled_keras_model()
      	return tff.learning.from_compiled_keras_model(keras_model, sample_batch)
  
      iterative_process = tff.learning.build_federated_averaging_process(model_fn)
      evaluation = tff.learning.build_federated_evaluation(model_fn)
      state = iterative_process.initialize()
      fd_test_accuracy = []
      fd_test_loss = []
      fd_train_loss = []
  
      for round_num in range(50):
          selected = np.random.choice(5, 4, replace=False)
          state, metrics = iterative_process.next(state, list(np.array(federated_train_data)[selected]))
          non_federated_model.set_weights(state.model.trainable)
          (loss, accuracy) = non_federated_model.evaluate(X_test, y_test)
          fd_train_loss.append(metrics[1])
          fd_test_accuracy.append(accuracy)
          fd_test_loss.append(loss)
  
      try:
      	with open('Log/Exp11/'+ EXP_CODE + '.txt', 'w') as log:
      		print(EXP_CODE + "Train = {}".format(fd_train_loss), file=log)
      		print(EXP_CODE + "Test = {}".format(fd_test_loss), file=log)
      		print(EXP_CODE + "Accuracy = {}".format(fd_test_accuracy), file=log)
  
      except IOError:
      	print('File Error')
  
  def get_indices_realistic(y, u):
      # split dataset into arrays of each class label
      all_indices = [i for i, d in enumerate(y)]
      shares_arr = [4000, 2000, 2000, 1000, 1000]
      user_indices = []
      for u in range(USERS):
          user_indices.append([all_indices.pop(0) for i in range(shares_arr[u])]) 
      return user_indices
  
  def get_indices_unbalanced(y):
      # split dataset into arrays of each class label
      indices_array = []
      for c in range(CLASSES):
          indices_array.append([i for i, d in enumerate(y) if d == c])
      # each user will have 2 classes excluded from their data sets, thus 250 examples * remaining 8 classes
      class_shares = 250
      # store indices for future use
      user_indices = []
      # auxilary index array to pop out pairs of classes missing at each user
      class_index = list(range(CLASSES))
      for u in range(USERS):
          columns_out = [class_index.pop(0) for i in range(2)]
          selected_columns = set(range(CLASSES)) - set(columns_out)
          starting_index = u*class_shares
          user_indices.append(
              np.array(indices_array)[list(selected_columns)].T[starting_index:starting_index + class_shares]
              .flatten())
      return user_indices
  
  def get_indices_unbalanced_completely(y):
      # split dataset into arrays of each class label
      indices_array = []
      for c in range(CLASSES):
          indices_array.append([i for i, d in enumerate(y) if d == c])
          class_shares = CLASSES // min(CLASSES, USERS)
      user_indices = []
      for u in range(USERS):
          user_indices.append(
              np.array(
                  [indices_array.pop(0)[:NUM_EXAMPLES_PER_USER//class_shares] for j in range(class_shares)])
              .flatten())
      return user_indices
  
  def get_indices_even(y):
      # split dataset into arrays of each class label
      indices_array = []
      for c in range(CLASSES):
          indices_array.append([i for i, d in enumerate(y) if d == c])
      
      user_indices = []
      class_shares = NUM_EXAMPLES_PER_USER // CLASSES
      
      # take even shares of each class for every user
      for u in range(USERS):
          starting_index = u*class_shares
          user_indices.append(np.array(indices_array).T[starting_index:starting_index + class_shares].flatten())   
      return user_indices
      
  def get_distributed(source, u, distribution):
      if distribution == 'i':
          indices = get_indices_even(source[1])[u]
      elif distribution == 'n':
          indices = get_indices_unbalanced(source[1])[u]
      elif distribution == 'r':
          indices = get_indices_realistic(source[1][:10000], u)[u]
      else:
          indices = get_indices_unbalanced_completely(source[1])[u]
      output_sequence = []
      for repeat in range(NUM_EPOCHS):
          for i in range(0, len(indices), BATCH_SIZE):
              batch_samples = indices[i:i + BATCH_SIZE]
              output_sequence.append({
                  'x': np.array([source[0][b] / 255.0 for b in batch_samples], dtype=np.float32),
                  'y': np.array([source[1][b] for b in batch_samples], dtype=np.int32)})
      return output_sequence
  
  def get_non_distributed(source):
      y = np.array(source[1][:10000], dtype=np.int32)
      X = np.array(source[0][:10000], dtype=np.float32) / 255.0
      return X, y
  
  def create_compiled_keras_model():
  	model = tf.keras.models.Sequential([
  		tf.keras.layers.Conv2D(32,(3, 3),
              activation="tanh",
              padding="same",
              input_shape=(WIDTH, HEIGHT, CHANNELS)),
          tf.keras.layers.MaxPooling2D(pool_size=(2,2)),
          tf.keras.layers.Conv2D(64, (3, 3), activation="tanh", padding="same"),
          tf.keras.layers.MaxPooling2D(pool_size=(2,2)),
          tf.keras.layers.Flatten(), 
          tf.keras.layers.Dense(128, activation="tanh"),
          tf.keras.layers.Dense(10, activation=tf.nn.softmax)])
  
  	def loss_fn(y_true, y_pred):
  		return tf.reduce_mean(tf.keras.losses.sparse_categorical_crossentropy(y_true, y_pred))
  
  	model.compile(loss=loss_fn, optimizer="adam", metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])
  	return model
  
  
  if __name__ == "__main__":
      
      mane()
  
  
  
  
  
  