# Source: content/notes/ml/probabilistic-and-instance-models.md
# Independent CPU example; use the curriculum environment.
# See /notes/ml/#example-environment or /notes/deep-learning/#example-environment.

import numpy as np
from sklearn.metrics import accuracy_score, log_loss
from sklearn.model_selection import train_test_split
from sklearn.naive_bayes import BernoulliNB, CategoricalNB, ComplementNB, MultinomialNB

rng = np.random.default_rng(9)
y = np.repeat([0, 1], 120)
rates = np.array([[4.0, 0.7, 2.0, 0.5], [0.6, 4.0, 0.7, 2.0]])
counts = rng.poisson(rates[y])
train, test = train_test_split(np.arange(len(y)), stratify=y, random_state=9)
for name, model, X in [
    ("Multinomial", MultinomialNB(alpha=1.0), counts),
    ("Complement", ComplementNB(alpha=1.0), counts),
    ("Bernoulli", BernoulliNB(alpha=1.0), (counts > 0).astype(int)),
    ("Categorical", CategoricalNB(alpha=1.0, min_categories=4), np.minimum(counts, 3)),
]:
    model.fit(X[train], y[train])
    pred = model.predict(X[test])
    prob = model.predict_proba(X[test])
    assert np.allclose(prob.sum(axis=1), 1)
    assert accuracy_score(y[test], pred) > 0.6
    print(name, "accuracy", accuracy_score(y[test], pred),
          "log loss", log_loss(y[test], prob))