# %%
import h2o

# %%
h2o.__version__

# %%
h2o.init()

# %%
from h2o.estimators import H2OGradientBoostingEstimator
h2o.init(jvm_custom_args=["sys.ai.h2o.debug.allowJavaVersions", "18"])

# Import the prostate dataset into H2O:
prostate = h2o.import_file("http://s3.amazonaws.com/h2o-public-test-data/smalldata/prostate/prostate.csv")

# Set the predictors and response; set the factors:
prostate["CAPSULE"] = prostate["CAPSULE"].asfactor()
predictors = ["ID","AGE","RACE","DPROS","DCAPS","PSA","VOL","GLEASON"]
response = "CAPSULE"

# Build and train the model:
pros_gbm = H2OGradientBoostingEstimator(nfolds=5,
                                        seed=1111,
                                        keep_cross_validation_predictions = True)
pros_gbm.train(x=predictors, y=response, training_frame=prostate)

# Eval performance:
perf = pros_gbm.model_performance()

# Generate predictions on a test set (if necessary):
pred = pros_gbm.predict(prostate)

# Extract feature interactions:
feature_interactions = pros_gbm.feature_interaction()

# %%
feature_interactions

# %%
#save model
h2o.save_model(model=pros_gbm, force=True)

# %%
pros_gbm.save_mojo('mojo')