"""==========================
Deep Ensemble on Two Moons
==========================

Train several independent copies of the same network from different initializations and average their predictions.
"""

from __future__ import annotations

from sklearn.datasets import make_moons
import torch
from torch import nn

from probly.representer import representer
from probly.transformation import ensemble

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()

# %%
# Model
# -----

base_model = MLPClassifier()

ensemble_model = ensemble(
    base_model,
    num_members=3,
    reset_params=True,  # fresh init per member maximizes diversity
    predictor_type="logit_classifier",
)

# %%
# Training
# --------
#
# Train each member independently with standard cross-entropy.
# The uncertainty signal comes from disagreement between members,
# not from the loss itself.

ensemble_model.train()
for member in ensemble_model:
    opt = torch.optim.Adam(member.parameters(), lr=1e-3)
    for epoch in range(250):
        opt.zero_grad()
        out = member(X_tensor)
        loss = nn.functional.cross_entropy(out, y_tensor)
        loss.backward()
        opt.step()

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

ensemble_model.eval()
rep = representer(ensemble_model)

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