# coding=utf-8
# Copyright 2023 DeepMind Technologies Limited.
#
# 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
#
#     http://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.

"""Propagate the relaxation through the network.

This is accomplished by traversing the JaxPR representation of the computation
and translating the computational graph.
"""
import abc
import functools
from typing import Callable, Tuple, List, Mapping, Optional, Union

import jax
from jax import lax
import jax.numpy as jnp
from jax_verify.src import activation_relaxation
from jax_verify.src import bound_propagation
from jax_verify.src import graph_traversal
from jax_verify.src import ibp
from jax_verify.src import synthetic_primitives
from jax_verify.src.types import Index, Primitive, Tensor  # pylint: disable=g-multiple-import
import numpy as np


class RelaxVariable(bound_propagation.Bound):
  """Variable used to build relaxation."""

  def __init__(self, idx: Index, base_bound):
    self.base_bound = base_bound
    self.name = 'rlxvar'
    for idx_cpt in idx:
      self.name += f'_{idx_cpt}'
    self.idx = idx
    self.constraints = None

  def set_constraints(self, constraints):
    self.constraints = constraints

  @property
  def lower(self):
    return self.base_bound.lower

  @property
  def upper(self):
    return self.base_bound.upper


class OptRelaxVariable(RelaxVariable):
  """Variant of RelaxVariable with lazy optimization tightening."""

  def __init__(self, base_relax_variable, optimize_transform):
    super(OptRelaxVariable, self).__init__(base_relax_variable.idx,
                                           base_relax_variable.base_bound)
    self._shape = base_relax_variable.shape
    self._opted_bounds = None
    self._optimize_transform = optimize_transform

  @property
  def shape(self):
    # Need to specify a shape method, because otherwise, a call to `.shape`
    # will attempt to read `.lower` and trigger optimization
    return self._shape

  @property
  def dtype(self):
    # Need to specify a type method, because otherwise, a call to `.dtype`
    # will attempt to read `.lower` and trigger optimization
    return self._dtype

  @property
  def lower(self):
    return self._optimized_bounds.lower

  @property
  def upper(self):
    return self._optimized_bounds.upper

  @property
  def _optimized_bounds(self):
    if self._opted_bounds is None:
      self._opted_bounds = self._optimize_transform.tight_bounds(self)
    return self._opted_bounds


class BinaryVariable():
  """Binary variable."""

  def __init__(self, idx, shape):
    self.shape = shape
    self.name = f'boolvar_{idx}'
    self.idx = idx


class MIPActivationConstraint:
  """MIP constraint to encode activation."""

  def __init__(self, outvar, invar, binvar, mask, binscale, scale, bias, sense):
    """Represents: outvar =(>)(<) scale * invar + binscale * binvar + bias."""
    self.outvar = outvar
    self.invar = invar
    self.binvar = binvar
    self.binscale = binscale
    self.mask = mask
    self.scale = scale
    self.bias = bias
    self.sense = sense

  def encode_into_solver(self, solver: 'MIPSolver', index: int):
    """Encode the linear constraints into the provided solver.

    Args:
      solver: MIPSolver to create the exact constraint into.
      index: Index in the batch for which to build the variable.
    """
    biases = np.reshape(self.bias[index, ...], [-1])
    slopes = np.reshape(self.scale[index, ...], [-1])
    binslopes = np.reshape(self.binscale[index, ...], [-1])
    mask = np.reshape(self.mask[index, ...], [-1])
    for act_index, (binslope, slope, bias) in enumerate(
        zip(binslopes, slopes, biases)):
      if mask[act_index]:
        solver.create_mip_activation_solver_constraint(
            self, act_index, binslope=binslope.item(),
            slope=slope.item(), bias=bias.item())


