# This file is generated automatically through:
#    d2lbook build lib
# Don't edit it directly

# Defined in file: ./chapter_preface/index.md
import collections
import hashlib
import math
import os
import random
import re
import shutil
import sys
import tarfile
import time
import zipfile
from collections import defaultdict

import pandas as pd
import requests
from IPython import display
from matplotlib import pyplot as plt

d2l = sys.modules[__name__]

# Defined in file: ./chapter_preface/index.md
import numpy as np
import tensorflow as tf

# Defined in file: ./chapter_preliminaries/calculus.md
def use_svg_display():
    """Use the svg format to display a plot in Jupyter."""
    display.set_matplotlib_formats('svg')

# Defined in file: ./chapter_preliminaries/calculus.md
def set_figsize(figsize=(3.5, 2.5)):
    """Set the figure size for matplotlib."""
    use_svg_display()
    d2l.plt.rcParams['figure.figsize'] = figsize

# Defined in file: ./chapter_preliminaries/calculus.md
def set_axes(axes, xlabel, ylabel, xlim, ylim, xscale, yscale, legend):
    """Set the axes for matplotlib."""
    axes.set_xlabel(xlabel)
    axes.set_ylabel(ylabel)
    axes.set_xscale(xscale)
    axes.set_yscale(yscale)
    axes.set_xlim(xlim)
    axes.set_ylim(ylim)
    if legend:
        axes.legend(legend)
    axes.grid()

# Defined in file: ./chapter_preliminaries/calculus.md
def plot(X, Y=None, xlabel=None, ylabel=None, legend=None, xlim=None,
         ylim=None, xscale='linear', yscale='linear',
         fmts=('-', 'm--', 'g-.', 'r:'), figsize=(3.5, 2.5), axes=None):
    """Plot data points."""
    if legend is None:
        legend = []

    set_figsize(figsize)
    axes = axes if axes else d2l.plt.gca()

    # Return True if `X` (tensor or list) has 1 axis
    def has_one_axis(X):
        return (hasattr(X, "ndim") and X.ndim == 1 or
                isinstance(X, list) and not hasattr(X[0], "__len__"))

    if has_one_axis(X):
        X = [X]
    if Y is None:
        X, Y = [[]] * len(X), X
    elif has_one_axis(Y):
        Y = [Y]
    if len(X) != len(Y):
        X = X * len(Y)
    axes.cla()
    for x, y, fmt in zip(X, Y, fmts):
        if len(x):
            axes.plot(x, y, fmt)
        else:
            axes.plot(y, fmt)
    set_axes(axes, xlabel, ylabel, xlim, ylim, xscale, yscale, legend)

# Defined in file: ./chapter_linear-networks/linear-regression.md
class Timer:
    """Record multiple running times."""
    def __init__(self):
        self.times = []
        self.start()

    def start(self):
        """Start the timer."""
        self.tik = time.time()

    def stop(self):
        """Stop the timer and record the time in a list."""
        self.times.append(time.time() - self.tik)
        return self.times[-1]

    def avg(self):
        """Return the average time."""
        return sum(self.times) / len(self.times)

    def sum(self):
        """Return the sum of time."""
        return sum(self.times)

    def cumsum(self):
        """Return the accumulated time."""
        return np.array(self.times).cumsum().tolist()

# Defined in file: ./chapter_linear-networks/linear-regression-scratch.md
def synthetic_data(w, b, num_examples):
    """Generate y = Xw + b + noise."""
    X = d2l.zeros((num_examples, w.shape[0]))
    X += tf.random.normal(shape=X.shape)
    y = d2l.matmul(X, tf.reshape(w, (-1, 1))) + b
    y += tf.random.normal(shape=y.shape, stddev=0.01)
    y = d2l.reshape(y, (-1, 1))
    return X, y

# Defined in file: ./chapter_linear-networks/linear-regression-scratch.md
def linreg(X, w, b):
    """The linear regression model."""
    return d2l.matmul(X, w) + b

# Defined in file: ./chapter_linear-networks/linear-regression-scratch.md
def squared_loss(y_hat, y):
    """Squared loss."""
    return (y_hat - d2l.reshape(y, y_hat.shape))**2 / 2

# Defined in file: ./chapter_linear-networks/linear-regression-scratch.md
def sgd(params, grads, lr, batch_size):
    """Minibatch stochastic gradient descent."""
    for param, grad in zip(params, grads):
        param.assign_sub(lr * grad / batch_size)

# Defined in file: ./chapter_linear-networks/linear-regression-concise.md
def load_array(data_arrays, batch_size, is_train=True):
    """Construct a TensorFlow data iterator."""
    dataset = tf.data.Dataset.from_tensor_slices(data_arrays)
    if is_train:
        dataset = dataset.shuffle(buffer_size=1000)
    dataset = dataset.batch(batch_size)
    return dataset

# Defined in file: ./chapter_linear-networks/image-classification-dataset.md
def get_fashion_mnist_labels(labels):
    """Return text labels for the Fashion-MNIST dataset."""
    text_labels = [
        't-shirt', 'trouser', 'pullover', 'dress', 'coat', 'sandal', 'shirt',
        'sneaker', 'bag', 'ankle boot']
    return [text_labels[int(i)] for i in labels]

# Defined in file: ./chapter_linear-networks/image-classification-dataset.md
def show_images(imgs, num_rows, num_cols, titles=None, scale=1.5):
    """Plot a list of images."""
    figsize = (num_cols * scale, num_rows * scale)
    _, axes = d2l.plt.subplots(num_rows, num_cols, figsize=figsize)
    axes = axes.flatten()
    for i, (ax, img) in enumerate(zip(axes, imgs)):
        ax.imshow(d2l.numpy(img))
        ax.axes.get_xaxis().set_visible(False)
        ax.axes.get_yaxis().set_visible(False)
        if titles:
            ax.set_title(titles[i])
    return axes

# Defined in file: ./chapter_linear-networks/image-classification-dataset.md
def load_data_fashion_mnist(batch_size, resize=None):
    """Download the Fashion-MNIST dataset and then load it into memory."""
    mnist_train, mnist_test = tf.keras.datasets.fashion_mnist.load_data()
    # Divide all numbers by 255 so that all pixel values are between
    # 0 and 1, add a batch dimension at the last. And cast label to int32
    process = lambda X, y: (tf.expand_dims(X, axis=3) / 255,
                            tf.cast(y, dtype='int32'))
    resize_fn = lambda X, y: (tf.image.resize_with_pad(X, resize, resize)
                              if resize else X, y)
    return (tf.data.Dataset.from_tensor_slices(
        process(*mnist_train)).batch(batch_size).shuffle(len(
            mnist_train[0])).map(resize_fn),
            tf.data.Dataset.from_tensor_slices(
                process(*mnist_test)).batch(batch_size).map(resize_fn))

# Defined in file: ./chapter_linear-networks/softmax-regression-scratch.md
def accuracy(y_hat, y):
    """Compute the number of correct predictions."""
    if len(y_hat.shape) > 1 and y_hat.shape[1] > 1:
        y_hat = d2l.argmax(y_hat, axis=1)
    cmp = d2l.astype(y_hat, y.dtype) == y
    return float(d2l.reduce_sum(d2l.astype(cmp, y.dtype)))

# Defined in file: ./chapter_linear-networks/softmax-regression-scratch.md
def evaluate_accuracy(net, data_iter):
    """Compute the accuracy for a model on a dataset."""
    metric = Accumulator(2)  # No. of correct predictions, no. of predictions
    for X, y in data_iter:
        metric.add(accuracy(net(X), y), d2l.size(y))
    return metric[0] / metric[1]

# Defined in file: ./chapter_linear-networks/softmax-regression-scratch.md
class Accumulator:
    """For accumulating sums over `n` variables."""
    def __init__(self, n):
        self.data = [0.0] * n

    def add(self, *args):
        self.data = [a + float(b) for a, b in zip(self.data, args)]

    def reset(self):
        self.data = [0.0] * len(self.data)

    def __getitem__(self, idx):
        return self.data[idx]

# Defined in file: ./chapter_linear-networks/softmax-regression-scratch.md
def train_epoch_ch3(net, train_iter, loss, updater):
    """The training loop defined in Chapter 3."""
    # Sum of training loss, sum of training accuracy, no. of examples
    metric = Accumulator(3)
    for X, y in train_iter:
        # Compute gradients and update parameters
        with tf.GradientTape() as tape:
            y_hat = net(X)
            # Keras implementations for loss takes (labels, predictions)
            # instead of (predictions, labels) that users might implement
            # in this book, e.g. `cross_entropy` that we implemented above
            if isinstance(loss, tf.keras.losses.Loss):
                l = loss(y, y_hat)
            else:
                l = loss(y_hat, y)
        if isinstance(updater, tf.keras.optimizers.Optimizer):
            params = net.trainable_variables
            grads = tape.gradient(l, params)
            updater.apply_gradients(zip(grads, params))
        else:
            updater(X.shape[0], tape.gradient(l, updater.params))
        # Keras loss by default returns the average loss in a batch
        l_sum = l * float(tf.size(y)) if isinstance(
            loss, tf.keras.losses.Loss) else tf.reduce_sum(l)
        metric.add(l_sum, accuracy(y_hat, y), tf.size(y))
    # Return training loss and training accuracy
    return metric[0] / metric[2], metric[1] / metric[2]

# Defined in file: ./chapter_linear-networks/softmax-regression-scratch.md
class Animator:
    """For plotting data in animation."""
    def __init__(self, xlabel=None, ylabel=None, legend=None, xlim=None,
                 ylim=None, xscale='linear', yscale='linear',
                 fmts=('-', 'm--', 'g-.', 'r:'), nrows=1, ncols=1,
                 figsize=(3.5, 2.5)):
        # Incrementally plot multiple lines
        if legend is None:
            legend = []
        d2l.use_svg_display()
        self.fig, self.axes = d2l.plt.subplots(nrows, ncols, figsize=figsize)
        if nrows * ncols == 1:
            self.axes = [self.axes,]
        # Use a lambda function to capture arguments
        self.config_axes = lambda: d2l.set_axes(self.axes[
            0], xlabel, ylabel, xlim, ylim, xscale, yscale, legend)
        self.X, self.Y, self.fmts = None, None, fmts

    def add(self, x, y):
        # Add multiple data points into the figure
        if not hasattr(y, "__len__"):
            y = [y]
        n = len(y)
        if not hasattr(x, "__len__"):
            x = [x] * n
        if not self.X:
            self.X = [[] for _ in range(n)]
        if not self.Y:
            self.Y = [[] for _ in range(n)]
        for i, (a, b) in enumerate(zip(x, y)):
            if a is not None and b is not None:
                self.X[i].append(a)
                self.Y[i].append(b)
        self.axes[0].cla()
        for x, y, fmt in zip(self.X, self.Y, self.fmts):
            self.axes[0].plot(x, y, fmt)
        self.config_axes()
        display.display(self.fig)
        display.clear_output(wait=True)

# Defined in file: ./chapter_linear-networks/softmax-regression-scratch.md
def train_ch3(net, train_iter, test_iter, loss, num_epochs, updater):
    """Train a model (defined in Chapter 3)."""
    animator = Animator(xlabel='epoch', xlim=[1, num_epochs], ylim=[0.3, 0.9],
                        legend=['train loss', 'train acc', 'test acc'])
    for epoch in range(num_epochs):
        train_metrics = train_epoch_ch3(net, train_iter, loss, updater)
        test_acc = evaluate_accuracy(net, test_iter)
        animator.add(epoch + 1, train_metrics + (test_acc,))
    train_loss, train_acc = train_metrics
    assert train_loss < 0.5, train_loss
    assert train_acc <= 1 and train_acc > 0.7, train_acc
    assert test_acc <= 1 and test_acc > 0.7, test_acc

# Defined in file: ./chapter_linear-networks/softmax-regression-scratch.md
class Updater():
    """For updating parameters using minibatch stochastic gradient descent."""
    def __init__(self, params, lr):
        self.params = params
        self.lr = lr

    def __call__(self, batch_size, grads):
        d2l.sgd(self.params, grads, self.lr, batch_size)

