import numpy as np
from scipy.linalg import lstsq
from scipy.special import expit
from sklearn.datasets import load_breast_cancer, load_iris
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis as skLinearDiscriminantAnalysis

class LinearDiscriminantAnalysis():
    def fit(self, X, y):
        self.classes_ = np.unique(y)
        n_features = X.shape[1]
        n_classes = len(self.classes_)
        self.priors_ = np.zeros(n_classes)
        self.means_ = np.zeros((n_classes, n_features))
        self.covariance_ = np.zeros((n_features, n_features))
        for i, c in enumerate(self.classes_):
            X_c = X[y == c]
            self.priors_[i] = X_c.shape[0] / X.shape[0]
            self.means_[i] = np.mean(X_c, axis=0)
            self.covariance_ += self.priors_[i] * np.cov(X_c.T, bias=True)
        self.coef_ = lstsq(self.covariance_, self.means_.T)[0].T
        self.intercept_ = (-0.5 * np.diag(np.dot(self.means_, self.coef_.T)) +
                           np.log(self.priors_))
        if len(self.classes_) == 2:
            self.coef_ = np.atleast_2d(self.coef_[1] - self.coef_[0])
            self.intercept_ = np.atleast_1d(self.intercept_[1] - self.intercept_[0])
        return self

    def decision_function(self, X):
        scores = np.dot(X, self.coef_.T) + self.intercept_
        if scores.shape[1] == 1:
            return scores.ravel()
        else:
            return scores

    def predict(self, X):
        scores = self.decision_function(X)
        if len(scores.shape) == 1:
            indices = (scores > 0).astype(int)
        else:
            indices = np.argmax(scores, axis=1)
        return self.classes_[indices]

    def predict_proba(self, X):
        scores = self.decision_function(X)
        if len(scores.shape) == 1:
            prob = expit(scores)
            prob = np.vstack((1 - prob, prob)).T
        else:
            scores -= np.max(scores, axis=1)[:, np.newaxis]
            prob = np.exp(scores)
            prob /= np.sum(prob, axis=1)[:, np.newaxis]
        return prob

X, y = load_breast_cancer(return_X_y=True)
clf1 = LinearDiscriminantAnalysis().fit(X, y)
clf2 = skLinearDiscriminantAnalysis(solver='lsqr').fit(X, y)
assert np.allclose(clf1.priors_, clf2.priors_)
assert np.allclose(clf1.means_, clf2.means_)
assert np.allclose(clf1.covariance_, clf2.covariance_)
assert np.allclose(clf1.coef_, clf2.coef_)
assert np.allclose(clf1.intercept_, clf2.intercept_)
prob1 = clf1.decision_function(X)
prob2 = clf2.decision_function(X)
assert np.allclose(prob1, prob2)
prob1 = clf1.predict_proba(X)
prob2 = clf2.predict_proba(X)
assert np.allclose(prob1, prob2)
pred1 = clf1.predict(X)
pred2 = clf2.predict(X)
assert np.array_equal(pred1, pred2)

X, y = load_iris(return_X_y=True)
clf1 = LinearDiscriminantAnalysis().fit(X, y)
clf2 = skLinearDiscriminantAnalysis(solver='lsqr').fit(X, y)
assert np.allclose(clf1.priors_, clf2.priors_)
assert np.allclose(clf1.means_, clf2.means_)
assert np.allclose(clf1.covariance_, clf2.covariance_)
assert np.allclose(clf1.coef_, clf2.coef_)
assert np.allclose(clf1.intercept_, clf2.intercept_)
prob1 = clf1.decision_function(X)
prob2 = clf2.decision_function(X)
assert np.allclose(prob1, prob2)
prob1 = clf1.predict_proba(X)
prob2 = clf2.predict_proba(X)
assert np.allclose(prob1, prob2)
pred1 = clf1.predict(X)
pred2 = clf2.predict(X)
assert np.array_equal(pred1, pred2)