# Uncomment these to install fedjax.
# !pip install fedjax
# !pip install --upgrade git+https://github.com/google/fedjax.git

import itertools

import jax
import jax.numpy as jnp
from jax.experimental import stax

import fedjax

"""
Explanation: Working with models in FedJAX

In this chapter, we will learn about fedjax.Model. This notebook assumes you already have finished the "Datasets" chapter. We first overview centralized training and evaluation with fedjax.Model and then describe how to add new neural architectures and specify additional evaluation metrics.
End of explanation
"""


# Load train/test splits of the EMNIST dataset.
train, test = fedjax.datasets.emnist.load_data()

# As a start, let's simply use a logistic regression model.
model = fedjax.models.emnist.create_logistic_model()

"""
Explanation: Centralized training & evaluation with fedjax.Model
Most federated learning algorithms are built upon common components from standard centralized learning. fedjax.Model holds these common components. In centralized learning, we are mostly concerned with two tasks:

Training: We want to optimize our model parameters on the training dataset.
Evaluation: We want to know the values of evaluation metrics (e.g. accuracy) of the current model parameters on a test dataset.

Let's first see how we can carry out these two tasks on the EMNIST dataset with fedjax.Model.
End of explanation
"""


params_rng = jax.random.PRNGKey(0)
params = model.init(params_rng)

"""
Explanation: Random initialization, the JAX way
To start training, we need some randomly initialized parameters. In JAX, pseudo random number generation works slightly differently. For now, it is sufficient to know we call jax.random.PRNGKey() to seed the random number generator. JAX has a detailed introduction on this topic, if you are interested.
To create the initial model parameters, we simply call fedjax.Model.init() with a PRNGKey.
End of explanation
"""


params

"""
Explanation: Here are our initial model parameters. With the same PRNGKey, we will always get the same random initialization. There are 2 parameters in our model, the weights w, and the bias b. They are organized into a FlapMapping, but in general any PyTree can be used to store model parameters.
End of explanation
"""


# We select first 16 batches using itertools.islice.
batched_test_data = list(itertools.islice(
    fedjax.padded_batch_federated_data(test, batch_size=128), 16))
batched_train_data = list(itertools.islice(
    fedjax.padded_batch_federated_data(train, batch_size=128), 16))

print('eval_test', fedjax.evaluate_model(model, params, batched_test_data))
print('eval_train', fedjax.evaluate_model(model, params, batched_train_data))

"""
Explanation: Evaluating model parameters
Before we start training, let's first see how our initial parameters fare on the train and test sets. Unsurprisingly, they do not do very well. We evaluate using the fedjax.evaluate_model() which takes in model, parameters, and datasets which are batched. As noted in the dataset tutorial, we batch using
fedjax.padded_batch_federated_data() for efficiency. fedjax.padded_batch_federated_data() is very similar to fedjax.ClientDataset.padded_batch() but operates over the entire federated dataset.
End of explanation
"""


model.eval_metrics

"""
Explanation: How does our model know what evaluation metrics to report? It is simply specified in the eval_metrics field. We will discuss evaluation metrics in more detail later.
End of explanation
"""


for client_id, dataset in itertools.islice(test.clients(), 4):
  print(
      client_id,
      fedjax.evaluate_model(model, params,
                            dataset.padded_batch(batch_size=128)))

"""
Explanation: Since fedjax.evaluate_model() simply takes a stream of batches, we can also use it to evaluate multiple clients.
End of explanation
"""


# train_batches is an infinite stream of shuffled batches of examples.
def train_batches():
  return fedjax.shuffle_repeat_batch_federated_data(
      train,
      batch_size=8,
      client_buffer_size=16,
      example_buffer_size=1024,
      seed=0)


# We obtain the first batch by using the `next` function.
example = next(train_batches())
output = model.apply_for_train(params, example, None)
per_example_loss = model.train_loss(example, output)