# Defined in file: ./chapter_linear-networks/softmax-regression-scratch.md
def predict_ch3(net, test_iter, n=6):
    """Predict labels (defined in Chapter 3)."""
    for X, y in test_iter:
        break
    trues = d2l.get_fashion_mnist_labels(y)
    preds = d2l.get_fashion_mnist_labels(d2l.argmax(net(X), axis=1))
    titles = [true + '\n' + pred for true, pred in zip(trues, preds)]
    d2l.show_images(d2l.reshape(X[0:n], (n, 28, 28)), 1, n,
                    titles=titles[0:n])

# Defined in file: ./chapter_multilayer-perceptrons/underfit-overfit.md
def evaluate_loss(net, data_iter, loss):
    """Evaluate the loss of a model on the given dataset."""
    metric = d2l.Accumulator(2)  # Sum of losses, no. of examples
    for X, y in data_iter:
        l = loss(net(X), y)
        metric.add(d2l.reduce_sum(l), d2l.size(l))
    return metric[0] / metric[1]

# Defined in file: ./chapter_multilayer-perceptrons/kaggle-house-price.md
DATA_HUB = dict()
DATA_URL = 'http://d2l-data.s3-accelerate.amazonaws.com/'

# Defined in file: ./chapter_multilayer-perceptrons/kaggle-house-price.md
def download(name, cache_dir=os.path.join('..', 'data')):
    """Download a file inserted into DATA_HUB, return the local filename."""
    assert name in DATA_HUB, f"{name} does not exist in {DATA_HUB}."
    url, sha1_hash = DATA_HUB[name]
    os.makedirs(cache_dir, exist_ok=True)
    fname = os.path.join(cache_dir, url.split('/')[-1])
    if os.path.exists(fname):
        sha1 = hashlib.sha1()
        with open(fname, 'rb') as f:
            while True:
                data = f.read(1048576)
                if not data:
                    break
                sha1.update(data)
        if sha1.hexdigest() == sha1_hash:
            return fname  # Hit cache
    print(f'Downloading {fname} from {url}...')
    r = requests.get(url, stream=True, verify=True)
    with open(fname, 'wb') as f:
        f.write(r.content)
    return fname

# Defined in file: ./chapter_multilayer-perceptrons/kaggle-house-price.md
def download_extract(name, folder=None):
    """Download and extract a zip/tar file."""
    fname = download(name)
    base_dir = os.path.dirname(fname)
    data_dir, ext = os.path.splitext(fname)
    if ext == '.zip':
        fp = zipfile.ZipFile(fname, 'r')
    elif ext in ('.tar', '.gz'):
        fp = tarfile.open(fname, 'r')
    else:
        assert False, 'Only zip/tar files can be extracted.'
    fp.extractall(base_dir)
    return os.path.join(base_dir, folder) if folder else data_dir

def download_all():
    """Download all files in the DATA_HUB."""
    for name in DATA_HUB:
        download(name)

# Defined in file: ./chapter_multilayer-perceptrons/kaggle-house-price.md
DATA_HUB['kaggle_house_train'] = (DATA_URL + 'kaggle_house_pred_train.csv',
                                  '585e9cc93e70b39160e7921475f9bcd7d31219ce')

DATA_HUB['kaggle_house_test'] = (DATA_URL + 'kaggle_house_pred_test.csv',
                                 'fa19780a7b011d9b009e8bff8e99922a8ee2eb90')

# Defined in file: ./chapter_deep-learning-computation/use-gpu.md
def try_gpu(i=0):
    """Return gpu(i) if exists, otherwise return cpu()."""
    if len(tf.config.experimental.list_physical_devices('GPU')) >= i + 1:
        return tf.device(f'/GPU:{i}')
    return tf.device('/CPU:0')

def try_all_gpus():
    """Return all available GPUs, or [cpu(),] if no GPU exists."""
    num_gpus = len(tf.config.experimental.list_physical_devices('GPU'))
    devices = [tf.device(f'/GPU:{i}') for i in range(num_gpus)]
    return devices if devices else [tf.device('/CPU:0')]

# Defined in file: ./chapter_convolutional-neural-networks/conv-layer.md
def corr2d(X, K):
    """Compute 2D cross-correlation."""
    h, w = K.shape
    Y = tf.Variable(tf.zeros((X.shape[0] - h + 1, X.shape[1] - w + 1)))
    for i in range(Y.shape[0]):
        for j in range(Y.shape[1]):
            Y[i, j].assign(tf.reduce_sum(X[i:i + h, j:j + w] * K))
    return Y

# Defined in file: ./chapter_convolutional-neural-networks/lenet.md
class TrainCallback(tf.keras.callbacks.Callback):
    """A callback to visiualize the training progress."""
    def __init__(self, net, train_iter, test_iter, num_epochs, device_name):
        self.timer = d2l.Timer()
        self.animator = d2l.Animator(
            xlabel='epoch', xlim=[1, num_epochs],
            legend=['train loss', 'train acc', 'test acc'])
        self.net = net
        self.train_iter = train_iter
        self.test_iter = test_iter
        self.num_epochs = num_epochs
        self.device_name = device_name

    def on_epoch_begin(self, epoch, logs=None):
        self.timer.start()

    def on_epoch_end(self, epoch, logs):
        self.timer.stop()
        test_acc = self.net.evaluate(self.test_iter, verbose=0,
                                     return_dict=True)['accuracy']
        metrics = (logs['loss'], logs['accuracy'], test_acc)
        self.animator.add(epoch + 1, metrics)
        if epoch == self.num_epochs - 1:
            batch_size = next(iter(self.train_iter))[0].shape[0]
            num_examples = batch_size * tf.data.experimental.cardinality(
                self.train_iter).numpy()
            print(f'loss {metrics[0]:.3f}, train acc {metrics[1]:.3f}, '
                  f'test acc {metrics[2]:.3f}')
            print(f'{num_examples / self.timer.avg():.1f} examples/sec on '
                  f'{str(self.device_name)}')

def train_ch6(net_fn, train_iter, test_iter, num_epochs, lr, device):
    """Train a model with a GPU (defined in Chapter 6)."""
    device_name = device._device_name
    strategy = tf.distribute.OneDeviceStrategy(device_name)
    with strategy.scope():
        optimizer = tf.keras.optimizers.SGD(learning_rate=lr)
        loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
        net = net_fn()
        net.compile(optimizer=optimizer, loss=loss, metrics=['accuracy'])
    callback = TrainCallback(net, train_iter, test_iter, num_epochs,
                             device_name)
    net.fit(train_iter, epochs=num_epochs, verbose=0, callbacks=[callback])
    return net

# Defined in file: ./chapter_convolutional-modern/resnet.md
class Residual(tf.keras.Model):
    """The Residual block of ResNet."""
    def __init__(self, num_channels, use_1x1conv=False, strides=1):
        super().__init__()
        self.conv1 = tf.keras.layers.Conv2D(num_channels, padding='same',
                                            kernel_size=3, strides=strides)
        self.conv2 = tf.keras.layers.Conv2D(num_channels, kernel_size=3,
                                            padding='same')
        self.conv3 = None
        if use_1x1conv:
            self.conv3 = tf.keras.layers.Conv2D(num_channels, kernel_size=1,
                                                strides=strides)
        self.bn1 = tf.keras.layers.BatchNormalization()
        self.bn2 = tf.keras.layers.BatchNormalization()

    def call(self, X):
        Y = tf.keras.activations.relu(self.bn1(self.conv1(X)))
        Y = self.bn2(self.conv2(Y))
        if self.conv3 is not None:
            X = self.conv3(X)
        Y += X
        return tf.keras.activations.relu(Y)

# Defined in file: ./chapter_recurrent-neural-networks/text-preprocessing.md
d2l.DATA_HUB['time_machine'] = (d2l.DATA_URL + 'timemachine.txt',
                                '090b5e7e70c295757f55df93cb0a180b9691891a')

def read_time_machine():
    """Load the time machine dataset into a list of text lines."""
    with open(d2l.download('time_machine'), 'r') as f:
        lines = f.readlines()
    return [re.sub('[^A-Za-z]+', ' ', line).strip().lower() for line in lines]

# Defined in file: ./chapter_recurrent-neural-networks/text-preprocessing.md
def tokenize(lines, token='word'):
    """Split text lines into word or character tokens."""
    if token == 'word':
        return [line.split() for line in lines]
    elif token == 'char':
        return [list(line) for line in lines]
    else:
        print('ERROR: unknown token type: ' + token)

# Defined in file: ./chapter_recurrent-neural-networks/text-preprocessing.md
class Vocab:
    """Vocabulary for text."""
    def __init__(self, tokens=None, min_freq=0, reserved_tokens=None):
        if tokens is None:
            tokens = []
        if reserved_tokens is None:
            reserved_tokens = []
        # Sort according to frequencies
        counter = count_corpus(tokens)
        self.token_freqs = sorted(counter.items(), key=lambda x: x[1],
                                  reverse=True)
        # The index for the unknown token is 0
        self.unk, uniq_tokens = 0, ['<unk>'] + reserved_tokens
        uniq_tokens += [
            token for token, freq in self.token_freqs
            if freq >= min_freq and token not in uniq_tokens]
        self.idx_to_token, self.token_to_idx = [], dict()
        for token in uniq_tokens:
            self.idx_to_token.append(token)
            self.token_to_idx[token] = len(self.idx_to_token) - 1

    def __len__(self):
        return len(self.idx_to_token)

    def __getitem__(self, tokens):
        if not isinstance(tokens, (list, tuple)):
            return self.token_to_idx.get(tokens, self.unk)
        return [self.__getitem__(token) for token in tokens]

    def to_tokens(self, indices):
        if not isinstance(indices, (list, tuple)):
            return self.idx_to_token[indices]
        return [self.idx_to_token[index] for index in indices]

def count_corpus(tokens):
    """Count token frequencies."""
    # Here `tokens` is a 1D list or 2D list
    if len(tokens) == 0 or isinstance(tokens[0], list):
        # Flatten a list of token lists into a list of tokens
        tokens = [token for line in tokens for token in line]
    return collections.Counter(tokens)

# Defined in file: ./chapter_recurrent-neural-networks/text-preprocessing.md
def load_corpus_time_machine(max_tokens=-1):
    """Return token indices and the vocabulary of the time machine dataset."""
    lines = read_time_machine()
    tokens = tokenize(lines, 'char')
    vocab = Vocab(tokens)
    # Since each text line in the time machine dataset is not necessarily a
    # sentence or a paragraph, flatten all the text lines into a single list
    corpus = [vocab[token] for line in tokens for token in line]
    if max_tokens > 0:
        corpus = corpus[:max_tokens]
    return corpus, vocab

# Defined in file: ./chapter_recurrent-neural-networks/language-models-and-dataset.md
def seq_data_iter_random(corpus, batch_size, num_steps):
    """Generate a minibatch of subsequences using random sampling."""
    # Start with a random offset (inclusive of `num_steps - 1`) to partition a
    # sequence
    corpus = corpus[random.randint(0, num_steps - 1):]
    # Subtract 1 since we need to account for labels
    num_subseqs = (len(corpus) - 1) // num_steps
    # The starting indices for subsequences of length `num_steps`
    initial_indices = list(range(0, num_subseqs * num_steps, num_steps))
    # In random sampling, the subsequences from two adjacent random
    # minibatches during iteration are not necessarily adjacent on the
    # original sequence
    random.shuffle(initial_indices)

    def data(pos):
        # Return a sequence of length `num_steps` starting from `pos`
        return corpus[pos:pos + num_steps]

    num_batches = num_subseqs // batch_size
    for i in range(0, batch_size * num_batches, batch_size):
        # Here, `initial_indices` contains randomized starting indices for
        # subsequences
        initial_indices_per_batch = initial_indices[i:i + batch_size]
        X = [data(j) for j in initial_indices_per_batch]
        Y = [data(j + 1) for j in initial_indices_per_batch]
        yield d2l.tensor(X), d2l.tensor(Y)

