import numpy as np
import pandas as pd

import time
import pathlib
from treefarms import TREEFARMS

def run_tree_farms(df, rashomon_parameter, depth):
    X, y = df.iloc[:, :-1], df.iloc[:, -1]
    h = df.columns[:-1]
    # train TREEFARMS model
    config = {
        "regularization": 0.01,  # regularization penalizes the tree with more leaves. We recommend to set it to relative high value to find a sparse tree.
        "rashomon_bound_multiplier": rashomon_parameter,  # rashomon bound multiplier indicates how large of a Rashomon set would you like to get
        "depth_budget": depth,
        "verbose": False
    }

    model = TREEFARMS(config)
    model.fit(X, y)
    return model.get_tree_count()

files = ['../datasets/monks3.csv']

result = {}

for file_id, file in enumerate(files):

    df = pd.read_csv(file)

    for rset_param in [0.03, 0.05, 0.07, 0.09, 0.11, 0.13, 0.15]:
        res = []
        for i in range(1,8):
            num_model = run_tree_farms(df, rset_param, i)
            res += [num_model]
        result[rset_param] = res
    
    
import pickle

# open a file, where you ant to store the data
file = open('tree_farms_rset_param', 'wb')

# dump information to that file
pickle.dump(result, file)

# close the file
file.close()