import numpy as np
from sklearn.datasets import load_iris
import matplotlib.pyplot as plt

class LDA:
    def __init__(self, X, y):
        self.X = list(X)
        self.y = list(y)
        self.dim = 2
        
    def fit(self):
        label_ = list(set(y))
    
        X_classify = {}
    
        for label in label_:
            X1 = np.array([X[i] for i in range(len(X)) if y[i] == label])
            X_classify[label] = X1
    
        mju = np.mean(X, axis=0)
        mju_classify = {}
    
        for label in label_:
            mju1 = np.mean(X_classify[label], axis=0)
            mju_classify[label] = mju1
    
        #St = np.dot((X - mju).T, X - mju)
    
        Sw = np.zeros((len(mju), len(mju)))  
        for i in label_:
            Sw += np.dot((X_classify[i] - mju_classify[i]).T,
                         X_classify[i] - mju_classify[i])
    
        # Sb=St-Sw
    
        Sb = np.zeros((len(mju), len(mju)))  
        for i in label_:
            Sb += len(X_classify[i]) * np.dot((mju_classify[i] - mju).reshape(
                (len(mju), 1)), (mju_classify[i] - mju).reshape((1, len(mju))))
    
        eig_vals, eig_vecs = np.linalg.eig(
            np.linalg.inv(Sw).dot(Sb))  
    
        sorted_indices = np.argsort(eig_vals)
        topk_eig_vecs = eig_vecs[:, sorted_indices[:-self.dim - 1:-1]]  
        return topk_eig_vecs

            


if '__main__' == __name__:

    iris = load_iris()
    X = iris.data
    y = iris.target
    L = LDA(X, y)
    W = L.fit()
    X_new = np.dot((X), W)
    plt.figure(1)
    plt.scatter(X_new[:, 0], X_new[:, 1], marker='o', c=y)
    plt.show()

    