output.shape, per_example_loss

"""
Explanation: The training objective
To train our model, we need two things: the objective function to minimize and an optimizer.
fedjax.Model contains two functions that can be used to arrive at the training objective:

apply_for_train(params, batch_example, rng) takes the current model parameters, a batch of examples, and a PRNGKey, and returns some output.
train_loss(batch_example, train_output) translates the output of apply_for_train() into a vector of per-example loss values.

In our example model, apply_for_train() produces a score for each class and train_loss() is simply the cross entropy loss. apply_for_train() in this case does not make use of a PRNGKey, so we can pass None instead for convenience. A different apply_for_train() might actually make use of the PRNGKey, for tasks such as dropout.
End of explanation
"""


per_example_loss_fn = fedjax.model_per_example_loss(model)
per_example_loss_fn(params, example, None)

"""
Explanation: Note that the output is per example predictions and has shape (8, 62), where 8 is the batch size and 62 is the number of classes. Alternatively, we can use model_per_example_loss() to get a function that gives us the same result. model_per_example_loss() is a convenience function that does exactly what we just did.
End of explanation
"""


def train_objective(params, example):
  return jnp.mean(per_example_loss_fn(params, example, None))

train_objective(params, example)

"""
Explanation: The training objective is a scalar, so why does train_loss() return a vector of per-example loss values? First of all, the training objective in most cases is just the average of the per-example loss values, so arriving at the final training objective isn't hard. Moreover, in certain algorithms, we not only use the train loss over a single batch of examples for a stochastic training step, but also need to estimate the average train loss over an entire (client) dataset. Having the per-example loss values there is instrumental in obtaining the correct estimate when the batch sizes may vary.
End of explanation
"""


optimizer = fedjax.optimizers.adam(1e-3)

"""
Explanation: Optimizers
With the training objective at hand, we just need an optimizer to find some good model parameters that minimize it.
There are many optimizer implementations in JAX out there, but FedJAX doesn't force one choice over any other. Instead, FedJAX provides a simple fedjax.optimizers.Optimizer interface so a new optimizer implementation can be wrapped. For convenience, FedJAX provides some common optimizers wrapped from optax.
End of explanation
"""


opt_state = optimizer.init(params)
grads = jax.grad(train_objective)(params, example)
opt_state, params = optimizer.apply(grads, opt_state, params)
train_objective(params, example)

"""
Explanation: An optimizer is simply a pair of two functions:

init(params) returns the initial optimizer state, such as initial values for accumulators of gradients.
apply(grads, opt_state, params) applies the gradients to update the current optimizer state and model parameters.

Instead of modifying opt_state or params, apply() returns a new pair of optimizer state and model parameters. In JAX, it is common to express computations in this stateless/mutation free style, often referred to as functional programming, or pure functions. The pureness of functions is crucial to many features in JAX, so it is always good practice to write functions that do not modify its inputs. You have probably also noticed that all the functions of fedjax.Model we have seen so far do not modify the model object itself (for example, init() returns model parameters instead of setting some attribute of model; apply_for_train() takes model parameters as an input argument, instead of getting it from model). FedJAX does this to keep all functions pure.
However, in the top level training loop, it is fine to mutate states since we are not in a function that may be transformed by JAX. Let's run our first training step, which resulted in a slight decrease in objective on the same batch of examples.
To obtain the gradients, we use jax.grad() which returns the gradient function. More details about jax.grad() can be found from the JAX documentation.
End of explanation
"""


model_grads = fedjax.model_grad(model)(params, example, None)
opt_state, params = optimizer.apply(grads, opt_state, params)
train_objective(params, example)

"""
Explanation: Instead of using jax.grad() directly, FedJAX also provides a convenient fedjax.model_grad() which computes the gradient of a model with respect to the averaged fedjax.model_per_example_loss().
End of explanation
"""


