import numpy as np
from scipy.special import logsumexp
from sklearn.datasets import load_iris
from sklearn.naive_bayes import GaussianNB as skGaussianNB

class GaussianNB():
    def fit(self, X, y):
        self.classes_ = np.unique(y)
        n_features = X.shape[1]
        n_classes = len(self.classes_)
        self.theta_ = np.zeros((n_classes, n_features))
        self.sigma_ = np.zeros((n_classes, n_features))
        self.class_count_ = np.zeros(n_classes)
        for i, c in enumerate(self.classes_):
            X_c = X[y == c]
            self.theta_[i] = np.mean(X_c, axis=0)
            self.sigma_[i] = np.var(X_c, axis=0)
            self.class_count_[i] = X_c.shape[0]
        self.class_prior_ = self.class_count_ / np.sum(self.class_count_)
        return self

    def _joint_log_likelihood(self, X):
        joint_log_likelihood = np.zeros((X.shape[0], len(self.classes_)))
        for i in range(len(self.classes_)):
            p1 = np.log(self.class_prior_[i])
            p2 = -0.5 * np.log(2 * np.pi * self.sigma_[i]) - 0.5 * (X - self.theta_[i]) ** 2 / self.sigma_[i]
            joint_log_likelihood[:, i] = p1 + np.sum(p2, axis=1)
        return joint_log_likelihood

    def predict(self, X):
        joint_log_likelihood = self._joint_log_likelihood(X)
        return self.classes_[np.argmax(joint_log_likelihood, axis=1)]

    def predict_proba(self, X):
        joint_log_likelihood = self._joint_log_likelihood(X)
        log_prob = joint_log_likelihood - logsumexp(joint_log_likelihood, axis=1)[:, np.newaxis]
        return np.exp(log_prob)

X, y = load_iris(return_X_y = True)
clf1 = GaussianNB().fit(X, y)
clf2 = skGaussianNB().fit(X, y)
assert np.allclose(clf1.theta_, clf2.theta_)
assert np.allclose(clf1.sigma_, clf2.sigma_)
assert np.allclose(clf1.class_count_, clf2.class_count_)
assert np.allclose(clf1.class_prior_, clf2.class_prior_)
prob1 = clf1._joint_log_likelihood(X)
prob2 = clf2._joint_log_likelihood(X)
assert np.allclose(prob1, prob2)
pred1 = clf1.predict(X)
pred2 = clf2.predict(X)
assert np.array_equal(pred1, pred2)
prob1 = clf1.predict_proba(X)
prob2 = clf2.predict_proba(X)
assert np.allclose(prob1, prob2)