# Standard Python libraries
from __future__ import absolute_import, division, print_function, unicode_literals

import os
import time
import numpy as np
import glob
import matplotlib.pyplot as plt
import PIL
import imageio

from IPython import display

import sklearn

import seaborn as sns

sns.set(style="ticks", color_codes=True)

import pandas as pd

pd.set_option("precision", 2)  # 2 decimal places
pd.set_option("display.max_rows", 20)
pd.set_option("display.max_columns", 30)
pd.set_option("display.width", 100)  # wide windows

"""
Explanation: <a href="https://colab.research.google.com/github/probml/pyprobml/blob/master/book1/intro/datasets.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>
Manipulating datasets
In this colab, we briefly discuss ways to access and manipulate common datasets that are used in the ML literature. Most of these are used for supervised learning experiments.
End of explanation
"""


from sklearn import datasets

iris = datasets.load_iris()
print(iris.keys())

X = iris["data"]
y = iris["target"]  # class labels
print(X.shape)
print(iris["feature_names"])  # meaning of each feature
print(iris["target_names"])  # meaning of each class

"""
Explanation: Tabular datasets
The UCI ML repository contains many smallish datasets, mostly tabular.
Kaggle also hosts many interesting datasets.
Sklearn has many small datasets builtin, making them easy to use for prototyping, as we illustrate below.
End of explanation
"""


# Standard Python libraries
from __future__ import absolute_import, division, print_function, unicode_literals

from typing import Any, Iterator, Mapping, NamedTuple, Sequence, Tuple

import os
import time
import numpy as np
import glob
import matplotlib.pyplot as plt
import PIL
import imageio
from IPython import display
import sklearn

# TensorFlow ≥2.0 is required
import tensorflow as tf
from tensorflow import keras

assert tf.__version__ >= "2.0"
import tensorflow_datasets as tfds

print("tf version {}".format(tf.__version__))

import jax
from typing import Any, Callable, Sequence, Optional, Dict, Tuple
import jax.numpy as jnp

rng = jax.random.PRNGKey(0)

# Useful type aliases

Array = jnp.ndarray
PRNGKey = Array
Batch = Mapping[str, np.ndarray]
OptState = Any

"""
Explanation: Tensorflow datasets
TFDS is a handy way to handle large datasets as a stream of minibatches, suitable for large scale training and parallel evaluation. It can be used by tensorflow and JAX code, as we illustrate below. (See the official colab for details.)
End of explanation
"""


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


def get_datasets_iris():
    iris = sklearn.datasets.load_iris()
    X = iris["data"]
    y = iris["target"]
    N, D = X.shape  # 150, 4
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.33, random_state=42)
    train_ds = {"X": X_train, "y": y_train}
    test_ds = {"X": X_test, "y": y_test}
    return train_ds, test_ds


train_ds, test_ds = get_datasets_iris()
print(train_ds["X"].shape)
print(train_ds["y"].shape)

iris = sklearn.datasets.load_iris()
print(iris.feature_names)
print(iris.target_names)

"""
Explanation: Minibatching without using TFDS
We first illustrate how to make streams of minibatches using vanilla numpy code. TFDS will then let us eliminate a lot of this boilerplate. As an example, let's package some small labeled datasets into two dictionaries, for train and test.
End of explanation
"""


def extract_batch(ds, ndx):
    batch = {k: v[ndx, ...] for k, v in ds.items()}
    # batch = {'X': ds['X'][ndx,:],  'y': ds['y'][ndx]}
    return batch


def process_epoch(train_ds, batch_size, rng):
    train_ds_size = len(train_ds["X"])
    steps_per_epoch = train_ds_size // batch_size
    perms = jax.random.permutation(rng, len(train_ds["X"]))
    perms = perms[: steps_per_epoch * batch_size]  # skip incomplete batch
    perms = perms.reshape((steps_per_epoch, batch_size))  # perms[i,:] is list of  data indices for step i
    for step, perm in enumerate(perms):
        batch = extract_batch(train_ds, perm)
        print("processing batch {} X shape {}, y shape {}".format(step, batch["X"].shape, batch["y"].shape))


batch_size = 30
process_epoch(train_ds, batch_size, rng)

"""
Explanation: Now we make one pass (epoch) over the data, computing random minibatches of size 30. There are 100 examples total, but with a batch size of 30,
we don't use all the data. We can solve such "boundary effects" later.
End of explanation
"""


def load_dataset_iris(split: str, batch_size: int) -> Iterator[Batch]:
    train_ds, test_ds = get_datasets_iris()
    if split == tfds.Split.TRAIN:
        ds = tf.data.Dataset.from_tensor_slices({"X": train_ds["X"], "y": train_ds["y"]})
    elif split == tfds.Split.TEST:
        ds = tf.data.Dataset.from_tensor_slices({"X": test_ds["X"], "y": test_ds["y"]})
    ds = ds.shuffle(buffer_size=1 * batch_size)
    ds = ds.batch(batch_size)
    ds = ds.cache()
    ds = ds.repeat()  # make infinite stream of batches
    return iter(tfds.as_numpy(ds))  # python iterator


