import time

import copt as cp
from sklearn.datasets import make_regression
from sklearn.linear_model import Ridge

X, y, c = make_regression(
    random_state=0, n_samples=600, n_features=5, coef=True)

print(X.shape, y.shape)

f = cp.SquaredLoss(X, y, alpha=1. / X.shape[0])

copt_start = time.time()
copt_ridge = cp.minimize_SAGA(f, trace=False, tol=1e-6)
copt_stop = time.time()

# note sklearn Ridge also performs some checks etc besides fitting the coefs
sklearn_ridge = Ridge(alpha=1., solver='saga', fit_intercept=False, tol=1e-6)
sklearn_start = time.time()
sklearn_ridge.fit(X, y)
sklearn_stop = time.time()

print('coef:', c)
print('copt coef:', copt_ridge.x)
print('sklearn coef:', sklearn_ridge.coef_)
print('copt time:', copt_stop - copt_start)
print('sklearn time:', sklearn_stop - sklearn_start)