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 = 'nE1'
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, 'n') for u in range(USERS)]
    federated_test_data = [get_distributed(cifar_test, u, 'n') for u in range(USERS)]
    sample_batch = federated_train_data[1][-2]
    
    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, 5, replace=False)
        state, metrics = iterative_process.next(state, list(np.array(federated_train_data)[selected]))
        test_metrics = evaluation(state.model, federated_test_data)
        fd_train_loss.append(metrics[1])
        fd_test_loss.append(test_metrics.loss)
        fd_test_accuracy.append(test_metrics.sparse_categorical_accuracy)

    try:
    	with open('Log/Exp10/'+ 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 = [5000, 3000, 1000, 750, 250]
    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 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()






tensorflow.compat.v1.enable_v2_behavior numpy.random.choice tensorflow.keras.layers.Dense tensorflow.keras.layers.Conv2D tensorflow.keras.losses.sparse_categorical_crossentropy tensorflow.keras.datasets.cifar10.load_data tensorflow.keras.metrics.SparseCategoricalAccuracy numpy.array tensorflow.keras.layers.MaxPooling2D tensorflow.keras.layers.Flatten
[(14, 'tensorflow.compat.v1.enable_v2_behavior', 'tf.compat.v1.enable_v2_behavior', ([], {}), True, 'import tensorflow as tf\n'), (29, 'tensorflow.keras.datasets.cifar10.load_data', 'tf.keras.datasets.cifar10.load_data', ([], {}), True, 'import tensorflow as tf\n'), (38, 'tensorflow_federated.python.learning.build_federated_averaging_process', 'tff.learning.build_federated_averaging_process', (['model_fn'], {}), True, 'from tensorflow_federated import python as tff\n'), (39, 'tensorflow_federated.python.learning.build_federated_evaluation', 'tff.learning.build_federated_evaluation', (['model_fn'], {}), True, 'from tensorflow_federated import python as tff\n'), (45, 'six.moves.range', 'range', (['(50)'], {}), False, 'from six.moves import range\n'), (67, 'six.moves.range', 'range', (['USERS'], {}), False, 'from six.moves import range\n'), (74, 'six.moves.range', 'range', (['CLASSES'], {}), False, 'from six.moves import range\n'), (82, 'six.moves.range', 'range', (['USERS'], {}), False, 'from six.moves import range\n'), (94, 'six.moves.range', 'range', (['CLASSES'], {}), False, 'from six.moves import range\n'), (98, 'six.moves.range', 'range', (['USERS'], {}), False, 'from six.moves import range\n'), (108, 'six.moves.range', 'range', (['CLASSES'], {}), False, 'from six.moves import range\n'), (115, 'six.moves.range', 'range', (['USERS'], {}), False, 'from six.moves import range\n'), (130, 'six.moves.range', 'range', (['NUM_EPOCHS'], {}), False, 'from six.moves import range\n'), (36, 'tensorflow_federated.python.learning.from_compiled_keras_model', 'tff.learning.from_compiled_keras_model', (['keras_model', 'sample_batch'], {}), True, 'from tensorflow_federated import python as tff\n'), (46, 'numpy.random.choice', 'np.random.choice', (['(5)', '(5)'], {'replace': '(False)'}), True, 'import numpy as np\n'), (81, 'six.moves.range', 'range', (['CLASSES'], {}), False, 'from six.moves import range\n'), (30, 'six.moves.range', 'range', (['USERS'], {}), False, 'from six.moves import range\n'), (31, 'six.moves.range', 'range', (['USERS'], {}), False, 'from six.moves import range\n'), (141, 'tensorflow.keras.layers.Conv2D', 'tf.keras.layers.Conv2D', (['(32)', '(3, 3)'], {'activation': '"""tanh"""', 'padding': '"""same"""', 'input_shape': '(WIDTH, HEIGHT, CHANNELS)'}), True, 'import tensorflow as tf\n'), (145, 'tensorflow.keras.layers.MaxPooling2D', 'tf.keras.layers.MaxPooling2D', ([], {'pool_size': '(2, 2)'}), True, 'import tensorflow as tf\n'), (146, 'tensorflow.keras.layers.Conv2D', 'tf.keras.layers.Conv2D', (['(64)', '(3, 3)'], {'activation': '"""tanh"""', 'padding': '"""same"""'}), True, 'import tensorflow as tf\n'), (147, 'tensorflow.keras.layers.MaxPooling2D', 'tf.keras.layers.MaxPooling2D', ([], {'pool_size': '(2, 2)'}), True, 'import tensorflow as tf\n'), (148, 'tensorflow.keras.layers.Flatten', 'tf.keras.layers.Flatten', ([], {}), True, 'import tensorflow as tf\n'), (149, 'tensorflow.keras.layers.Dense', 'tf.keras.layers.Dense', (['(128)'], {'activation': '"""tanh"""'}), True, 'import tensorflow as tf\n'), (150, 'tensorflow.keras.layers.Dense', 'tf.keras.layers.Dense', (['(10)'], {'activation': 'tf.nn.softmax'}), True, 'import tensorflow as tf\n'), (153, 'tensorflow.keras.losses.sparse_categorical_crossentropy', 'tf.keras.losses.sparse_categorical_crossentropy', (['y_true', 'y_pred'], {}), True, 'import tensorflow as tf\n'), (83, 'six.moves.range', 'range', (['(2)'], {}), False, 'from six.moves import range\n'), (84, 'six.moves.range', 'range', (['CLASSES'], {}), False, 'from six.moves import range\n'), (155, 'tensorflow.keras.metrics.SparseCategoricalAccuracy', 'tf.keras.metrics.SparseCategoricalAccuracy', ([], {}), True, 'import tensorflow as tf\n'), (47, 'numpy.array', 'np.array', (['federated_train_data'], {}), True, 'import numpy as np\n'), (68, 'six.moves.range', 'range', (['shares_arr[u]'], {}), False, 'from six.moves import range\n'), (134, 'numpy.array', 'np.array', (['[(source[0][b] / 255.0) for b in batch_samples]'], {'dtype': 'np.float32'}), True, 'import numpy as np\n'), (135, 'numpy.array', 'np.array', (['[source[1][b] for b in batch_samples]'], {'dtype': 'np.int32'}), True, 'import numpy as np\n'), (117, 'numpy.array', 'np.array', (['indices_array'], {}), True, 'import numpy as np\n'), (87, 'numpy.array', 'np.array', (['indices_array'], {}), True, 'import numpy as np\n'), (101, 'six.moves.range', 'range', (['class_shares'], {}), False, 'from six.moves import range\n')]