class LinearConstraint:
  """Linear constraint, to be encoded into a solver."""

  def __init__(self, vars_and_coeffs, bias, sense):
    self._vars_and_coeffs = vars_and_coeffs
    self._bias = bias
    self.sense = sense
    self.sample_dependent = bool(self._bias.shape)

  def bias(self, index: int):
    """Get the bias corresponding to the minibatch sample `index`.

    If bias has a dimension, it means it is sample dependent. Otherwise,
    the bias is the same for all samples.

    Args:
      index: Index in the batch for which to build the variable.
    Returns:
      bias: value of the bias.
    """
    return self._bias[index] if self.sample_dependent else self._bias

  def vars_and_coeffs(self, index: int):
    """Get the variable and coefficients corresponding to the sample `index`.

    If coeffs has a dimension, the coefficients are sample dependent. Otherwise,
    the coefficients are the same for all samples.

    Args:
      index: Index in the batch for which to build the variable.
    Returns:
      vars_and_coeffs: vars_and_coeffs list where the coefficients are the one
        corresponding to the sample `index`
    """
    if self.sample_dependent:
      return [(var, (cpts, coeffs[index]))
              for (var, (cpts, coeffs)) in self._vars_and_coeffs]
    else:
      return self._vars_and_coeffs

  def encode_into_solver(self, solver: 'RelaxationSolver', index: int):
    """Encode the linear constraints into the provided solver.

    Args:
      solver: RelaxationSolver to create the linear constraint into.
      index: Index in the batch for which to build the variable.
    """
    solver.create_linear_solver_constraint(self, index)


class RelaxActivationConstraint:
  """Linear constraint involved in the relaxation of an activation."""

  def __init__(self, outvar, invar, mask, scale, bias, sense):
    """Represents the constraint outvar =(>)(<) scale * invar + bias."""
    self.outvar = outvar
    self.invar = invar
    self.mask = mask
    self.scale = scale
    self.bias = bias
    self.sense = sense

  def encode_into_solver(self, solver: 'RelaxationSolver', index: int):
    """Encode the linear constraints into the provided solver.

    Args:
      solver: RelaxationSolver to create the linear constraint into.
      index: Index in the batch for which to build the variable.
    """
    biases = np.reshape(self.bias[index, ...], [-1])
    slopes = np.reshape(self.scale[index, ...], [-1])
    mask = np.reshape(self.mask[index, ...], [-1])
    for act_index, (slope, bias) in enumerate(zip(slopes, biases)):
      if mask[act_index]:
        solver.create_activation_solver_constraint(
            self, act_index, slope.item(), bias.item())


class RelaxationSolver(metaclass=abc.ABCMeta):
  """Abstract solver for the relaxation."""

  def maybe_create_solver_variable(
      self,
      var: RelaxVariable,
      index: int):
    """Create a new solver variable for var if it has not been created yet.

    Args:
      var: Variable generated by the relaxation bound propagation.
      index: Index in the batch for which to build the variable.
    """
    if not self._variable_already_created(var):
      self._create_solver_relax_variable(var, index)

  @abc.abstractmethod
  def _create_solver_relax_variable(
      self,
      var: RelaxVariable,
      index: int):
    """Create a new bound-constrained variable based on a RelaxVariable.

    Args:
      var: Variable generated by the relaxation bound propagation.
      index: Index in the batch for which to build the variable.
    """

  @abc.abstractmethod
  def _variable_already_created(
      self,
      var: Union[RelaxVariable, BinaryVariable],
  ) -> bool:
    """Check whether the solver has already created a variable for var.

    Args:
      var: Variable generated by the relaxation bound propagation.
    """

  @abc.abstractmethod
  def create_linear_solver_constraint(
      self,
      constraint: LinearConstraint,
      index: int):
    """Create a new solver linear constraint.

    Args:
      constraint: Constraint generated by the relaxation bound propagation.
      index: Index in the batch for which to build the variable.
    """

  @abc.abstractmethod
  def create_activation_solver_constraint(
      self,
      constraint: RelaxActivationConstraint,
      act_index: int,
      slope: float,
      bias: float):
    """Create the linear constraint involved in the activation relaxation.

    Args:
      constraint: Constraint generated by the relaxation bound propagation.
      act_index: Index of the activation to encode (in the variables involved
        in constraint)
      slope : Slope coefficients of the linear inequality.
      bias: Bias of the linear inequality
    """

  @abc.abstractmethod
  def minimize_objective(
      self,
      var_name: str,
      objective: Tensor,
      objective_bias: float,
      time_limit_millis: Optional[int],
  ) -> Tuple[float, Mapping[str, Tensor], bool]:
    """Minimize a linear function.

    Args:
      var_name: Index of the variable to define a linear function over the
        components.
      objective: Coefficients of the linear function.
      objective_bias: Bias of the linear function.
      time_limit_millis: Maximum solve time in ms. Use None for unbounded.
    Returns:
      val: Value of the minimum.
      solution: Solution found.
      status: Status of the optimization function.
    """


