# -*- coding: utf-8 -*-
"""Custom Federated Algorithms, Part 2: Implementing Federated Averaging

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/github/tensorflow/federated/blob/master/docs/tutorials/custom_federated_algorithms_2.ipynb

##### Copyright 2019 The TensorFlow Authors.
"""

#@title 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.

"""# Custom Federated Algorithms, Part 2: Implementing Federated Averaging

    View on TensorFlow.org

    Run in Google Colab

    View source on GitHub

    Download notebook

This tutorial is the second part of a two-part series that demonstrates how to
implement custom types of federated algorithms in TFF using the
[Federated Core (FC)](../federated_core.md), which serves as a foundation for
the [Federated Learning (FL)](../federated_learning.md) layer (`tff.learning`).

We encourage you to first read the
[first part of this series](custom_federated_algorithms_1.ipynb), which
introduce some of the key concepts and programming abstractions used here.

This second part of the series uses the mechanisms introduced in the first part
to implement a simple version of federated training and evaluation algorithms.

We encourage you to review the
[image classification](federated_learning_for_image_classification.ipynb) and
[text generation](federated_learning_for_text_generation.ipynb) tutorials for a
higher-level and more gentle introduction to TFF's Federated Learning APIs, as
they will help you put the concepts we describe here in context.

## Before we start

Before we start, try to run the following "Hello World" example to make sure
your environment is correctly setup. If it doesn't work, please refer to the
[Installation](../install.md) guide for instructions.
"""

#@test {"skip": true}
# !pip install --quiet --upgrade tensorflow-federated
# !pip install --quiet --upgrade nest-asyncio

import nest_asyncio
nest_asyncio.apply()

import collections

import numpy as np
import tensorflow as tf
import tensorflow_federated as tff

# TODO(b/148678573,b/148685415): must use the reference context because it
# supports unbounded references and tff.sequence_* intrinsics.
tff.backends.reference.set_reference_context()

@tff.federated_computation
def hello_world():
  return 'Hello, World!'

hello_world()

"""## Implementing Federated Averaging

As in
[Federated Learning for Image Classification](federated_learning_for_image_classification.ipynb),
we are going to use the MNIST example, but since this is intended as a low-level
tutorial, we are going to bypass the Keras API and `tff.simulation`, write raw
model code, and construct a federated data set from scratch.

### Preparing federated data sets

For the sake of a demonstration, we're going to simulate a scenario in which we
have data from 10 users, and each of the users contributes knowledge how to
recognize a different digit. This is about as
non-[i.i.d.](https://en.wikipedia.org/wiki/Independent_and_identically_distributed_random_variables)
as it gets.

First, let's load the standard MNIST data:
"""

mnist_train, mnist_test = tf.keras.datasets.mnist.load_data()

[(x.dtype, x.shape) for x in mnist_train]

"""The data comes as Numpy arrays, one with images and another with digit labels, both
with the first dimension going over the individual examples. Let's write a
helper function that formats it in a way compatible with how we feed federated
sequences into TFF computations, i.e., as a list of lists - the outer list
ranging over the users (digits), the inner ones ranging over batches of data in
each client's sequence. As is customary, we will structure each batch as a pair
of tensors named `x` and `y`, each with the leading batch dimension. While at
it, we'll also flatten each image into a 784-element vector and rescale the
pixels in it into the `0..1` range, so that we don't have to clutter the model
logic with data conversions.
"""

NUM_EXAMPLES_PER_USER = 1000
BATCH_SIZE = 100

def get_data_for_digit(source, digit):
  output_sequence = []
  all_samples = [i for i, d in enumerate(source[1]) if d == digit]
  for i in range(0, min(len(all_samples), NUM_EXAMPLES_PER_USER), BATCH_SIZE):
    batch_samples = all_samples[i:i + BATCH_SIZE]
    output_sequence.append({
        'x':
            np.array([source[0][i].flatten() / 255.0 for i in batch_samples],
                     dtype=np.float32),
        'y':
            np.array([source[1][i] for i in batch_samples], dtype=np.int32)
    })
  return output_sequence

federated_train_data = [get_data_for_digit(mnist_train, d) for d in range(10)]

federated_test_data = [get_data_for_digit(mnist_test, d) for d in range(10)]

"""As a quick sanity check, let's look at the `Y` tensor in the last batch of data
contributed by the fifth client (the one corresponding to the digit `5`).
"""

federated_train_data[5][-1]['y']

"""Just to be sure, let's also look at the image corresponding to the last element of that batch."""

from matplotlib import pyplot as plt

plt.imshow(federated_train_data[5][-1]['x'][-1].reshape(28, 28), cmap='gray')
plt.grid(False)
plt.show()

"""### On combining TensorFlow and TFF

In this tutorial, for compactness we immediately decorate functions that
introduce TensorFlow logic with `tff.tf_computation`. However, for more complex
logic, this is not the pattern we recommend. Debugging TensorFlow can already be
a challenge, and debugging TensorFlow after it has been fully serialized and
then re-imported necessarily loses some metadata and limits interactivity,
making debugging even more of a challenge.

Therefore, **we strongly recommend writing complex TF logic as stand-alone
Python functions** (that is, without `tff.tf_computation` decoration). This way
the TensorFlow logic can be developed and tested using TF best practices and
tools (like eager mode), before serializing the computation for TFF (e.g., by invoking `tff.tf_computation` with a Python function as the argument).

### Defining a loss function

Now that we have the data, let's define a loss function that we can use for
training. First, let's define the type of input as a TFF named tuple. Since the
size of data batches may vary, we set the batch dimension to `None` to indicate
that the size of this dimension is unknown.
"""

BATCH_SPEC = collections.OrderedDict(
    x=tf.TensorSpec(shape=[None, 784], dtype=tf.float32),
    y=tf.TensorSpec(shape=[None], dtype=tf.int32))
BATCH_TYPE = tff.to_type(BATCH_SPEC)

str(BATCH_TYPE)

"""You may be wondering why we can't just define an ordinary Python type. Recall
the discussion in [part 1](custom_federated_algorithms_1.ipynb), where we
explained that while we can express the logic of TFF computations using Python,
under the hood TFF computations *are not* Python. The symbol `BATCH_TYPE`
defined above represents an abstract TFF type specification. It is important to
distinguish this *abstract* TFF type from concrete Python *representation*
types, e.g., containers such as `dict` or `collections.namedtuple` that may be
used to represent the TFF type in the body of a Python function. Unlike Python,
TFF has a single abstract type constructor `tff.StructType` for tuple-like
containers, with elements that can be individually named or left unnamed. This
type is also used to model formal parameters of computations, as TFF
computations can formally only declare one parameter and one result - you will
see examples of this shortly.

Let's now define the TFF type of model parameters, again as a TFF named tuple of
*weights* and *bias*.
"""

MODEL_SPEC = collections.OrderedDict(
    weights=tf.TensorSpec(shape=[784, 10], dtype=tf.float32),
    bias=tf.TensorSpec(shape=[10], dtype=tf.float32))
MODEL_TYPE = tff.to_type(MODEL_SPEC)

print(MODEL_TYPE)

"""With those definitions in place, now we can define the loss for a given model, over a single batch. Note the usage of `@tf.function` decorator inside the `@tff.tf_computation` decorator. This allows us to write TF using Python like semantics even though were inside a `tf.Graph` context created by the `tff.tf_computation` decorator."""

# NOTE: `forward_pass` is defined separately from `batch_loss` so that it can
# be later called from within another tf.function. Necessary because a
# @tf.function  decorated method cannot invoke a @tff.tf_computation.

@tf.function
def forward_pass(model, batch):
  predicted_y = tf.nn.softmax(
      tf.matmul(batch['x'], model['weights']) + model['bias'])
  return -tf.reduce_mean(
      tf.reduce_sum(
          tf.one_hot(batch['y'], 10) * tf.math.log(predicted_y), axis=[1]))

@tff.tf_computation(MODEL_TYPE, BATCH_TYPE)
def batch_loss(model, batch):
  return forward_pass(model, batch)

"""As expected, computation `batch_loss` returns `float32` loss given the model and
a single data batch. Note how the `MODEL_TYPE` and `BATCH_TYPE` have been lumped
together into a 2-tuple of formal parameters; you can recognize the type of
`batch_loss` as `(<MODEL_TYPE,BATCH_TYPE> -> float32)`.
"""

str(batch_loss.type_signature)

"""As a sanity check, let's construct an initial model filled with zeros and
compute the loss over the batch of data we visualized above.
"""

initial_model = collections.OrderedDict(
    weights=np.zeros([784, 10], dtype=np.float32),
    bias=np.zeros([10], dtype=np.float32))

sample_batch = federated_train_data[5][-1]

batch_loss(initial_model, sample_batch)

