"""==================================
Laplace Approximation on Two Moons
==================================

The Laplace Approximation is a post-hoc method that turns a deterministically
trained neural network into a Bayesian Neural Network by approximating the
posterior over weights with a Gaussian.  Uncertainty concentrates along the
decision boundary and grows away from the training manifold.
"""

from __future__ import annotations

from laplace import Laplace
from sklearn.datasets import make_moons
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

from probly.representer import representer

from examples.utils.model import MLPClassifier
from examples.utils.plotting import plot_example_uncertainty

# %%
# Setup
# -----

X, y = make_moons(n_samples=500, noise=0.05, random_state=0)
X_tensor = torch.from_numpy(X).float()
y_tensor = torch.from_numpy(y).long()

dataset = TensorDataset(X_tensor, y_tensor)
train_loader = DataLoader(dataset, batch_size=32, shuffle=True)

# %%
# Model
# -----
#
# Train a standard MLP classifier with mini-batched cross-entropy; the
# Laplace approximation is applied afterwards as a post-hoc uncertainty wrapper.

base_model = MLPClassifier()

opt = torch.optim.Adam(base_model.parameters(), lr=1e-3)

base_model.train()
for _epoch in range(300):
    for X_batch, y_batch in train_loader:
        opt.zero_grad()
        out = base_model(X_batch)
        loss = nn.functional.cross_entropy(out, y_batch)
        loss.backward()
        opt.step()

# %%
# Laplace Approximation
# ---------------------
#
# Fit a Kronecker-factored (KFAC) Laplace approximation over the last layer of
# the trained model.  No retraining is needed.

base_model.eval()

fit_loader = DataLoader(dataset, batch_size=32)
laplace_model = Laplace(
    base_model,
    "classification",
    subset_of_weights="last_layer",
    hessian_structure="kron",
)
laplace_model.fit(fit_loader)

# %%
# Uncertainty Evaluation
# ----------------------

rep = representer(laplace_model, num_samples=200)

plot = plot_example_uncertainty(X, y, rep, title="Laplace Predictive Uncertainty", notion="total")
plot.show()