from time import time

import numpy as np
import torch as xp
from tqdm import tqdm

from sklearn._config import config_context
from sklearn.preprocessing._polynomial import PolynomialFeatures

X_np = np.random.rand(100000, 100)
X_xp_cuda = xp.asarray(X_np, device="cuda")

# Numpy benchmarks
fit_times = []
transform_times = []
n_iter = 10
for _ in tqdm(range(n_iter), desc="Numpy Flow"):
    start = time()
    pf_np = PolynomialFeatures(degree=2)
    pf_np.fit(X_np)
    fit_times.append(time() - start)

    start = time()
    pf_np.transform(X_np)
    transform_times.append(time() - start)

avg_fit_time_numpy = sum(fit_times) / n_iter
avg_transform_time_numpy = sum(transform_times) / n_iter
print(f"Avg fit time for numpy: {avg_fit_time_numpy:.3f}")
print(f"Avg transform time for numpy: {avg_transform_time_numpy:.3f}")


# Torch cuda benchmarks
fit_times = []
transform_times = []
for _ in tqdm(range(n_iter), desc="Torch cuda Flow"):
    with config_context(array_api_dispatch=True):
        start = time()
        pf_xp = PolynomialFeatures(degree=2)
        pf_xp.fit(X_xp_cuda)
        fit_times.append(time() - start)

        start = time()
        float(pf_xp.transform(X_xp_cuda)[0, 0])
        transform_times.append(time() - start)

avg_fit_time_cuda = sum(fit_times) / n_iter
avg_transform_time_cuda = sum(transform_times) / n_iter
print(
    f"Avg fit time for torch cuda: {avg_fit_time_cuda:.3f}, "
    f"speed-up: {avg_fit_time_numpy / avg_fit_time_cuda:.1f}x"
)
print(
    f"Avg transform time for torch cuda: {avg_transform_time_cuda:.3f} "
    f"speed-up: {avg_transform_time_numpy / avg_transform_time_cuda:.1f}x"
)