# Defined in file: ./chapter_recurrent-neural-networks/language-models-and-dataset.md
def seq_data_iter_sequential(corpus, batch_size, num_steps):
    """Generate a minibatch of subsequences using sequential partitioning."""
    # Start with a random offset to partition a sequence
    offset = random.randint(0, num_steps)
    num_tokens = ((len(corpus) - offset - 1) // batch_size) * batch_size
    Xs = d2l.tensor(corpus[offset:offset + num_tokens])
    Ys = d2l.tensor(corpus[offset + 1:offset + 1 + num_tokens])
    Xs = d2l.reshape(Xs, (batch_size, -1))
    Ys = d2l.reshape(Ys, (batch_size, -1))
    num_batches = Xs.shape[1] // num_steps
    for i in range(0, num_batches * num_steps, num_steps):
        X = Xs[:, i:i + num_steps]
        Y = Ys[:, i:i + num_steps]
        yield X, Y

# Defined in file: ./chapter_recurrent-neural-networks/language-models-and-dataset.md
class SeqDataLoader:
    """An iterator to load sequence data."""
    def __init__(self, batch_size, num_steps, use_random_iter, max_tokens):
        if use_random_iter:
            self.data_iter_fn = d2l.seq_data_iter_random
        else:
            self.data_iter_fn = d2l.seq_data_iter_sequential
        self.corpus, self.vocab = d2l.load_corpus_time_machine(max_tokens)
        self.batch_size, self.num_steps = batch_size, num_steps

    def __iter__(self):
        return self.data_iter_fn(self.corpus, self.batch_size, self.num_steps)

# Defined in file: ./chapter_recurrent-neural-networks/language-models-and-dataset.md
def load_data_time_machine(batch_size, num_steps, use_random_iter=False,
                           max_tokens=10000):
    """Return the iterator and the vocabulary of the time machine dataset."""
    data_iter = SeqDataLoader(batch_size, num_steps, use_random_iter,
                              max_tokens)
    return data_iter, data_iter.vocab

# Defined in file: ./chapter_recurrent-neural-networks/rnn-scratch.md
class RNNModelScratch:
    """A RNN Model implemented from scratch."""
    def __init__(self, vocab_size, num_hiddens, init_state, forward_fn,
                 get_params):
        self.vocab_size, self.num_hiddens = vocab_size, num_hiddens
        self.init_state, self.forward_fn = init_state, forward_fn
        self.trainable_variables = get_params(vocab_size, num_hiddens)

    def __call__(self, X, state):
        X = tf.one_hot(tf.transpose(X), self.vocab_size)
        X = tf.cast(X, tf.float32)
        return self.forward_fn(X, state, self.trainable_variables)

    def begin_state(self, batch_size, *args, **kwargs):
        return self.init_state(batch_size, self.num_hiddens)

# Defined in file: ./chapter_recurrent-neural-networks/rnn-scratch.md
def predict_ch8(prefix, num_preds, net, vocab):
    """Generate new characters following the `prefix`."""
    state = net.begin_state(batch_size=1, dtype=tf.float32)
    outputs = [vocab[prefix[0]]]
    get_input = lambda: d2l.reshape(d2l.tensor([outputs[-1]]), (1, 1)).numpy()
    for y in prefix[1:]:  # Warm-up period
        _, state = net(get_input(), state)
        outputs.append(vocab[y])
    for _ in range(num_preds):  # Predict `num_preds` steps
        y, state = net(get_input(), state)
        outputs.append(int(y.numpy().argmax(axis=1).reshape(1)))
    return ''.join([vocab.idx_to_token[i] for i in outputs])

# Defined in file: ./chapter_recurrent-neural-networks/rnn-scratch.md
def grad_clipping(grads, theta):
    """Clip the gradient."""
    theta = tf.constant(theta, dtype=tf.float32)
    norm = tf.math.sqrt(
        sum((tf.reduce_sum(grad**2)).numpy() for grad in grads))
    norm = tf.cast(norm, tf.float32)
    new_grad = []
    if tf.greater(norm, theta):
        for grad in grads:
            new_grad.append(grad * theta / norm)
    else:
        for grad in grads:
            new_grad.append(grad)
    return new_grad

# Defined in file: ./chapter_recurrent-neural-networks/rnn-scratch.md
def train_epoch_ch8(net, train_iter, loss, updater, use_random_iter):
    """Train a model within one epoch (defined in Chapter 8)."""
    state, timer = None, d2l.Timer()
    metric = d2l.Accumulator(2)  # Sum of training loss, no. of tokens
    for X, Y in train_iter:
        if state is None or use_random_iter:
            # Initialize `state` when either it is the first iteration or
            # using random sampling
            state = net.begin_state(batch_size=X.shape[0], dtype=tf.float32)
        with tf.GradientTape(persistent=True) as g:
            y_hat, state = net(X, state)
            y = d2l.reshape(tf.transpose(Y), (-1))
            l = loss(y, y_hat)
        params = net.trainable_variables
        grads = g.gradient(l, params)
        grads = grad_clipping(grads, 1)
        updater.apply_gradients(zip(grads, params))

        # Keras loss by default returns the average loss in a batch
        # l_sum = l * float(d2l.size(y)) if isinstance(
        #     loss, tf.keras.losses.Loss) else tf.reduce_sum(l)
        metric.add(l * d2l.size(y), d2l.size(y))
    return math.exp(metric[0] / metric[1]), metric[1] / timer.stop()

# Defined in file: ./chapter_recurrent-neural-networks/rnn-scratch.md
def train_ch8(net, train_iter, vocab, lr, num_epochs, strategy,
              use_random_iter=False):
    """Train a model (defined in Chapter 8)."""
    with strategy.scope():
        loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
        updater = tf.keras.optimizers.SGD(lr)
    animator = d2l.Animator(xlabel='epoch', ylabel='perplexity',
                            legend=['train'], xlim=[10, num_epochs])
    predict = lambda prefix: predict_ch8(prefix, 50, net, vocab)
    # Train and predict
    for epoch in range(num_epochs):
        ppl, speed = train_epoch_ch8(net, train_iter, loss, updater,
                                     use_random_iter)
        if (epoch + 1) % 10 == 0:
            print(predict('time traveller'))
            animator.add(epoch + 1, [ppl])
    device = d2l.try_gpu()._device_name
    print(f'perplexity {ppl:.1f}, {speed:.1f} tokens/sec on {str(device)}')
    print(predict('time traveller'))
    print(predict('traveller'))

# Defined in file: ./chapter_recurrent-neural-networks/rnn-concise.md
class RNNModel(tf.keras.layers.Layer):
    def __init__(self, rnn_layer, vocab_size, **kwargs):
        super(RNNModel, self).__init__(**kwargs)
        self.rnn = rnn_layer
        self.vocab_size = vocab_size
        self.dense = tf.keras.layers.Dense(vocab_size)

    def call(self, inputs, state):
        X = tf.one_hot(tf.transpose(inputs), self.vocab_size)
        # Later RNN like `tf.keras.layers.LSTMCell` return more than two values
        Y, *state = self.rnn(X, state)
        output = self.dense(tf.reshape(Y, (-1, Y.shape[-1])))
        return output, state

    def begin_state(self, *args, **kwargs):
        return self.rnn.cell.get_initial_state(*args, **kwargs)

# Defined in file: ./chapter_recurrent-modern/machine-translation-and-dataset.md
d2l.DATA_HUB['fra-eng'] = (d2l.DATA_URL + 'fra-eng.zip',
                           '94646ad1522d915e7b0f9296181140edcf86a4f5')

def read_data_nmt():
    """Load the English-French dataset."""
    data_dir = d2l.download_extract('fra-eng')
    with open(os.path.join(data_dir, 'fra.txt'), 'r') as f:
        return f.read()

# Defined in file: ./chapter_recurrent-modern/machine-translation-and-dataset.md
def preprocess_nmt(text):
    """Preprocess the English-French dataset."""
    def no_space(char, prev_char):
        return char in set(',.!?') and prev_char != ' '

    # Replace non-breaking space with space, and convert uppercase letters to
    # lowercase ones
    text = text.replace('\u202f', ' ').replace('\xa0', ' ').lower()
    # Insert space between words and punctuation marks
    out = [
        ' ' + char if i > 0 and no_space(char, text[i - 1]) else char
        for i, char in enumerate(text)]
    return ''.join(out)

# Defined in file: ./chapter_recurrent-modern/machine-translation-and-dataset.md
def tokenize_nmt(text, num_examples=None):
    """Tokenize the English-French dataset."""
    source, target = [], []
    for i, line in enumerate(text.split('\n')):
        if num_examples and i > num_examples:
            break
        parts = line.split('\t')
        if len(parts) == 2:
            source.append(parts[0].split(' '))
            target.append(parts[1].split(' '))
    return source, target

# Defined in file: ./chapter_recurrent-modern/machine-translation-and-dataset.md
def truncate_pad(line, num_steps, padding_token):
    """Truncate or pad sequences."""
    if len(line) > num_steps:
        return line[:num_steps]  # Truncate
    return line + [padding_token] * (num_steps - len(line))  # Pad

# Defined in file: ./chapter_recurrent-modern/machine-translation-and-dataset.md
def build_array_nmt(lines, vocab, num_steps):
    """Transform text sequences of machine translation into minibatches."""
    lines = [vocab[l] for l in lines]
    lines = [l + [vocab['<eos>']] for l in lines]
    array = d2l.tensor([
        truncate_pad(l, num_steps, vocab['<pad>']) for l in lines])
    valid_len = d2l.reduce_sum(d2l.astype(array != vocab['<pad>'], d2l.int32),
                               1)
    return array, valid_len

# Defined in file: ./chapter_recurrent-modern/machine-translation-and-dataset.md
def load_data_nmt(batch_size, num_steps, num_examples=600):
    """Return the iterator and the vocabularies of the translation dataset."""
    text = preprocess_nmt(read_data_nmt())
    source, target = tokenize_nmt(text, num_examples)
    src_vocab = d2l.Vocab(source, min_freq=2,
                          reserved_tokens=['<pad>', '<bos>', '<eos>'])
    tgt_vocab = d2l.Vocab(target, min_freq=2,
                          reserved_tokens=['<pad>', '<bos>', '<eos>'])
    src_array, src_valid_len = build_array_nmt(source, src_vocab, num_steps)
    tgt_array, tgt_valid_len = build_array_nmt(target, tgt_vocab, num_steps)
    data_arrays = (src_array, src_valid_len, tgt_array, tgt_valid_len)
    data_iter = d2l.load_array(data_arrays, batch_size)
    return data_iter, src_vocab, tgt_vocab

# Defined in file: ./chapter_attention-mechanisms/attention-cues.md
def show_heatmaps(matrices, xlabel, ylabel, titles=None, figsize=(2.5, 2.5),
                  cmap='Reds'):
    d2l.use_svg_display()
    num_rows, num_cols = matrices.shape[0], matrices.shape[1]
    fig, axes = d2l.plt.subplots(num_rows, num_cols, figsize=figsize,
                                 sharex=True, sharey=True, squeeze=False)
    for i, (row_axes, row_matrices) in enumerate(zip(axes, matrices)):
        for j, (ax, matrix) in enumerate(zip(row_axes, row_matrices)):
            pcm = ax.imshow(d2l.numpy(matrix), cmap=cmap)
            if i == num_rows - 1:
                ax.set_xlabel(xlabel)
            if j == 0:
                ax.set_ylabel(ylabel)
            if titles:
                ax.set_title(titles[j])
    fig.colorbar(pcm, ax=axes, shrink=0.6)

# Defined in file: ./chapter_optimization/optimization-intro.md
def annotate(text, xy, xytext):
    d2l.plt.gca().annotate(text, xy=xy, xytext=xytext,
                           arrowprops=dict(arrowstyle='->'))

# Defined in file: ./chapter_optimization/gd.md
def train_2d(trainer, steps=20, f_grad=None):
    """Optimize a 2D objective function with a customized trainer."""
    # `s1` and `s2` are internal state variables that will be used later
    x1, x2, s1, s2 = -5, -2, 0, 0
    results = [(x1, x2)]
    for i in range(steps):
        if f_grad:
            x1, x2, s1, s2 = trainer(x1, x2, s1, s2, f_grad)
        else:
            x1, x2, s1, s2 = trainer(x1, x2, s1, s2)
        results.append((x1, x2))
    print(f'epoch {i + 1}, x1: {float(x1):f}, x2: {float(x2):f}')
    return results

def show_trace_2d(f, results):
    """Show the trace of 2D variables during optimization."""
    d2l.set_figsize()
    d2l.plt.plot(*zip(*results), '-o', color='#ff7f0e')
    x1, x2 = d2l.meshgrid(d2l.arange(-5.5, 1.0, 0.1),
                          d2l.arange(-3.0, 1.0, 0.1))
    d2l.plt.contour(x1, x2, f(x1, x2), colors='#1f77b4')
    d2l.plt.xlabel('x1')
    d2l.plt.ylabel('x2')

# Defined in file: ./chapter_optimization/minibatch-sgd.md
d2l.DATA_HUB['airfoil'] = (d2l.DATA_URL + 'airfoil_self_noise.dat',
                           '76e5be1548fd8222e5074cf0faae75edff8cf93f')

def get_data_ch11(batch_size=10, n=1500):
    data = np.genfromtxt(d2l.download('airfoil'), dtype=np.float32,
                         delimiter='\t')
    data = (data - data.mean(axis=0)) / data.std(axis=0)
    data_iter = d2l.load_array((data[:n, :-1], data[:n, -1]), batch_size,
                               is_train=True)
    return data_iter, data.shape[1] - 1

# Defined in file: ./chapter_optimization/minibatch-sgd.md
def train_ch11(trainer_fn, states, hyperparams, data_iter, feature_dim,
               num_epochs=2):
    # Initialization
    w = tf.Variable(
        tf.random.normal(shape=(feature_dim, 1), mean=0, stddev=0.01),
        trainable=True)
    b = tf.Variable(tf.zeros(1), trainable=True)

    # Train
    net, loss = lambda X: d2l.linreg(X, w, b), d2l.squared_loss
    animator = d2l.Animator(xlabel='epoch', ylabel='loss',
                            xlim=[0, num_epochs], ylim=[0.22, 0.35])
    n, timer = 0, d2l.Timer()

    for _ in range(num_epochs):
        for X, y in data_iter:
            with tf.GradientTape() as g:
                l = tf.math.reduce_mean(loss(net(X), y))

            dw, db = g.gradient(l, [w, b])
            trainer_fn([w, b], [dw, db], states, hyperparams)
            n += X.shape[0]
            if n % 200 == 0:
                timer.stop()
                p = n / X.shape[0]
                q = p / tf.data.experimental.cardinality(data_iter).numpy()
                r = (d2l.evaluate_loss(net, data_iter, loss),)
                animator.add(q, r)
                timer.start()
    print(f'loss: {animator.Y[0][-1]:.3f}, {timer.avg():.3f} sec/epoch')
    return timer.cumsum(), animator.Y[0]

# Defined in file: ./chapter_optimization/minibatch-sgd.md
def train_concise_ch11(trainer_fn, hyperparams, data_iter, num_epochs=2):
    # Initialization
    net = tf.keras.Sequential()
    net.add(
        tf.keras.layers.Dense(
            1, kernel_initializer=tf.random_normal_initializer(stddev=0.01)))
    optimizer = trainer_fn(**hyperparams)
    loss = tf.keras.losses.MeanSquaredError()
    # Note: L2 Loss = 1/2 * MSE Loss. TensorFlow has MSE Loss which is
    # slightly different from MXNet's L2Loss by a factor of 2. Hence we halve
    # the loss value to get L2Loss in TensorFlow
    animator = d2l.Animator(xlabel='epoch', ylabel='loss',
                            xlim=[0, num_epochs], ylim=[0.22, 0.35])
    n, timer = 0, d2l.Timer()
    for _ in range(num_epochs):
        for X, y in data_iter:
            with tf.GradientTape() as g:
                out = net(X)
                l = loss(y, out) / 2
                params = net.trainable_variables
                grads = g.gradient(l, params)
            optimizer.apply_gradients(zip(grads, params))
            n += X.shape[0]
            if n % 200 == 0:
                timer.stop()
                p = n / X.shape[0]
                q = p / tf.data.experimental.cardinality(data_iter).numpy()
                r = (d2l.evaluate_loss(net, data_iter, loss) / 2,)
                animator.add(q, r)
                timer.start()
    print(f'loss: {animator.Y[0][-1]:.3f}, {timer.avg():.3f} sec/epoch')

# Defined in file: ./chapter_computational-performance/hybridize.md
class Benchmark:
    def __init__(self, description='Done'):
        self.description = description

    def __enter__(self):
        self.timer = d2l.Timer()
        return self

    def __exit__(self, *args):
        print(f'{self.description}: {self.timer.stop():.4f} sec')

# Defined in file: ./chapter_computer-vision/bounding-box.md
def box_corner_to_center(boxes):
    """Convert from (upper_left, bottom_right) to (center, width, height)"""
    x1, y1, x2, y2 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]
    cx = (x1 + x2) / 2
    cy = (y1 + y2) / 2
    w = x2 - x1
    h = y2 - y1
    boxes = d2l.stack((cx, cy, w, h), axis=-1)
    return boxes