batch_size = 30
train_ds = load_dataset_iris(tfds.Split.TRAIN, batch_size)
valid_ds = load_dataset_iris(tfds.Split.TEST, batch_size)

print(train_ds)

training_steps = 5
for step in range(training_steps):
    batch = next(train_ds)
    print("processing batch {} X shape {}, y shape {}".format(step, batch["X"].shape, batch["y"].shape))

"""
Explanation: Minibatching with TFDS
Below we show how to convert a numpy array into a TFDS.
We shuffle the records and convert to minibatches, and then repeat these batches indefinitely to create an infinite stream,
which we can convert to a python iterator. We pass this iterator of batches to our training loop.
End of explanation
"""


def process_record(batch):
    image = batch["image"]
    label = batch["label"]
    # reshape image to standard size, just for fun
    image = tf.image.resize(image, (32, 32))
    # flatten image to vector
    shape = image.get_shape().as_list()
    D = np.prod(shape)  # no batch dimension
    image = tf.reshape(image, (D,))
    # rescale to -1..+1
    image = tf.cast(image, dtype=tf.float32)
    image = ((image / 255.0) - 0.5) * 2.0
    # convert to standard names
    return {"X": image, "y": label}


def load_mnist(split, batch_size):
    dataset, info = tfds.load("mnist", split=split, with_info=True)
    dataset = dataset.map(process_record)
    if split == "train":
        dataset = dataset.shuffle(10 * batch_size, seed=0)
    dataset = dataset.batch(batch_size)
    dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
    dataset = dataset.cache()
    dataset = dataset.repeat()
    dataset = tfds.as_numpy(dataset)  # leave TF behind
    num_examples = info.splits[split].num_examples
    return iter(dataset), num_examples


batch_size = 100
train_iter, num_train = load_mnist("train", batch_size)
test_iter, num_test = load_mnist("test", batch_size)

num_epochs = 3
num_steps = num_train // batch_size
print(f"{num_epochs} epochs with batch size {batch_size} will take {num_steps} steps")

batch = next(train_iter)
print(batch["X"].shape)
print(batch["y"].shape)

"""
Explanation: Preprocessing the data
We can process the data before creating minibatches.
We can also use pre-fetching to speed things up (see
this TF tutorial for details.)
We illustrate this below for MNIST.
End of explanation
"""


ds, info = tfds.load("binarized_mnist", split=tfds.Split.TRAIN, shuffle_files=True, with_info=True)
print(ds)
print(info)

train_ds, info = tfds.load("mnist", split=tfds.Split.TRAIN, shuffle_files=True, with_info=True)
print(train_ds)
print(info)

ds = tfds.load("mnist", split="train")
print(type(ds))
ds = ds.take(1)  # Only take a single example
print(type(ds))

for example in ds:  # example is `{'image': tf.Tensor, 'label': tf.Tensor}`
    print(list(example.keys()))
    image = example["image"]
    label = example["label"]
    print(image.shape, label)

ds, info = tfds.load("mnist", split="train", with_info=True)
fig = tfds.show_examples(ds, info, rows=2, cols=5)

# This function is not well documented. But source code for show_examples is here:
# https://github.com/tensorflow/datasets/blob/v4.2.0/tensorflow_datasets/core/visualization/image_visualizer.py

"""
Explanation: Vision datasets
MNIST
There are many standard versions of MNIST,
some of which are available from https://www.tensorflow.org/datasets. We give some examples below.
End of explanation
"""


ds, info = tfds.load("cifar10", split="train", with_info=True)
fig = tfds.show_examples(ds, info, rows=2, cols=5)

"""
Explanation: CIFAR
The CIFAR dataset is commonly used for prototyping.
The CIFAR-10 version consists of 60000 32x32 colour images in 10 classes, with 6000 images per class. There are 50000 training images and 10000 test images. The dataset is divided into five training batches and one test batch, each with 10000 images. There is also a 100 class version.
An easy way to get this data is to use TFDS, as we show below.
End of explanation
"""


import tensorflow_datasets as tfds

imagenette_builder = tfds.builder("imagenette/full-size")
imagenette_info = imagenette_builder.info
print(imagenette_info)

imagenette_builder.download_and_prepare()

datasets = imagenette_builder.as_dataset(as_supervised=True)

train_examples = imagenette_info.splits["train"].num_examples
validation_examples = imagenette_info.splits["validation"].num_examples
print("ntrain", train_examples, "nvalidation", validation_examples)

train, test = datasets["train"], datasets["validation"]

import tensorflow as tf

batch_size = 32

train_batch = (
    train.map(lambda image, label: (tf.image.resize(image, (448, 448)), label)).shuffle(100).batch(batch_size).repeat()
)

