#@title Python imports
import collections
import datetime
from functools import partial
import math

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
from scipy import stats
import seaborn as sns

from sklearn.datasets import make_regression
from sklearn.model_selection import train_test_split

from IPython.display import display, Math
import sympy

# Make colab plots larger
plt.rcParams['figure.figsize'] = [8, 6]
plt.rcParams['figure.dpi'] = 100

# Functions to display matrices nicely

def matrix_to_latex(a):
  return sympy.latex(sympy.Matrix(np.round(a, 3)))

def display_latex(s, newline=True):
  display(Math(s))
  if newline:
    print()

"""
Explanation: Copyright 2021 Google LLC
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.
Make a copy of this notebook!
End of explanation
"""


import jax
import jax.numpy as jnp
from jax import jit, grad, vmap, api, random, jacfwd, jacrev
from jax.experimental import optimizers, stax

"""
Explanation: Jax
Autodiff cookbook
Jax is a python library that provides the ability to differentiate many python
functions.
End of explanation
"""


# Define our functions

def f(x):
  return jnp.sqrt(x)

def g(x):
  return x*x + 1

def y(x):
  return f(g(x))

# We'll use x0 as our point so as not to over use the letter x
x0 = 1.

# Compute the gradient functions.
dy_dx = jacfwd(y)
dg_dx = jacfwd(g)
df_dg = jacfwd(f)

# Compute the values and print
display_latex(r"\frac{{d y}}{{d x}}({}) = {}".format(x0, dy_dx(x0)))
display_latex(r"\frac{{dg}}{{dx}}({}) = {}".format(x0, dg_dx(x0)))
display_latex(r"\frac{{\partial f}}{{dx}} ({}) = {}".format(g(x0), df_dg(g(x0))))
display_latex(r"{} * {} = {}".format(dg_dx(x0), df_dg(g(x0)), np.dot(dg_dx(x0), df_dg(g(x0)))))

"""
Explanation: Section 1: Multivariate Chain Rules
In this section we'll cover:

Using Jax
Multivariate chain rule examples

Jacobians and Chain rules
Let's verify the chain rule for some examples similar to those from the slides.
Single variable
First let's simply try a single variable example.
$$ y = \sqrt{x^2+1} = f(g(x))$$
where
$$ f(x) = \sqrt{x}$$
and 
$$g(x) = x^2+1$$
Let's compute the derivatives with Jax. The chain rule is:
$$ \frac{dy}{dx} = \frac{df}{dg} \frac{dg}{dx}$$
In this case all the inputs and outputs are scalars so the multiplication of the derivatives is just ordinary scalar multiplication.
End of explanation
"""


# Experiments here!

def f(x):
  # Your outer function
  pass

def g(x):
  # Your inner function
  pass

def y(x):
  # The composite
  return f(g(x))

# Copy the derivatives code from the above cells


# @title Solution (double click to show)

def f(x):
  return 1. / (1 + jnp.exp(x))
  pass

def g(x):
  # Your inner function
  return -2 * x

def y(x):
  # The composite
  return f(g(x))

# We'll use x0 as our point so as not to over use the letter x
x0 = 0.

# Compute the gradient functions.
dy_dx = jacfwd(y)
dg_dx = jacfwd(g)
df_dg = jacfwd(f)

# Compute the values and print
display_latex(r"\frac{{d y}}{{d x}}({}) = {}".format(x0, dy_dx(x0)))
display_latex(r"\frac{{dg}}{{dx}}({}) = {}".format(x0, dg_dx(x0)))
display_latex(r"\frac{{\partial f}}{{dx}} ({}) = {}".format(g(x0), df_dg(g(x0))))
display_latex(r"{} * {} = {}".format(dg_dx(x0), df_dg(g(x0)), np.dot(dg_dx(x0), df_dg(g(x0)))))


"""
Explanation: So we see that JAX gets the same answer by computing the derivative of the composite or by multiplying together the two component derivatives.
Try it out
Feel free to experiment with other functions and compositions.
Need an idea? Try a sigmoid like
$$f(x) = \frac{1}{1 + e^{-2x}}$$
End of explanation
"""


def f(x, y):
  return jnp.exp((x*x - y*y)/4)

def g(s, t):
  return s + t

def h(s, t):
  return s - t