def box_center_to_corner(boxes):
    """Convert from (center, width, height) to (upper_left, bottom_right)"""
    cx, cy, w, h = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]
    x1 = cx - 0.5 * w
    y1 = cy - 0.5 * h
    x2 = cx + 0.5 * w
    y2 = cy + 0.5 * h
    boxes = d2l.stack((x1, y1, x2, y2), axis=-1)
    return boxes

# Defined in file: ./chapter_computer-vision/bounding-box.md
def bbox_to_rect(bbox, color):
    """Convert bounding box to matplotlib format."""
    # Convert the bounding box (top-left x, top-left y, bottom-right x,
    # bottom-right y) format to matplotlib format: ((upper-left x,
    # upper-left y), width, height)
    return d2l.plt.Rectangle(xy=(bbox[0], bbox[1]), width=bbox[2] - bbox[0],
                             height=bbox[3] - bbox[1], fill=False,
                             edgecolor=color, linewidth=2)

# Alias defined in config.ini
size = lambda a: tf.size(a).numpy()

reshape = tf.reshape
ones = tf.ones
zeros = tf.zeros
meshgrid = tf.meshgrid
sin = tf.sin
sinh = tf.sinh
cos = tf.cos
cosh = tf.cosh
tanh = tf.tanh
linspace = tf.linspace
exp = tf.exp
normal = tf.random.normal
rand = tf.random.uniform
matmul = tf.matmul
reduce_sum = tf.reduce_sum
argmax = tf.argmax
tensor = tf.constant
arange = tf.range
astype = tf.cast
int32 = tf.int32
float32 = tf.float32
transpose = tf.transpose
concat = tf.concat
stack = tf.stack
abs = tf.abs
eye = tf.eye
numpy = lambda x, *args, **kwargs: x.numpy(*args, **kwargs)


# To add a new cell, type '# %%'
# To add a new markdown cell, type '# %% [markdown]'
# %%
from IPython import get_ipython

# %%
from IPython import get_ipython

# %% [markdown]
#  # COVID-19 Epidemiology Models
# %% [markdown]
#  First, some preliminary imports. You may need to pip install `holoviews` and `GitPython`. For `holoviews`, some extras might be needed (see https://holoviews.org/install.html).

# %%
import pandas as pd
import numpy as np
import scipy.integrate
import matplotlib.pyplot as plt

import bokeh.io
import bokeh.application
import bokeh.application.handlers
import bokeh.models

import holoviews as hv

bokeh.io.output_notebook()
hv.extension('bokeh')

# %% [markdown]
#  Let's load the data from the relevant folder. If this data doesn't exist for you, you'll need to run the `processing/raw_data_processing/daily_refresh.sh` script (which may require `pip install us`).

# %%
import git

repo = git.Repo("./", search_parent_directories=True)
homedir = repo.working_dir
datadir = f"{homedir}/data/us/"

# %% [markdown]
#  Load the US data by state (50 states, but we also make a bunch of other countries pay taxes without getting the right to vote (Puerto Rico, Guam, etc.), so we include their data too).
#  This results in more than 50 "regions".

# %%
df = pd.read_csv(datadir + 'covid/nyt_us_states.csv')

# %% [markdown]
# We want to count days as integers from some starting point.

# %%
df['date_processed'] = pd.to_datetime(df['date'].values)
df['date_processed'] = (df['date_processed'] - df['date_processed'].min()) / np.timedelta64(1, 'D')

# %% [markdown]
# Now we import the population for each state, and add that information correctly to the Pandas "dataframe".

# %%
populations = pd.read_csv(datadir + 'demographics/state_populations_augmented.csv')
def get_population(region):
    return populations[populations['state'] == region]['population'].values[0]

def get_ticker(state):
    state_to_ticker = dict([('CA','California'),('TX','Texas'),('FL','Florida'),('NY','New York'),('PA','Pennsylvania'),('IL','Illinois'),('OH','Ohio'),('GA','Georgia'),('NC','North Carolina'),
                           ('MI','Michigan'),('NJ','New Jersey'),('VA','Virginia'),('WA','Washington'),('AZ','Arizona'),('MA','Massachusetts'),('TN','Tennessee'),('IN','Indiana'),('MO','Missouri'),
                           ('MD','Maryland'),('WI','Wisconsin'),('CO','Colorado'),('MN','Minnesota'),('SC','South Carolina'),('AL','Alabama'),('LA','Louisiana'),('KY','Kentucky'),('OR','Oregon'),
                           ('OK','Oklahoma'),('CT','Connecticut'),('UT','Utah'),('IA','Iowa'),('NV','Nevada'),('AR','Arkansas'),('MS','Mississippi'),('KS','Kansas'),('NM','New Mexico'),('NE','Nebraska'),
                           ('WV','West Virginia'),('ID','Idaho'),('HI','Hawaii'),('NH','New Hampshire'),('ME','Maine'),('MT','Montana'),('RI','Rhode Island'),('DE','Delaware'),('SD','South Dakota'),
                           ('ND','North Dakota'),('AK','Alaska'),('DC','District of Columbia'),('VT','Vermont'),('WY','Wyoming'),('PR','Puerto Rico'),('VI','Virgin Islands'),('GU','Guam'),
                           ('NMI','Northern Mariana Islands')])
    state_to_ticker = {v: k for k, v in state_to_ticker.items()} # I accidentally typed the dictionary backwards (:
    return state_to_ticker[state]

# %% [markdown]
# Let's also include the state ticker (i.e. Minnesota --> MN) as a column in the Pandas dataframe. (Because the US population info is given by state tickers).

# %%
df['state_ticker'] = df.apply(lambda row: get_ticker(row.state),axis = 1)
df['Population'] = df.apply(lambda row: get_population(row.state_ticker), axis=1)

# %% [markdown]
#  Just checking to make sure the data is here...

# %%
df.head()

# %% [markdown]
#  Great! Let's also make a helper function to select data from a state, starting when the pandemic hit to be able to fit models.
#
#  We could change the parameter "min_deaths" here if we want... it seems like an arbitrary choice.

# %%
# return data ever since first min_cases cases
def select_region(df, region, min_deaths=50):
    d = df.loc[df['state'] == region]
    start = np.where(d['deaths'].values > min_deaths)[0][0]
    d = d[start:]
    return d

# %% [markdown]
#  ## MODEL 1 of 3: erf model
# %% [markdown]
#  Let's start with a simple model used in this major paper by IHME: https://www.medrxiv.org/content/10.1101/2020.03.27.20043752v1. Just fit an erf function, nothing else.

