# Copyright 2021 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Benchmark JAX implementation vs. NumPy implementation of proximal gradient."""

import time

from typing import NamedTuple
from typing import Sequence

from absl import app
from absl import flags

from sklearn import datasets
from sklearn import preprocessing

import numpy as onp

import jax
import jax.numpy as jnp

from jaxopt import base
from jaxopt import proximal_gradient2 as pg
from jaxopt.prox import prox_lasso


FLAGS = flags.FLAGS

flags.DEFINE_string("dataset", default="boston", help=("Dataset to use."))
flags.DEFINE_bool("float64", default=False, help=("Enable double precision."))
flags.DEFINE_float("lam", default=1.0, help=("Regularization value."))
flags.DEFINE_integer("maxiter", default=200, help=("Max # of iterations."))
flags.DEFINE_integer("n_samples", default=100000, help=("Number of samples."))
flags.DEFINE_integer("n_features", default=200, help=("Number of features."))
flags.DEFINE_bool("verbose", default=False, help=("Enable verbose output."))


def _make_linesearch(fun, prox, maxls):
  def linesearch(curr_x, curr_x_fun_val, curr_x_fun_grad, curr_stepsize):
    """A pure NumPy re-implementation of linesearch for benchmarking reasons."""
    for _ in range(maxls):
      next_x = prox(curr_x - curr_stepsize * curr_x_fun_grad, curr_stepsize)
      diff = next_x - curr_x
      sqdist = onp.vdot(diff, diff)
      value_F = fun(next_x)
      value_Q = (curr_x_fun_val + onp.vdot(diff, curr_x_fun_grad) +
                 0.5 * sqdist / curr_stepsize)
      if value_F <= value_Q:
        return next_x, curr_stepsize
      curr_stepsize *= 0.5

    # Undo the last decrase when `maxls` is reached.
    curr_stepsize *= 2

    return next_x, curr_stepsize
  return linesearch


def proximal_gradient_onp(fun, init, prox, stepsize, maxiter=200, maxls=15,