class MIPSolver(RelaxationSolver):
  """Abstract solver for the MIP encoding."""

  def maybe_create_solver_variable(
      self,
      var: Union[RelaxVariable, BinaryVariable],
      index: int):
    """Create a new solver variable for var if it has not been created yet.

    Args:
      var: Variable generated by the relaxation bound propagation.
      index: Index in the batch for which to build the variable.
    """
    if not self._variable_already_created(var):
      if isinstance(var, BinaryVariable):
        self._create_solver_bool_variable(var, index)
      else:
        self._create_solver_relax_variable(var, index)

  @abc.abstractmethod
  def _create_solver_bool_variable(
      self,
      var: BinaryVariable,
      index: int):
    """Create a new bound-constrained variable based on a BinaryVariable.

    Args:
      var: Variable generated by the relaxation bound propagation.
      index: Index in the batch for which to build the variable.
    """

  @abc.abstractmethod
  def create_mip_activation_solver_constraint(
      self,
      constraint: MIPActivationConstraint,
      act_index: int,
      binslope: float,
      slope: float,
      bias: float):
    """Create the linear constraint involved in the activation relaxation.

    Args:
      constraint: Constraint generated by the relaxation bound propagation.
      act_index: Index of the activation to encode (in the variables involved
        in constraint)
      binslope: Slope coefficients for the binary variable.
      slope : Slope coefficients for the input variable.
      bias: Bias of the linear inequality
    """


def encode_relaxation(
    solver_ctor: Callable[[], RelaxationSolver],
    env: Mapping[jax.core.Var, Union[RelaxVariable, Tensor]],
    index: int,
) -> RelaxationSolver:
  """Creates a solver and encodes the relaxation into it.

  Args:
    solver_ctor: Constructor for the solver.
    env: Environment created by applying boundprop with relaxation.py
    index: The index in the minibatch for which the relaxation should be
      encoded.
  Returns:
    solver: Solver containing the relaxation of the network encoded.
  """
  solver = solver_ctor()
  for key in env.keys():
    if isinstance(env[key], RelaxVariable):
      variable = env[key]
      # Create the variable in the solver.
      solver.maybe_create_solver_variable(variable, index)
      # Create the constraints in the solver.
      if variable.constraints:
        for constraint in variable.constraints:
          if isinstance(constraint, MIPActivationConstraint):
            # create manually the binary variable because it is not collected
            # automatically by the graph propagation.
            solver.maybe_create_solver_variable(constraint.binvar, index)
          constraint.encode_into_solver(solver, index)
  return solver


def solve_relaxation(
    solver_ctor: Callable[[], RelaxationSolver],
    objective: Tensor,
    objective_bias: float,
    variable_opt: RelaxVariable,
    env: Mapping[jax.core.Var, Union[RelaxVariable, Tensor]],
    index: int,
    time_limit_millis: Optional[int] = None,
) -> Tuple[float, Mapping[str, Tensor], bool]:
  """Solves the relaxation using the provided LP solver.

  Args:
    solver_ctor: Constructor for the solver.
    objective: Objective to optimize, given as an array of coefficients to be
      applied to the variable to form a linear objective function
    objective_bias: Bias to add to objective
    variable_opt: RelaxVariable over which the linear function to optimize
      is defined.
    env: Environment created by applying boundprop with relaxation.py
    index: The index in the minibatch for which the LP should be solved
    time_limit_millis: Time limit on solver. None if unbounded.
  Returns:
    opt: Value of the solution found.
    solution: Solution found by the solver.
    status: Whether the optimal solution has been achieved.
  """
  solver = encode_relaxation(solver_ctor, env, index)
  return solver.minimize_objective(variable_opt.name,
                                   objective, objective_bias, time_limit_millis)