"""Note that we feed the TFF computation with the initial model defined as a
`dict`, even though the body of the Python function that defines it consumes
model parameters as `model['weight']` and `model['bias']`. The arguments of the call
to `batch_loss` aren't simply passed to the body of that function.

What happens when we invoke `batch_loss`?
The Python body of `batch_loss` has already been traced and serialized  in the above cell where it was defined.  TFF acts as the caller to `batch_loss`
at the computation definition time, and as the target of invocation at the time
`batch_loss` is invoked. In both roles, TFF serves as the bridge between TFF's
abstract type system and Python representation types. At the invocation time,
TFF will accept most standard Python container types (`dict`, `list`, `tuple`,
`collections.namedtuple`, etc.) as concrete representations of abstract TFF
tuples. Also, although as noted above, TFF computations formally only accept a
single parameter, you can use the familiar Python call syntax with positional
and/or keyword arguments in case where the type of the parameter is a tuple - it
works as expected.

### Gradient descent on a single batch

Now, let's define a computation that uses this loss function to perform a single
step of gradient descent. Note how in defining this function, we use
`batch_loss` as a subcomponent. You can invoke a computation constructed with
`tff.tf_computation` inside the body of another computation, though typically
this is not necessary - as noted above, because serialization looses some
debugging information, it is often preferable for more complex computations to
write and test all the TensorFlow without the `tff.tf_computation` decorator.
"""

@tff.tf_computation(MODEL_TYPE, BATCH_TYPE, tf.float32)
def batch_train(initial_model, batch, learning_rate):
  # Define a group of model variables and set them to `initial_model`. Must
  # be defined outside the @tf.function.
  model_vars = collections.OrderedDict([
      (name, tf.Variable(name=name, initial_value=value))
      for name, value in initial_model.items()
  ])
  optimizer = tf.keras.optimizers.SGD(learning_rate)

  @tf.function
  def _train_on_batch(model_vars, batch):
    # Perform one step of gradient descent using loss from `batch_loss`.
    with tf.GradientTape() as tape:
      loss = forward_pass(model_vars, batch)
    grads = tape.gradient(loss, model_vars)
    optimizer.apply_gradients(
        zip(tf.nest.flatten(grads), tf.nest.flatten(model_vars)))
    return model_vars

  return _train_on_batch(model_vars, batch)

str(batch_train.type_signature)

"""When you invoke a Python function decorated with `tff.tf_computation` within the
body of another such function, the logic of the inner TFF computation is
embedded (essentially, inlined) in the logic of the outer one. As noted above,
if you are writing both computations, it is likely preferable to make the inner
function (`batch_loss` in this case) a regular Python or `tf.function` rather
than a `tff.tf_computation`. However, here we illustrate that calling one
`tff.tf_computation` inside another basically works as expected. This may be
necessary if, for example, you do not have the Python code defining
`batch_loss`, but only its serialized TFF representation.

Now, let's apply this function a few times to the initial model to see whether
the loss decreases.
"""

model = initial_model
losses = []
for _ in range(5):
  model = batch_train(model, sample_batch, 0.1)
  losses.append(batch_loss(model, sample_batch))

losses

"""### Gradient descent on a sequence of local data

Now, since `batch_train` appears to work, let's write a similar training
function `local_train` that consumes the entire sequence of all batches from one
user instead of just a single batch. The new computation will need to now
consume `tff.SequenceType(BATCH_TYPE)` instead of `BATCH_TYPE`.
"""

LOCAL_DATA_TYPE = tff.SequenceType(BATCH_TYPE)

@tff.federated_computation(MODEL_TYPE, tf.float32, LOCAL_DATA_TYPE)
def local_train(initial_model, learning_rate, all_batches):

  # Mapping function to apply to each batch.
  @tff.federated_computation(MODEL_TYPE, BATCH_TYPE)
  def batch_fn(model, batch):
    return batch_train(model, batch, learning_rate)

  return tff.sequence_reduce(all_batches, initial_model, batch_fn)

str(local_train.type_signature)

"""There are quite a few details buried in this short section of code, let's go
over them one by one.

First, while we could have implemented this logic entirely in TensorFlow,
relying on `tf.data.Dataset.reduce` to process the sequence similarly to how
we've done it earlier, we've opted this time to express the logic in the glue
language, as a `tff.federated_computation`. We've used the federated operator
`tff.sequence_reduce` to perform the reduction.

The operator `tff.sequence_reduce` is used similarly to
`tf.data.Dataset.reduce`. You can think of it as essentially the same as
`tf.data.Dataset.reduce`, but for use inside federated computations, which as
you may remember, cannot contain TensorFlow code. It is a template operator with
a formal parameter 3-tuple that consists of a *sequence* of `T`-typed elements,
the initial state of the reduction (we'll refer to it abstractly as *zero*) of
some type `U`, and the *reduction operator* of type `(<U,T> -> U)` that alters the
state of the reduction by processing a single element. The result is the final
state of the reduction, after processing all elements in a sequential order. In
our example, the state of the reduction is the model trained on a prefix of the
data, and the elements are data batches.

Second, note that we have again used one computation (`batch_train`) as a
component within another (`local_train`), but not directly. We can't use it as a
reduction operator because it takes an additional parameter - the learning rate.
To resolve this, we define an embedded federated computation `batch_fn` that
binds to the `local_train`'s parameter `learning_rate` in its body. It is
allowed for a child computation defined this way to capture a formal parameter
of its parent as long as the child computation is not invoked outside the body
of its parent. You can think of this pattern as an equivalent of
`functools.partial` in Python.

The practical implication of capturing `learning_rate` this way is, of course,
that the same learning rate value is used across all batches.

Now, let's try the newly defined local training function on the entire sequence
of data from the same user who contributed the sample batch (digit `5`).
"""

locally_trained_model = local_train(initial_model, 0.1, federated_train_data[5])

"""Did it work? To answer this question, we need to implement evaluation.

### Local evaluation

Here's one way to implement local evaluation by adding up the losses across all data
batches (we could have just as well computed the average; we'll leave it as an
exercise for the reader).
"""

@tff.federated_computation(MODEL_TYPE, LOCAL_DATA_TYPE)
def local_eval(model, all_batches):
  # TODO(b/120157713): Replace with `tff.sequence_average()` once implemented.
  return tff.sequence_sum(
      tff.sequence_map(
          tff.federated_computation(lambda b: batch_loss(model, b), BATCH_TYPE),
          all_batches))

str(local_eval.type_signature)

"""Again, there are a few new elements illustrated by this code, let's go over them
one by one.

First, we have used two new federated operators for processing sequences:
`tff.sequence_map` that takes a *mapping function* `T->U` and a *sequence* of
`T`, and emits a sequence of `U` obtained by applying the mapping function
pointwise, and `tff.sequence_sum` that just adds all the elements. Here, we map
each data batch to a loss value, and then add the resulting loss values to
compute the total loss.

Note that we could have again used `tff.sequence_reduce`, but this wouldn't be
the best choice - the reduction process is, by definition, sequential, whereas
the mapping and sum can be computed in parallel. When given a choice, it's best
to stick with operators that don't constrain implementation choices, so that
when our TFF computation is compiled in the future to be deployed to a specific
environment, one can take full advantage of all potential opportunities for a
faster, more scalable, more resource-efficient execution.

Second, note that just as in `local_train`, the component function we need
(`batch_loss`) takes more parameters than what the federated operator
(`tff.sequence_map`) expects, so we again define a partial, this time inline by
directly wrapping a `lambda` as a `tff.federated_computation`. Using wrappers
inline with a function as an argument is the recommended way to use
`tff.tf_computation` to embed TensorFlow logic in TFF.

Now, let's see whether our training worked.
"""

print('initial_model loss =', local_eval(initial_model,
                                         federated_train_data[5]))
print('locally_trained_model loss =',
      local_eval(locally_trained_model, federated_train_data[5]))

"""Indeed, the loss decreased. But what happens if we evaluated it on another
user's data?
"""

print('initial_model loss =', local_eval(initial_model,
                                         federated_train_data[0]))
print('locally_trained_model loss =',
      local_eval(locally_trained_model, federated_train_data[0]))

"""As expected, things got worse. The model was trained to recognize `5`, and has
never seen a `0`. This brings the question - how did the local training impact
the quality of the model from the global perspective?

### Federated evaluation

This is the point in our journey where we finally circle back to federated types
and federated computations - the topic that we started with. Here's a pair of
TFF types definitions for the model that originates at the server, and the data
that remains on the clients.
"""

SERVER_MODEL_TYPE = tff.type_at_server(MODEL_TYPE)
CLIENT_DATA_TYPE = tff.type_at_clients(LOCAL_DATA_TYPE)

"""With all the definitions introduced so far, expressing federated evaluation in
TFF is a one-liner - we distribute the model to clients, let each client invoke
local evaluation on its local portion of data, and then average out the loss.
Here's one way to write this.
"""

@tff.federated_computation(SERVER_MODEL_TYPE, CLIENT_DATA_TYPE)
def federated_eval(model, data):
  return tff.federated_mean(
      tff.federated_map(local_eval, [tff.federated_broadcast(model), data]))

