from sklearn.model_selection import KFold,StratifiedKFold
import numpy as np
import pandas as pd
from openpyxl import Workbook
import time
start = time.clock()
def loadDataSet(filename):
    E = pd.read_excel(filename,sheet_name=0,header=None)
    E = np.array(E)
    return E
def MGD_P(X,A,B,row,column,mini_batch,lamda_P):
    descentP_up = np.zeros((F, 1))
    descentP_down = np.zeros((F, 1))
    X_head = np.dot(A, B)
    for f in range(F):
        a_up = 0
        a_down = 0
        for n in range(mini_batch):
            hang = np.uint8(row[n][0])
            lie = np.uint8(column[n][0])
            a_up += X[hang][lie] * B[f][lie]
            a_down += B[f][lie] * X_head[hang][lie] + lamda_P * A[hang][f]
        descentP_up[f][0] = a_up
        descentP_down[f][0] = a_down
    for h in range(len(X)):
        for f2 in range(F):
            P[h][f2] = A[h][f2] * (descentP_up[f2][0] / descentP_down[f2][0])
    return P
def MGD_Q(X,A,B,row,column,mini_batch,lamda_Q):
    X_head1 = np.dot(A, B)
    descentQ_up = np.zeros((F, 1))
    descentQ_down = np.zeros((F, 1))
    for f3 in range(F):
        b_up = 0
        b_down = 0
        for n1 in range(mini_batch):
            hang1 = np.uint8(row[n1][0])
            lie1 = np.uint8(column[n1][0])
            b_up += A[hang1][f3] * X[hang1][lie1]
            b_down += A[hang1][f3] * X_head1[hang1][lie1] + lamda_Q * B[f3][lie1]
        descentQ_up[f3][0] = b_up
        descentQ_down[f3][0] = b_down
    for f5 in range(F):
        for t in range(len(X[0])):
            Q[f5][t] = B[f5][t] * (descentQ_up[f5][0] / descentQ_down[f5][0])
    return Q
def fill(nX,X,XX1,Xm):
    X_orginal = XX1
    X_estimate = nX
    X2 = np.zeros((len(X_orginal), len(X_orginal[0])))
    # X3 = np.zeros((len(X_orginal), len(X_orginal[0])))
    for i in range(len(X_orginal)):
        for j in range(len(X_orginal[0])):
            if Xm[i][j] == 11:
                X2[i][j] = X_estimate[i][j]
            else:
                X2[i][j] = X_orginal[i][j]
    return X2

def MSE(X_fill,X,XX,mun,Xmm):
    Y = X_fill
    Q = X
    target = XX
    prediction = Y
    error = np.zeros((len(target), len(target[0])))
    for i in range(len(Q)):
        for j in range(len(Q[0])):
            if Xmm[i][j] == 11:
                if Q[i][j] == 0:
                    error[i][j] = target[i][j] - prediction[i][j]
                else:
                    error[i][j] = 0
    squaredError = np.zeros((len(target), len(target[0])))
    # absError = np.zeros((len(target), len(target[0])))
    for i in range(len(target)):
        for j in range(len(target[0])):
            squaredError[i][j] = error[i][j] * error[i][j]
            # absError[i][j] = abs(error[i][j])
    s = np.sum(squaredError)
    MSE = s / mun
    # print("MSE = ", s / 35379)  # 均方误差MSE
    return MSE
