"""====================
Het-Net on Two Moons
====================

Het-Net augments a standard classifier with a learnable heteroscedastic noise
head that draws Monte Carlo samples in logit space.  The representer reuses the
same noise mechanism at inference time to estimate per-sample aleatoric
uncertainty, which is helpful on data with input-dependent label noise such as
this noisy Two Moons setup.
"""

from __future__ import annotations

from sklearn.datasets import make_moons
import torch
import torch.nn.functional as F
import numpy as np

from probly.layers.torch import HeteroscedasticLayer
from probly.method.het_net import het_net
from probly.representer import representer

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

# %%
# Setup
# -----
#
# Two Moons with input-dependent (heteroscedastic) Gaussian noise: points
# further from each class' core get a larger noise scale, mimicking
# spatially varying label/observation noise.

X, y = make_moons(n_samples=500, noise=0.0, random_state=0)
rng = np.random.default_rng(0)

noise_scale = np.zeros_like(X[:, 0])

x0 = X[y == 0, 0]
noise_scale[y == 0] = 0.05 + 0.4 * (np.max(x0) - x0) / (np.max(x0) - np.min(x0))

x1 = X[y == 1, 0]
noise_scale[y == 1] = 0.05 + 0.4 * (x1 - np.min(x1)) / (np.max(x1) - np.min(x1))

X += rng.normal(scale=np.expand_dims(noise_scale, 1), size=X.shape)

X_tensor = torch.from_numpy(X).float()
y_tensor = torch.from_numpy(y).long()

# %%
# Model
# -----
#
# Het-Net wraps the backbone with a ``HeteroscedasticLayer`` that learns a
# per-sample noise distribution over the logits.  The noise head is co-trained
# with the backbone and captures input-dependent (aleatoric) uncertainty.

base_model = SequentialModel()
het_net_model = het_net(base_model, predictor_type="logit_classifier")

# %%
# Training
# --------
#
# Setting ``training_samples = S`` on every ``HeteroscedasticLayer`` makes the
# head draw S noise samples per input in a single vectorized forward pass and
# return the log of the softmax-averaged probabilities, optimized with NLL.

opt = torch.optim.Adam(het_net_model.parameters(), lr=1e-3)
training_samples = 4

het_layers = [m for m in het_net_model.modules() if isinstance(m, HeteroscedasticLayer)]
for layer in het_layers:
    layer.training_samples = training_samples

het_net_model.train()
try:
    for _epoch in range(500):
        opt.zero_grad()
        log_probs = het_net_model(X_tensor)
        loss = F.nll_loss(log_probs, y_tensor)
        loss.backward()
        opt.step()
finally:
    for layer in het_layers:
        layer.training_samples = 1

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

het_net_model.eval()
rep = representer(het_net_model, num_samples=800)

plot = plot_example_uncertainty(X, y, rep, title="HET-Net Predictive Uncertainty", notion="aleatoric")
plot.show()