# %%
import sklearn
sklearn.__version__

# %%
import time
from collections import defaultdict
import numpy as np
from scipy import sparse
from sklearn.neighbors import KNeighborsRegressor

n_samples, n_features = 1_000, [500, 1_000, 2_500, 5_000, 7_500, 10_000, 25_000, 50_000]

results = defaultdict(list)
for nf in n_features:
    X = sparse.random(n_samples, nf, density=0.01, format="csr", random_state=0)
    rng = np.random.default_rng(0)
    y = rng.integers(0, 2, size=n_samples)
    knn = KNeighborsRegressor(n_neighbors=5).fit(X, y)
    start = time.time()
    knn.kneighbors(X, return_distance=False)
    elapsed_time = time.time() - start
    results["version"].append(sklearn.__version__)
    results["n_features"].append(nf)
    results["elapsed_time"].append(elapsed_time)

# %%
import pandas as pd

results = pd.DataFrame(results)
results.to_csv(f"bench_{sklearn.__version__}.csv", index=False)

# %%
import pandas as pd

results_main = pd.read_csv(f"bench_1.5.dev0.csv")
results_1_1 = pd.read_csv(f"bench_1.1.3.csv")
results = pd.concat([results_main, results_1_1], axis=0)
results

# %%
import seaborn as sns

sns.set_context("talk")

# %%
sns.relplot(
    data=results,
    x="n_features",
    y="elapsed_time",
    hue="version",
    kind="line",
    height=6,
    aspect=1.5,
    legend=True,
)
sns.despine(offset=10, trim=True)