def _get_linear(primitive, outval, *eqn_invars, **params):
  """Get linear expressions corresponding to an affine layer.

  Args:
    primitive: jax primitive
    outval: dummy tensor shaped according to a single example's outputs
    *eqn_invars: Arguments of the primitive, wrapped as RelaxVariables
    **params: Keyword Arguments of the primitive.

  Returns:
    For each output component, a pair `(bias, coefficients)`, where
    `coefficients` is a list of `(component, coefficient)` pairs.
  """
  def funx(x):
    if isinstance(x, RelaxVariable):
      return jnp.zeros(x.shape)
    else:
      return x
  def fungrad(i, args):
    return jnp.reshape(primitive.bind(*args, **params)[0, ...], [-1])[i]
  results = []
  # Loop over output dimensions one at a time to avoid creating a large
  # materialized tensor
  # TODO: Replace with something more efficient
  for i in range(outval.size):
    fung = functools.partial(fungrad, i)
    bias, current_grad = jax.value_and_grad(fung, allow_int=True)(
        [funx(x) for x in eqn_invars])
    coefficients = []
    for res, in_var in zip(current_grad, eqn_invars):
      if isinstance(in_var, RelaxVariable):
        components = jnp.flatnonzero(res)
        coefficients.append((components, res.ravel()[components]))
      else:
        coefficients.append(None)
    results.append((bias, coefficients))
  return results


def _relax_input(
    index: Index, in_bounds: bound_propagation.Bound,
) -> RelaxVariable:
  """Generates the initial inputs for the relaxation.

  Args:
    index: Integer identifying the input node.
    in_bounds: Concrete bounds on the input node.
  Returns:
    `RelaxVariable` for the initial inputs.
  """
  # Wrap initial bound as RelaxVariable bound.
  in_variable = RelaxVariable(index, in_bounds)
  return in_variable


_order_preserving_reshapes = [lax.reshape_p, lax.squeeze_p]
_affine_primitives_list = [
    *bound_propagation.AFFINE_PRIMITIVES,
    *bound_propagation.RESHAPE_PRIMITIVES,
    lax.div_p,
]