"""We've already seen examples of `tff.federated_mean` and `tff.federated_map`
in simpler scenarios, and at the intuitive level, they work as expected, but
there's more in this section of code than meets the eye, so let's go over it
carefully.

First, let's break down the *let each client invoke local evaluation on its
local portion of data* part. As you may recall from the preceding sections,
`local_eval` has a type signature of the form `(<MODEL_TYPE, LOCAL_DATA_TYPE> ->
float32)`.

The federated operator `tff.federated_map` is a template that accepts as a
parameter a 2-tuple that consists of the *mapping function* of some type `T->U`
and a federated value of type `{T}@CLIENTS` (i.e., with member constituents of
the same type as the parameter of the mapping function), and returns a result of
type `{U}@CLIENTS`.

Since we're feeding `local_eval` as a mapping function to apply on a per-client
basis, the second argument should be of a federated type `{<MODEL_TYPE,
LOCAL_DATA_TYPE>}@CLIENTS`, i.e., in the nomenclature of the preceding sections,
it should be a federated tuple. Each client should hold a full set of arguments
for `local_eval` as a member consituent. Instead, we're feeding it a 2-element
Python `list`. What's happening here?

Indeed, this is an example of an *implicit type cast* in TFF, similar to
implicit type casts you may have encountered elsewhere, e.g., when you feed an
`int` to a function that accepts a `float`. Implicit casting is used scarcily at
this point, but we plan to make it more pervasive in TFF as a way to minimize
boilerplate.

The implicit cast that's applied in this case is the equivalence between
federated tuples of the form `{<X,Y>}@Z`, and tuples of federated values
`<{X}@Z,{Y}@Z>`. While formally, these two are different type signatures,
looking at it from the programmers's perspective, each device in `Z` holds two
units of data `X` and `Y`. What happens here is not unlike `zip` in Python, and
indeed, we offer an operator `tff.federated_zip` that allows you to perform such
conversions explicity. When the `tff.federated_map` encounters a tuple as a
second argument, it simply invokes `tff.federated_zip` for you.

Given the above, you should now be able to recognize the expression
`tff.federated_broadcast(model)` as representing a value of TFF type
`{MODEL_TYPE}@CLIENTS`, and `data` as a value of TFF type
`{LOCAL_DATA_TYPE}@CLIENTS` (or simply `CLIENT_DATA_TYPE`), the two getting
filtered together through an implicit `tff.federated_zip` to form the second
argument to `tff.federated_map`.

The operator `tff.federated_broadcast`, as you'd expect, simply transfers data
from the server to the clients.

Now, let's see how our local training affected the average loss in the system.
"""

print('initial_model loss =', federated_eval(initial_model,
                                             federated_train_data))
print('locally_trained_model loss =',
      federated_eval(locally_trained_model, federated_train_data))

"""Indeed, as expected, the loss has increased. In order to improve the model for
all users, we'll need to train in on everyone's data.

### Federated training

The simplest way to implement federated training is to locally train, and then
average the models. This uses the same building blocks and patters we've already
discussed, as you can see below.
"""

SERVER_FLOAT_TYPE = tff.type_at_server(tf.float32)

@tff.federated_computation(SERVER_MODEL_TYPE, SERVER_FLOAT_TYPE,
                           CLIENT_DATA_TYPE)
def federated_train(model, learning_rate, data):
  return tff.federated_mean(
      tff.federated_map(local_train, [
          tff.federated_broadcast(model),
          tff.federated_broadcast(learning_rate), data
      ]))

"""Note that in the full-featured implementation of Federated Averaging provided by
`tff.learning`, rather than averaging the models, we prefer to average model
deltas, for a number of reasons, e.g., the ability to clip the update norms,
for compression, etc.

Let's see whether the training works by running a few rounds of training and
comparing the average loss before and after.
"""

model = initial_model
learning_rate = 0.1
for round_num in range(5):
  model = federated_train(model, learning_rate, federated_train_data)
  learning_rate = learning_rate * 0.9
  loss = federated_eval(model, federated_train_data)
  print('round {}, loss={}'.format(round_num, loss))

"""For completeness, let's now also run on the test data to confirm that our model
generalizes well.
"""

print('initial_model test loss =',
      federated_eval(initial_model, federated_test_data))
print('trained_model test loss =', federated_eval(model, federated_test_data))

"""This concludes our tutorial.

Of course, our simplified example doesn't reflect a number of things you'd need
to do in a more realistic scenario - for example, we haven't computed metrics
other than loss. We encourage you to study
[the implementation](https://github.com/tensorflow/federated/blob/master/tensorflow_federated/python/learning/federated_averaging.py)
of federated averaging in `tff.learning` as a more complete example, and as a
way to demonstrate some of the coding practices we'd like to encourage.
"""


"""
Collection of algorithms (geared towards abstract math and numerical problems)

Picked up from MA and now continued from Cooling Singapore. Oct 2019
"""

# This control flag is used to skip sections if their prerequisite modules are not available, so that DK_Numerical
# can be imported more flexibly with regards to the environment.

_control = {section: True for section in ('NetworkX', 'Pandas', 'numpy')}

try:
	import networkx as nx
except ModuleNotFoundError:
	_control['NetworkX'] = False

try:
	import pandas as pd
except ModuleNotFoundError:
	_control['Pandas'] = False

try:
	import numpy as np
except ModuleNotFoundError:
	_control['numpy'] = False

import warnings

# -------------------------------------------- NetworkX ------------------------------------------------ #
if _control['NetworkX']:
	def get_treeroot(tree, guess=None):
		"""
		Finds the root node of a NetworkX DiGraph tree.

		Args:
			tree: NetworkX DiGraph tree
			guess: Best bet for which node is the root. The closer this node is to the root, the faster the search.

		Returns:
			Root node of tree.
		"""

		# Step 1: Param checks
		if not isinstance(tree, nx.DiGraph):
			raise TypeError("Must pass a NetworkX Digraph to parameter 'tree.'")

		if not nx.is_tree(tree):
			raise ValueError("Passed graph is not a tree.")

		if guess is None:
			guess = list(tree.nodes)[0]
		else:
			if guess not in tree:
				warnings.warn("The guessed root node is not in the tree.")
				guess = list(tree.nodes)[0]

		# Step 2: Follow predecessors until you find the root
		root = guess
		counter = 1
		n_nodes = len(tree)

		# This is a redundant loop, but Iterator Protocol + NetworkX guarantees that this is unnecessary.
		while counter <= n_nodes:
			try:
				root = next(tree.predecessors(root))
				counter += 1
			except StopIteration: # raised for the root node (no predecessors)
				break

		return root

	def compare_trees(T1, T2):
		"""
		Returns true if the trees are IDENTICAL (isomorphism is necessary but insufficient), and False otherwise.

		Two trees are identical iff:
			1) same set of nodes
			2) same set of branches
			3) same branch params
			(for now, assumed no node params)

		DEVELOPER'S CORNER:
			I learned that dictionaries can be tested for equality via ==. There's also talk on StackOverflow about this
			equality test holding for nested dicts (and other nesting combinations with other data structures).

		"""
		if not isinstance(T1, nx.DiGraph) or not isinstance(T2, nx.DiGraph):
			raise TypeError("This function only accepts NetworkX DiGraphs.")

		if not nx.faster_could_be_isomorphic(T1, T2):
			return False

		# Check 1: Compare nodes
		if set(T1) != set(T2):
			return False

		# Check 2: Compare branches
		br1 = set(br for br in T1.edges())
		br2 = set(br for br in T2.edges())

		if br1 != br2:
			return False

		# Check 3: Compare branch parameters (final test)
		return all([T1.edges[up, down] == T2.edges[up, down] for up, down in br1])

# -------------------------------------------- Pandas ------------------------------------------------ #
if _control['Pandas']:
	def empty_df_with_dtypes(col_dtype: dict):
		"""Initializes a pandas DataFrame with NO ROWS, but with col labels AND datatypes. The DataFrame constructor only
		lets you specify one datatype for all the columns, so this function fills that gap.

		Args:
			col_dtype: as dictionary of 'column label':'numpy dtype' pairs

		Note: This solution is adapted from the one proposed in stackoverflow,
		https://stackoverflow.com/questions/36462257/create-empty-dataframe-in-pandas-specifying-column-types/48374031#48374031

		Returns:
			Empty DataFrame, with columns + dtype specified.

		"""
		if not isinstance(col_dtype, dict):
			raise TypeError("Pls. pass a dictionary of 'column label':'numpy dtype' pairs.")

		if len(col_dtype) == 0:
			raise RuntimeError("Requested table has no columns.")

		df = pd.DataFrame()

		for c, d in col_dtype.items():
			df[c] = pd.Series(dtype=d)

		return df

	def is_xymonotonic(x, y, slope='pos', getdf=False, as_assertion=False):
		"""Checks if the given relationship <x, y> is monotonic in the specified direction ('pos' or 'neg' slope). A
		monotonic line is one whose slope is consistently either non-neg or non-pos. Pass vectors (iterables) to x
		and y. x has to be strictly increasing, whereas y can have repeating subsequent points (flat slope).

		Returns a tuple (bool, list):
			Bool is True if monotonic in the specified direction.

			List is empty if bool is True. Otherwise, contains the indices (0-indexed range is assigned to the data)

			if getdf is True, then a third item is returned as a pandas DataFrame of vectors x and y, with column
			names as 'x' and 'y'.

		To assert monotonicity:
			assert is_xymonotonic(x, y, slope='pos')[0]

		To assert monotinicity internally
			is_xymonotonic(x, y, slope='pos', as_assertion=True)

		"""

		# ----------------------------------------------------------------------- Checks
		length = len(x)
		if length != len(y):
			raise ValueError("Vectors x and y have to be of the same length.")
		if length < 3:
			raise ValueError("Vectors x and y must have at least 3 points.")

		# ----------------------------------------------------------------------- Convert to Series
		# x2, y2 is the same data but shifted one index forward wrt x1 and y1  -- old implementation
		x = pd.Series(data=(pt for pt in x))
		# x2 = pd.Series(data=x1.values, index=pd.RangeIndex(-1, length-1))

		y = pd.Series(data=(pt for pt in y))
		# y2 = pd.Series(data=y1.values, index=pd.RangeIndex(-1, length-1))

		# ----------------------------------------------------------------------- Calc dx and dy
		# dx = x2-x1
		dx = x.diff()
		dx = dx.loc[dx.notna()]

		# dy = y2-y1
		dy = y.diff()
		dy = dy.loc[dy.notna()]

		# Check that all dx > 0
		Lf_invalid_dx = dx <= 0
		if Lf_invalid_dx.any():
			# Collect indeces of points (+next point) where dx <= 0
			invalid_idx = Lf_invalid_dx.loc[Lf_invalid_dx].index
			invalid_idx_withadj = pd.Index(set(invalid_idx).union(invalid_idx+1))

			print("x vector can only be increasing. Pls. view the slice: \n{}".format(x.loc[invalid_idx_withadj]))
			raise ValueError(invalid_idx_withadj)

		# ----------------------------------------------------------------------- Check if monotonic
		sign = {'pos': 1}.get(slope, -1)
		# sign * dy must be >= 0 for all entries to fulfill the specified monotonicity
		Lf_failed = sign * dy < 0

		returned_value = (not any(Lf_failed), list(Lf_failed.loc[Lf_failed].index))

		if as_assertion:
			assert returned_value[0], returned_value[1]

		if getdf:
			returned_value += (pd.DataFrame({'x': x, 'y': y}), )
		return returned_value

	def sortxy(x, y, ascending=True):
		"""Given vectors x and y (as iterables that are ordered and 1:1), sortxy() sorts these vectors in ascending
		or descending order (via param 'ascending'), and returns them in a DataFrame with a range index. No check is
		done if x has repeating values."""
		# todo consider keeping the original index, if ever necessary

		# ----------------------------------------------------------------------- Checks
		length = len(x)
		if length != len(y):
			raise ValueError("Vectors x and y have to be of the same length.")

		# ----------------------------------------------------------------------- Convert to Series
		# Range index here stores the mapping x --> y
		Serx = pd.Series(data=(pt for pt in x))
		Sery = pd.Series(data=(pt for pt in y))

		# ----------------------------------------------------------------------- Sort and return df
		Serx_sorted = Serx.sort_values(ascending=ascending)

		return pd.DataFrame({
			'x': Serx_sorted.values,
			'y': Sery.loc[Serx_sorted.index].values,
		})

