from sklearn.metrics import average_precision_score

from onad.model.iforest.mondrian import MondrianForest
from onad.stream.dataset import Dataset, load

model = MondrianForest(n_estimators=250, subspace_size=500, seed=1)

labels, scores = [], []
dataset = load(Dataset.SHUTTLE)

for i, (x, y) in enumerate(dataset.stream()):
    if i < 10_000:
        if y == 0:
            model.learn_one(x)
        continue
    model.learn_one(x)
    score = model.score_one(x)

    labels.append(y)
    scores.append(score)

print(f"PR_AUC: {round(average_precision_score(labels, scores), 3)}")  # 0.329