# %%
from scipy.special import erf
from scipy.optimize import curve_fit

def erf_curve(t, logp, a, b):
    p = 10**logp
    deaths = p/2*(1+erf(a*(t-b)))
    return deaths

def erf_model(params, data, future=0):
    # initial conditions
    p = 10**params[0]
    a = params[1]
    b = params[2]

    t = data['date_processed'].values
    if future > 0:
        extrapolation = np.arange(future)
        t = np.concatenate((t, extrapolation + t[-1] + 1))
    deaths = p/2*(1+erf(a*(t-b)))

    return t, deaths

def erf_fit(params, data):
    Draw = data['deaths'].values
    Draw[np.where(Draw < 1)] = 1
    Ddata = Draw
    t, Draw = erf_model(params, data)
    Draw[np.where(Draw < 1)] = 1
    D = Draw

    error = mse(D, Ddata)

    return error

# %% [markdown]
#  Now let's try to have something to run the fit and plot the results, including a prediction. To estimate errors, we'll use `curve_fit` and the covariance matrix it returns. Sampling parameters from a Gaussian distribution around the fit, we then bootstrap errors corresponding to 25th and 75th percentile predictions.
#
#  Maybe we can change the way that the quantile predictions are made. The quantile predictions get crazy later (~1e150, which is way to fucking big).

# %%
from bokeh.models import Span

def plot_erf_with_errors_sample(df, region, extrapolate=1, boundary=None):
    '''
    df: main italy data frame
    region: string (region name)
    extrapolate: run model on this * length of data (ex. extrapolate=1 does no prediction, extrapolate=2 predicts a time interval equal in length to the data length)
    boundary: train model on data until the day demarcated by boundary; if boundary is None, train on all available data
    '''
    data = select_region(df, region)
    keys = ['deaths', 'cases']
    all_out = []
    proper = []
    for k in keys:
        erf_params0 = [np.log10(np.max(data[k])), .1, 30]
        y = data[k].values[:boundary]
        popt, pcov = curve_fit(erf_curve, data['date_processed'].values[:boundary], y, p0=erf_params0)
        errors = np.sqrt(np.diag(pcov))
        #======BEGIN=======DEBUGGING================================================================================
        if k == 'deaths':
            print('popt = ', popt)
            print('errors = ', errors)
        #========END=========DEBUGGING================================================================================
        all_s = []
        samples = 100
        for i in range(samples):
            sample = np.random.normal(loc=popt, scale=errors)
            t, y = erf_model(sample, data, future=len(data)*(extrapolate-1))
            all_s.append(y)

        all_s = np.array(all_s)
        all_out.append(all_s)

        t, y = erf_model(popt, data, future=len(data)*(extrapolate-1))
        proper.append(y)

    t = np.arange(0, len(data))
    tp = np.arange(0, len(data)*extrapolate)

    p = bokeh.plotting.figure(plot_width=600,
                              plot_height=400,
                             title = region + ' erf Model',
                             x_axis_label = 't (days)',
                             y_axis_label = '# people')

    p1 = 25
    p2 = 75
    p.varea(x=tp, y1=np.percentile(all_out[1], p1, axis=0), y2=np.percentile(all_out[1], p2, axis=0), color='purple', fill_alpha=0.2)
    p.varea(x=tp, y1=np.percentile(all_out[0], p1, axis=0), y2=np.percentile(all_out[0], p2, axis=0), color='black', fill_alpha=0.2)

#     s1 = np.percentile(all_s, 40, axis=0)
#     s2 = np.percentile(all_s, 60, axis=0)
#     p.varea(x=tp, y1=s1[:, 2], y2=s2[:, 2], color='red', fill_alpha=0.2)
#     p.varea(x=tp, y1=s1[:, 3], y2=s2[:, 3], color='purple', fill_alpha=0.2)
#     p.varea(x=tp, y1=s1[:, 5], y2=s2[:, 5], color='black', fill_alpha=0.2)

    p.line(tp, proper[0], color = 'black', line_width = 1, legend_label = 'Deceased')
    p.line(tp, proper[1] , color = 'purple', line_width = 1, legend_label = 'Symptomatic infected')

    # deaths
    p.circle(t, data[keys[0]], color ='black')

    # cases
    p.circle(t, data[keys[1]], color ='purple')

    if boundary is not None and boundary < len(data):
        vline = Span(location=boundary, dimension='height', line_color='black', line_width=3)
        p.renderers.extend([vline])

    p.legend.location = 'top_left'
    bokeh.io.show(p)

# %% [markdown]
#  And how well does this model do at prediction? We train it on days 0-16 (left of the vertical line), and we try to predict days afterwards.

# %%
plot_erf_with_errors_sample(df, 'Washington', 2, 10)

# %% [markdown]
#  Hm. Let's check how it performs more systematically. Let's try adjusting how long we train and predict for.
#
#  There is obviously some kind of fitting error that's happening here. But I don't care that much because the erf model is not even a great model.

# %%
start = 12
step = 4
ind = 0
results = []
one_more = False
while start + ind*step <= 18:
    boundary = start + ind*step
    plot_erf_with_errors_sample(df, 'Washington', 2, boundary)
    ind += 1

# %% [markdown]
#  The model seems to think that the pandemic will flatten out in the next couple days. Given the earlier plots, I doubt that is right. We can check it on another region and see what happens! What regions can we choose from? We need a region with enough data.

# %%
regions = sorted(np.unique(df['state'].values)) # this is just alphabetical
print(regions)
for r in regions:
    print(r + ' has ',[np.amax(df.loc[df['state'] == r]['cases'].values)], ' cases')

# %% [markdown]
#  The predictions look a little better for Emilia Romagna, but it appears that the error bars for anything earlier than 16-day prediction are enormous. By the time you reach 20 days, however, we only have 22 days of data, so it seems a little trivial to predict.

# %%
# start = 16
# step = 4
# ind = 0
# results = []
# one_more = False
# while start + ind*step <= 24:
#     boundary = start + ind*step
#     plot_erf_with_errors_sample(df, 'Alabama', 2, boundary)
#     ind += 1

# %% [markdown]
#  Time to move on to a more standard epidemiological model.
# %% [markdown]
#  ## MODEL 2 of 3: SEIR-QD model
# %% [markdown]
#  Ok, will this model isn't too standard. It's selected as one of the best-performing models from "Rational evaluation of various epidemic models based on the COVID-19 data of China" (https://www.medrxiv.org/content/10.1101/2020.03.12.20034595v1.full.pdf), where it turned out to be one of the better models in terms of early-stage data collection. It takes into effect quarantine and self-protection, but not from data; it just makes assumptions about how people behave. The differential equations can be found on page 8 of the supplement: https://www.medrxiv.org/content/medrxiv/suppl/2020/03/16/2020.03.12.20034595.DC1/2020.03.12.20034595-1.pdf. Thanks @Liana Merk for doing the initial work on this model!
#
#  Parameters:
#  * $\beta$ = infection rate, from earlier plotting, $10^{(-8)} - 10^{(-6)}$ seem reasonable.
#  * $\delta$ = recovery rate, which we think is on the order of 10-40 days.
#  * $\gamma$ = transition of exposed individuals to infected, which we aren't sure of, especially with the unknown number of asymptomatics.
#  * $\alpha$ = protection rate of susceptible individuals, which we also don't know, and is most likely dynamic over the course of the outbreak.
#  * $\lambda$ = transition rate of infected to quarantined with infection, same as above.
#  * $\kappa$ = death rate, which we think is around 0.01-0.06. We will leave a range between 0.01 and 0.1.

# %%
def seirqd(dat, t, params, N):
    beta = params[0] / N
    delta = params[1]
    gamma = params[2]
    alpha = params[3]
    lambda_ = params[4]
    kappa = params[5]

    s = dat[0]
    e = dat[1]
    i = dat[2]
    q = dat[3]
    r = dat[4]
    d = dat[5]
    sa = dat[6]

    dsdt = - beta * s * i - alpha * s
    dedt = beta * s * i - gamma * e
    didt = gamma * e - lambda_ * i
    dqdt = lambda_ * i - delta * q - kappa * q
    drdt = delta * q
    dddt = kappa * q
    dsadt = alpha * s

    # susceptible, exposed, infected, quarantined, recovered, died, unsusceptible
    return [dsdt, dedt, didt, dqdt, drdt, dddt, dsadt]

# %%
from sklearn.metrics import mean_squared_error

def mse_qd(A, B):
    Ap = np.nan_to_num(A)
    Bp = np.nan_to_num(B)
    Ap[A == -np.inf] = 0
    Bp[B == -np.inf] = 0
    Ap[A == np.inf] = 0
    Bp[B == np.inf] = 0
    return mean_squared_error(Ap, Bp)

def model_qd(params, data, tmax=-1):
    # initial conditions
    N = data['Population'].values[0] # total population
    initial_conditions = N * np.array(params[-5:]) # the parameters are a fraction of the population so multiply by the population

    # initial conditions
    e0 = initial_conditions[0]
    i0 = initial_conditions[1]
    q0 = initial_conditions[2]
    r0 = initial_conditions[3]
    sa0 = initial_conditions[4]

    d0 = data['deaths'].values[0]
    s0 = N - np.sum(initial_conditions) - d0

    yz_0 = np.array([s0, e0, i0, q0, r0, d0, sa0])

    # Package parameters into a tuple
    args = (params, N)

    n = len(data)
    if tmax > 0:
        n = tmax

    # Integrate ODEs
    s = scipy.integrate.odeint(seirqd, yz_0, np.arange(0, n), args=args)

    return s

def fit_leastsq_qd(params, data):
    Ddata = (data['deaths'].values)
    Idata = (data['cases'].values)
    s = model_qd(params, data)

    S = s[:,0]
    E = s[:,1]
    I = s[:,2]
    Q = s[:,3]
    R = s[:,4]
    D = s[:,5]
    SA = s[:,6]

    error = np.concatenate((D-Ddata, I - Idata))
    return error

# %% [markdown]
#  Let's check if this is reasonable with a manual fit to the US state of Washington. Anyway, we need to provide the optimization method an initial guess of the parameters, so this gives us a good chance to make a reasonable guess and speed up the optimization. Some good guesses are provided by the SEIR model here: https://gabgoh.github.io/COVID/index.html.
#
#  NOT SURE WHY THIS FIT DOESN'T WORK WELL

# %%
get_ipython().run_line_magic('matplotlib', 'notebook')
get_ipython().run_line_magic('matplotlib', 'inline')

plt.figure()
d = select_region(df, 'Washington')
# parameters: beta, delta, gamma, alpha, lambda, kappa
params = [2.0, 0.3, 0.2, 0.05, 0.2, 0.03]
# parameters: beta, sigma, ra, shift
# params = [1.1, 0.8, 0.2, 0.3]
# conditions: E, I, Q, R, SA
initial_conditions = [0.5e-3, 0.5e-3, 0.3e-3, 0.1e-4, 0.5]
s = model_qd(params + initial_conditions, d, False)
plt.scatter(d['date_processed'], d['deaths'])
plt.plot(d['date_processed'], s[:, 5])
plt.show()

print('Do the initial conditions sum to one? If not, what do they physically mean?')
print('They sum to ',sum(initial_conditions))

# %% [markdown]
#  Looks good! How about actually fitting it now? We need a plotting function!
#
#  ^ DOES NOT LOOK GOOD! LOOKS BAD! Do we need to change the initial conditions? This model ends up working later, so something must change to make it work for the state of Washington.

# %%
import itertools

