import numpy as np
import matplotlib.pyplot as plt

np.random.seed(seed=1)
import math
import collections
import os
import random

import json
import multiprocessing

try:
    import torch
except ModuleNotFoundError:
    %pip install -qq torch
    import torch
from torch import nn
from torch.nn import functional as F

try:
    import torchvision
except ModuleNotFoundError:
    %pip install -qq torchvision
    import torchvision

!mkdir figures # for saving plots

!wget https://raw.githubusercontent.com/d2l-ai/d2l-en/master/d2l/torch.py -q -O d2l.py
import d2l

"""
Explanation: Please find jax implementation of this notebook here: https://colab.research.google.com/github/probml/pyprobml/blob/master/notebooks/book1/09/naive_bayes_mnist_jax.ipynb
<a href="https://colab.research.google.com/github/probml/probml-notebooks/blob/main/notebooks-d2l/naive_bayes_mnist.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>
Naive Bayes classifiers
We show how to implement Naive Bayes classifiers from scratch.
We use binary features, and 2 classes.
Based on sec 18.9 of http://d2l.ai/chapter_appendix-mathematics-for-deep-learning/naive-bayes.html.
End of explanation
"""


data_transform = torchvision.transforms.Compose([torchvision.transforms.ToTensor()])

mnist_train = torchvision.datasets.MNIST(root="./temp", train=True, transform=data_transform, download=True)
mnist_test = torchvision.datasets.MNIST(root="./temp", train=False, transform=data_transform, download=True)

print(mnist_train)

image, label = mnist_train[2]
print(type(image))
print(image.shape)
print(type(label))
print(label)

image[0, 15:20, 15:20]  # not binary (pytorch rescales to 0:1)

[torch.min(image), torch.max(image)]

print(mnist_train[0][0].shape)  # (1,28,28)
indices = [0, 1]
xx = torch.stack([mnist_train[i][0] for i in indices])
print(xx.shape)
xx = torch.stack([mnist_train[i][0] for i in indices], dim=1)
print(xx.shape)
xx = torch.stack([mnist_train[i][0] for i in indices], dim=1).squeeze(0)
print(xx.shape)

# convert from torch.tensor to numpy, extract subset of indices, optionally binarize
def get_data(data, indices=None, binarize=True):
    N = len(data)
    if indices is None:
        indices = range(0, N)
    # X = torch.stack([data[i][0] for i in indices], dim=1).squeeze(0) # (N,28,28)
    X = np.stack([data[i][0].numpy() for i in indices], axis=1).squeeze(0)  # (N,28,28)
    if binarize:
        X = X > 0.5
    # y = torch.tensor([data[i][1] for i in indices])
    y = np.array([data[i][1] for i in indices])
    return X, y

indices = range(0, 10)
images, labels = get_data(mnist_train, indices, False)
print([images.shape, labels.shape])
print(images[0, 15:20, 15:20])  # not binary
d2l.show_images(images, 1, 10);

indices = range(0, 10)
images, labels = get_data(mnist_train, indices, True)
print([images.shape, labels.shape])
print(images[0, 15:20, 15:20])  # binary
d2l.show_images(images, 1, 10);

X_train, y_train = get_data(mnist_train)
X_test, y_test = get_data(mnist_test)

print(X_train.shape)
print(type(X_train))

print(X_train[0, 15:20, 15:20])

"""
Explanation: Get data
We use a binarized version of MNIST.
End of explanation
"""


n_y = np.zeros(10)
for y in range(10):
    n_y[y] = (y_train == y).sum()
P_y = n_y / n_y.sum()
P_y

# Training set is not equally balanced across classes...
print(np.unique(y_train))
from collections import Counter

cnt = Counter(y_train)
print(cnt.keys())
print(cnt.values())

"""
Explanation: Training
End of explanation
"""


n_x = np.zeros((10, 28, 28))
for y in range(10):
    n_x[y] = X_train[y_train == y].sum(axis=0)
# using pseudo counts of 1
# P_xy = (n_x + 1) / (n_y + 1).reshape(10, 1, 1)
P_xy = (n_x + 1) / (n_y + 2).reshape(10, 1, 1)
print(P_xy.shape)
print(type(P_xy))

d2l.show_images(P_xy, 1, 10)
plt.tight_layout()
plt.savefig("nbc_mnist_centroids.pdf", dpi=300)

"""
Explanation: We use add-one smoothing for class conditional Bernoulli distributions.
End of explanation
"""


log_P_xy = np.log(P_xy)
log_P_xy_neg = np.log(1 - P_xy)
log_P_y = np.log(P_y)


def bayes_pred_stable(x):
    # x = x.unsqueeze(0)  # (28, 28) -> (1, 28, 28)
    x = np.expand_dims(x, 0)  # (28, 28) -> (1, 28, 28)
    p_xy = log_P_xy * x + log_P_xy_neg * (1 - x)  # select the 0 and 1 pixels
    p_xy = p_xy.reshape(10, -1).sum(axis=1)  # p(x|y)
    return p_xy + log_P_y


def predict(X):
    return [np.argmax(bayes_pred_stable(x)) for x in X]


# image, label = mnist_test[0]
image = X_test[0]
label = y_test[0]
py = bayes_pred_stable(image)
print(py)
print("ytrue ", label, "yhat ", np.argmax(py))
print(predict([image]))

indices = range(0, 10)
X1, y1 = get_data(mnist_test, indices, True)
preds = predict(X1)
d2l.show_images(X1, 1, 10, titles=[str(d) for d in preds])
plt.tight_layout()
plt.savefig("nbc_mnist_preds.pdf", dpi=300)

indices = range(5, 10)
X1, y1 = get_data(mnist_test, indices, True)
preds = predict(X1)
d2l.show_images(X1, 1, 5, titles=[str(d) for d in preds])
plt.tight_layout()
plt.savefig("nbc_mnist_preds.pdf", dpi=300)

indices = range(30, 40)
X1, y1 = get_data(mnist_test, indices, True)
preds = predict(X1)
d2l.show_images(X1, 1, 10, titles=[str(d) for d in preds]);

preds = predict(X_test)
float((preds == y_test).sum()) / len(y_test)  # test accuracy

"""
Explanation: Testing
End of explanation
"""


