import argparse
from time import perf_counter
import json
from statistics import mean, stdev

from sklearn.tree import DecisionTreeClassifier
from sklearn.datasets import make_classification
from collections import defaultdict


N_SAMPLES = [1_000, 5_000, 10_000, 20_000, 50_000]
N_REPEATS = 20

parser = argparse.ArgumentParser()
parser.add_argument("results", type=argparse.FileType("w"))
args = parser.parse_args()

results = defaultdict(list)

for n_samples in N_SAMPLES:
    for n_repeat in range(N_REPEATS):
        X, y = make_classification(
            random_state=n_repeat, n_samples=n_samples, n_features=100
        )
        tree = DecisionTreeClassifier(random_state=n_repeat)
        start = perf_counter()
        tree.fit(X, y)
        duration = perf_counter() - start
        results[n_samples].append(duration)
    results_mean, results_stdev = mean(results[n_samples]), stdev(results[n_samples])
    print(f"n_samples={n_samples} with {results_mean:.3f} +/- {results_stdev:.3f}")

json.dump(results, args.results)