# -------------------------------------------- numpy ------------------------------------------------ #
if _control['numpy']:
	def clip(f, lb=None, ub=None):
		"""Return a clipped version of callable f(). This is the counterpart of numpy.clip(),
		which only works on explicit arrays, for callables.

		lb and ub follow the same requirements as in their counterparts in numpy.clip().
		"""
		if ub < lb:
			raise ValueError("ub < lb")
		return lambda x: np.clip(f(x), lb, ub)

	def apply_f_every_n(arr, n, func=np.sum, res_dtype='f8'):
		"""Apply function func() to every n elements of array. The length of array arr must be a multiple of n,
		but this is not checked (last operation would fall short of elements w/o raising an exception).

		sample usage:
			# Downscale the resolution by half, while getting the mean of the data points
			apply_f_every_n(ser.values, 2, np.mean)

			# You can use this via DataFrame.apply(). Here, we are changing a 24-h, half-hourly dataset into an hourly one.
			df_24h_hh.apply(apply_f_every_n, axis=0, args=(2, np.mean))
		"""
		start_idxs = np.arange(0, arr.shape[0], n)

		return np.fromiter((func(arr[idx:idx + n]) for idx in start_idxs), dtype=res_dtype)

# -------------------------------------------- Uncategorized Python ------------------------------------------------ #
def get_dictvals(mydict):
	"""This function recursively gets all non-dict items in a hierarchical dict data structure. The order of the items have no intended meaning.

	e.g.
		myheir = {
			0: '0.0',
			1: '0.1',
			2: {0: '1.0', 1: '1.1'},
			3: {0: '1.2', 1: {0: '2.0', 1: '2.1', 2: {0: '3.0'}, 3:'2.2'}},
			4: {0: '1.3', 1: '1.4'},
			5: '0.2',
		}

	get_dictvals(myheir)
	>> ['0.0', '0.1', '1.0', '1.1', '1.2', '2.0', '2.1', '3.0', '2.2', '1.3', '1.4', '0.2']

	"""
	values = []

	for val in mydict.values():
		if isinstance(val, dict):
			values.extend(get_dictvals(val))
		else:
			values.append(val)
	return values

# -------------------------------------------- Tuple Arithmetic ------------------------------------------------ #
# FOR DEVELOPMENT:
#    1) Sequences of tuples should be contained in sets.
#

def totuple(item):
	if isinstance(item, (str, int)):
		return tuple([item])
	else:
		return tuple(item)

def tupsum(*addends, strAsOne=True):
	"""
	Performs a tuple addition on an indefinite number of tuples. Arguments must be tuple-coercibles (they must be
	iterables or iterators), or scalars (i.e. tuple([arg]) succeeds).

	Returns:
		A tuple of the tuple sum of the addends.

	DEVELOPER'S CORNER:
		Tuples support concatenation, ie (1) + (0) = (1,0), which is really what tupple addition is in the context of
		the Cartesian product.

		tuple() returns an empty tuple, which guarantees that 'sum' will support concatenation.

		If an addend is not a tuple, tuple() coerces it (same elements).
		If an addend cannot be coerced directly, then it is expected to be a scalar (usu. numeric or string),
		and thus containing it via tuple([a]) should work.
	"""
	sum = tuple()

	for a in addends:
		if type(a) is str and strAsOne:
			a = tuple([a])
		try:
			sum = sum + tuple(a)
		except TypeError:
			sum = sum + tuple([a])

	return sum

def tupadd(a, b):
	"""
	Performs a tuple addition of a and b.

	Args: a and b can either be 1) scalars (i.e. single coordinate) or 2) elementary tuples

	Any argument that is a tuple (assumed elementary; i.e. all its elements are scalars) is unpacked as can be seen
	below.

	As far as the author knows, this function cannot be extended to an arbitrary number of addends (by taking a
	var-positional parameter), because the statement a = *a is not allowed. To sum an indefinite number of addends,
	one way to achieve this is by looping every addend pair, much like how tuppi works.

	Returns:
		The tuple sum of a and b (GUARANTEED elementary tuple if a and b are elementary tuples or scalars).
	"""
	if type(a) == tuple and type(b) == tuple:
		return tuple([*a, *b])
	elif type(a) == tuple:
		return tuple([*a, b])
	elif type(b) == tuple:
		return tuple([a, *b])
	else:
		return tuple([a, b])

def tuppi(*factors):
	"""
	Performs a Cartesian product of all of its factors.

	Args:
		*factors: An iterable of tuple-coercibles. If you want a 1-dim tuple, pass it as (a,) and not a.

	Returns:
		The Cartesian product as a tuple.

	DEVELOPER'S CORNER:
		If the factors are all elementary tuples or scalars, the product is GUARANTEED to be an elementary tuple.
		Even if after every iteration of the for factor .. loop, product is contained in a tuple, this container is
		"opened up" by the same generator expression come next iteration. The container is opened and new (and more)
		elements are repacked in it.

	"""
	# Require that factors is an iterable of tuple-coercibles
	tup_factors = [tuple(i) for i in factors]

	# init product to 1st factor * 1
	product = tup_factors[0]

	for factor in tup_factors[1:]:
		product = tuple(tupadd(a, b) for a in product for b in factor)
		#         RHS is precisely the def'n of the Cartesian Product, product x factor

	return product

def tupsumproduct(tree):
	"""
	This function performs a sum of products.

	Args:
		tree: dictionary of {iterable of parents : iterable of children}

	Returns:
		The tuple sum of products as a tuple, i.e., sum( keys x values )

	DEVELOPER'S CORNER:
		This generator expression will iterate through all its for clauses, before stringing them together.
		As tupadd() is guaranteed to return elementary tuples (so long as arguments are elementary tuples or scalars),
		the generator produces a sequence of elementary tuples, which are then finally packed in the tuple() constructor.
		I could have used a list(), but using tuple() is safe as tuple(any_tuple) == any_tuple. The external
		container type can be any sequence. I just chose tuples for both this function and tuppi so that the
		containers are immutable, if you want to use them directlty.

	"""
	prm_tuples = {totuple(key): totuple(val) for key, val in tree.items()}
	return tuple(tupadd(parent, child) for key, val in prm_tuples.items() for parent in key for child in val)
	#                                  SUM (of heterogenous key:val mappings)     PRODUCTS (key x val)


"""
=====================
Visualize Evoked data
=====================

In this tutorial we focus on the plotting functions of :class:`mne.Evoked`.
"""
import os.path as op
import numpy as np
import matplotlib.pyplot as plt

import mne

# sphinx_gallery_thumbnail_number = 9

###############################################################################
# First we read the evoked object from a file. Check out
# :ref:`tut_epoching_and_averaging` to get to this stage from raw data.
data_path = mne.datasets.sample.data_path()
fname = op.join(data_path, 'MEG', 'sample', 'sample_audvis-ave.fif')
evoked = mne.read_evokeds(fname, baseline=(None, 0), proj=True)
print(evoked)

###############################################################################
# Notice that ``evoked`` is a list of :class:`evoked <mne.Evoked>` instances.
# You can read only one of the categories by passing the argument ``condition``
# to :func:`mne.read_evokeds`. To make things more simple for this tutorial, we
# read each instance to a variable.
evoked_l_aud = evoked[0]
evoked_r_aud = evoked[1]
evoked_l_vis = evoked[2]
evoked_r_vis = evoked[3]

