import os
import time
from multiprocessing import Process
from typing import Tuple

import flwr as fl
import numpy as np
import tensorflow as tf
from flwr.server.strategy import FedAvg

import dataset

# generate random integer values
from random import seed
from random import randint

# Make TensorFlow log less verbose
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"

# K: Prevent TF from using GPU (not enough memory)
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"


DATASET = Tuple[Tuple[np.ndarray, np.ndarray], Tuple[np.ndarray, np.ndarray]]


def start_server(num_rounds: int, num_clients: int, fraction_fit: float):
    """Start the server with a slightly adjusted FedAvg strategy."""
    strategy = FedAvg(min_available_clients=num_clients, fraction_fit=fraction_fit)
    # Exposes the server by default on port 8080
    fl.server.start_server(strategy=strategy, config={"num_rounds": num_rounds})


def start_client(dataset: DATASET) -> None:
    """Start a single client with the provided dataset."""

    # Load and compile a Keras model for CIFAR-10
    #model = tf.keras.applications.MobileNetV2((32, 32, 3), classes=10, weights=None)
    model = tf.keras.Sequential(
    [
        tf.keras.Input(shape=(32, 32, 3)),
        tf.keras.layers.Conv2D(32, kernel_size=(3, 3), activation="relu"),
        tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
        tf.keras.layers.Conv2D(64, kernel_size=(3, 3), activation="relu"),
        tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
        tf.keras.layers.Flatten(),
        tf.keras.layers.Dropout(0.5),
        tf.keras.layers.Dense(10, activation="softmax"),
    ]
    )
    model.compile("adam", "sparse_categorical_crossentropy", metrics=[tf.keras.metrics.CategoricalAccuracy(), tf.keras.metrics.MeanSquaredError()])
    ### @TODO: check if "accuracy" and tf.keras.metrics.CategoricalAccuracy() return the same results

    # Unpack the CIFAR-10 dataset partition
    (x_train, y_train), (x_test, y_test) = dataset

    # Define a Flower client
    class CifarClient(fl.client.NumPyClient):
        def get_parameters(self):
            """Return current weights."""
            return model.get_weights()

        def fit(self, parameters, config):
            """Fit model and return new weights as well as number of training
            examples."""
            model.set_weights(parameters)
            # Remove steps_per_epoch if you want to train over the full dataset
            # https://keras.io/api/models/model_training_apis/#fit-method
            #nap_time = randint (0, 5)
            #time.sleep (nap_time)
            #print ("Slept for", nap_time,  "seconds.")
            model.fit(x_train, y_train, epochs=10, batch_size=256, steps_per_epoch=10)
            return model.get_weights(), len(x_train), {}

        def evaluate(self, parameters, config):
            """Evaluate using provided parameters."""
            model.set_weights(parameters)
            loss, accuracy, mse = model.evaluate(x_test, y_test)
            print ('"Loss:', loss, ". Accuracy:", accuracy, ". MSE:", mse, ".")
            return loss, len(x_test), {"accuracy": accuracy}

    # Start Flower client
    fl.client.start_numpy_client("0.0.0.0:8080", client=CifarClient())


def run_simulation(num_rounds: int, num_clients: int, fraction_fit: float):
    """Start a FL simulation."""

    # This will hold all the processes which we are going to create
    processes = []

    # Start the server
    server_process = Process(
        target=start_server, args=(num_rounds, num_clients, fraction_fit)
    )
    server_process.start()
    processes.append(server_process)

    # Optionally block the script here for a second or two so the server has time to start
    time.sleep(2)

    # Load the dataset partitions
    partitions = dataset.load(num_partitions=num_clients)

    # Start all the clients
    for partition in partitions:
        client_process = Process(target=start_client, args=(partition,))
        client_process.start()
        processes.append(client_process)

    # Block until all processes are finished
    for p in processes:
        p.join()


if __name__ == "__main__":
    run_simulation(num_rounds=100, num_clients=5, fraction_fit=0.5)

tensorflow.keras.Input tensorflow.keras.layers.Dense tensorflow.keras.metrics.MeanSquaredError tensorflow.keras.layers.Conv2D tensorflow.keras.metrics.CategoricalAccuracy tensorflow.keras.layers.Dropout tensorflow.keras.layers.MaxPooling2D tensorflow.keras.layers.Flatten
[(29, 'flwr.server.strategy.FedAvg', 'FedAvg', ([], {'min_available_clients': 'num_clients', 'fraction_fit': 'fraction_fit'}), False, 'from flwr.server.strategy import FedAvg\n'), (31, 'flwr.server.start_server', 'fl.server.start_server', ([], {'strategy': 'strategy', 'config': "{'num_rounds': num_rounds}"}), True, 'import flwr as fl\n'), (93, 'multiprocessing.Process', 'Process', ([], {'target': 'start_server', 'args': '(num_rounds, num_clients, fraction_fit)'}), False, 'from multiprocessing import Process\n'), (100, 'time.sleep', 'time.sleep', (['(2)'], {}), False, 'import time\n'), (103, 'dataset.load', 'dataset.load', ([], {'num_partitions': 'num_clients'}), False, 'import dataset\n'), (107, 'multiprocessing.Process', 'Process', ([], {'target': 'start_client', 'args': '(partition,)'}), False, 'from multiprocessing import Process\n'), (41, 'tensorflow.keras.Input', 'tf.keras.Input', ([], {'shape': '(32, 32, 3)'}), True, 'import tensorflow as tf\n'), (42, 'tensorflow.keras.layers.Conv2D', 'tf.keras.layers.Conv2D', (['(32)'], {'kernel_size': '(3, 3)', 'activation': '"""relu"""'}), True, 'import tensorflow as tf\n'), (43, 'tensorflow.keras.layers.MaxPooling2D', 'tf.keras.layers.MaxPooling2D', ([], {'pool_size': '(2, 2)'}), True, 'import tensorflow as tf\n'), (44, 'tensorflow.keras.layers.Conv2D', 'tf.keras.layers.Conv2D', (['(64)'], {'kernel_size': '(3, 3)', 'activation': '"""relu"""'}), True, 'import tensorflow as tf\n'), (45, 'tensorflow.keras.layers.MaxPooling2D', 'tf.keras.layers.MaxPooling2D', ([], {'pool_size': '(2, 2)'}), True, 'import tensorflow as tf\n'), (46, 'tensorflow.keras.layers.Flatten', 'tf.keras.layers.Flatten', ([], {}), True, 'import tensorflow as tf\n'), (47, 'tensorflow.keras.layers.Dropout', 'tf.keras.layers.Dropout', (['(0.5)'], {}), True, 'import tensorflow as tf\n'), (48, 'tensorflow.keras.layers.Dense', 'tf.keras.layers.Dense', (['(10)'], {'activation': '"""softmax"""'}), True, 'import tensorflow as tf\n'), (51, 'tensorflow.keras.metrics.CategoricalAccuracy', 'tf.keras.metrics.CategoricalAccuracy', ([], {}), True, 'import tensorflow as tf\n'), (51, 'tensorflow.keras.metrics.MeanSquaredError', 'tf.keras.metrics.MeanSquaredError', ([], {}), True, 'import tensorflow as tf\n')]