"""
=====================================================
MixedNB Equivalence with BernoulliNB
=====================================================

This example demonstrates that :class:`skbn.mixed_nb.MixedNB`
produces identical results to :class:`sklearn.naive_bayes.BernoulliNB`
when all features are binary (Bernoulli).

The plot shows the decision boundaries for both classifiers. As expected,
the boundaries are identical, and the predicted probabilities for the
dataset are all-close.
"""

# Author: The scikit-bayes Developers
# SPDX-License-Identifier: BSD-3-Clause

import matplotlib.pyplot as plt
import numpy as np
from numpy.testing import assert_allclose
from sklearn.naive_bayes import BernoulliNB

from skbn.mixed_nb import MixedNB

# 1. Generate a 2D Bernoulli dataset (0s and 1s)
np.random.seed(42)
n_samples = 100
X = np.random.randint(0, 2, size=(n_samples, 2), dtype=int)
# Logic: AND gate (Class 1 only if both are 1)
y = (X[:, 0] & X[:, 1]).astype(int)

# 2. Fit both classifiers
bnb = BernoulliNB(alpha=1.0)
bnb.fit(X, y)
probs_bnb = bnb.predict_proba(X)

# MixedNB will auto-detect both features as 'bernoulli' (unique_values=2)
mnb = MixedNB(alpha=1.0)
mnb.fit(X, y)
probs_mnb = mnb.predict_proba(X)

# 3. Assert equivalence
try:
    assert_allclose(probs_bnb, probs_mnb, rtol=1e-7, atol=1e-7)
    equivalence_message = "Probabilities are identical."
except AssertionError as e:
    equivalence_message = f"Probabilities are NOT identical:\n{e}"

print(f"BernoulliNB vs MixedNB Equivalence Check: {equivalence_message}")

# 4. Plot decision boundaries
fig, axes = plt.subplots(1, 2, figsize=(14, 6), sharey=True)

models = [bnb, mnb]
titles = ["1. scikit-learn BernoulliNB", "2. skbn MixedNB (auto-detected)"]

# Define grid (Centers for prediction, Edges for plotting)
# Binary features: 0, 1
centers = np.array([0, 1])
edges = np.array([-0.5, 0.5, 1.5])

# Create prediction grid
xx, yy = np.meshgrid(centers, centers)
grid_pred = np.c_[xx.ravel(), yy.ravel()]

for ax, model, title in zip(axes, models, titles):
    # Predict probabilities on binary grid
    probs = model.predict_proba(grid_pred)[:, 1]
    Z = probs.reshape(xx.shape)

    # Plot Heatmap - VIRIDIS
    # Using pcolormesh with edges defines the 4 quadrants perfectly
    ax.pcolormesh(
        edges,
        edges,
        Z,
        cmap="viridis",
        vmin=0,
        vmax=1,
        shading="flat",
        alpha=0.8,
        edgecolors="none",
    )

    # Overlay real data points with consistent style
    # Jitter points slightly to show density (since many points overlap on 0/1)
    x_jit = X[:, 0] + np.random.uniform(-0.15, 0.15, size=n_samples)
    y_jit = X[:, 1] + np.random.uniform(-0.15, 0.15, size=n_samples)

    # Class 0 -> Indigo Circle
    ax.scatter(
        x_jit[y == 0],
        y_jit[y == 0],
        c="indigo",
        marker="o",
        s=50,
        alpha=0.8,
        edgecolors="w",
        linewidth=0.8,
        label="Class 0",
    )

    # Class 1 -> Gold Triangle
    ax.scatter(
        x_jit[y == 1],
        y_jit[y == 1],
        c="gold",
        marker="^",
        s=50,
        alpha=0.9,
        edgecolors="k",
        linewidth=0.5,
        label="Class 1",
    )

    ax.set_title(title, fontsize=12)
    ax.set_xticks([0, 1])
    ax.set_yticks([0, 1])
    ax.grid(True, alpha=0.2, linestyle="--")

axes[0].set_ylabel("Feature 1 (Bernoulli)")
axes[0].set_xlabel("Feature 0 (Bernoulli)")
axes[1].set_xlabel("Feature 0 (Bernoulli)")
axes[0].legend(loc="lower left")

fig.suptitle("Equivalence of MixedNB and BernoulliNB on Binary Data", fontsize=16)
plt.tight_layout()
plt.subplots_adjust(top=0.85)
plt.show()
