from sklearn.linear_model import Ridge
def stacking_model_input(models, model_name, data):
    X_test, y_test, X_train, y_train = data
    X_data = [X_test, X_train]
    stack_Xs = [pd.DataFrame(), pd.DataFrame()]
    for i, model in enumerate(models):
        model.fit(X_train, y_train)
        for j, x in enumerate(X_data):
            preds = model.predict(x)
            preds = preds.reshape(len(preds))
            stack_Xs[j][model_name[i]] = preds
    return stack_Xs[0], stack_Xs[1]


def ensemble_stacking_R2(models, model_name, Data, alpha, k_fold = 5):
    R_squares = []
    kf = KFold(n_splits = k_fold)

    for train_index, test_index in kf.split(X):
        split_data = split_test_train(Data, (train_index, test_index))
        y_test, y_train = split_data[1], split_data[3]

        stack_X_test, stack_X_train = stacking_model_input(models, model_name, split_data)
#         print(stack_X_test, y_test)
        reg = Ridge(alpha = alpha).fit(stack_X_train, y_train)
        y_pred = reg.predict(stack_X_test)
        R_squares.append(r2_score(y_test, y_pred))
    return np.array([R_squares]).squeeze(axis=0)


Data = (X, Y)
models = [SVR_model, KRR_model, RFG_model, GBR_model]
# models = [SVR_model, KRR_model]
model_name = ['Support Vector Regression', 'Kernel Ridge Regression'
              , 'Random Forest Regression', 'Gradient Boosting Regression']

alphas = [10, 20, 30, 50]
for alpha in alphas:
    stacking_scores = ensemble_stacking_R2(models, model_name, Data, alpha)
    print(alpha)
    print(stacking_scores)
    print('the R2 of stacking method is %.2f with std %.2f' %(stacking_scores.mean(), stacking_scores.std()))