"""=================
DDU on Two Moons
=================

Deep Deterministic Uncertainty (DDU) applies spectral normalization to a
feature extractor, then fits a class-conditional Gaussian density model on
the training features.  A single deterministic forward pass yields both
a class prediction and a feature-density score for epistemic uncertainty.
"""

from __future__ import annotations

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

import torch
from torch import nn

from probly.representer import representer
from probly.method.ddu import ddu

from examples.utils.model import ResFFN
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=64, shuffle=True)



# %%
# Model
# -----
#
# DDU wraps the backbone with spectral normalization (controlled by ``sn_coeff``)
# to smooth the Lipschitz constant of the feature map, which is required for
# the density score to be a reliable distance proxy.

base_model = ResFFN()

ddu_model = ddu(base_model, sn_coeff = 3.0, predictor_type="logit_classifier")



# %%
# Training
# --------

opt = torch.optim.Adam(ddu_model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()

ddu_model.train()
for epoch in range(200):
    opt.zero_grad()

    features = ddu_model.encoder(X_tensor)
    logits = ddu_model.classification_head(features)

    loss = criterion(logits, y_tensor)

    loss.backward()
    opt.step()

# %%
# Fit Density Head
# ----------------
#
# Collect all training features in one pass and fit the class-conditional
# Gaussians.  This only needs to happen once after training.

ddu_model.eval()

all_features = []
all_labels = []

with torch.no_grad():
    for inputs, targets in train_loader:
        features = ddu_model.encoder(inputs)
        all_features.append(features.detach().cpu())
        all_labels.append(targets.detach().cpu())

features_cat = torch.cat(all_features)
labels_cat = torch.cat(all_labels)

density_head = ddu_model.density_head
density_head.fit(features_cat, labels_cat)

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

ddu_model.eval()

rep = representer(ddu_model)

plot = plot_example_uncertainty(X, y, rep, title="DDU Predictive Uncertainty", notion="epistemic", log_scale=True)

plot.show()