# If running in TPU mode
import jax.tools.colab_tpu

jax.tools.colab_tpu.setup_tpu()

%%capture
!pip install sgmcmcjax

import jax
import jax.numpy as jnp
from jax import jit, grad, vmap, random

print(jax.__version__)
print(jax.devices())

from sgmcmcjax.samplers import build_sgld_sampler
from sgmcmcjax.kernels import build_sgld_kernel, build_psgld_kernel, build_sgldAdam_kernel, build_sghmc_kernel

from tqdm.auto import tqdm
import matplotlib.pyplot as plt
import numpy as np

key = random.PRNGKey(0)

"""
Explanation: <a href="https://colab.research.google.com/github/probml/probml-notebooks/blob/main/notebooks/sg_mcmc_jax.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>
SG-MCMC-JAX library
https://github.com/jeremiecoullon/SGMCMCJax
Setup
End of explanation
"""


# define model in JAX
def loglikelihood(theta, x):
    return -0.5 * jnp.dot(x - theta, x - theta)


def logprior(theta):
    return -0.5 * jnp.dot(theta, theta) * 0.01


# generate dataset
N, D = 10_000, 100
key = random.PRNGKey(0)
mu_true = random.normal(key, shape=(D,))
X_data = random.normal(key, shape=(N, D)) + mu_true

# build sampler
batch_size = int(0.1 * N)
dt = 1e-5
my_sampler = build_sgld_sampler(dt, loglikelihood, logprior, (X_data,), batch_size)

%%time 
Nsamples = 10_000
samples = my_sampler(key, Nsamples, jnp.zeros(D))




print(samples.shape)
mu_est = jnp.mean(samples, axis=0)
print(jnp.allclose(mu_true, mu_est, atol=1e-1))
print(mu_true[:10])
print(mu_est[:10])

data = (X_data,)
init_fn, sgld_kernel, get_params = build_sgld_kernel(dt, loglikelihood, logprior, data, batch_size)


key, subkey = random.split(key)
params = random.normal(subkey, shape=(D,))

key, subkey = random.split(key)
state = init_fn(subkey, params)
print(state)

"""
Explanation: Gaussian posterior
https://sgmcmcjax.readthedocs.io/en/latest/nbs/gaussian.html
End of explanation
"""


from models.logistic_regression import gen_data, loglikelihood, logprior

key = random.PRNGKey(42)
dim = 10
Ndata = 100000

theta_true, X, y_data = gen_data(key, dim, Ndata)

"""
Explanation: Logistic regression
https://sgmcmcjax.readthedocs.io/en/latest/nbs/logistic_regression.html
End of explanation
"""


from models.bayesian_NN.NN_data import X_train, X_test, y_train, y_test
from models.bayesian_NN.NN_model import init_network, loglikelihood, logprior, accuracy

from sgmcmcjax.kernels import build_sgld_kernel
from tqdm.auto import tqdm

"""
Explanation: Bayesian neural network
https://sgmcmcjax.readthedocs.io/en/latest/nbs/BNN.html
End of explanation
"""


%%capture
!pip install --upgrade git+https://github.com/google/flax.git

import tensorflow_datasets as tfds

from flax import linen as nn

class CNN(nn.Module):
    """A simple CNN model."""

    @nn.compact
    def __call__(self, x):
        x = nn.Conv(features=32, kernel_size=(3, 3))(x)
        x = nn.relu(x)
        x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
        x = nn.Conv(features=64, kernel_size=(3, 3))(x)
        x = nn.relu(x)
        x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
        x = x.reshape((x.shape[0], -1))  # flatten
        x = nn.Dense(features=256)(x)
        x = nn.relu(x)
        x = nn.Dense(features=10)(x)
        x = nn.log_softmax(x)
        return x

cnn = CNN()


def loglikelihood(params, x, y):
    x = x[jnp.newaxis]  # add an axis so that it works for a single data point
    logits = cnn.apply({"params": (params)}, x)
    label = jax.nn.one_hot(y, num_classes=10)
    return jnp.sum(logits * label)