###############################################################################
# Let's start with a simple one. We plot event related potentials / fields
# (ERP/ERF). The bad channels are not plotted by default. Here we explicitly
# set the ``exclude`` parameter to show the bad channels in red. All plotting
# functions of MNE-python return a handle to the figure instance. When we have
# the handle, we can customise the plots to our liking.
fig = evoked_l_aud.plot(exclude=(), time_unit='s')

###############################################################################
# All plotting functions of MNE-python return a handle to the figure instance.
# When we have the handle, we can customise the plots to our liking. For
# example, we can get rid of the empty space with a simple function call.
fig.tight_layout()

###############################################################################
# Now we will make it a bit fancier and only use MEG channels. Many of the
# MNE-functions include a ``picks`` parameter to include a selection of
# channels. ``picks`` is simply a list of channel indices that you can easily
# construct with :func:`mne.pick_types`, :func:`mne.pick_channels`,
# :func:`mne.pick_channels_regexp`, or a list of strings that can be
# interpreted as channel names or channel types.
#
# Using ``spatial_colors=True``, the individual channel lines are color coded
# to show the sensor positions - specifically, the x, y, and z locations of
# the sensors are transformed into R, G and B values.
evoked_l_aud.plot(spatial_colors=True, gfp=True, picks='meg')

###############################################################################
# Notice the legend on the left. The colors would suggest that there may be two
# separate sources for the signals. This wasn't obvious from the first figure.
# Try painting the slopes with left mouse button. It should open a new window
# with topomaps (scalp plots) of the average over the painted area. There is
# also a function for drawing topomaps separately.
evoked_l_aud.plot_topomap(time_unit='s')

###############################################################################
# By default the topomaps are drawn from evenly spread out points of time over
# the evoked data. We can also define the times ourselves.
times = np.arange(0.05, 0.151, 0.05)
evoked_r_aud.plot_topomap(times=times, ch_type='mag', time_unit='s')

###############################################################################
# Or we can select automatically the peaks.
evoked_r_aud.plot_topomap(times='peaks', ch_type='mag', time_unit='s')

###############################################################################
# See :ref:`sphx_glr_auto_examples_visualization_plot_evoked_topomap.py` for
# more advanced topomap plotting options. You can also take a look at the
# documentation of :func:`mne.Evoked.plot_topomap` or simply write
# ``evoked_r_aud.plot_topomap?`` in your Python console to see the different
# parameters you can pass to this function. Most of the plotting functions also
# accept ``axes`` parameter. With that, you can customise your plots even
# further. First we create a set of matplotlib axes in a single figure and plot
# all of our evoked categories next to each other.
fig, ax = plt.subplots(1, 5, figsize=(8, 2))
kwargs = dict(times=0.1, show=False, vmin=-300, vmax=300, time_unit='s')
evoked_l_aud.plot_topomap(axes=ax[0], colorbar=True, **kwargs)
evoked_r_aud.plot_topomap(axes=ax[1], colorbar=False, **kwargs)
evoked_l_vis.plot_topomap(axes=ax[2], colorbar=False, **kwargs)
evoked_r_vis.plot_topomap(axes=ax[3], colorbar=False, **kwargs)
for ax, title in zip(ax[:4], ['Aud/L', 'Aud/R', 'Vis/L', 'Vis/R']):
    ax.set_title(title)
plt.show()

###############################################################################
# Notice that we created five axes, but had only four categories. The fifth
# axes was used for drawing the colorbar. You must provide room for it when you
# create this kind of custom plots or turn the colorbar off with
# ``colorbar=False``. That's what the warnings are trying to tell you. Also, we
# used ``show=False`` for the three first function calls. This prevents the
# showing of the figure prematurely. The behavior depends on the mode you are
# using for your Python session. See https://matplotlib.org/users/shell.html
# for more information.
#
# We can combine the two kinds of plots in one figure using the
# :func:`mne.Evoked.plot_joint` method of Evoked objects. Called as-is
# (``evoked.plot_joint()``), this function should give an informative display
# of spatio-temporal dynamics.
# You can directly style the time series part and the topomap part of the plot
# using the ``topomap_args`` and ``ts_args`` parameters. You can pass key-value
# pairs as a Python dictionary. These are then passed as parameters to the
# topomaps (:func:`mne.Evoked.plot_topomap`) and time series
# (:func:`mne.Evoked.plot`) of the joint plot.
# For an example of specific styling using these ``topomap_args`` and
# ``ts_args`` arguments, here, topomaps at specific time points
# (90 and 200 ms) are shown, sensors are not plotted (via an argument
# forwarded to `plot_topomap`), and the Global Field Power is shown:
ts_args = dict(gfp=True, time_unit='s')
topomap_args = dict(sensors=False, time_unit='s')
evoked_r_aud.plot_joint(title='right auditory', times=[.09, .20],
                        ts_args=ts_args, topomap_args=topomap_args)

###############################################################################
# Sometimes, you may want to compare two or more conditions at a selection of
# sensors, or e.g. for the Global Field Power. For this, you can use the
# function :func:`mne.viz.plot_compare_evokeds`. The easiest way is to create
# a  Python dictionary, where the keys are condition names and the values are
# :class:`mne.Evoked` objects. If you provide lists of :class:`mne.Evoked`
# objects, such as those for multiple subjects, the grand average is plotted,
# along with a confidence interval band - this can be used to contrast
# conditions for a whole experiment.
# First, we load in the evoked objects into a dictionary, setting the keys to
# '/'-separated tags (as we can do with event_ids for epochs). Then, we plot
# with :func:`mne.viz.plot_compare_evokeds`.
# The plot is styled with dict arguments, again using "/"-separated tags.
# We plot a MEG channel with a strong auditory response.
#
# For move advanced plotting using :func:`mne.viz.plot_compare_evokeds`.
# See also :ref:`sphx_glr_auto_tutorials_plot_metadata_epochs.py`.
conditions = ["Left Auditory", "Right Auditory", "Left visual", "Right visual"]
evoked_dict = dict()
for condition in conditions:
    evoked_dict[condition.replace(" ", "/")] = mne.read_evokeds(
        fname, baseline=(None, 0), proj=True, condition=condition)
print(evoked_dict)

colors = dict(Left="Crimson", Right="CornFlowerBlue")
linestyles = dict(Auditory='-', visual='--')
pick = evoked_dict["Left/Auditory"].ch_names.index('MEG 1811')

mne.viz.plot_compare_evokeds(evoked_dict, picks=pick, colors=colors,
                             linestyles=linestyles, split_legend=True)

###############################################################################
# We can also plot the activations as images. The time runs along the x-axis
# and the channels along the y-axis. The amplitudes are color coded so that
# the amplitudes from negative to positive translates to shift from blue to
# red. White means zero amplitude. You can use the ``cmap`` parameter to define
# the color map yourself. The accepted values include all matplotlib colormaps.
evoked_r_aud.plot_image(picks='meg')

###############################################################################
# Finally we plot the sensor data as a topographical view. In the simple case
# we plot only left auditory responses, and then we plot them all in the same
# figure for comparison. Click on the individual plots to open them bigger.
title = 'MNE sample data\n(condition : %s)'
evoked_l_aud.plot_topo(title=title % evoked_l_aud.comment,
                       background_color='k', color=['white'])
mne.viz.plot_evoked_topo(evoked, title=title % 'Left/Right Auditory/Visual',
                         background_color='w')

###############################################################################
# We can also plot the activations as arrow maps on top of the topoplot.
# The arrows represent an estimation of the current flow underneath the MEG
# sensors. Here, sample number 175 corresponds to the time of the maximum
# sensor space activity.
evoked_l_aud_mag = evoked_l_aud.copy().pick_types(meg='mag')
mne.viz.plot_arrowmap(evoked_l_aud_mag.data[:, 175], evoked_l_aud_mag.info)

###############################################################################
# Visualizing field lines in 3D
# -----------------------------
# We now compute the field maps to project MEG and EEG data to the MEG helmet
# and scalp surface.
#
# To do this, we need coregistration information. See
# :ref:`tut_forward` for more details. Here we just illustrate usage.

subjects_dir = data_path + '/subjects'
trans_fname = data_path + '/MEG/sample/sample_audvis_raw-trans.fif'

maps = mne.make_field_map(evoked_l_aud, trans=trans_fname, subject='sample',
                          subjects_dir=subjects_dir, n_jobs=1)

# Finally, explore several points in time
field_map = evoked_l_aud.plot_field(maps, time=.1)

###############################################################################
# .. note::
#     If trans_fname is set to None then only MEG estimates can be visualized.


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: notebooks//ipynb,markdown_files//md,python_scripts//py:percent
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.2'
#       jupytext_version: 1.2.1
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# %% [markdown] {"deletable": true, "editable": true}
# # Out-of-core Learning - Large Scale Text Classification for Sentiment Analysis

# %% [markdown] {"deletable": true, "editable": true}
# ## Scalability Issues