if __name__ =='__main__':
    KF = KFold(n_splits=2)
    F = 5
    steps = 1200
    mini_batch = 2000
    list_mse = []
    list_lam =[]
    list_mseoldtrain = []
    list_lamda = []
    list_mse_train_value = []
    list_mse_test = []
    list_lam_test = []
    list_mseoldtest = []
    list_mse_test_value = []
    X = loadDataSet('A(10%_33416).xlsx')
    X_m = loadDataSet('missing_index_A(10).xlsx')
    # X1 = loadDataSet('stand_A.xlsx')
    XX1 = loadDataSet('standA.xlsx')
    # list_lamda = [0.00001, 0.0001, 0.001, 0.01, 0.1]
    # list_lamda = [0.01, 0.02, 0.03, 0.04, 0.05, 0.06, 0.07, 0.08, 0.09, 0.1]
    for a in np.arange(0,1,0.02):
        lamda_P = a
        lamda_Q = a
        print('lamda', lamda_Q)
        for train_index, test_index in KF.split(X,XX1,X_m):
            print('TRAIN:', train_index, 'TEST:', test_index)
            X_train, X_test = X[train_index], X[test_index]
            X1_train, X1_test = XX1[train_index], XX1[test_index]
            X_m_train, X_m_test = X_m[train_index], X_m[test_index]
            # train
            number_train_0 = 0
            for i in range(len(X_train)):
                for j in range(len(X_train[0])):
                    if X_m_train[i][j] == 11:
                        number_train_0 += 1
            print('11',number_train_0)
            list_mseoldtrain = []
            P1 = np.random.rand(len(X_train), F)
            Q1 = np.random.rand(F, len(X_train[0]))
            A = P1.copy()
            B = Q1.copy()
            P = np.zeros((len(X_train), F))
            Q = np.zeros((F, len(X_train[0])))
            for step in range(steps):
                # print('train_step', step)
                number = 0
                column = np.zeros((mini_batch, 1))
                row = np.zeros((mini_batch, 1))
                while number < mini_batch:
                    u = np.random.randint(0, high=len(X_train))
                    i = np.random.randint(0, high=len(X_train[0]))
                    if X_train[u][i] != 0:
                        row[number][0] = u
                        column[number][0] = i
                        number += 1
                nP = MGD_P(X_train, A, B, row, column, mini_batch, lamda_P)
                nQ = MGD_Q(X_train, A, B, row, column, mini_batch, lamda_Q)
                nX = np.dot(nP, nQ)
                X_fill = fill(nX, X_train,X1_train,X_m_train)
                mse_train = MSE(X_fill, X_train, X1_train,number_train_0,X_m_train)
                list_mseoldtrain.append(mse_train)
                print('train_mse', step,mse_train)
                A = nP.copy()
                B = nQ.copy()
            mse_min_train = min(list_mseoldtrain)
            list_mse.append(mse_min_train)  # 存储在某一lambda下每次交叉验证每一折的最小的mse值
            # test
            number_test_0 = 0
            for i in range(len(X_test)):
                for j in range(len(X_test[0])):
                    if X_m_test[i][j] == 11:
                        number_test_0 += 1
            list_mseoldtest = []
            P1_test = np.random.rand(len(X_test), F)
            Q1_test = np.random.rand(F, len(X_test[0]))
            A_test = P1_test.copy()
            B_test = Q1_test.copy()
            P = np.zeros((len(X_test), F))
            Q = np.zeros((F, len(X_test[0])))
            nE_old_test = np.linalg.norm(np.dot(P1_test, Q1_test))
            for step_test in range(steps):
                # print('step_test', step_test)
                number_test = 0
                column_test = np.zeros((mini_batch, 1))
                row_test = np.zeros((mini_batch, 1))
                while number_test < mini_batch:
                    u_test = np.random.randint(0, high=len(X_test))
                    i_test = np.random.randint(0, high=len(X_test[0]))
                    if X_test[u_test][i_test] != 0:
                        row_test[number_test][0] = u_test
                        column_test[number_test][0] = i_test
                        number_test += 1
                nP_test = MGD_P(X_test, A_test, B_test, row_test, column_test, mini_batch, lamda_P)
                nQ_test = MGD_Q(X_test, A_test, B_test, row_test, column_test, mini_batch, lamda_Q)
                nX_test = np.dot(nP_test, nQ_test)
                X_fill_test = fill(nX_test, X_test,X1_test,X_m_test)
                mse_test = MSE(X_fill_test, X_test, X1_test,number_test_0,X_m_test)
                list_mseoldtest.append(mse_test)
                print('test_mse', step_test,mse_test)
                A_test = nP_test.copy()
                B_test = nQ_test.copy()
            mse_min_test = min(list_mseoldtest)
            list_mse_test.append(mse_min_test)
    list_mse = np.matrix(list_mse)
    list_mse = np.array(list_mse).T
    print(list_mse)# 所有lambda值的训练集每一折最小的lambda值的矩阵
    list_mse_test = np.matrix(list_mse_test)
    list_mse_test = np.array(list_mse_test).T
    list_all = np.hstack((list_mse,list_mse_test))
    list_abs = np.zeros((len(list_all), 1))
    for i in range(len(list_all)):
        for j in range(len(list_all[0])):
            list_abs[i][0] = abs(list_all[i][0] - list_all[i][1])
    list_all = np.hstack((list_all, list_abs))
    book = Workbook()
    sheet = book.create_sheet('Sheet', index=0)
    [w, r] = list_all.shape
    for i in range(w):
        for j in range(r):
            sheet.cell(row=i + 1, column=j + 1, value=list_all[i][j])
    book.save('lambda-A_10.xlsx')
    # book = Workbook()
    # sheet = book.create_sheet('Sheet', index=0)
    # [w, r] = list_mse_test.shape
    # for i in range(w):
    #     for j in range(r):
    #         sheet.cell(row=i + 1, column=j + 1, value=list_mse_test[i][j])
    # book.save('A(10)_test(lambdalist).xlsx')
    end = time.clock()
    t = end - start
    print('Run time is :', t)




#         mse_tra = np.matrix(list_mse)
#         mse_tra = np.array(mse_tra).T
#         mse_tra = list(mse_tra)
#         lam = mse_tra.index(min(mse_tra))
#         mse_tra = np.array(mse_tra)
#         mse_train_value = min(mse_tra)
#         list_mse_train_value.append(mse_train_value)
#
#
#     mse_tes = np.matrix(list_mse_test)
#     mse_tes = np.array(mse_tes).T
#     mse_test_value = min(mse_tes)
#     list_mse_test_value.append(mse_test_value)
# lis_mse_tra = np.matrix(list_mse_train_value)
# mse_tr = np.array(lis_mse_tra).T
# lis_mse_tes = np.matrix(list_mse_test_value)
# mse_te = np.array(lis_mse_tes).T








