from uplift_kit.trees import UpliftRandomForestModel
import pandas as pd

model = UpliftRandomForestModel(
    n_estimators=10,      # number of uplift trees
    max_features=10,      # maximum number of features considered in one split
    max_depth=10,         # maximum depth of one single tree
    min_sample_leaf=100,  # minumum number of samples classified to a leaf
    eval_func="ED",       # split evaluation function, support `ED, KL, CHI`
    max_bins=10,          # maximum bins considered when calculating best split
    balance=False,        # whether to use weighted average to calculate score, False mean not
    regularization=True,  # whether to add regularization term
    alpha=0.9,            # param for the regularization term
)

data = pd.read_parquet("../train.parquet")
x_names = list(data.columns[:-2])

# model will use columns of `x_names` as features
# treatment_col should contains 0,1,2,...k, where 0 indicates control sample, 1~k means treatment 1~k.
# outcome_col should only contains 0,1 as integer values, i.e. binary outcome.
model.fit(
    data,
    x_names=x_names,
    treatment_col="treats",
    outcome_col="outcome",
    n_threads=8,
)

# In prediction, model will automatically choose the feature columns (x_names) from input dataframe. 
# It returns a numpy array, where k columns per sample indicate the uplift value for treatment k.
test = pd.read_parquet("../test.parquet")
res = model.predict(data=test[x_names], n_threads=8)
print(res[:10])

model.fit(...)
model.save("saved_model.json")

new_model = UpliftRandomForestModel()
new_model.load("saved_model.json")
new_model.predict(...)

res = model.predict_row([1,2,"ASIA",...]) # input a list of features, consistent with `x_names`