def _relax_primitive(
    index: Index, out_bounds: bound_propagation.Bound,
    primitive: Primitive,
    *args, use_mip: bool = False, **kwargs
) -> RelaxVariable:
  """Generates the relaxation for a given primitive op.

  Args:
    index: Integer identifying the computation node.
    out_bounds: Concrete bounds on the outputs of the primitive op.
    primitive: jax primitive.
    *args: Arguments of the primitive, wrapped as RelaxVariables
    use_mip: whether to use mixed integer programming for the activation
      constraints.
    **kwargs: Keyword Arguments of the primitive.
  Returns:
    `RelaxVariable` that contains the output of this primitive for the
    relaxation, with all the constraints linking the output to inputs.
  """
  # Create variable for output of this primitive
  out_variable = RelaxVariable(index, out_bounds)
  # Create constraints linking output and input of primitive
  constraints = []
  if primitive in _order_preserving_reshapes:
    invar = args[0]
    constraints = [RelaxActivationConstraint(outvar=out_variable,
                                             invar=invar,
                                             mask=jnp.ones(invar.shape),
                                             scale=jnp.ones(invar.shape),
                                             bias=jnp.zeros(invar.shape),
                                             sense=0)]
  elif primitive in _affine_primitives_list:
    if primitive == lax.div_p and isinstance(args[1], RelaxVariable):
      raise NotImplementedError(
          'Division with non-constant divisor is not supported')
    results = _get_linear(primitive, out_bounds.lower[0, ...],
                          *args, **kwargs)
    for i, (bias, coeffs) in enumerate(results):
      # Coefficients of the input variable(s).
      vars_and_coeffs = [
          (arg, coeff) for arg, coeff in zip(args, coeffs)
          if isinstance(arg, RelaxVariable)]
      # Equate with the output variable, by using a coefficient of -1.
      out_coeff = (np.array([i], dtype=np.int64), np.array([-1.]))
      vars_and_coeffs.append((out_variable, out_coeff))
      constraints.append(LinearConstraint(vars_and_coeffs, bias, 0))
  elif primitive in activation_relaxation.relaxation_fns:
    # Generate convex relaxation.
    safe_kwargs = synthetic_primitives.filter_jaxverify_kwargs(kwargs)
    activation = activation_relaxation.relaxation_fns[primitive]
    lb_funs, ub_funs = activation.piecewise_linear_relaxation_fn(*args,
                                                                 **safe_kwargs)
    invar, = args
    zeros = jnp.zeros_like(invar.lower)
    ones = jnp.ones_like(invar.lower)
    if activation.pos_neg_linear and (len(lb_funs) == 1 or len(ub_funs) == 1):
      # Use equality constraints if linear over the entire interval.
      ambiguous = (invar.lower < 0) & (invar.upper > 0)
      chord, = lb_funs if len(lb_funs) == 1 else ub_funs
      constraints.append(RelaxActivationConstraint(
          outvar=out_variable,
          invar=invar,
          mask=(~ambiguous),
          scale=(chord(ones) - chord(zeros)),
          bias=chord(zeros),
          sense=0))
    else:
      ambiguous = ones

    for lb_fun in lb_funs:
      # act(x) >= lb(x)
      constraints.append(RelaxActivationConstraint(
          outvar=out_variable,
          invar=invar,
          mask=ambiguous,
          scale=(lb_fun(ones) - lb_fun(zeros)),
          bias=lb_fun(zeros),
          sense=1))

    for ub_fun in ub_funs:
      # act(x) <= ub(x)
      constraints.append(RelaxActivationConstraint(
          outvar=out_variable,
          invar=invar,
          mask=ambiguous,
          scale=(ub_fun(ones) - ub_fun(zeros)),
          bias=ub_fun(zeros),
          sense=-1))

    if use_mip:
      if primitive is not synthetic_primitives.relu_p:
        raise ValueError(
            f'Only ReLU activations supported. Encountered {primitive}')
      binvar = BinaryVariable(index, out_bounds.lower.shape)
      constraints += [
          # outvar <= upper_bound * binvar
          MIPActivationConstraint(outvar=out_variable,
                                  invar=invar,
                                  binvar=binvar,
                                  binscale=invar.upper,
                                  mask=ambiguous,
                                  scale=zeros,
                                  bias=zeros,
                                  sense=-1),
          # outvar <= invar - lower_bound * (1. - binvar)
          MIPActivationConstraint(outvar=out_variable,
                                  invar=invar,
                                  binvar=binvar,
                                  binscale=invar.lower,
                                  mask=ambiguous,
                                  scale=ones,
                                  bias=-invar.lower,
                                  sense=-1),
      ]
  else:
    raise NotImplementedError(f'Unsupported primitive: {primitive}')

  out_variable.set_constraints(constraints)
  return out_variable


class RelaxationTransform(graph_traversal.GraphTransform[RelaxVariable]):
  """Transform to produce `RelaxVariable`s for each op."""

  def __init__(
      self,
      boundprop_transform: bound_propagation.BoundTransform,
      use_mip: bool = False,
  ):
    """Defines relaxation constraint propagation.

    Args:
      boundprop_transform: Basic Jax primitive ops' equivalents for
        the underlying bound propagation method.
      use_mip: whether to use mixed integer programming for the activation
        constraints.
    """
    self._boundprop_transform = boundprop_transform
    self._use_mip = use_mip

  def input_transform(self, context, input_bound):
    in_bounds = self._boundprop_transform.input_transform(
        context, input_bound)
    return _relax_input(context.index, in_bounds)

  def primitive_transform(self, context, primitive, *args, **params):
    interval_args = [arg.base_bound if isinstance(arg, RelaxVariable) else arg
                     for arg in args]
    out_bounds, = self._boundprop_transform.equation_transform(
        context, primitive, *interval_args, **params)
    return _relax_primitive(
        context.index, out_bounds, primitive, *args,
        use_mip=self._use_mip, **params)


