import h2o
from h2o.estimators import H2ORuleFitEstimator, H2ORandomForestEstimator

# init h2o cluster
h2o.init()

df = h2o.import_file("https://s3.amazonaws.com/h2o-public-test-data/smalldata/gbm_test/titanic.csv",
                       col_types={'pclass': "enum", 'survived': "enum"})
x =  ["age", "sibsp", "parch", "sex", "pclass"]

# Split the dataset into train and test
train, test = df.split_frame(ratios=[.8], seed=1234)

"""
Explanation: Rulefit demo - Titanic Dataset
H2O Rulefit algorithm
Rulefit algorithm combines tree ensembles and linear models to take advantage of both methods: a tree ensemble accuracy and a linear model interpretability. The general algorithm fits a tree ensebmle to the data, builds a rule ensemble by traversing each tree, evaluates the rules on the data to build a rule feature set and fits a sparse linear model (LASSO) to the rule feature set joined with the original feature set.
For more information, refer to: http://statweb.stanford.edu/~jhf/ftp/RuleFit.pdf by Jerome H. Friedman and Bogden E. Popescu.
Demo example
We will train a rulefit model to predict the rules defining whether or not someone will survive:
End of explanation
"""


rfit = H2ORuleFitEstimator(algorithm="drf", 
                               min_rule_length=1, 
                               max_rule_length=10, 
                               max_num_rules=100, 
                               model_type="rules_and_linear",
                               rule_generation_ntrees=50,
                               seed=1234)
rfit.train(training_frame=train, x=x, y="survived")

"""
Explanation: Using the algorithm parameter, a user can set whether algorithm will use DRF or GBM to fit a tree enseble. 
Using the min_rule_length and max_rule_length parameters, a user can set interval of tree enseble depths to be fitted. The bigger this interval is, the more tree ensembles will be fitted (1 per each depth) and the bigger the rule feature set will be.
Using the max_num_rules parameter, the maximum number of rules to return can be set.
Using the model_type parameter, the type of base learners in the enseble can be set.
Using the rule_generation_ntrees parameter, the number of trees for tree enseble can be set.
End of explanation
"""


from IPython.display import display
display(rfit.rule_importance())

"""
Explanation: The output for the Rulefit model includes:
    - model parameters
    - rule importences in tabular form
    - training and validation metrics of the underlying linear model
End of explanation
"""


