import numpy as np
from sklearn.decomposition import PCA
from nn import nn

class LDA:
    dataset = []
    labels = []
    label_ = []
    W=[]
    k=2

    def __init__(self, k):
        '''
            lda=LDA(k)
            lda.fit_transform(dataset,labels)
            lda.W
        '''
        self.k=k

    def fit(self,dataset,labels):
        self.dataset = np.array(dataset,dtype="float")
        self.labels = labels
        self.label_ = list(set(labels))
        X_classify = {}
        for label in self.label_:
            x = np.array([self.dataset[i] for i in range(
                len(self.dataset)) if self.labels[i] == label])
            X_classify[label] = x

        mju = np.mean(self.dataset, axis=0)
        mju_classify = {}

        for label in self.label_:
            mju1 = np.mean(X_classify[label], axis=0)
            mju_classify[label] = mju1

        # St=np.dot((dataset-mju).T,(dataset-mju))

        Sw = np.zeros((len(mju), len(mju)))
        for i in self.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 self.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))  # 计算Sw-1*Sb的特征值和特征矩阵

        sorted_indices = np.argsort(eig_vals)
        topk_eig_vecs = eig_vecs[:, sorted_indices[:-self.k - 1:-1]]  # 提取前k个特征向量
        self.W=topk_eig_vecs
        return topk_eig_vecs

    def transform(self,X,personalW=None):
        if personalW is None: personalW=self.W
        return np.array(np.dot(X,personalW)).astype(float)

    def fit_transform(self,dataset,labels):
        self.fit(dataset,labels)
        return self.transform(dataset)


# %% get PCA+LDA new dataset
def handledData(train_dat,train_lab,all_dat,all_lab,pca_n,lda_n):
    pca = PCA(n_components=pca_n)
    x_train = pca.fit_transform(train_dat)
    x_train=np.array(x_train).astype(float)
    y_train = train_lab

    x_all=pca.transform(all_dat)
    x_all=np.array(x_all).astype(float)
    y_all=all_lab

    lda=LDA(lda_n)
    lda_train=lda.fit_transform(x_train,y_train)
    lda_all=lda.transform(x_all)
    return lda_train,y_train,lda_all,y_all

# %% test LDA
def testLDA(train_dat,train_lab,all_dat,all_lab,pca_n,lda_n,knn_k):
    pca = PCA(n_components=pca_n)
    x_train = pca.fit_transform(train_dat)
    x_train=np.array(x_train).astype(float)
    y_train = train_lab

    x_all=pca.transform(all_dat)
    x_all=np.array(x_all).astype(float)
    y_all=all_lab

    lda=LDA(lda_n)
    lda_train=lda.fit_transform(x_train,y_train)
    lda_all=lda.transform(x_all)

    return nn(knn_k,lda_train,y_train,lda_all,y_all,tips="",console=False)

def bestLDA(train_dat,train_lab,all_dat,all_lab,pca_n,knn_k=1):
    '''
        return [best lda_n,success_rate]
    '''
    pca = PCA(n_components=pca_n)
    x_train = pca.fit_transform(train_dat)
    x_train=np.array(x_train).astype(float)
    y_train = train_lab

    x_all=pca.transform(all_dat)
    x_all=np.array(x_all).astype(float)
    y_all=all_lab

    # 优化
    res=[]
    lda=LDA(pca.n_components_-1)
    W=lda.fit(x_train,y_train)
    for i in range(1,pca.n_components_):
        lda_train=lda.transform(x_train,W[:,0:i])
        lda_all=lda.transform(x_all,W[:,0:i])
        res.append([i,nn(knn_k,lda_train,y_train,lda_all,y_all,tips="",console=False)])
    return sorted(res,key=lambda x: x[1])[-1]

    # 原始
    # res=[]
    # for i in range(1,pca.n_components_):
    #     lda=LDA(i)
    #     lda_train=lda.fit_transform(x_train,y_train)
    #     lda_all=lda.transform(x_all)
    #     res.append([i,nn(knn_k,lda_train,y_train,lda_all,y_all,tips="",console=False)])
    # return sorted(res,key=lambda x: x[1])[-1]