def plot_qd(res, p0_params, p0_initial_conditions, df, region, extrapolate=1, boundary=None, plot_infectious=False):
    data = select_region(df, region)

    s = model_qd(res.x, data, len(data)*extrapolate)
    S = s[:,0]
    E = s[:,1]
    I = s[:,2]
    Q = s[:,3]
    R = s[:,4]
    D = s[:,5]
    SA = s[:,6]

    t = np.arange(0, len(data))
    tp = np.arange(0, len(data)*extrapolate)

    p = bokeh.plotting.figure(plot_width=600,
                              plot_height=400,
                             title = region + ' SEIIRD+Q Model',
                             x_axis_label = 't (days)',
                             y_axis_label = '# people')

    if plot_infectious:
        p.line(tp, I, color = 'red', line_width = 1, legend_label = 'All infected')
    p.line(tp, D, color = 'black', line_width = 1, legend_label = 'Deceased')

    # death
    p.circle(t, data['deaths'], color ='black')

    # quarantined
    p.circle(t, data['cases'], color ='purple', legend_label='Tested infected')

    if boundary is not None:
        vline = Span(location=boundary, dimension='height', line_color='black', line_width=3)
        p.renderers.extend([vline])

    p.legend.location = 'top_left'
    bokeh.io.show(p)

# %% [markdown]
#  We'll need to estimate errors from the optimization procedure. I transferred some of the relevant code from `curve_fit` to work for this kind of model:

# %%
from scipy.linalg import svd
# return params, 1 standard deviation errors
def get_errors(res, p0):
    p0 = np.array(p0)
    ysize = len(res.fun)
    cost = 2 * res.cost  # res.cost is half sum of squares!
    popt = res.x
    # Do Moore-Penrose inverse (Pseudoinverse) discarding zero singular values.
    _, s, VT = svd(res.jac, full_matrices=False)
    threshold = np.finfo(float).eps * max(res.jac.shape) * s[0]
    s = s[s > threshold]
    VT = VT[:s.size]
    pcov = np.dot(VT.T / s**2, VT)

    warn_cov = False
    absolute_sigma = False
    if pcov is None:
        # indeterminate covariance
        pcov = zeros((len(popt), len(popt)), dtype=float)
        pcov.fill(inf)
        warn_cov = True
    elif not absolute_sigma:
        if ysize > p0.size:
            s_sq = cost / (ysize - p0.size)
            pcov = pcov * s_sq
        else:
            pcov.fill(inf)
            warn_cov = True

    if warn_cov:
        print('cannot estimate variance')
        return None

    perr = np.sqrt(np.diag(pcov))
    return perr

# %% [markdown]
#  To do the optimization, we should define reasonable ranges for the parameters as well.

# %%
param_ranges = [(0.5, 3.0), (0.0, 0.5), (0.0, 0.5), (0.01, 0.5), (0.0, 0.5), (0.005, 0.1)]
initial_ranges = [(1.0e-7, 0.01), (1.0e-7, 0.01), (1.0e-7, 0.01), (1.0e-7, 0.01), (1.0e-7, 0.9)]
guesses = params + initial_conditions
ranges = param_ranges + initial_ranges

# %%
from scipy.optimize import least_squares
res = least_squares(fit_leastsq_qd, guesses, args=(select_region(df, 'Wisconsin'),), bounds=np.transpose(np.array(ranges)))

# %%
plot_qd(res, params, initial_conditions, df, 'Wisconsin', extrapolate=2, plot_infectious=True)

# %% [markdown]
#  This is quite misleading! We don't actually know how many people are infected, so this model isn't quite what we're looking for. Instead, we only know the people who were tested as infected; they usually have symptoms. I won't bother with the error analysis for this model, but we can see how it holds up to prediction.
#
#  HERE THE MODEL FINALLY WORKS FOR US DATA (Washington). Note that initial conditions are given here. I'm not sure where they got the initial conditions, but in the code, they are defined 5 or 6 cells above.

# %%
start = 10
step = 4
ind = 0
results = []
one_more = False
while start + ind*step <= 18:
    boundary = start + ind*step
    res = least_squares(fit_leastsq_qd, guesses, args=(select_region(df, 'New York')[:boundary],), bounds=np.transpose(np.array(ranges)))
    plot_qd(res, params, initial_conditions, df, 'New York', extrapolate=2, boundary=boundary, plot_infectious=True)
    ind += 1

# %% [markdown]
#  Despite being conceptually wrong, this produces much more satisfactory predictions of fatality rates than the erf model. The infectious predictions should be taken with a (very large) grain of salt anyway, since we know it isn't really predicting the infected population.
# %% [markdown]
#  ## MODEL 3 of 3: $\mathbf{SEI_A I_S R}$ with empirical quarantine
# %% [markdown]
#  The motivation for this model is to add two crucial ingredients: quarantine data and asymptomatic cases. For quarantine analysis, we find an effective population size based on what fraction of the population is moving according to https://citymapper.com/cmi/milan. (Since Milan is the capital of Lombardy, we perform the analysis for that region.) To make the quarantine more realistic, we model a "leaky" quarantine, where the susceptible population is given by the mobility from above plus some offset. To treat asymptomatic cases, we introduce states $I_A$ (asymptomatic) and $I_S$ (symptomatic) according to the following sketch and differential equations:
# %% [markdown]
#  ![SEIIR + quarantine](images/overview.png)
# %% [markdown]
#  Since this is prototyping the model, we manually enter the chart above (raw data is at the link above) and implement it. We also have a testing function $T(t)$ (called `tau(t)` in the code) that allows us to try out different testing strategies for asymptomatic populations. Sorry about the confusing variable names below: parameters are renamed as $\sigma\to$ `alpha`, $s\to$ `sigma`, and $d\to$ `delta`. Not shown in the equations above but included in the diagram is a fixed offset (`offset`) for the leaky quarantine model.
#
#  ***Regarding "TODO fix data imputation", we need to somehow get this data based on US states. Currently, the mobility data (and therefore quarantine data) is hard coded into the vector 'moving'. I'm not sure what shift is. Some arbitrary parameter that probably also needs to be fit...

# %%
# TODO fix data imputation
def q(t, N, shift):
    moving = np.array([57, 54, 52, 51, 49, 47, 46, 45, 44, 43, 39, 37, 34, 23, 19, 13, 10, 7, 6, 5, 5, 7, 6, 5, 4, 4, 3, 3, 4, 4, 3, 3, 3, 3, 2])/100
    q = N*(1-moving) - shift*N
    if np.round(t) >= len(q):
        return q[-1]
    return q[int(np.round(t))]

def tau(t):
    return 0

def seiirq(dat, t, params, N, max_t, offset):
    if t >= max_t:
        return [0]*8
    beta = params[0]
    alpha = params[1] # rate from e to ia
    sigma = params[2] # rate of asymptomatic people becoming symptotic
    ra = params[3] # rate of asymptomatic recovery
    rs = params[4] # rate of symptomatic recovery
    delta = params[5] # death rate
    shift = params[6] # shift quarantine rate vertically from CityMapper data

    s = dat[0]
    e = dat[1]
    i_a = dat[2]
    i_s = dat[3]

    Qind = (q(t + offset, N, shift) - tau(t + offset)*i_a)/(s + e + i_a - tau(t + offset)*i_a)
    Qia = Qind + (1-Qind)*tau(t + offset)

    dsdt = - beta * s * i_a * (1 - Qind) * (1 - Qia) / N
    dedt = beta * s * i_a* (1 - Qind) * (1 - Qia) / N  - alpha * e
    diadt = alpha * e - (sigma + ra) * i_a
    disdt = sigma * i_a - (delta + rs) * i_s
    dddt = delta * i_s
    drdt = ra * i_a + rs * i_s

    # susceptible, exposed, infected, quarantined, recovered, died, unsusceptible
    out = [dsdt, dedt, diadt, disdt, drdt, dddt]
    return out

# %%
from sklearn.metrics import mean_squared_error

def mse(A, B):
    Ap = np.nan_to_num(A)
    Bp = np.nan_to_num(B)
    Ap[A == -np.inf] = 0
    Bp[B == -np.inf] = 0
    Ap[A == np.inf] = 0
    Bp[B == np.inf] = 0
    return mean_squared_error(Ap, Bp)

def model_z(params, data, tmax=-1):
    # initial conditions
    N = data['Population'].values[0] # total population
    initial_conditions = N * np.array(params[-4:]) # the parameters are a fraction of the population so multiply by the population

    e0 = initial_conditions[0]
    ia0 = initial_conditions[1]
    is0 = initial_conditions[2]
    r0 = initial_conditions[3]

    d0 = data['deaths'].values[0]
    s0 = N - np.sum(initial_conditions) - d0

    offset = data['date_processed'].min()
    yz_0 = np.array([s0, e0, ia0, is0, r0, d0])

    n = len(data)
    if tmax > 0:
        n = tmax

    # Package parameters into a tuple
    args = (params, N, n, offset)

    # Integrate ODEs
    try:
        s = scipy.integrate.odeint(seiirq, yz_0, np.arange(0, n), args=args)
    except RuntimeError:
#         print('RuntimeError', params)
        return np.zeros((n, len(yz_0)))

    return s

def fit_leastsq_z(params, data):
    Ddata = (data['deaths'].values)
    Idata = (data['cases'].values)
    s = model_z(params, data)

    S = s[:,0]
    E = s[:,1]
    I_A = s[:,2]
    I_S = s[:,3]
    R = s[:,4]
    D = s[:,5]

    error = np.concatenate(((D-Ddata), I_S - Idata))
    return error

# %% [markdown]
#  Again, we find some good initial parameters.
#
#  *** We need to find some good initial parameters. ):

# %%
get_ipython().run_line_magic('matplotlib', 'notebook')
get_ipython().run_line_magic('matplotlib', 'inline')

plt.figure()
d = select_region(df, 'Washington')
# parameters: beta, alpha, sigma, ra, rs, delta, shift
params = [1.8, 0.35, 0.1, 0.15, 0.34, 0.015, 0.5]
# conditions: E, IA, IS, R
initial_conditions = [4e-6, 0.0009, 0.0005, 0.0002]
s = model_z(params + initial_conditions, d)
plt.scatter(d['date_processed'], d['deaths'])
plt.plot(d['date_processed'], s[:, 5])
plt.show()

# %%
import itertools

def plot_with_errors_sample_z(res, p0_params, p0_initial_conditions, df, region, extrapolate=1, boundary=None, plot_infectious=False):
    data = select_region(df, region)
    errors = get_errors(res, list(p0_params) + list(p0_initial_conditions))
    errors[len(p0_params):] = 0

    all_s = []
    samples = 100
    for i in range(samples):
        sample = np.random.normal(loc=res.x, scale=errors)
        s = model_z(sample, data, len(data)*extrapolate)
        all_s.append(s)

    all_s = np.array(all_s)

    s = model_z(res.x, data, len(data)*extrapolate)
    S = s[:,0]
    E = s[:,1]
    I_A = s[:,2]
    I_S = s[:,3]
    R = s[:,4]
    D = s[:,5]

    t = np.arange(0, len(data))
    tp = np.arange(0, len(data)*extrapolate)

    p = bokeh.plotting.figure(plot_width=600,
                              plot_height=400,
                             title = region + ' SEIIRD+Q Model',
                             x_axis_label = 't (days)',
                             y_axis_label = '# people')

    s1 = np.percentile(all_s, 25, axis=0)
    s2 = np.percentile(all_s, 75, axis=0)
    if plot_infectious:
        p.varea(x=tp, y1=s1[:, 2], y2=s2[:, 2], color='red', fill_alpha=0.2)
    p.varea(x=tp, y1=s1[:, 3], y2=s2[:, 3], color='purple', fill_alpha=0.2)
    p.varea(x=tp, y1=s1[:, 5], y2=s2[:, 5], color='black', fill_alpha=0.2)

    if plot_infectious:
        p.line(tp, I_A, color = 'red', line_width = 1, legend_label = 'Asymptomatic infected')
    p.line(tp, D, color = 'black', line_width = 1, legend_label = 'Deceased')
    p.line(tp, I_S , color = 'purple', line_width = 1, legend_label = 'Symptomatic infected')

    # death
    p.circle(t, data['deaths'], color ='black')

    # quarantined
    p.circle(t, data['cases'], color ='purple')

    if boundary is not None:
        vline = Span(location=boundary, dimension='height', line_color='black', line_width=3)
        p.renderers.extend([vline])

    p.legend.location = 'top_left'
    bokeh.io.show(p)
    return all_s,s1,s2

