# Path: jax_verify/extensions/functional_lagrangian/inner_solvers/pga/utils.py
# def linf_project_fn(epsilon: float, bounds: Tuple[float, float]) -> ProjectFn:
#   def project_fn(x, origin_x):
# def bounded_initialize_fn(
#     bounds: Optional[Tuple[chex.Array, chex.Array]] = None,) -> InitializeFn:
#     def _initialize_fn(rng, x):
# def noop_initialize_fn() -> InitializeFn:
#   def _initialize_fn(rng, x):


from typing import Callable, List, Tuple
from jax_verify.extensions.functional_lagrangian.inner_solvers.pga import utils
import chex
import jax
import jax.numpy as jnp


def _schedule(values: List[float],
              boundaries: List[int],
              dtype=jnp.float32) -> Callable[[chex.Array], chex.Numeric]:
  """Schedule the value of p, the proportion of elements to be modified."""
  large_step = max(boundaries) + 1
  boundaries = boundaries + [large_step, large_step + 1]
  num_values = len(values)
  values = jnp.array(values, dtype=jnp.float32)
  large_step = jnp.array([large_step] * len(boundaries), dtype=jnp.int32)
  boundaries = jnp.array(boundaries, dtype=jnp.int32)

  def _get(step):
    """Returns the value according to the current step and schedule."""
    b = boundaries - jnp.minimum(step + 1, large_step + 1)
    b = jnp.where(b < 0, large_step, b)
    idx = jnp.minimum(jnp.argmin(b), num_values - 1)
    return values[idx].astype(dtype)

  return _get


class Square:
  """Performs a blackbox optimization as in https://arxiv.org/pdf/1912.00049."""

  def __init__(
      self,
      num_steps: int,
      epsilon: chex.Numeric,

# Path: jax_verify/extensions/functional_lagrangian/inner_solvers/pga/utils.py
# def linf_project_fn(epsilon: float, bounds: Tuple[float, float]) -> ProjectFn:
#   def project_fn(x, origin_x):
# def bounded_initialize_fn(
#     bounds: Optional[Tuple[chex.Array, chex.Array]] = None,) -> InitializeFn:
#     def _initialize_fn(rng, x):
# def noop_initialize_fn() -> InitializeFn:
#   def _initialize_fn(rng, x):


# Path: jax_verify/extensions/functional_lagrangian/inner_solvers/pga/square.py
from typing import Callable, List, Tuple
from jax_verify.extensions.functional_lagrangian.inner_solvers.pga import utils
import chex
import jax
import jax.numpy as jnp



def _schedule(values: List[float],
              boundaries: List[int],
              dtype=jnp.float32) -> Callable[[chex.Array], chex.Numeric]:
  """Schedule the value of p, the proportion of elements to be modified."""
  large_step = max(boundaries) + 1
  boundaries = boundaries + [large_step, large_step + 1]
  num_values = len(values)
  values = jnp.array(values, dtype=jnp.float32)
  large_step = jnp.array([large_step] * len(boundaries), dtype=jnp.int32)
  boundaries = jnp.array(boundaries, dtype=jnp.int32)

  def _get(step):
    """Returns the value according to the current step and schedule."""
    b = boundaries - jnp.minimum(step + 1, large_step + 1)
    b = jnp.where(b < 0, large_step, b)
    idx = jnp.minimum(jnp.argmin(b), num_values - 1)
    return values[idx].astype(dtype)

  return _get


class Square:
  """Performs a blackbox optimization as in https://arxiv.org/pdf/1912.00049."""

  def __init__(
      self,
      num_steps: int,
      epsilon: chex.Numeric,

      initialize_fn: utils.InitializeFn,