# We'll use z for the composite so as not to define f twice in code

def z(s, t):
  return f(g(s, t), h(s, t))


"""
Explanation: Multivariate
Now let's try a multivariate example. Consider the function
$$z = f(x, y) = e^{(x^2 - y^2) / 4}$$
and suppose that $x$ and $y$ are themselves functions of $s$ and $t$:
$$x = g(s, t) = s + t$$
$$y = h(s, t) = s - t$$
Let's compute the two partial derivatives
$$\frac{\partial f}{\partial s} = \frac{\partial f}{\partial x} \frac{\partial x}{\partial s} + \frac{\partial f}{\partial y} \frac{\partial y}{\partial s}$$
$$\frac{\partial f}{\partial t} = \frac{\partial f}{\partial x} \frac{\partial x}{\partial t} + \frac{\partial f}{\partial y} \frac{\partial y}{\partial t}$$
First let's define the functions and compute the partial derivative with respect to $t$.
End of explanation
"""


# Use JAX to get the Jacobian, which in this case is again a function that returns
# a scalar value
# argnums=0 gives us the partial with respect to the first variable
dz_ds = jacfwd(z, argnums=0)

for s, t in [(1., 0.), (0., 1.), (1., 1.)]:
  display_latex(r"\frac{{\partial z}}{{\partial s}}({}, {}) = {}".format(s, t, dz_ds(s, t)))


"""
Explanation: Now we can compute the partial
$$\frac{\partial z}{\partial s}$$
at some example values.
End of explanation
"""


dz_dx = jacfwd(f, argnums=0)
dx_ds = jacfwd(g, argnums=0)

dz_dy = jacfwd(f, argnums=1)
dy_ds = jacfwd(h, argnums=0)

for s, t in [(1., 0.), (0., 1.), (1., 1.)]:
  x_ = g(s, t)
  y_ = h(s, t)
  dz_ds = dz_dx(x_, y_) *  dx_ds(s, t) + dz_dy(x_, y_) * dy_ds(s, t)
  display_latex(r"\frac{{\partial z}}{{\partial s}}({}, {}) = {}".format(s, t, dz_ds))

"""
Explanation: Let's compare to the right-hand-side of the partial formula:
$$\frac{\partial f}{\partial s} = \frac{\partial f}{\partial x} \frac{\partial x}{\partial s} + \frac{\partial f}{\partial y} \frac{\partial y}{\partial s}$$
End of explanation
"""


# Your code here (copy and modify the above code cells)

# @title Solution (double click to show)

# Use JAX to get the Jacobian, which in this case is a function that returns
# a scalar value
# argnums=0 gives us the partial with respect to the first variable
dz_dt = jacfwd(z, argnums=1)

for s, t in [(1., 0.), (0., 1.), (1., 1.)]:
  display_latex(r"\frac{{\partial z}}{{\partial s}}({}, {}) = {}".format(t, s, dz_dt(s, t)))


dz_dx = jacfwd(f, argnums=0)
dx_dt = jacfwd(g, argnums=1)

dz_dy = jacfwd(f, argnums=1)
dy_dt = jacfwd(h, argnums=1)

for s, t in [(1., 0.), (0., 1.), (1., 1.)]:
  x = g(s, t)
  y = h(s, t)
  dz_dt = dz_dx(x, y) *  dx_dt(s, t) + dz_dy(x, y) * dy_dt(s, t)
  display_latex(r"\frac{{\partial z}}{{\partial s}}({}, {}) = {}".format(s, t, dz_dt))

"""
Explanation: Same results!
Exercise
Compute $\frac{\partial z}{\partial t}$ both ways as above.
End of explanation
"""


# Fix the dimension
n = 4

def f(x):
  return jnp.sum(x * x)

def g(t):
  return t * np.array([i for i in range(1, n+1)])

def h(t):
  return f(g(t))

# Compute the partials
t = 1.
dz_dx = jacfwd(f)
dx_dt = jacfwd(g)
dh_dt = jacfwd(h)

