from sklearn.discriminant_analysis import LinearDiscriminantAnalysis
import numpy as np


def cul_w(X_, y_):
    y_ = np.array(y_.T)
    pos = y_ == 1
    neg = y_ == 0

    X0 = X_[neg[0]]
    X1 = X_[pos[0]]

    u0 = X0.mean(0)
    u1 = X1.mean(0)
    u = X_.mean(0)

    # calc Sw
    sw0 = 0
    for i in range(0, X0.shape[0]):
        sw0 = sw0 + (X0[i, :] - u0).T * (X0[i, :] - u0)

    sw1 = 0
    for i in range(0, X1.shape[0]):
        sw1 = sw1 + (X1[i, :] - u1).T * (X1[i, :] - u1)
    sw = sw0 + sw1

    # calc sb
    sb = (u - u0).T * (u - u0) + (u - u1).T * (u - u1)

    eig_value, vec = np.linalg.eig(np.mat(sw).I * sb)
    print ("eig is :",eig_value)
    print ("vec is :", vec)


if '__main__' == __name__:
    X=[[2, 50, 12500, 98], [0, 13, 3250, 28], [1, 16, 4000, 35], [2, 20, 5000, 45], [1, 24, 6000, 77], [4, 4, 1000, 4], [2, 7, 1750, 14]]
    y=[[1], [1], [1], [1], [0], [0], [1]]
    X=np.mat(X)
    y=np.mat(y).astype(int)

    cul_w(X, y)