# %% [markdown]
#  Let's define the initial ranges of the constants for the ODE.

# %%
# beta, alpha, sigma, ra, rs, delta, shift
param_ranges = [(1.0, 2.0), (0.1, 0.5), (0.1, 0.5), (0.05, 0.5), (0.32, 0.36), (0.005, 0.05), (0.1, 0.6)]
initial_ranges = [(1.0e-7, 0.001), (1.0e-7, 0.001), (1.0e-7, 0.001), (1.0e-7, 0.001)]

guesses = params + initial_conditions
ranges = param_ranges + initial_ranges

# %%
# beta, alpha, sigma, ra, rs, delta, shift
param_ranges = [(1.0, 2.0), (0.1, 0.5), (0.1, 0.5), (0.05, 0.5), (0.32, 0.36), (0.005, 0.05), (0.1, 0.6)]
initial_ranges = [(1.0e-7, 0.001), (1.0e-7, 0.001), (1.0e-7, 0.001), (1.0e-7, 0.001)]

guesses = params + initial_conditions
ranges = param_ranges + initial_ranges

boundary = 16
res = least_squares(fit_leastsq_z, guesses, args=(select_region(df, 'Washington')[:boundary],), bounds=np.transpose(np.array(ranges)),jac = '3-point')
all_s,s1,s2 = plot_with_errors_sample_z(res, params, initial_conditions, df, 'Washington', extrapolate=1, boundary=boundary, plot_infectious=True)
perr = get_errors(res,np.zeros((11,1)))
print('parameter standard deviations = ',perr)

# %%
start = 8
step = 4
ind = 0
results = []
one_more = False
while start + ind*step <= 18:
    boundary = start + ind*step
    res = least_squares(fit_leastsq_z, guesses, args=(select_region(df, 'Washington')[:boundary],), bounds=np.transpose(np.array(ranges)),jac = '3-point')
    plot_with_errors_sample_z(res, params, initial_conditions, df, 'Washington', extrapolate=1, boundary=boundary, plot_infectious=True)
    ind += 1

# %% [markdown]
# ***HUUUUGE uncertainty bars. Need to figure out why. Maybe because we don't have good fit parameters, so the covariance is huge.
#
#
#  Just like SEIR-QD, the predictions of cases aren't so great, although the predictions of fatalaties (which matters more and has less data bias) is reasonably accurate. The model also produces prediction that around 2/3 of the cases are asymptomatic (or at least not tested). This corresponds roughly to some recent studies, such as the 50-75% number reported after testing an entire town of 3,300 in Italy (https://www.repubblica.it/salute/medicina-e-ricerca/2020/03/16/news/coronavirus_studio_il_50-75_dei_casi_a_vo_sono_asintomatici_e_molto_contagiosi-251474302/?ref=RHPPTP-BH-I251454518-C12-P3-S2.4-T1) and similar results in Iceland (https://www.government.is/news/article/2020/03/15/Large-scale-testing-of-general-population-in-Iceland-underway/).

# %%


import numpy as np

import matplotlib
import matplotlib.pyplot as plt
from matplotlib import rcParams
from mpl_toolkits.mplot3d import Axes3D

import ipywidgets as widgets
from ipywidgets import AppLayout, FloatSlider
from ipywidgets import GridspecLayout
from matplotlib import rc

from copy import deepcopy

try:
    from addict import Dict
except:
    import subprocess
    import sys
    subprocess.check_call([sys.executable, "-m", "pip", "install", "addict"])
    from addict import Dict

try:
    import ipympl
except:
    import subprocess
    import sys
    subprocess.check_call([sys.executable, "-m", "pip", "install", "ipympl"])
    install("ipympl")

def dict_merge(destination, source):
    out = deepcopy(destination)
    for key, value in source.items():
        if isinstance(value, dict):
            # get node or create one
            node = out.setdefault(key, {})
            dict_merge(node, value)
        else:
            out[key] = value
    return out

PARAM = Dict()
PARAM.background_color = (0.0, 0.0, 0.0, 0.0)
PARAM.zlim = (-4,4)
PARAM.grid_color = (0.0, 0.0, 0.0, 0.25)
PARAM.grid_linewidth = 0.25
PARAM.grid_size = 20
PARAM.border_color = (0.0, 0.0, 0.0, 0.25)
PARAM.border_linewidth = 0.25
PARAM.tick_color = (0.0, 0.0, 0.0, 0.25)
PARAM.tick_linewidth = 0.25
PARAM.label_size = 7
PARAM.label_color = (0,0,0,1)
PARAM.azimuth = -60
PARAM.elevation = 30
PARAM.header_visible = False
PARAM.style = ""
PARAM.plot_box = [[-2,2],[-2,2]]
PARAM.cmap = plt.cm.jet
PARAM.levels_number = 15
PARAM.density = 0.5
PARAM.iter_max = 40

def set_axis3d(ax=None, **param):
    # api for 3D plot : https://matplotlib.org/3.1.1/api/_as_gen/mpl_toolkits.mplot3d.axes3d.Axes3D.html
    # last component is a level of gray : 0.0 for white to 1.0 for black
    # https://matplotlib.org/_modules/mpl_toolkits/mplot3d/axes3d.html

    param = Dict(dict_merge(PARAM, param))
    if ax is None:
        ax = plt.axes(projection='3d')

    if param.style == "raw":
        param.grid_color = (0.0, 0.0, 0.0, 0)
        param.border_color = (0.0, 0.0, 0.0, 0)
        param.tick_color = (0.0, 0.0, 0.0, 0)
        ax.set_xticklabels([])
        ax.set_yticklabels([])
        ax.set_zticklabels([])

    # Now we apply the parameters:

    # Background of the plot
    ax.w_xaxis.set_pane_color(param.background_color)
    ax.w_yaxis.set_pane_color(param.background_color)
    ax.w_zaxis.set_pane_color(param.background_color)

    # Grid of the background
    ax.xaxis._axinfo["grid"]['color'] = param.grid_color
    ax.yaxis._axinfo["grid"]['color'] = param.grid_color
    ax.zaxis._axinfo["grid"]['color'] = param.grid_color
    ax.xaxis._axinfo["grid"]['linewidth'] = param.grid_linewidth
    ax.yaxis._axinfo["grid"]['linewidth'] = param.grid_linewidth
    ax.zaxis._axinfo["grid"]['linewidth'] = param.grid_linewidth

    # Main axes (with the ticks labels etc) :
    ax.w_xaxis.line.set_color(param.border_color)
    ax.w_yaxis.line.set_color(param.border_color)
    ax.w_zaxis.line.set_color(param.border_color)
    ax.w_xaxis.line.set_linewidth(param.border_linewidth)
    ax.w_yaxis.line.set_linewidth(param.border_linewidth)
    ax.w_zaxis.line.set_linewidth(param.border_linewidth)

    # Ticks
    ax.tick_params(color=param.tick_color,
                   size=param.tick_linewidth, # this is broken..
                   axis='both',
                   which='both',
                   labelsize=param.label_size,
                   labelcolor=param.label_color)
    # other options : ['size', 'width', 'color', 'tickdir', 'pad', 'labelsize', 'labelcolor', 'zorder', 'gridOn', 'tick1On', 'tick2On', 'label1On', 'label2On', 'length', 'direction', 'left', 'bottom', 'right', 'top', 'labelleft', 'labelbottom', 'labelright', 'labeltop', 'labelrotation', 'grid_agg_filter', 'grid_alpha', 'grid_animated', 'grid_antialiased', 'grid_clip_box', 'grid_clip_on', 'grid_clip_path', 'grid_color', 'grid_contains', 'grid_dash_capstyle', 'grid_dash_joinstyle', 'grid_dashes', 'grid_data', 'grid_drawstyle', 'grid_figure', 'grid_fillstyle', 'grid_gid', 'grid_in_layout', 'grid_label', 'grid_linestyle', 'grid_linewidth', 'grid_marker', 'grid_markeredgecolor', 'grid_markeredgewidth', 'grid_markerfacecolor', 'grid_markerfacecoloralt', 'grid_markersize', 'grid_markevery', 'grid_path_effects', 'grid_picker', 'grid_pickradius', 'grid_rasterized', 'grid_sketch_params', 'grid_snap', 'grid_solid_capstyle', 'grid_solid_joinstyle', 'grid_transform', 'grid_url', 'grid_visible', 'grid_xdata', 'grid_ydata', 'grid_zorder', 'grid_aa', 'grid_c', 'grid_ds', 'grid_ls', 'grid_lw', 'grid_mec', 'grid_mew', 'grid_mfc', 'grid_mfcalt', 'grid_ms']

    # title and stuff around the plot
    if "title" in param.keys():
        ax.title.set_text(param["title"])
    ax.view_init(param.elevation, param.azimuth)

    return ax

def function_values(fun, **param):
    # fun is a function taking two floats as an argument, returning one
    if "plot_box" not in param.keys():
        param["plot_box"] = [[-2,2],[-2,2]]
    if "grid_size" not in param.keys():
        param["grid_size"] = 20

    x = np.outer(np.linspace(param["plot_box"][0][0], param["plot_box"][0][1], param["grid_size"]), np.ones(param["grid_size"]))
    y = np.outer(np.linspace(param["plot_box"][1][0], param["plot_box"][1][1], param["grid_size"]), np.ones(param["grid_size"])).T

    z = np.zeros((param["grid_size"],param["grid_size"]))
    for idx, _ in np.ndenumerate(x):
        z[idx] = fun(x[idx], y[idx])
    return x, y, z

def field_values(fun, **param):
    # fun is a function taking two floats as an argument, returning two
    if "plot_box" not in param.keys():
        param["plot_box"] = [[-2,2],[-2,2]]
    if "grid_size" not in param.keys():
        param["grid_size"] = 20

    x = np.outer(np.linspace(param["plot_box"][0][0], param["plot_box"][0][1], param["grid_size"]), np.ones(param["grid_size"]))
    y = np.outer(np.linspace(param["plot_box"][1][0], param["plot_box"][1][1], param["grid_size"]), np.ones(param["grid_size"])).T

    zx = np.zeros((param["grid_size"],param["grid_size"]))
    zy = np.zeros((param["grid_size"],param["grid_size"]))
    for idx, _ in np.ndenumerate(x):
        zx[idx], zy[idx] = fun(x[idx], y[idx])
    return x, y, zx, zy

def matrix_definiteness(A):
    if np.linalg.norm(A-A.T)> 1e-8 : # not symmetric
        return matrix_definiteness(0.5*(A+A.T))
    spec = np.linalg.eig(A)[0]
    if (spec>0).all():
        return "Définie Positive"
    elif (spec>=0).all():
        return "Positive"
    elif (spec<0).all():
        return "Définie Négative"
    elif (spec<=0).all():
        return "Négative"
    else:
        return "Non Définie"

def matrix_symmetricness(A):
    if np.linalg.norm(A-A.T) < 1e-8:
        return "symétrique"
    elif np.linalg.norm(A+A.T) < 1e-8:
        return "antisymétrique"
    else:
        return "non symétrique"

def print_info(A):
    return "La matrice A est " + matrix_definiteness(A) + ", " +  matrix_symmetricness(A)

def get_plots_we_want(**plot_param):
    # what do we want to plot? Let's look at the parameters
    possible_plots = ["graph", "levelset", "gradient", "flow"]
    have_to_plot = {plot : ((plot in plot_param.keys()) and plot_param[plot]) for plot in possible_plots}# a dict with bool values
    number_have_to_plot = sum(list(have_to_plot.values())) # number of Trues in the dict

    return have_to_plot, number_have_to_plot

def is_in_box(x, box):
    return box[0][0] < x[0] < box[0][1] and box[1][0] < x[1] < box[1][1]