class OptimizedRelaxationTransform(
    graph_traversal.GraphTransform[OptRelaxVariable]):
  """Wraps a RelaxVariable-producing BoundTransform to add optimization."""

  def __init__(
      self,
      transform: graph_traversal.GraphTransform[RelaxVariable],
      solver_ctor: Callable[[], RelaxationSolver],
      time_limit_millis: Optional[int] = None):
    self._transform = transform
    self.solver_ctor = solver_ctor
    self.solvers: List[RelaxationSolver] = []
    self._time_limit_millis = time_limit_millis

  def tight_bounds(self, variable: RelaxVariable) -> ibp.IntervalBound:
    """Compute tighter bounds based on the LP relaxation.

    Args:
      variable: Variable as created by the base boundprop transform. This is a
        RelaxVariable that has already been encoded into the solvers.
    Returns:
      tightened_base_bound: Bounds tightened by optimizing with the LP solver.
    """
    lbs = []
    ubs = []
    for solver in self.solvers:
      nb_targets = np.prod(variable.shape[1:])
      sample_lbs = []
      sample_ubs = []
      for target_idx in range(nb_targets):
        objective = (jnp.arange(nb_targets) == target_idx).astype(jnp.float32)
        lb, _, optimal_lb = solver.minimize_objective(
            variable.name, objective, 0., self._time_limit_millis)
        assert optimal_lb
        neg_ub, _, optimal_ub = solver.minimize_objective(
            variable.name, -objective, 0., self._time_limit_millis)
        assert optimal_ub
        sample_lbs.append(lb)
        sample_ubs.append(-neg_ub)
      lbs.append(sample_lbs)
      ubs.append(sample_ubs)

    tightened_base_bound = ibp.IntervalBound(
        jnp.reshape(jnp.array(lbs), variable.shape),
        jnp.reshape(jnp.array(ubs), variable.shape))
    return tightened_base_bound

  def input_transform(
      self,
      context: graph_traversal.TransformContext[OptRelaxVariable],
      input_bound: graph_traversal.InputBound,
  ) -> RelaxVariable:
    in_bounds = self._transform.input_transform(
        context, input_bound)
    for minibatch_index in range(in_bounds.shape[0]):
      # Create one solver instance for each problem in the batch because they
      # will have different constraints.
      if minibatch_index >= len(self.solvers):
        self.solvers.append(self.solver_ctor())
      solver = self.solvers[minibatch_index]
      solver.maybe_create_solver_variable(in_bounds, minibatch_index)
    return in_bounds

  def primitive_transform(
      self,
      context: graph_traversal.TransformContext[OptRelaxVariable],
      primitive: Primitive,
      *args: Union[RelaxVariable, Tensor],
      **params,
  ) -> RelaxVariable:
    basic_relax_var, = self._transform.equation_transform(
        context, primitive, *args, **params)
    opt_relax_var = OptRelaxVariable(basic_relax_var, self)

    for minibatch_index, solver in enumerate(self.solvers):
      # Encode the new variable and the associated constraints. We encode the
      # basic variable, not the optimized one. This way, the optimization is
      # only performed if the lower / upper bounds are required and the bounds
      # on that variable are accessed "from outside".
      solver.maybe_create_solver_variable(basic_relax_var, minibatch_index)
      if basic_relax_var.constraints:
        for constraint in basic_relax_var.constraints:
          if isinstance(constraint, MIPActivationConstraint):
            # create manually the binary variable because it is not collected
            # automatically by the graph propagation.
            solver.maybe_create_solver_variable(
                constraint.binvar, minibatch_index)
          constraint.encode_into_solver(solver, minibatch_index)
    return opt_relax_var
