# Path: jax_verify/extensions/sdp_verify/utils.py
# class SdpDualVerifInstance(_SdpDualVerifInstance):
# class VerifInstanceTypes(enum.Enum):
# class DualVarTypes(enum.Enum):
# def mlp_layer_sizes(params):
# def nn_layer_sizes(params):
# def layer_sizes_from_bounds(bounds):
# def predict_mlp(params, inputs):
# def fwd(inputs, layer_params):
# def predict_cnn(params, inputs, include_preactivations=False):
# def get_network_activs(params, x):
# def get_layer_params(fun_to_extract, example_input):
#   def _get_const_input_arg(eqn):
# def boundprop(params, bounds_in):
# def init_bound(x, epsilon, input_bounds=(0., 1.), add_batch_dim=True):
# def ibp_bound_elided(verif_instance):
# def ibp_bound_nonelided(verif_instance):
# def make_relu_robust_verif_instance(
#     params, bounds=None, target_label=1, label=2, input_bounds=None):
# def make_relu_robust_verif_instance_elided(
#     params, bounds=None, input_bounds=None):
# def output_size_and_verif_type(params, bounds):
# def elide_params(params, label, target_label, op_size):
# def preprocess_cifar(image, inception_preprocess=False, perturbation=False):
# def preprocessed_cifar_eps_and_input_bounds(
#     shape=(32, 32, 3), epsilon=2/255, inception_preprocess=False):
# def adv_objective(model_fn, x, label, target_label):
# def fgsm_single(model_fn, x, label, target_label, epsilon, num_steps,
#                 step_size, input_bounds=(0., 1.)):
# def untargeted_margin_loss(logits, labels):
# def pgd_default(model_fn, x, label, epsilon, num_steps, step_size,
#                 input_bounds=(0., 1.)):
# def pgd(adv_loss, x_init, epsilon, num_steps, step_size, input_bounds=(0., 1.)):
# def scale_by_variable_opt(multipliers):
#   def init_fn(params):
#   def update_fn(updates, _, params=None):
# def flatten(pytree, backend=np):
# def unflatten_like(a, pytree):
# def structure_like(tree1, tree2):
#   MLP_ELIDED = 'mlp_elided'
#   CNN_ELIDED = 'cnn_elided'
#   EQUALITY = 'equality'
#   INEQUALITY = 'inequality'


import collections
import jax.numpy as jnp
import numpy as np
from jax_verify.extensions.sdp_verify import utils
#   nu_quad: IBP quadratic matrix constraint: (x_i - l_i)(x_j - u_j) <= 0
#   muminus: x'>=0
#   muminus2: Triangle linear Relu relaxation - u(Wx+b) - ul - (u-l)x' >= 0
#       where l = min(l, 0), u = max(u, 0)
#   muplus: x'>=Wx+b
DualVar = collections.namedtuple(
    'DualVar', ['lam', 'nu', 'nu_quad', 'muminus', 'muplus', 'muminus2'])
DualVarFin = collections.namedtuple('DualVarFin', ['nu', 'nu_quad'])
DEFAULT_DISABLED_DUAL_VARS = ('nu_quad', 'muminus2')
NECESSARY_DUAL_VARS = ('lam', 'muplus', 'muminus')


def make_relu_network_lagrangian(dual_vars, params, bounds, obj):
  """Returns a function that computes the Lagrangian for a ReLU network.

  This function assumes `params` represent a feedforward ReLU network i.e.
  x_{i+1} = relu(W_i x_i + b_i). It defines the Lagrangian by applying the
  objective `obj` to the final layer activations, and encoding the Lagrangian
  terms for each of the constraints defining the ReLU network. It then returns
  this function.

  Args:
    dual_vars: A length L+1 list of dual variables at each layer
    params: A length L list of (W, b) pairs, elided network weights
    bounds: A length L+1 list of `IntBound`s, elided bounds at each layer
    obj: function, taking final layer activations as input

  Returns:
    Function that computes Lagrangian L(x) with fixed `dual_vars`.
  """