def logprior(params):
    return 1.0


@jit
def accuracy_cnn(params, X, y):
    target_class = y
    predicted_class = jnp.argmax(cnn.apply({"params": (params)}, X), axis=1)
    return jnp.mean(predicted_class == target_class)


def get_datasets():
    """Load MNIST train and test datasets into memory."""
    ds_builder = tfds.builder("mnist")
    ds_builder.download_and_prepare()
    train_ds = tfds.as_numpy(ds_builder.as_dataset(split="train", batch_size=-1))
    test_ds = tfds.as_numpy(ds_builder.as_dataset(split="test", batch_size=-1))
    train_ds["image"] = jnp.float32(train_ds["image"]) / 255.0
    test_ds["image"] = jnp.float32(test_ds["image"]) / 255.0
    return train_ds, test_ds


train_ds, test_ds = get_datasets()

X_train_s = train_ds["image"]
y_train_s = jnp.array(train_ds["label"])
X_test_s = test_ds["image"]
y_test_s = jnp.array(test_ds["label"])


data = (X_train_s, y_train_s)
batch_size = int(0.01 * data[0].shape[0])
print(batch_size)

params = cnn.init(key, jnp.ones([1, 28, 28, 1]))
print(params.keys())
params = params["params"]
print(params["Dense_0"]["kernel"].shape)
print(params["Dense_1"]["kernel"].shape)

def run_sgmcmc(key, Nsamples, init_fn, my_kernel, get_params, record_accuracy_every=50):
    "Run SGMCMC sampler and return the test accuracy list"
    accuracy_list = []
    params = cnn.init(key, jnp.ones([1, 28, 28, 1]))["params"]
    key, subkey = random.split(key)
    state = init_fn(subkey, params)

    for i in tqdm(range(Nsamples)):
        key, subkey = random.split(key)
        state = my_kernel(i, subkey, state)
        if i % record_accuracy_every == 0:
            test_acc = accuracy_cnn(get_params(state), X_test_s, y_test_s)
            accuracy_list.append((i, test_acc))

    return accuracy_list

kernel_dict = {}

# SGLD
init_fn, kernel, get_params = build_sgld_kernel(5e-6, loglikelihood, logprior, data, batch_size)
kernel = jit(kernel)
kernel_dict["sgld"] = (init_fn, kernel, get_params)

# SG-HMC
init_fn, kernel, get_params = build_sghmc_kernel(
    1e-6, 4, loglikelihood, logprior, data, batch_size, compiled_leapfrog=False
)
kernel = jit(kernel)
kernel_dict["sghmc"] = (init_fn, kernel, get_params)

# PSGLD
init_fn, kernel, get_params = build_psgld_kernel(1e-3, loglikelihood, logprior, data, batch_size)
kernel = jit(kernel)
kernel_dict["psgld"] = (init_fn, kernel, get_params)

# SGLD-Adam
init_fn, kernel, get_params = build_sgldAdam_kernel(1e-2, loglikelihood, logprior, data, batch_size)
kernel = jit(kernel)
kernel_dict["sgldAdam"] = (init_fn, kernel, get_params)

Nsamples = 500
acc_dict = {}

for (name, fns) in kernel_dict.items():
    print(name)
    init_fn, kernel, get_params = fns
    accuracy_list = run_sgmcmc(random.PRNGKey(0), Nsamples, init_fn, kernel, get_params)
    print(accuracy_list)
    acc_dict[name] = accuracy_list

plt.figure()
for name, acc_list in acc_dict.items():
    steps = [s for (s, a) in acc_list]
    accs = [a for (s, a) in acc_list]
    plt.plot(steps, accs, "-o", label=name)

plt.legend(fontsize=16)
plt.title("Test accuracy of SGMCMC samplers on MNIST with a CNN")
plt.xlabel("Iterations", size=20)
plt.ylabel("Test accuracy", size=20)

"""
Explanation: FLAX CNN
https://sgmcmcjax.readthedocs.io/en/latest/nbs/Flax_MNIST.html
End of explanation
"""


