# %%
import scipy.sparse as sp
import numpy as np
from sklearn.linear_model import ElasticNet
from time import perf_counter


rng = np.random.default_rng(42)
X = rng.uniform(size=(1_000, 100_000))
X[X < 0.9] = 0.0  # Sparsify the matrix
X_sparse = sp.csc_array(X)

w_true = np.zeros(X.shape[1])
w_true[0:300] = 10.0

assert X.mean(axis=0)[0:5].min() > 0.01

y = X_sparse @ w_true + 1 + rng.normal(size=X.shape[0]) * 0.1
# %%

reg = ElasticNet(
    alpha=0.1, l1_ratio=0.5, fit_intercept=True, selection="random", random_state=42
)

coef_dense = reg.fit(X, y).coef_
intercept_dense = reg.intercept_
print("10 first dense coefficients:", coef_dense[:10])
print("dense intercept:", intercept_dense)
# %%
for i in range(5):
    tic = perf_counter()
    coef_sparse = reg.fit(X_sparse, y).coef_
    toc = perf_counter()
    print(f"Time for sparse fit: {toc - tic:.3f} seconds")

intercept_sparse = reg.intercept_
print("10 first sparse coefficients:", coef_sparse[:10])
print("sparse intercept:", intercept_sparse)
assert np.allclose(coef_dense, coef_sparse)
assert np.allclose(intercept_dense, intercept_sparse)