@jax.jit
def train_step(example, opt_state, params):
  grads = jax.grad(train_objective)(params, example)
  return optimizer.apply(grads, opt_state, params)

for example in itertools.islice(train_batches(), 5000):
  opt_state, params = train_step(example, opt_state, params)

print('eval_test', fedjax.evaluate_model(model, params, batched_test_data))
print('eval_train', fedjax.evaluate_model(model, params, batched_train_data))

"""
Explanation: Let's wrap everything into a single JIT compiled function and train a few more steps, and evaluate again.
End of explanation
"""


def hinge_loss(example, output):
  label = example['y']
  num_classes = output.shape[-1]
  mask = jax.nn.one_hot(label, num_classes)
  label_score = jnp.sum(output * mask, axis=-1)
  best_score = jnp.max(output + 1 - mask, axis=-1)
  return best_score - label_score


hinge_model = model.replace(train_loss=hinge_loss)
fedjax.model_per_example_loss(hinge_model)(params, example, None)

"""
Explanation: Building a custom model
fedjax.Model was designed with customization in mind. We have already seen how to switch to a different training loss. In this section, we will discuss how the rest of a fedjax.Model can be customized.
Training loss
Because train_loss() is separate from apply_for_train(), it is easy to switch to a different loss function.
End of explanation
"""


only_accuracy = model.replace(
    eval_metrics={'accuracy': fedjax.metrics.Accuracy()})
fedjax.evaluate_model(only_accuracy, params, batched_test_data)

"""
Explanation: Evaluation metrics
We have already seen that the eval_metrics field of a fedjax.Model tells the model what metrics to evaluate. eval_metrics is a mapping from metric names to fedjax.metrics.Metric objects. A fedjax.metrics.Metric object tells us how to calculate a metric's value from multiple batches of examples. Like fedjax.Model, a fedjax.metrics.Metric is stateless.
To customize the metrics to evaluate on, or what names to give to each, simply specify a different mapping.
End of explanation
"""


jnp.all(
    model.apply_for_train(params, example, None) == model.apply_for_eval(
        params, example))

"""
Explanation: There are already some concrete Metrics in fedjax.metrics. It is also easy to implement a new one. You can read more about how to implement a Metric in its own introduction.
The bit of fedjax.Model that is directly relevant to evaluation is apply_for_eval(). The relation between apply_for_eval() and an evaluation metric is similar to that between apply_for_train() and train_loss(): apply_for_eval(params, example) takes the model parameters and a batch of examples (notice there is no randomness in evaluation so we don't need a PRNGKey), and produces some prediction that evaluation metrics can consume. In our example, the outputs from apply_for_eval() and apply_for_train() are identical, but they don't have to be.
End of explanation
"""


fedjax.metrics.Accuracy()

"""
Explanation: What apply_for_eval() needs to produce really just depends on what evaluation fedjax.metrics.Metrics will be used. In our case, we are using fedjax.metrics.Accuracy, and fedjax.metrics.CrossEntropyLoss. They are similar in their requirements on the inputs:

They both need to know the true label from the example, using a target_key that defaults to "y".
They both need to know the predicted scores from apply_for_eval(), customizable as pred_key. If pred_key is None, apply_for_eval() should return just a vector of per-class scores; otherwise pred_key can be a string key, and apply_for_eval() should return a mapping (e.g. dict) that maps the key to a vector of per-class scores.
End of explanation
"""


def cross_entropy_loss(example, output):
  label = example['y']
  num_classes = output.shape[-1]
  mask = jax.nn.one_hot(label, num_classes)
  return -jnp.sum(jax.nn.log_softmax(output) * mask, axis=-1)

