import sklearn
from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor, GradientBoostingClassifier, GradientBoostingRegressor
import numpy as np
import os
import matplotlib.pyplot as plt

# read data
bin_file = './feature000.npz'
data = np.load(bin_file)
day = data['day']
grade = data['grade']
print "shape of day: {}; shape of grade: {}".format(day.shape, grade.shape)


# week: sID x weeks x features
week = np.concatenate([np.expand_dims(np.sum(day[:, i:(i+7), :], axis=1), axis=1) for i in range(0, day.shape[1], 7)], axis = 1)
print "shape of whole data: {}".format(week.shape)


""" 
Regression based on the random forest and gradient boosting
"""
regression_models = {"random_forest":     RandomForestRegressor(n_estimators=100,n_jobs = 4),
                     "gradient_boosting": GradientBoostingRegressor(learning_rate=0.1, max_depth=5)}

def model_evaluator(model, w=10):
    cur_week = week[:,:w,:]
    data = cur_week.reshape(week.shape[0],-1) # expected array with dim 2
    model.fit(data, grade)
    grade_pred = model.predict(data)
    plt.plot(grade, grade_pred,'.')
    plt.show()


def key_features(feature_importance, w):
    importance = feature_importance.reshape([w, feature_importance.shape[0]/w])
    for i in range(importance.shape[0]):
        for j in range(importance.shape[1]):
            print "%.4f"%importance[i,j],
        print

for key, model in regression_models.items():
    print "train the model: " + key
    model_evaluator(model, w=15)
    print

for key, model in regression_models.items():
    print "the importance of features in "+key+" :"
    key_features(model.feature_importances_, w=15)
    print

"""
Regression with the default cross validation in sklearn
"""
from sklearn.model_selection import cross_val_score, KFold
# sklearn.cross_validation will be removed in 0.20.
# Two sets of [cross_val_score, KFold] are slightly different.
cur_week = week[:,:10,:]
data = cur_week.reshape(week.shape[0],-1)

k_folder = KFold(n_splits =5, shuffle=True, random_state=1)
k_folder.get_n_splits(data)

rf_reg = RandomForestRegressor(n_estimators=200,n_jobs = 4)
rf_cross_score = cross_val_score(rf_reg, data, grade, scoring = 'r2', cv=k_folder.split(data), n_jobs=1)
print "cross r2 score of random forest: {:6.4f}".format(np.mean(rf_cross_score))

gb_reg = GradientBoostingRegressor(learning_rate=0.1, max_depth=5)
gb_cross_score = cross_val_score(gb_reg, data, grade, scoring = 'r2', cv=k_folder.split(data), n_jobs=2)
print "cross r2 score of gradient boosting: {:6.4f}".format(np.mean(gb_cross_score))

"""
Customized cross validation and learning monitoring
"""

def evaluate(y, y_pred, score='r2'):
    if score=="r2":
        mu = np.mean(y)
        return 1-np.sum((y-y_pred)**2)/np.sum((y-mu)**2)
    

def train(model, x_inp, y_inp):
    score = {"train": list(), "test": list()}
    for train_index, test_index in k_folder.split(x_inp):
        x_train, x_test = x_inp[train_index], x_inp[test_index]
        y_train, y_test = y_inp[train_index], y_inp[test_index]
        
        model.fit(x_train, y_train)
        
        y_train_pred = model.predict(x_train)
        train_score = evaluate(y_train, y_train_pred)
        
        y_test_pred  = model.predict(x_test)
        test_score  = evaluate(y_test, y_test_pred)
        print "train score: {:6.4f}; test score: {:6.4f}".format(train_score, test_score)
        
        score["train"] += [train_score]
        score["test"]  += [test_score]
    return score

def list_mean(x):
    return sum(x)/float(len(x))

rf_score = train(rf_reg, x_inp=data, y_inp=grade)
print "train cross score: {:6.4f}; test score: {:6.4f}".format(list_mean(rf_score["train"]), list_mean(rf_score["test"]))

gb_score = train(gb_reg, x_inp=data, y_inp=grade)
print "train cross score: {:6.4f}; test score: {:6.4f}".format(list_mean(gb_score["train"]), list_mean(gb_score["test"]))
