import numpy as np
from sklearn.datasets import load_iris
from sklearn.linear_model import Lasso as Lasso_sklearn
from skglm import Lasso as Lasso_skglm
from benchopt import dataset, solver, objective, run_benchmark


@dataset(X=data.data, y=data.target)
def get_data():
    data = load_iris()
    return data.data, data.target


@solver(X=x, y=y, lmdb=lmdb, coef=lasso.coef_)
def solve_lasso_sklearn(X, y, lmbd):
    lasso = Lasso_sklearn(alpha=lmbd)
    lasso.fit(X, y)
    return lasso.coef_


@solver(X=x, y=y, lmdb=lmdb, coef=lasso.coef_)
def solve_lasso_skglm(X, y, lmbd):
    lasso = Lasso_skglm(alpha=lmbd)
    lasso.fit(X, y)
    return lasso.coef_


@objective(X=x, y=y, lmdb=lmdb, coef=coef)
def objective_lasso(X, y, lmbd, coef):
    return 0.5 * np.linalg.norm(y - X @ coef) ** 2 + lmbd * np.linalg.norm(coef, ord=1)


if __name__ == "__main__":
    run_benchmark(
        datasets=get_data,
        solvers=[solve_lasso_sklearn, solve_lasso_skglm],
        objectives=objective_lasso,
    )