validation_batch = (
    test.map(lambda image, label: (tf.image.resize(image, (448, 448)), label)).shuffle(100).batch(batch_size).repeat()
)

i = 0
for X, y in train_batch:
    # print(b)
    # X = b['image']
    # y = b['label']
    print("image {}, X shape {}, y shape {}".format(i, X.shape, y.shape))
    i += 1
    if i > 1:
        break

fig = tfds.show_examples(train, imagenette_info, rows=2, cols=5)

"""
Explanation: Imagenet
A lot of vision experiments use the Imagenet dataset, with 1000 classes and ~1M images.
However, this takes a long time to download and process.
The FastAI team made a smaller version called ImageNette, that only has 10 classes of size 160 or 320 pixels (largest dimension). This is good for prototyping, and the images tend to be easier to interpret that CIFAR. A version of the raw data, in a more convenient format (all images 224x224, no dependence on FastAI library) can be found here. It is also bundled into TFDS, as we show below.
End of explanation
"""


def get_datasets_mnist():
    ds_builder = tfds.builder("mnist")
    ds_builder.download_and_prepare()
    train_ds_all = tfds.as_numpy(ds_builder.as_dataset(split="train", batch_size=-1))
    test_ds_all = tfds.as_numpy(ds_builder.as_dataset(split="test", batch_size=-1))

    num_train = len(train_ds_all["image"])
    train_ds["X"] = jnp.reshape(jnp.float32(train_ds_all["image"]) / 255.0, (num_train, -1))
    train_ds["y"] = train_ds_all["label"]

    num_test = len(test_ds_all["image"])
    test_ds["X"] = jnp.reshape(jnp.float32(test_ds["image"]) / 255.0, (num_test, -1))
    test_ds["y"] = test_ds_all["label"]

    return train_ds, test_ds

dataset = load_dataset_iris(tfds.Split.TRAIN, 30)
batches = dataset.repeat().batch(batch_size)

step = 0
num_minibatches = 5
for batch in batches:
    if step >= num_minibatches:
        break
    X, y = batch["image"], batch["label"]
    print("processing batch {} X shape {}, y shape {}".format(step, X.shape, y.shape))
    step = step + 1

print("batchified version v2")
batch_stream = batches.as_numpy_iterator()
for step in range(num_minibatches):
    batch = batch_stream.next()
    X, y = batch["image"], batch["label"]  # convert to canonical names
    print("processing batch {} X shape {}, y shape {}".format(step, X.shape, y.shape))
    step = step + 1

ds = tfds.as_numpy(train_ds)
print(ds)
for i, batch in enumerate(ds):
    print(type(batch))
    X = batch["image"]
    y = batch["label"]
    print(X.shape)
    print(y.shape)
    i += 1
    if i > 2:
        break


ds = tfds.load("mnist", split="train")
ds = ds.take(100)
# ds = tfds.as_numpy(ds)

batches = ds.repeat(2).batch(batch_size)
print(type(batches))
print(batches)

batch_stream = batches.as_numpy_iterator()
print(type(batch_stream))
print(batch_stream)

b = next(batch_stream)
print(type(b))
print(b["image"].shape)

b = batch_stream.next()
print(type(b))
print(b["image"].shape)


ds = tfds.load("mnist", split="train")
batches = ds.repeat().batch(batch_size)
batch_stream = batches.as_numpy_iterator()


def process_stream(stream):
    b = next(stream)
    X = b["image"]
    y = b["label"]
    d = {"X": X, "y": y}
    yield d


my_stream = process_stream(batch_stream)

b = next(my_stream)
print(type(b))
print(b["X"].shape)

b = my_stream.next()
print(type(b))
print(b["X"].shape)

def sample_categorical(N, C):
    p = (1 / C) * np.ones(C)
    y = np.random.choice(C, size=N, p=p)
    return y


def get_datasets_rnd():
    Ntrain = 1000
    Ntest = 1000
    D = 5
    C = 10
    train_ds = {"X": np.random.randn(Ntrain, D), "y": sample_categorical(Ntrain, C)}
    test_ds = {"X": np.random.randn(Ntest, D), "y": sample_categorical(Ntest, C)}
    return train_ds, test_ds


def get_datasets_logreg(key):
    Ntrain = 1000
    Ntest = 1000
    D = 5
    C = 10
    W = jax.random.normal(key, (D, C))

    Xtrain = jax.random.normal(key, (Ntrain, D))
    logits = jnp.dot(Xtrain, W)
    ytrain = jax.random.categorical(key, logits)

    Xtest = jax.random.normal(key, (Ntest, D))
    logits = jnp.dot(Xtest, W)
    ytest = jax.random.categorical(key, logits)

    train_ds = {"X": Xtrain, "y": ytrain}
    test_ds = {"X": Xtest, "y": ytest}
    return train_ds, test_ds

"""
Explanation: Language datasets
Various datasets are used in the natural language processing (NLP) communities.
TODO: fill in.
Graveyard
Here we store some scratch code that you can ignore,
End of explanation
"""