# %% [markdown] {"deletable": true, "editable": true}
# The `sklearn.feature_extraction.text.CountVectorizer` and `sklearn.feature_extraction.text.TfidfVectorizer` classes suffer from a number of scalability issues that all stem from the internal usage of the `vocabulary_` attribute (a Python dictionary) used to map the unicode string feature names to the integer feature indices.
#
# The main scalability issues are:
#
# - **Memory usage of the text vectorizer**: all the string representations of the features are loaded in memory
# - **Parallelization problems for text feature extraction**: the `vocabulary_` would be a shared state: complex synchronization and overhead
# - **Impossibility to do online or out-of-core / streaming learning**: the `vocabulary_` needs to be learned from the data: its size cannot be known before making one pass over the full dataset
#
#
# To better understand the issue let's have a look at how the `vocabulary_` attribute work. At `fit` time the tokens of the corpus are uniquely indentified by a integer index and this mapping stored in the vocabulary:

# %% {"deletable": true, "editable": true}
from sklearn.feature_extraction.text import CountVectorizer

vectorizer = CountVectorizer(min_df=1)

vectorizer.fit([
    "The cat sat on the mat.",
])
vectorizer.vocabulary_

# %% [markdown] {"deletable": true, "editable": true}
# The vocabulary is used at `transform` time to build the occurrence matrix:

# %% {"deletable": true, "editable": true}
X = vectorizer.transform([
    "The cat sat on the mat.",
    "This cat is a nice cat.",
]).toarray()

print(len(vectorizer.vocabulary_))
print(vectorizer.get_feature_names())
print(X)

# %% [markdown] {"deletable": true, "editable": true}
# Let's refit with a slightly larger corpus:

# %% {"deletable": true, "editable": true}
vectorizer = CountVectorizer(min_df=1)

vectorizer.fit([
    "The cat sat on the mat.",
    "The quick brown fox jumps over the lazy dog.",
])
vectorizer.vocabulary_

# %% [markdown] {"deletable": true, "editable": true}
# The `vocabulary_` is the (logarithmically) growing with the size of the training corpus. Note that we could not have built the vocabularies in parallel on the 2 text documents as they share some words hence would require some kind of shared datastructure or synchronization barrier which is complicated to setup, especially if we want to distribute the processing on a cluster.
#
# With this new vocabulary, the dimensionality of the output space is now larger:

# %% {"deletable": true, "editable": true}
X = vectorizer.transform([
    "The cat sat on the mat.",
    "This cat is a nice cat.",
]).toarray()

print(len(vectorizer.vocabulary_))
print(vectorizer.get_feature_names())
print(X)

# %% [markdown] {"deletable": true, "editable": true}
# ## The IMDb movie dataset

# %% [markdown] {"deletable": true, "editable": true}
# To illustrate the scalability issues of the vocabulary-based vectorizers, let's load a more realistic dataset for a classical text classification task: sentiment analysis on text documents. The goal is to tell apart negative from positive movie reviews from the [Internet Movie Database](http://www.imdb.com) (IMDb).
#
# In the following sections, with a [large subset](http://ai.stanford.edu/~amaas/data/sentiment/) of movie reviews from the IMDb that has been collected by Maas et al.
#
# - A. L. Maas, R. E. Daly, P. T. Pham, D. Huang, A. Y. Ng, and C. Potts. Learning Word Vectors for Sentiment Analysis. In the proceedings of the 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies, pages 142–150, Portland, Oregon, USA, June 2011. Association for Computational Linguistics.
#
# This dataset contains 50,000 movie reviews, which were split into 25,000 training samples and 25,000 test samples. The reviews are labeled as either negative (neg) or positive (pos). Moreover, *positive* means that a movie received >6 stars on IMDb; negative means that a movie received <5 stars, respectively.
#
#
# Assuming that the `../fetch_data.py` script was run successfully the following files should be available:

# %% {"deletable": true, "editable": true}
import os

train_path = os.path.join('datasets', 'IMDb', 'aclImdb', 'train')
test_path = os.path.join('datasets', 'IMDb', 'aclImdb', 'test')

# %% [markdown] {"deletable": true, "editable": true}
# Now, let's load them into our active session via scikit-learn's `load_files` function

# %% {"deletable": true, "editable": true}
from sklearn.datasets import load_files

train = load_files(container_path=(train_path),
                   categories=['pos', 'neg'])

test = load_files(container_path=(test_path),
                  categories=['pos', 'neg'])

# %% [markdown] {"deletable": true, "editable": true}
#

#     NOTE:
#

#

#       Since the movie datasets consists of 50,000 individual text files, executing the code snippet above may take ~20 sec or longer.
#

#

#

# %% [markdown] {"deletable": true, "editable": true}
# The `load_files` function loaded the datasets into `sklearn.datasets.base.Bunch` objects, which are Python dictionaries:

# %% {"deletable": true, "editable": true}
train.keys()

# %% [markdown] {"deletable": true, "editable": true}
# In particular, we are only interested in the `data` and `target` arrays.

# %% {"deletable": true, "editable": true}
import numpy as np

for label, data in zip(('TRAINING', 'TEST'), (train, test)):
    print('\n\n%s' % label)
    print('Number of documents:', len(data['data']))
    print('\n1st document:\n', data['data'][0])
    print('\n1st label:', data['target'][0])
    print('\nClass names:', data['target_names'])
    print('Class count:',
          np.unique(data['target']), ' -> ',
          np.bincount(data['target']))

# %% [markdown] {"deletable": true, "editable": true}
# As we can see above the `'target'` array consists of integers `0` and `1`, where `0` stands for negative and `1` stands for positive.

# %% [markdown] {"deletable": true, "editable": true}
# ## The Hashing Trick

# %% [markdown] {"deletable": true, "editable": true}
# Remember the bag of word representation using a vocabulary based vectorizer:
#
#

# %% [markdown] {"deletable": true, "editable": true}
# To workaround the limitations of the vocabulary-based vectorizers, one can use the hashing trick. Instead of building and storing an explicit mapping from the feature names to the feature indices in a Python dict, we can just use a hash function and a modulus operation:

# %% [markdown] {"deletable": true, "editable": true}
#

# %% [markdown] {"deletable": true, "editable": true}
# More info and reference for the original papers on the Hashing Trick in the [following site](http://www.hunch.net/~jl/projects/hash_reps/index.html) as well as a description specific to language [here](http://blog.someben.com/2013/01/hashing-lang/).

# %% {"deletable": true, "editable": true}
from sklearn.utils.murmurhash import murmurhash3_bytes_u32

# encode for python 3 compatibility
for word in "the cat sat on the mat".encode("utf-8").split():
    print("{0} => {1}".format(
        word, murmurhash3_bytes_u32(word, 0) % 2 ** 20))

# %% [markdown] {"deletable": true, "editable": true}
# This mapping is completely stateless and the dimensionality of the output space is explicitly fixed in advance (here we use a modulo `2 ** 20` which means roughly 1M dimensions). The makes it possible to workaround the limitations of the vocabulary based vectorizer both for parallelizability and online / out-of-core learning.

# %% [markdown] {"deletable": true, "editable": true}
# The `HashingVectorizer` class is an alternative to the `CountVectorizer` (or `TfidfVectorizer` class with `use_idf=False`) that internally uses the murmurhash hash function:

# %% {"deletable": true, "editable": true}
from sklearn.feature_extraction.text import HashingVectorizer

h_vectorizer = HashingVectorizer(encoding='latin-1')
h_vectorizer

# %% [markdown] {"deletable": true, "editable": true}
# It shares the same "preprocessor", "tokenizer" and "analyzer" infrastructure:

# %% {"deletable": true, "editable": true}
analyzer = h_vectorizer.build_analyzer()
analyzer('This is a test sentence.')

# %% [markdown] {"deletable": true, "editable": true}
# We can vectorize our datasets into a scipy sparse matrix exactly as we would have done with the `CountVectorizer` or `TfidfVectorizer`, except that we can directly call the `transform` method: there is no need to `fit` as `HashingVectorizer` is a stateless transformer:

# %% {"deletable": true, "editable": true}
docs_train, y_train = train['data'], train['target']
docs_valid, y_valid = test['data'][:12500], test['target'][:12500]
docs_test, y_test = test['data'][12500:], test['target'][12500:]

# %% [markdown] {"deletable": true, "editable": true}
# The dimension of the output is fixed ahead of time to `n_features=2 ** 20` by default (nearly 1M features) to minimize the rate of collision on most classification problem while having reasonably sized linear models (1M weights in the `coef_` attribute):

# %% {"deletable": true, "editable": true}
h_vectorizer.transform(docs_train)

# %% [markdown] {"deletable": true, "editable": true}
# Now, let's compare the computational efficiency of the `HashingVectorizer` to the `CountVectorizer`:

# %% {"deletable": true, "editable": true}
h_vec = HashingVectorizer(encoding='latin-1')
# %timeit -n 1 -r 3 h_vec.fit(docs_train, y_train)

# %% {"deletable": true, "editable": true}
count_vec =  CountVectorizer(encoding='latin-1')
# %timeit -n 1 -r 3 count_vec.fit(docs_train, y_train)

# %% [markdown] {"deletable": true, "editable": true}
# As we can see, the HashingVectorizer is much faster than the Countvectorizer in this case.

# %% [markdown] {"deletable": true, "editable": true}
# Finally, let us train a LogisticRegression classifier on the IMDb training subset:

# %% {"deletable": true, "editable": true}
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import Pipeline

h_pipeline = Pipeline([
    ('vec', HashingVectorizer(encoding='latin-1')),
    ('clf', LogisticRegression(random_state=1)),
])

h_pipeline.fit(docs_train, y_train)

# %% {"deletable": true, "editable": true}
print('Train accuracy', h_pipeline.score(docs_train, y_train))
print('Validation accuracy', h_pipeline.score(docs_valid, y_valid))