# Path: jax_verify/extensions/sdp_verify/utils.py
# class SdpDualVerifInstance(_SdpDualVerifInstance):
# class VerifInstanceTypes(enum.Enum):
# class DualVarTypes(enum.Enum):
# def mlp_layer_sizes(params):
# def nn_layer_sizes(params):
# def layer_sizes_from_bounds(bounds):
# def predict_mlp(params, inputs):
# def fwd(inputs, layer_params):
# def predict_cnn(params, inputs, include_preactivations=False):
# def get_network_activs(params, x):
# def get_layer_params(fun_to_extract, example_input):
#   def _get_const_input_arg(eqn):
# def boundprop(params, bounds_in):
# def init_bound(x, epsilon, input_bounds=(0., 1.), add_batch_dim=True):
# def ibp_bound_elided(verif_instance):
# def ibp_bound_nonelided(verif_instance):
# def make_relu_robust_verif_instance(
#     params, bounds=None, target_label=1, label=2, input_bounds=None):
# def make_relu_robust_verif_instance_elided(
#     params, bounds=None, input_bounds=None):
# def output_size_and_verif_type(params, bounds):
# def elide_params(params, label, target_label, op_size):
# def preprocess_cifar(image, inception_preprocess=False, perturbation=False):
# def preprocessed_cifar_eps_and_input_bounds(
#     shape=(32, 32, 3), epsilon=2/255, inception_preprocess=False):
# def adv_objective(model_fn, x, label, target_label):
# def fgsm_single(model_fn, x, label, target_label, epsilon, num_steps,
#                 step_size, input_bounds=(0., 1.)):
# def untargeted_margin_loss(logits, labels):
# def pgd_default(model_fn, x, label, epsilon, num_steps, step_size,
#                 input_bounds=(0., 1.)):
# def pgd(adv_loss, x_init, epsilon, num_steps, step_size, input_bounds=(0., 1.)):
# def scale_by_variable_opt(multipliers):
#   def init_fn(params):
#   def update_fn(updates, _, params=None):
# def flatten(pytree, backend=np):
# def unflatten_like(a, pytree):
# def structure_like(tree1, tree2):
#   MLP_ELIDED = 'mlp_elided'
#   CNN_ELIDED = 'cnn_elided'
#   EQUALITY = 'equality'
#   INEQUALITY = 'inequality'


# Path: jax_verify/extensions/sdp_verify/problem.py
import collections
import jax.numpy as jnp
import numpy as np
from jax_verify.extensions.sdp_verify import utils

#   nu_quad: IBP quadratic matrix constraint: (x_i - l_i)(x_j - u_j) <= 0
#   muminus: x'>=0
#   muminus2: Triangle linear Relu relaxation - u(Wx+b) - ul - (u-l)x' >= 0
#       where l = min(l, 0), u = max(u, 0)
#   muplus: x'>=Wx+b
DualVar = collections.namedtuple(
    'DualVar', ['lam', 'nu', 'nu_quad', 'muminus', 'muplus', 'muminus2'])
DualVarFin = collections.namedtuple('DualVarFin', ['nu', 'nu_quad'])
DEFAULT_DISABLED_DUAL_VARS = ('nu_quad', 'muminus2')
NECESSARY_DUAL_VARS = ('lam', 'muplus', 'muminus')


def make_relu_network_lagrangian(dual_vars, params, bounds, obj):
  """Returns a function that computes the Lagrangian for a ReLU network.

  This function assumes `params` represent a feedforward ReLU network i.e.
  x_{i+1} = relu(W_i x_i + b_i). It defines the Lagrangian by applying the
  objective `obj` to the final layer activations, and encoding the Lagrangian
  terms for each of the constraints defining the ReLU network. It then returns
  this function.

  Args:
    dual_vars: A length L+1 list of dual variables at each layer
    params: A length L list of (W, b) pairs, elided network weights
    bounds: A length L+1 list of `IntBound`s, elided bounds at each layer
    obj: function, taking final layer activations as input

  Returns:
    Function that computes Lagrangian L(x) with fixed `dual_vars`.
  """

  layer_sizes = utils.layer_sizes_from_bounds(bounds)