display_latex(r"\frac{{d f}}{{d t}}({}) = {}".format(t, dh_dt((t))))
display_latex(r"\frac{{\partial \mathbf{{x}}}}{{\partial t}}({}) = {}".format(t, dx_dt(t)))
display_latex(r"\frac{{\partial f}}{{\partial \mathbf{{x}} }} ({}) = {}".format(g(t), dz_dx(g(t))))
display_latex(r"{} * {} = {}".format(dx_dt(t), dz_dx(g(t)), np.dot(dx_dt(t), dz_dx(g(t)))))

"""
Explanation: Example
Let's do another example, this time with more inputs.
$$ z = f(\textbf{x}) = x_1^2 + \cdots + x_n^2$$
with
$$ x = g(t) = (t, 2t, \ldots, nt)$$
Now we compute both sides of
$$\frac{df}{dt} = \frac{\partial f}{\partial \mathbf{x}} \frac{\partial x}{\partial t}$$
On an actual point $t_0$, note carefully the arguments on which the Jacobians are evaluated:
$$\frac{df}{dt}(t_0) = \frac{\partial f}{\partial \mathbf{x}}(x(t_0)) \frac{\partial x}{\partial t}(t_0)$$
End of explanation
"""


# Dimensions of our matrices, n x m and m x k
n = 3
m = 4
k = 2

# Two random matrices
A_1 = np.random.rand(n, m)
A_2 = np.random.rand(m, k)

display_latex(r"A_1 = {}".format(matrix_to_latex(A_1)))
display_latex(r"A_2 = {}".format(matrix_to_latex(A_2)))

def f(x):
  return jnp.dot(A_1, x)

def g(x):
  return jnp.dot(A_2, x)

def h(x):
  return f(g(x))

test_point = jnp.array([1.]*k)

display_latex(r"\mathbf{{x}} = {}".format(matrix_to_latex(test_point)))
display_latex(r"g(\mathbf{{x}}) = {}".format(matrix_to_latex(g(test_point))))
display_latex(r"f(g(\mathbf{{x}})) = {}".format(matrix_to_latex(f(g(test_point)))))
display_latex(r"h(\mathbf{{x}}) = {}".format(matrix_to_latex(h(test_point))))

"""
Explanation: Exercise: Try different values of $n$ and $t$ in the above example
Composition of vector-valued functions
Consider two functions $f(x) = A_1x$ and $g(x) = A_2x$ defined by an $n \times m$ matrix $A$ and an $m \times k$ matrix $A_2$.
Since the dimensions of the
matrices $A$ and $B$ match up, we can form the composition $h(x) = f(g(x)) = A_1 A_2 x$.
First let's define the functions.
End of explanation
"""


# This time we'll go ahead and evaluate 
dh_dx = jacfwd(h)(test_point)
dg_dx = jacfwd(g)(test_point)
df_dg = jacfwd(f)(g(test_point))

display_latex(r"\frac{{\partial f}}{{\partial g}} = {}".format(matrix_to_latex(df_dg)))
display_latex(r"\frac{{\partial g}}{{\partial x}} = {}".format(matrix_to_latex(dg_dx)))
display_latex(r"\frac{{\partial h}}{{\partial x}} = {}".format(matrix_to_latex(dh_dx)))

display_latex(r"\frac{{\partial f}}{{\partial g}} \frac{{\partial g}}{{\partial x}} = {} {} = {}".format(
    matrix_to_latex(df_dg),
    matrix_to_latex(dg_dx),
    matrix_to_latex(np.dot(df_dg, dg_dx))))

"""
Explanation: Now let's compare the derivatives. We want to see that the chain rule holds
$$ \frac{\partial \mathbf{h}}{\partial \mathbf{x}} = \frac{\partial \mathbf{f}}{\partial \mathbf{g}} \frac{\partial \mathbf{g}}{\partial \mathbf{x}}$$
Again note carefully on which arguments we evaluate the Jacobians:
$$ \frac{\partial \mathbf{h}}{\partial \mathbf{x}}(\mathbf{x}) = \frac{\partial \mathbf{f}}{\partial \mathbf{g}}(g(\mathbf{x})) \frac{\partial \mathbf{g}}{\partial \mathbf{x}}(\mathbf{x})$$
End of explanation
"""


test_point = jnp.array(np.identity(k))

## Take the 0-th index to collapse the outer array layer
dh_dx = jacfwd(h)(test_point[0])
dg_dx = jacfwd(g)(test_point[0])
df_dg = jacfwd(f)(jnp.array(np.identity(m))[0])