# %% {"deletable": true, "editable": true}
import gc

del count_vec
del h_pipeline

gc.collect()

# %% [markdown] {"deletable": true, "editable": true}
# # Out-of-Core learning

# %% [markdown] {"deletable": true, "editable": true}
# Out-of-Core learning is the task of training a machine learning model on a dataset that does not fit into memory or RAM. This requires the following conditions:
#
# - a **feature extraction** layer with **fixed output dimensionality**
# - knowing the list of all classes in advance (in this case we only have positive and negative reviews)
# - a machine learning **algorithm that supports incremental learning** (the `partial_fit` method in scikit-learn).
#
# In the following sections, we will set up a simple batch-training function to train an `SGDClassifier` iteratively.

# %% [markdown] {"deletable": true, "editable": true}
# But first, let us load the file names into a Python list:

# %% {"deletable": true, "editable": true}
train_path = os.path.join('datasets', 'IMDb', 'aclImdb', 'train')
train_pos = os.path.join(train_path, 'pos')
train_neg = os.path.join(train_path, 'neg')

fnames = [os.path.join(train_pos, f) for f in os.listdir(train_pos)] +\
         [os.path.join(train_neg, f) for f in os.listdir(train_neg)]

fnames[:3]

# %% [markdown] {"deletable": true, "editable": true}
# Next, let us create the target label array:

# %% {"deletable": true, "editable": true}
y_train = np.zeros((len(fnames), ), dtype=int)
y_train[:12500] = 1
np.bincount(y_train)

# %% [markdown] {"deletable": true, "editable": true}
# Now, we implement the `batch_train function` as follows:

# %% {"deletable": true, "editable": true}
from sklearn.base import clone

def batch_train(clf, fnames, labels, iterations=25, batchsize=1000, random_seed=1):
    vec = HashingVectorizer(encoding='latin-1')
    idx = np.arange(labels.shape[0])
    c_clf = clone(clf)
    rng = np.random.RandomState(seed=random_seed)

    for i in range(iterations):
        rnd_idx = rng.choice(idx, size=batchsize)
        documents = []
        for i in rnd_idx:
            with open(fnames[i], 'r', encoding='latin-1') as f:
                documents.append(f.read())
        X_batch = vec.transform(documents)
        batch_labels = labels[rnd_idx]
        c_clf.partial_fit(X=X_batch,
                          y=batch_labels,
                          classes=[0, 1])

    return c_clf

# %% [markdown] {"deletable": true, "editable": true}
# Note that we are not using `LogisticRegression` as in the previous section, but we will use a `SGDClassifier` with a logistic cost function instead. SGD stands for `stochastic gradient descent`, an optimization alrogithm that optimizes the weight coefficients iteratively sample by sample, which allows us to feed the data to the classifier chunk by chuck.

# %% [markdown] {"deletable": true, "editable": true}
# And we train the `SGDClassifier`; using the default settings of the `batch_train` function, it will train the classifier on 25*1000=25000 documents. (Depending on your machine, this may take >2 min)

# %% {"deletable": true, "editable": true}
from sklearn.linear_model import SGDClassifier

sgd = SGDClassifier(loss='log', random_state=1, max_iter=1000)

sgd = batch_train(clf=sgd,
                  fnames=fnames,
                  labels=y_train)

# %% [markdown] {"deletable": true, "editable": true}
# Eventually, let us evaluate its performance:

# %% {"deletable": true, "editable": true}
vec = HashingVectorizer(encoding='latin-1')
sgd.score(vec.transform(docs_test), y_test)

# %% [markdown] {"deletable": true, "editable": true}
# ### Limitations of the Hashing Vectorizer

# %% [markdown] {"deletable": true, "editable": true}
# Using the Hashing Vectorizer makes it possible to implement streaming and parallel text classification but can also introduce some issues:
#
# - The collisions can introduce too much noise in the data and degrade prediction quality,
# - The `HashingVectorizer` does not provide "Inverse Document Frequency" reweighting (lack of a `use_idf=True` option).
# - There is no easy way to inverse the mapping and find the feature names from the feature index.
#
# The collision issues can be controlled by increasing the `n_features` parameters.
#
# The IDF weighting might be reintroduced by appending a `TfidfTransformer` instance on the output of the vectorizer. However computing the `idf_` statistic used for the feature reweighting will require to do at least one additional pass over the training set before being able to start training the classifier: this breaks the online learning scheme.
#
# The lack of inverse mapping (the `get_feature_names()` method of `TfidfVectorizer`) is even harder to workaround. That would require extending the `HashingVectorizer` class to add a "trace" mode to record the mapping of the most important features to provide statistical debugging information.
#
# In the mean time to debug feature extraction issues, it is recommended to use `TfidfVectorizer(use_idf=False)` on a small-ish subset of the dataset to simulate a `HashingVectorizer()` instance that have the `get_feature_names()` method and no collision issues.

# %% [markdown] {"deletable": true, "editable": true}
#

#     EXERCISE:
#

#

#       In our implementation of the batch_train function above, we randomly draw *k* training samples as a batch in each iteration, which can be considered as a random subsampling ***with*** replacement. Can you modify the `batch_train` function so that it iterates over the documents ***without*** replacement, i.e., that it uses each document ***exactly once*** per iteration?
#

#

#

# %% {"deletable": true, "editable": true}
# # %load solutions/23_batchtrain.py


"""
.. _tut-fnirs-vis-brain:

Utilising Anatomical Information
================================

This example demonstrates how you can utilise anatomical and sensor position
information in your analysis pipeline. This information can be used to
verify measurement/analysis and also improve analysis accuracy
:footcite:`novi2020integration`.

This example demonstrates how to plot your data on a 3D brain
and overlay the sensor locations and regions of interest.

This tutorial glosses over the processing details, see the
:ref:`GLM tutorial <tut-fnirs-hrf>` for details on the preprocessing.

.. contents:: Page contents
   :local:
   :depth: 2

"""
# sphinx_gallery_thumbnail_number = 5

# Authors: Robert Luke <mail@robertluke.net>
#
# License: BSD (3-clause)

import numpy as np
import pandas as pd

import mne
from mne.preprocessing.nirs import optical_density, beer_lambert_law

import statsmodels.formula.api as smf

from mne_bids import BIDSPath, read_raw_bids, get_entity_vals
import mne_nirs

from mne_nirs.experimental_design import make_first_level_design_matrix
from mne_nirs.statistics import run_glm, statsmodels_to_results
from mne_nirs.channels import get_long_channels, get_short_channels
from mne_nirs.io.fold import fold_landmark_specificity
from mne_nirs.visualisation import plot_nirs_source_detector, plot_glm_surface_projection
from mne_nirs.datasets import fnirs_motor_group

# %%
# Download example data
# -------------------------------
#
# First, the data required data for this tutorial is downloaded.

# %%
# Download example fNIRS data
# ^^^^^^^^^^^^^^^^^^^^^^^^^^^
#
# Download the ``audio_or_visual_speech`` dataset and load the first measurement.

root = mne_nirs.datasets.audio_or_visual_speech.data_path()
dataset = BIDSPath(root=root, suffix="nirs", extension=".snirf", subject="04",
                   task="AudioVisualBroadVsRestricted", datatype="nirs", session="01")
raw = read_raw_bids(dataset)

# %%
# Download annotation information
# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
#
# Download the HCP-MMP parcellation.

# Download anatomical locations
subjects_dir = mne.datasets.sample.data_path() + '/subjects'
mne.datasets.fetch_hcp_mmp_parcellation(subjects_dir=subjects_dir, accept=True)
labels = mne.read_labels_from_annot('fsaverage', 'HCPMMP1', 'lh', subjects_dir=subjects_dir)
labels_combined = mne.read_labels_from_annot('fsaverage', 'HCPMMP1_combined', 'lh', subjects_dir=subjects_dir)

# %%
# Verify placement of sensors
# ---------------------------
#
# The first thing we can do is plot the location of the optodes and channels
# over an average brain surface to verify the data, specifically the 3D coordinates,
# have been loaded correctly. The sources are represented as red dots,
# the detectors are represented as black dots, the whit lines represent source-detector
# pairs, and the orange dots represent channel locations.
# In this example we can see channels over the left inferior frontal gyrus,
# auditory cortex, planum temporale, and occipital lobe.

brain = mne.viz.Brain('fsaverage', subjects_dir=subjects_dir, background='w', cortex='0.5')
brain.add_sensors(raw.info, trans='fsaverage', fnirs=['channels', 'pairs', 'sources', 'detectors'])
brain.show_view(azimuth=180, elevation=80, distance=450)

# %%
# .. _tut-fnirs-vis-brain-plot-3d-montage:
#
# Plot sensor channel numbers
# ---------------------------
# Often for publications and sanity checking, it's convenient to create an
# image showing the channel numbers along with the (typically) 10-20 location
# in the correct locations in a 3D view. The function
# :func:`mne_nirs.visualisation.plot_3d_montage` gives us this once we
# specify which views to use to show each channel pair:

view_map = {
    'left-lat': np.r_[np.arange(1, 27), 28],
    'caudal': np.r_[27, np.arange(43, 53)],
    'right-lat': np.r_[np.arange(29, 43), 44],
}

fig_montage = mne_nirs.visualisation.plot_3d_montage(
    raw.info, view_map=view_map, subjects_dir=subjects_dir)