def sequence_gradient(func, x0=np.ones(2),stepsize=0.1,iter_max=100, **param):
    # here func is a function with a gradient parameter
    seq = []
    x = x0
    if 'plot_box' in param.keys() and is_in_box(x, param['plot_box']):
        seq.append(x)
    for t in range(iter_max):
        grad = func(x[0], x[1], gradient=True) # grad is a tuple, we can compine it with arrays
        x = x - stepsize*np.array(grad)
        if 'plot_box' in param.keys():
            # we check that the sequence is still within the bounds of our plot.
            if is_in_box(x, param['plot_box']):
                seq.append(x)
        else:
            seq.append(x)
    X = [x[0] for x in seq]
    Y = [x[1] for x in seq]
    return X, Y

def quadratic(A, b=np.zeros(2), c=0):
    # returns the function 0.5*<AX,X> + <b,X> + c or its gradient
    def func(x,y, gradient=None):
        if gradient is None:
            return 0.5*( x*(A[0,0]*x + A[0,1]*y) + y*(A[1,0]*x + A[1,1]*y) ) + b[0]*x + b[1]*y + c
        else:
            AA = 0.5*(A + A.T)
            return ( AA[0,0]*x + AA[0,1]*y + b[0], AA[1,0]*x + AA[1,1]*y + b[1] )
    return func

def plot2d_function(func, fig=None, **plot_param):
    """ Reprensents a func : (x,y) ---> z in various ways
        If we ask for more than one representation,
        they are all displayed next to each other in a subplot

    INPUT:
      - fig : the handle of the figure in which drawing
      - func : a function with the signature func(x,y,gradient)
            func(x,y) returns the value of the function at (x,y)
            func(x,y,gradient=True) returns the value of the
            gradient at (x,y), as a 2-tuple

    OUTPUT: None

    OPTIONAL PARAMETERS: there is a *lot* of them, here are the main ones. default values
        are stored in the PARAM dictionary
      - cmap, colormap (plt.cm.jet) : the same colormap is used for every plot
      - style, str (None) : "raw" plots naked graphs, without axis, background or ticks
      - plot_box, 2D array ([[-2,2],[-2,2]]) : The syntax is [[xmin, xmax],[ymin,ymax]].
            Defines the box in which all the objects will be plotted

      - graph, bool (False) : tells if plotting the 3D graph of the function. Displayed by default.
          - grid_size, int (20) : Controls the precision of the mesh for plotting the graph.
          - title, str (None) : A title to display on top of the graph

      - levelset, bool (False) : tells if plotting the 2D levelsets of the function
          - levels, int (15) : number of levelsets to be displayed

      - gradient, bool (False) : tells if plotting the 2D vector field of the gradient of the function

      - flow, bool (False) : tells if plotting the descent gradient flow curves of the function
          - density, float (0.5) : density of the flow curves. Impacts the performance.

      - sequence, (list,list) (None) : if a sequence is specified, it will be plotted on top of the 2D
            graphs (levelset, gradient, flow). Impossible to plot it on 3D surfaces so far.
            The syntax is : a 2-tuple, each component being a list of floats (same size) representing
            the x- and y-coordinates of the sequence.
      - algo, str (None) : instead of specifying a sequence, you can just pass an argument to tell
            which algorithm you want to run on the function. Currently supported:
          - "gradient" : Runs the gradient descent algorithm. You can pass options:
              - x0, 1D arrray (np.ones(2)) : the initialization of the algorithm
              - iter_max, int (100) : the maximum number of iterations allowed
              - stepsize, float (0.1) : the stepsize in the iteration x <-- x - stepsize*gradient
    """
    # deal with optional parameters
    plot_param = Dict(dict_merge(PARAM, plot_param))
    have_to_plot, number_plots = get_plots_we_want(**plot_param)
    if fig is None:
        fig = plt.gcf()
    if 'algo' in plot_param.keys() and plot_param.algo is not None:
        if plot_param.algo == 'gradient':
            seq = sequence_gradient(func, **plot_param)
            if len(seq[0]) == 0: # we are out ouf bounds
                plot_param.sequence = None
            else:
                plot_param.sequence = seq
                s = np.ones(len(plot_param.sequence[0]))*40
                s[0] = 120
                plot_param.scatter_size = s
    fig.canvas.header_visible = plot_param.header_visible

    k = 0
    if have_to_plot['graph']:
        k = k+1
        ax = fig.add_subplot(1, number_plots, k, projection='3d')
        ax = set_axis3d(ax, **plot_param)
        ax.plot_surface(*function_values(func, **plot_param), cmap=plot_param.cmap, rstride=1, cstride=1);
        ax.set_zlim(plot_param.zlim);
    if have_to_plot['levelset']:
        k = k+1
        ax = fig.add_subplot(1, number_plots, k)
        contour = ax.contour(*function_values(func, **plot_param), cmap=plot_param.cmap, levels=plot_param.levels_number)
        if plot_param["style"] != "raw":
            ax.clabel(contour, inline=1, fontsize=7);
        else:
            plt.axis('off')
        if 'sequence' in plot_param.keys() and plot_param['sequence'] is not None:
            ax.scatter(*plot_param['sequence'], s=plot_param.scatter_size)
    if have_to_plot['gradient']:
        k = k+1
        ax = fig.add_subplot(1, number_plots, k)
        X,Y,U,V = field_values(lambda x,y : func(x,y,gradient=True), **plot_param)
        _,_,COLOR = function_values(func, **plot_param)
        contour = ax.quiver(X,Y,U,V,COLOR, cmap=plot_param.cmap)
        if plot_param["style"] == "raw":
            plt.axis('off')
        if 'sequence' in plot_param.keys() and plot_param['sequence'] is not None:
            ax.scatter(*plot_param['sequence'], s=plot_param.scatter_size)
    if have_to_plot['flow']:
        k = k+1
        ax = fig.add_subplot(1, number_plots, k)
        X,Y,U,V = field_values(lambda x,y : func(x,y,gradient=True), **plot_param)
        _,_,COLOR = function_values(func, **plot_param)
        # Here we need to TRANSPOSE the vector field because??? streamplot is stupid? arrays are indexed in a different way??
        contour = ax.streamplot(X.T[0],Y[0],-U.T,-V.T, color=COLOR, density=plot_param.density, cmap=plot_param.cmap)
        if plot_param["style"] == "raw":
            plt.axis('off')
        if 'sequence' in plot_param.keys() and plot_param['sequence'] is not None:
            ax.scatter(*plot_param['sequence'], s=plot_param.scatter_size)
    #print(plot_param)
    #return fig

def widget_quadratic(**plot_param):
    # given a function handle (taking two arguments, returning one)
    # plots an interactive widget with graph and level set

    # parameters for the plot
    default_param = {
        'plot_box' : [[-2,2],[-2,2]],
        'grid_size' : 30,
        'style' : "normal",
        'dpi' : 80,
        'sequence' : None,
        'algo' : None
    }
    if get_plots_we_want(**plot_param)[1] == 0:
        default_param['graph'] = True
    plot_param = { **default_param, **plot_param }
    dpi = plot_param['dpi']

    # define all the sliders we want to manipulate
    slider_param = {
        'min' : -2,
        'max' : 2,
        'orientation' : 'horizontal',
        #'layout' : Layout(height='auto', width='auto')#Layout(width='40%', margin='0px 30% 0px 30%'),
    }
    slider_a11 = FloatSlider(description='$A_{11}$', value=2.0, **slider_param)
    slider_a12 = FloatSlider(description='$A_{12}$', value=0.0, **slider_param)
    slider_a21 = FloatSlider(description='$A_{21}$', value=0.0, **slider_param)
    slider_a22 = FloatSlider(description='$A_{22}$', value=1.0, **slider_param)
    slider_b1  = FloatSlider(description='$b_1$', value=0.0, **slider_param)
    slider_b2  = FloatSlider(description='$b_2$', value=0.0, **slider_param)
    if plot_param['algo'] is not None:
        slider_x0  = FloatSlider(description='$x_0$', value=1.0, **slider_param)
        slider_y0  = FloatSlider(description='$y_0$', value=1.0, **slider_param)
        slider_stepsize = FloatSlider(description='stepsize', value=0.1, min=0.01, max=1, step=0.05)

    # we initialize everything
    A_slider = np.array([[slider_a11.value, slider_a12.value], [slider_a21.value, slider_a22.value]])
    b_slider = np.array([slider_b1.value, slider_b2.value])
    plot_param["title"] = print_info(A_slider)
    if plot_param['algo'] is not None:
        x0_slider = np.array([slider_x0.value, slider_y0.value])
        plot_param['x0']= np.array([slider_x0.value, slider_y0.value])
        stepsize_slider = slider_stepsize.value
        plot_param['stepsize'] = slider_stepsize.value
        #plot_param['sequence'] = sequence_gradient(A_slider, b_slider, x0_slider, stepsize_slider, **plot_param)

    # open the figure
    plt.ioff() # turn off interactive mode to be able to display widget. I honestly don't understand why. see https://github.com/matplotlib/ipympl/issues/220 or https://github.com/matplotlib/matplotlib/pull/17371
    fig = plt.figure(dpi=dpi)
    plt.ion()
    fig.subplots_adjust(left=0.1, right=0.9, bottom=0.1, top=0.9)
    _, number_plots = get_plots_we_want(**plot_param)
    plt.gcf().set_size_inches(plt.figaspect(1/number_plots))
    # plot a first draw
    plot2d_function(quadratic(A_slider,b_slider), fig,  **plot_param)

    # the function handling the change of parameters
    def update_plot(change, slider, idx):
        # store the camera 3D view, then clears the figure
        if 'graph' in plot_param.keys() and plot_param['graph']:
            plot_param['azimuth'] = fig.get_axes()[0].azim # remember there are two subplots
            plot_param['elevation'] = fig.get_axes()[0].elev
        fig.clear()
        # update the parameters
        A = A_slider
        b = b_slider
        if slider == 'A':
            A[idx] = change.new
        elif slider == 'b':
            b[idx] = change.new
        plot_param["title"] = print_info(A)
        if plot_param['algo'] is not None:
            plot_param['x0'] = x0_slider
            plot_param['stepsize'] = stepsize_slider
            if slider == 'x0':
                plot_param['x0'][idx] = change.new
            elif slider == 'stepsize':
                plot_param['stepsize'] = change.new
        # plot
        plot2d_function(quadratic(A,b), fig, **plot_param)
        _=fig.canvas.draw()
        _=fig.canvas.flush_events()

    # keep track of changes
    slider_a11.observe(lambda change : update_plot(change, 'A', (0,0)), names='value')
    slider_a21.observe(lambda change : update_plot(change, 'A', (1,0)), names='value')
    slider_a12.observe(lambda change : update_plot(change, 'A', (0,1)), names='value')
    slider_a22.observe(lambda change : update_plot(change, 'A', (1,1)), names='value')
    slider_b1.observe(lambda change : update_plot(change, 'b', (0,)), names='value')
    slider_b2.observe(lambda change : update_plot(change, 'b', (1,)), names='value')
    if plot_param['algo'] is not None:
        slider_x0.observe(lambda change : update_plot(change, 'x0', (0,)), names='value')
        slider_y0.observe(lambda change : update_plot(change, 'x0', (1,)), names='value')
        slider_stepsize.observe(lambda change : update_plot(change, 'stepsize', None), names='value')

    # now we display all of this
    # a grid of sliders for the parameters
    grid = GridspecLayout(2, 3)
    grid[0,0] = slider_a11
    grid[1,0] = slider_a21
    grid[0,1] = slider_a12
    grid[1,1] = slider_a22
    grid[0,2] = slider_b1
    grid[1,2] = slider_b2
    header=grid
    # optional sliders for the sequence
    if plot_param['algo'] is not None:
        grid2 = GridspecLayout(1, 3)
        grid2[0,0] = slider_x0
        grid2[0,1] = slider_y0
        grid2[0,2] = slider_stepsize
        footer=grid2
    else:
        footer=None

    # we gather everything
    return AppLayout(
        header=header,
        left_sidebar=None,
        center=fig.canvas,
        right_sidebar=None,
        footer=footer
    )