display_latex(r"\frac{{\partial f}}{{\partial g}} = {} = A_1".format(matrix_to_latex(df_dg)))
display_latex(r"\frac{{\partial g}}{{\partial x}} = {} = A_2".format(matrix_to_latex(dg_dx)))
display_latex(r"\frac{{\partial h}}{{\partial x}} = {} = A_1 A_2".format(matrix_to_latex(dh_dx)))
display_latex(r"A_1 A_2 = {}".format(matrix_to_latex(np.dot(A_1, A_2))))

"""
Explanation: So we see that the chain rule holds for our test point.
If we test against identity matrices, we should get back $A_1$, $A_2$, and $A_1 A_2$.
End of explanation
"""


def activation(x):
  return jax.nn.sigmoid(x)
  # return 1. / (1. + jnp.exp(-x))

def g(W, b, x):
  # vmap applies the activation component-wise, in case
  # the output is a vector
  return vmap(activation)(jnp.dot(W, x) + b)

# Initialize W and b
# Choose their sizes
n = 1
m = 3

# Pick two random matrices
W = np.random.rand(n, m)
b = np.random.rand(n, 1)

print("Random initial parameter values")

display_latex(r"W = {}".format(matrix_to_latex(W)))
display_latex(r"b = {}".format(matrix_to_latex(b)))

"""
Explanation: Section 2: Neural Networks
Now we'll look at multivariate derivatives and gradient descent,
again using Jax.
Basic Neural Network

Now let's look at a basic one-layer neural network, which is a function of the form:
$$ g(\mathbf{x}) = f(W \mathbf{x} + \mathbf{b})$$
where the activation function is sigmoid:
$$f(x) = \frac{1}{1 + e^{-x}}$$
Let's see how to compute the derivatives for the weights and bias on a given data point.
End of explanation
"""


def loss(y, W, b, x):
  return (y - g(W, b, x))**2

# Compute the derivative of the loss function at an initial point
# from our "data"
y = np.random.rand(n, 1)
x_0 = np.random.rand(m, 1)

# Uncomment if you want to see the values
# print(x_0)
# print(y)

print("Partials for weights ({} x {} values):\n".format(n, m))
dloss_W = jacrev(loss, argnums=1)(y, W, b, x_0)
print(dloss_W)

print("\nPartials for bias ({} values):\n".format(n))
dloss_b = jacrev(loss, argnums=2)(y, W, b, x_0)
print(dloss_b)

"""
Explanation: Let's use squared error as the loss for a data point $(x, y)$:
$$C(y, x) = (y - g(x))^2$$
End of explanation
"""


## Data representing XOR
data = np.array([[0, 0], [0, 1], [1, 0], [1, 1]], dtype=np.float32)
labels = np.array([0, 1, 1, 0], dtype=np.float32)

plt.scatter(data[:,0], data[:,1], c=labels, cmap='viridis', s=100)
plt.show()

"""
Explanation: We could then use these partial derivatives to update our weights, and iterate
over additional data points. More examples here.
A single layer network like the one above just isn't very powerful. More complicated networks could be build up as a composition of such functions. But it's more common to specify the network architecture layers as in the following example.
XOR and Neural Networks
(Example from safyan@)
XOR is a classic example of a function that can't be well-fit with a single layer neural network. It can, however, be fit with a two layer network.
Here we'll use JAX to fit some candidate neural networks. See here for non-JAX example that also explains gradient descent and backpropagation.
End of explanation
"""


# Initialize a fully-connected network
def init_params(key, in_dim, out_dims: list):
  params = {}
  num_layers = len(out_dims)
  # Build up the layers iteratively
  for layer in range(num_layers):
    out_dim = out_dims[layer]
    key, subkey1, subkey2 = random.split(key, num=3)
    W = jax.nn.initializers.glorot_normal()(subkey1, (in_dim, out_dim))
    b = jax.nn.initializers.normal()(subkey2, (out_dim,))
    params["layer_{}".format(layer)] = {'W': W, 'b': b}
    in_dim = out_dim
  return params