# %%
# Plot sensor channels and anatomical region of interest
# ------------------------------------------------------
#
# Once the data has been loaded we can highlight anatomical regions of interest
# to ensure that the sensors are appropriately placed to measure from
# the relevant brain structures.
# In this example we highlight the primary auditory cortex in blue,
# and we can see that a number of channels are placed over this structure.

brain = mne.viz.Brain('fsaverage', subjects_dir=subjects_dir, background='w', cortex='0.5')
brain.add_sensors(raw.info, trans='fsaverage', fnirs=['channels', 'pairs', 'sources', 'detectors'])

aud_label = [label for label in labels if label.name == 'L_A1_ROI-lh'][0]
brain.add_label(aud_label, borders=False, color='blue')
brain.show_view(azimuth=180, elevation=80, distance=450)

# %%
# Plot channels sensitive to anatomical region of interest
# --------------------------------------------------------
#
# .. sidebar:: fOLD Toolbox
#
#    You should use the fOLD toolbox to pick your optode locations
#    when designing your experiment.
#    The tool is very intuitive and easy to use.
#    Be sure to cite the authors if you use their tool or data:
#
#    Morais, Guilherme Augusto Zimeo, Joana Bisol Balardin, and João Ricardo Sato. "fNIRS optodes’ location decider (fOLD): a toolbox for probe arrangement guided by brain regions-of-interest." Scientific reports 8.1 (2018): 1-11.
#
# Rather than simply eye balling the sensor and ROIs of interest, we can
# quantify the specificity of each channel to the anatomical region of interest
# and select channels that are sufficiently sensitive for further analysis.
# In this example we highlight the left inferior frontal gyrus (IFG) and
# use data from the fOLD toolbox :footcite:`morais2018fnirs`.
# To see more details about how to use the fOLD data see
# :ref:`this tutorial <tut-fnirs-group-relating>`.

# Return specificity of each channel to the Left IFG
specificity = fold_landmark_specificity(raw, 'L IFG (p. Triangularis)')

# Retain only channels with specificity to left IFG of greater than 50%
raw_IFG = raw.copy().pick(picks=np.where(specificity > 50)[0])

brain = mne.viz.Brain('fsaverage', subjects_dir=subjects_dir, background='w', cortex='0.5')
brain.add_sensors(raw_IFG.info, trans='fsaverage', fnirs=['channels', 'pairs'])

ifg_label = [label for label in labels_combined if label.name == 'Inferior Frontal Cortex-lh'][0]
brain.add_label(ifg_label, borders=False, color='green')

brain.show_view(azimuth=140, elevation=95, distance=360)

# %%
#
# Alternatively, we can retain all channels and visualise the specificity of each channel the ROI
# by encoding the specificty in the color of the line between each source and detector.
# In this example we see that several channels have substantial specificity to
# the region of interest.
#
# Note: this function currently doesn't support the new MNE brain API, so does
# not allow the same behaviour as above (adding sensors, highlighting ROIs etc).
# It should be updated in the near future.

fig = plot_nirs_source_detector(specificity, raw.info, surfaces='brain',
                                subject='fsaverage', subjects_dir=subjects_dir, trans='fsaverage')
mne.viz.set_3d_view(fig, azimuth=140, elevation=95)

# %%
# Anatomically informed weighting in region of interest analysis
# --------------------------------------------------------------
#
# As observed above, some channels have greater specificity to the desired
# brain region than other channels.
# Thus, when doing a region of interest analysis you may wish to give extra
# weight to channels with greater sensitivity to the desired ROI.
# This can be done by manually specifying the weights used in the region of
# interest function call.
# The details of the GLM analysis will not be described here, instead view the
# :ref:`fNIRS GLM tutorial <tut-fnirs-hrf>`. Instead, comments are provided
# for the weighted region of interest function call.

# Basic pipeline, simplified for example
raw_od = optical_density(raw)
raw_haemo = beer_lambert_law(raw_od)
raw_haemo.resample(0.3).pick("hbo")  # Speed increase for web server
sht_chans = get_short_channels(raw_haemo)
raw_haemo = get_long_channels(raw_haemo)
design_matrix = make_first_level_design_matrix(raw_haemo, stim_dur=13.0)
design_matrix["ShortHbO"] = np.mean(sht_chans.copy().pick(picks="hbo").get_data(), axis=0)
glm_est = run_glm(raw_haemo, design_matrix)

# First we create a dictionary for each region of interest.
# Here we include all channels in each ROI, as we will later be applying
# weights based on their specificity to the brain regions of interest.
rois = dict()
rois["Audio_weighted"] = range(len(glm_est.ch_names))
rois["Visual_weighted"] = range(len(glm_est.ch_names))

# Next we compute the specificity for each channel to the auditory and visual cortex.
spec_aud = fold_landmark_specificity(raw_haemo, '42 - Primary and Auditory Association Cortex', atlas="Brodmann")
spec_vis = fold_landmark_specificity(raw_haemo, '17 - Primary Visual Cortex (V1)', atlas="Brodmann")

# Next we create a dictionary to store the weights for each channel in the ROI.
# The weights will be the specificity to the ROI.
# The keys and length of each dictionary entry must match the ROI dictionary.
weights = dict()
weights["Audio_weighted"] = spec_aud
weights["Visual_weighted"] = spec_vis

# Finally we compute region of interest results using the weights specified above
out = glm_est.to_dataframe_region_of_interest(rois, ["Video", "Control"], weighted=weights)
out["Significant"] = out["p"] < 0.05
out

# %%
# In the table above we observe that the response to the visual condition
# is only present in the visual region of interest. You can use this
# technique to load any custom weighting, including weights exported from
# other software.

# %%
# Preprocess fNIRS data
# ---------------------
#
# We can also use the 3D information to project the results on to the cortical surface.
# First, we process the fNIRS data. This is a duplication of the GLM tutorial
# analysis. The details will not be described here, instead view the
# :ref:`fNIRS GLM tutorial <tut-fnirs-hrf>`.

def individual_analysis(bids_path, ID):

    raw_intensity = read_raw_bids(bids_path=bids_path, verbose=False)
     # sanitize event names
    raw_intensity.annotations.description[:] = [
        d.replace('/', '_') for d in raw_intensity.annotations.description]

    # Convert signal to haemoglobin and resample
    raw_od = optical_density(raw_intensity)
    raw_haemo = beer_lambert_law(raw_od, ppf=0.1)
    raw_haemo.resample(0.3)

    # Cut out just the short channels for creating a GLM repressor
    sht_chans = get_short_channels(raw_haemo)
    raw_haemo = get_long_channels(raw_haemo)

    # Create a design matrix
    design_matrix = make_first_level_design_matrix(raw_haemo, stim_dur=5.0)

    # Append short channels mean to design matrix
    design_matrix["ShortHbO"] = np.mean(sht_chans.copy().pick(picks="hbo").get_data(), axis=0)
    design_matrix["ShortHbR"] = np.mean(sht_chans.copy().pick(picks="hbr").get_data(), axis=0)

    # Run GLM
    glm_est = run_glm(raw_haemo, design_matrix)

    # Extract channel metrics
    cha = glm_est.to_dataframe()

    # Add the participant ID to the dataframes
    cha["ID"] = ID

    # Convert to uM for nicer plotting below.
    cha["theta"] = [t * 1.e6 for t in cha["theta"]]

    return raw_haemo, cha

# Get dataset details
root = fnirs_motor_group.data_path()
dataset = BIDSPath(root=root, task="tapping",
                   datatype="nirs", suffix="nirs", extension=".snirf")
subjects = get_entity_vals(root, 'subject')

df_cha = pd.DataFrame()  # To store channel level results
for sub in subjects:  # Loop from first to fifth subject

    # Create path to file based on experiment info
    bids_path = dataset.update(subject=sub)

    # Analyse data and return both ROI and channel results
    raw_haemo, channel = individual_analysis(bids_path, sub)

    # Append individual results to all participants
    df_cha = pd.concat([df_cha, channel], ignore_index=True)

ch_summary = df_cha.query("Condition in ['Tapping_Right']")
assert len(ch_summary)
ch_summary = ch_summary.query("Chroma in ['hbo']")
ch_model = smf.mixedlm("theta ~ -1 + ch_name", ch_summary,
                       groups=ch_summary["ID"]).fit(method='nm')
model_df = statsmodels_to_results(ch_model, order=raw_haemo.copy().pick("hbo").ch_names)

# %%
# Plot surface projection of GLM results
# --------------------------------------
#
# Finally, we can project the GLM results from each channel to the nearest cortical surface
# and overlay the sensor positions and two different regions of interest.
# In this example we also highlight the premotor cortex and auditory association cortex
# in green and blue respectively.

# Plot the projection and sensor locations
brain = plot_glm_surface_projection(raw_haemo.copy().pick("hbo"), model_df, colorbar=True)
brain.add_sensors(raw_haemo.info, trans='fsaverage', fnirs=['channels', 'pairs', 'sources', 'detectors'])

# mark the premotor cortex in green
aud_label = [label for label in labels_combined if label.name == 'Premotor Cortex-lh'][0]
brain.add_label(aud_label, borders=True, color='green')

# mark the auditory association cortex in blue
aud_label = [label for label in labels_combined if label.name == 'Auditory Association Cortex-lh'][0]
brain.add_label(aud_label, borders=True, color='blue')

brain.show_view(azimuth=160, elevation=60, distance=400)

# %%
# Bibliography
# -----------------------------------------------
#
# .. footbibliography::