def mlp_model(num_input_units, num_units, num_classes):

  def mlp_init(rng):
    w0_rng, w1_rng = jax.random.split(rng)
    w0 = jax.random.uniform(w0_rng, [num_input_units, num_units])
    b0 = jnp.zeros([num_units])
    w1 = jax.random.uniform(w1_rng, [num_units, num_classes])
    b1 = jnp.zeros([num_classes])
    return w0, b0, w1, b1

  def mlp_apply(params, batch, rng=None):
    w0, b0, w1, b1 = params
    x = batch['x']
    batch_size = x.shape[0]
    h = jax.nn.relu(x.reshape([batch_size, -1]) @ w0 + b0)
    return h @ w1 + b1

  return fedjax.Model(
      init=mlp_init,
      apply_for_train=mlp_apply,
      apply_for_eval=mlp_apply,
      train_loss=cross_entropy_loss,
      eval_metrics={'accuracy': fedjax.metrics.Accuracy()})


# There are 28*28 input pixels, and 62 classes in EMNIST.
mlp = mlp_model(28 * 28, 128, 62)

@jax.jit
def mlp_train_step(example, opt_state, params):

  @jax.grad
  def grad_fn(params, example):
    return jnp.mean(fedjax.model_per_example_loss(mlp)(params, example, None))

  grads = grad_fn(params, example)
  return optimizer.apply(grads, opt_state, params)


params = mlp.init(jax.random.PRNGKey(0))
opt_state = optimizer.init(params)
print('eval_test before training:',
      fedjax.evaluate_model(mlp, params, batched_test_data))
for example in itertools.islice(train_batches(), 5000):
  opt_state, params = mlp_train_step(example, opt_state, params)
print('eval_test after training:',
      fedjax.evaluate_model(mlp, params, batched_test_data))

"""
Explanation: Neural network architectures
We have now covered all five parts of a fedjax.Model, namely init(), apply_for_train(), apply_for_eval(), train_loss(), and eval_metrics. train_loss() and eval_metrics are easy to customize since they are mostly agnostic to the actual neural network architecture of the model. init(), apply_for_train(), and apply_for_eval() on the other hand, are closely related.
In principle, as long as these three functions meet the interface we have seen so far, they can be used to build a custom model. Let's try to build a model that uses multi-layer perceptron and hinge loss.
End of explanation
"""


def stax_cnn_model(input_shape, num_classes):
  stax_init, stax_apply = stax.serial(
      stax.Conv(
          out_chan=64, filter_shape=(3, 3), strides=(1, 1), padding='SAME'),
      stax.Relu,
      stax.Flatten,
      stax.Dense(256),
      stax.Relu,
      stax.Dense(num_classes),
  )
  return fedjax.create_model_from_stax(
      stax_init=stax_init,
      stax_apply=stax_apply,
      sample_shape=input_shape,
      train_loss=cross_entropy_loss,
      eval_metrics={'accuracy': fedjax.metrics.Accuracy()})


stax_cnn = stax_cnn_model([-1, 28, 28, 1], 62)


@jax.jit
def stax_cnn_train_step(example, opt_state, params):

  @jax.grad
  def grad_fn(params, example):
    return jnp.mean(
        fedjax.model_per_example_loss(stax_cnn)(params, example, None))

  grads = grad_fn(params, example)
  return optimizer.apply(grads, opt_state, params)


params = stax_cnn.init(jax.random.PRNGKey(0))
opt_state = optimizer.init(params)
print('eval_test before training:',
      fedjax.evaluate_model(stax_cnn, params, batched_test_data))
for example in itertools.islice(train_batches(), 1000):
  opt_state, params = stax_cnn_train_step(example, opt_state, params)
print('eval_test after training:',
      fedjax.evaluate_model(stax_cnn, params, batched_test_data))

"""
Explanation: While writing custom neural network architectures from scratch is possible, most of the time, it is much more convenient to use a neural network library such as Haiku or jax.experimental.stax. The two functions fedjax.create_model_from_haiku and fedjax.create_model_from_stax can convert a neural network expressed in the respective framework into a fedjax.Model. Let's build a convolutional network using jax.experimental.stax this time.
End of explanation
"""