@jit
def predict(params, inputs):
  # This is the forward pass where we calculate the outputs of each layer
  n_layers = len(params.keys())
  for i, layer in enumerate(params.keys()):
    W, b = params[layer].values()
    if i == n_layers:
      outputs = jnp.dot(inputs, W) + b
    # Apply this activation to the hidden layers
    else:
      outputs = jax.nn.tanh(jnp.dot(inputs, W) + b)
    inputs = outputs
  return jax.nn.sigmoid(outputs)

"""
Explanation: Now we construct neural networks to fit to this data.
End of explanation
"""


@jit
def cross_entropy(y_hat, y):
  return -y * jnp.log(y_hat) - (1 - y) * jnp.log(1 - y_hat)

@jit
def loss_fn(params, data, labels):
  predictions = predict(params, data)
  return jnp.mean(cross_entropy(predictions, labels))

@jit
def update(params, grads, alpha=0.1):
  # Gradient descent update
  return jax.tree_multimap(lambda p, g: p - alpha * g, params, grads)

"""
Explanation: We'll use cross entropy for the loss function.
$$L = -y \log {\hat{y}} - (1-y) \log{(1-\hat{y}}) $$
End of explanation
"""


# Network with no hidden layers
key = random.PRNGKey(0)
params = init_params(key, in_dim=2, out_dims=[1])
params

# Parameters for a network with one hidden layer with 10 neurons
key = random.PRNGKey(0)
params = init_params(key, in_dim=2, out_dims=[10, 1])
params

"""
Explanation: For different choices of parameters we obtain different network architectures. For example we may want one or more hidden inner layers, or more neurons per hidden layer.
End of explanation
"""


# per example gradients
key = random.PRNGKey(0)
params = init_params(key, in_dim=2, out_dims=[1])
vmap(partial(grad(loss_fn), params))(data, labels)

"""
Explanation: We can use grad to compute the derivatives.
End of explanation
"""


tkey = random.PRNGKey(0)
params = init_params(key, in_dim=2, out_dims=[1])
loss_grad_fn = grad(loss_fn)
losses = []
iterations = 500
for i in range(iterations):
  grads = loss_grad_fn(params, data, labels)
  losses.append(loss_fn(params, data, labels))
  params = update(params, grads, alpha=.1)

plt.plot(losses)
plt.xlabel("Iteration")
plt.ylabel("Cross entropy loss")

"""
Explanation: First, we'll try to fit XOR with a one-layer network (i.e, no hidden layers) with gradient descent.
End of explanation
"""


fig, ax = plt.subplots()
xx, yy = np.mgrid[-.1:1.1:.01, -.1:1.1:.01]
grid = np.c_[xx.ravel(), yy.ravel()]
probs = predict(params, grid).reshape(xx.shape)
contour = ax.contourf(xx, yy, probs, levels=50)
fig.colorbar(contour)
ax.scatter(data[:,0], data[:,1], c=labels, s=100, cmap='viridis')
plt.show()

"""
Explanation: Plotting the predictions of the network on a grid, we find that we can't fit XOR. In hindsight, this is forseeable because the classes (0 and 1) are not linearly separable.
End of explanation
"""


key = random.PRNGKey(9)
params = init_params(key, in_dim=2, out_dims=[10, 1])
loss_grad_fn = grad(loss_fn)
losses = []
iterations = 100
for i in range(iterations):
  grads = loss_grad_fn(params, data, labels)
  losses.append(loss_fn(params, data, labels))
  params = update(params, grads, alpha=.8)

plt.plot(losses)
plt.xlabel("Iteration")
plt.ylabel("Cross entropy loss")

"""
Explanation: What about a two-layer network? With a two-layer neural network, we add a "hidden" layer between the input and output. The hidden layer applies an affine function to the input followed by a non-linear "activation" function (in this case, hyperbolic tangent).
End of explanation
"""


fig, ax = plt.subplots()
xx, yy = np.mgrid[-.1:1.1:.01, -.1:1.1:.01]
grid = np.c_[xx.ravel(), yy.ravel()]
probs = predict(params, grid).reshape(xx.shape)
contour = ax.contourf(xx, yy, probs, levels=50)
fig.colorbar(contour)
ax.scatter(data[:,0], data[:,1], c=labels, s=100, cmap='viridis')
plt.show()

"""
Explanation: Again, we plot the predictions of the network and find that this time we're able to correctly classify the points.
End of explanation
"""


