"""
Baseline random forest script.
Leon Zheng.

ToDo: choose the right feature and fine tune.
"""


import preprocessing
from sksurv.ensemble import RandomSurvivalForest
from sklearn.model_selection import cross_validate

"""
Reading data
"""

def load_data(features=None):
    X_df, y_df, _ = preprocessing.load_owkin_data()
    if features != None:
        X_df = X_df[features]
    X = X_df.to_numpy()
    y = preprocessing.y_dataframe_to_rsf_input(y_df)
    return X_df, y_df, X, y

X_df, y_df, X, y = load_data()
feature_name = list(X_df.columns.values)

"""
Train model
"""
params = {'min_samples_leaf': 1, 'min_samples_split': 10, 'n_estimators': 10}
rsf = RandomSurvivalForest(n_estimators = params['n_estimators'],
                           min_samples_split = params['min_samples_split'],
                           min_samples_leaf = params['min_samples_leaf'],
                           max_features="sqrt",
                           n_jobs=-1
                           )
print(cross_validate(rsf, X, y, cv=5))