# 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.

"""Routines to solve the subproblem during verification."""
from typing import Dict, List, Mapping, Optional, Tuple, Union

from absl import logging
import cvxpy as cp
from cvxpy.reductions.solvers.defines import INSTALLED_MI_SOLVERS as MIP_SOLVERS
import jax
import jax.numpy as jnp
from jax_verify.src import utils
from jax_verify.src.mip_solver import relaxation
import numpy as np

CvxpyConstraint = cp.constraints.constraint.Constraint
Tensor = np.ndarray

Variable = Union[relaxation.RelaxVariable,
                 relaxation.BinaryVariable]
Solution = Mapping[str, jnp.ndarray]


class CvxpySolver(relaxation.RelaxationSolver):
  """Holder class to represent problem being built."""

  def __init__(self):
    """Build and solve the relaxation using CVXPY."""
    self.solver_variables: Dict[str, cp.Variable] = {}
    self.constraints: List[CvxpyConstraint] = []
    self.problem_kwargs = {}

  def _variable_already_created(self, var: Variable) -> bool:
    return var.name in self.solver_variables

  def _create_solver_relax_variable(
      self,
      relax_var: relaxation.RelaxVariable,
      index: int):
    """Create a new bound-constrained CVXPY variable based on a RelaxVariable.

    Args:
      relax_var: Variable generated by the relaxation bound propagation.
      index: Index in the batch for which to build the variable.
    """
    lower = np.reshape(relax_var.lower[index, ...], [-1])
    upper = np.reshape(relax_var.upper[index, ...], [-1])
    var = cp.Variable(lower.shape)
    self.constraints += [lower <= var]
    self.constraints += [var <= upper]
    self.solver_variables[relax_var.name] = var

  def create_linear_solver_constraint(self,
                                      constraint: relaxation.LinearConstraint,
                                      index: int):
    """Create a new CVXPY linear constraint.

    Args:
      constraint: Constraint generated by the relaxation bound propagation.
      index: Index in the batch for which to build the variable.
    """
    rhs = cp.expressions.constants.Constant(constraint.bias(index))
    for (constraint_variable,
         (cpts, coeffs)) in constraint.vars_and_coeffs(index):
      if isinstance(constraint_variable, relaxation.RelaxVariable):
        current_variable = self.solver_variables[constraint_variable.name]
        for component, coeff in zip(cpts, coeffs):
          if np.abs(coeff) > utils.EPSILON:
            rhs += current_variable[component] * coeff.item()
    if constraint.sense == 0:
      self.constraints += [rhs == 0]
    if constraint.sense == 1:
      self.constraints += [rhs <= 0]
    if constraint.sense == -1:
      self.constraints += [rhs >= 0]

  def create_activation_solver_constraint(
      self, constraint: relaxation.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
    """
    outvar = self.solver_variables[constraint.outvar.name]
    invar = self.solver_variables[constraint.invar.name]
    expression = invar[act_index] * slope + bias - outvar[act_index]
    if constraint.sense == 0:
      self.constraints += [expression == 0]
    if constraint.sense == 1:
      self.constraints += [expression <= 0]
    if constraint.sense == -1:
      self.constraints += [expression >= 0]

  def minimize_objective(
      self,
      var_name: str,
      objective: Tensor,
      objective_bias: float,
      time_limit: Optional[int] = None,
  ) -> Tuple[float, Solution, 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: Maximum solve time in ms. Use None for unbounded. CVXPY does
        not support time_limit so any other value will raise a ValueError.
    Returns:
      val: Value of the minimum.
      solution: solution of the problem
      status: Status of the optimization function.
    """
    if time_limit is not None:
      raise ValueError('Cvxpy Solver does not support time limit.')
    # Define the objective function
    obj = cp.expressions.constants.Constant(0.)
    for i, coeff in enumerate(objective):
      if np.abs(coeff) > utils.EPSILON:
        obj += self.solver_variables[var_name][i] * coeff.item()
    objective = cp.Minimize(obj)
    prob = cp.Problem(objective, self.constraints)
    logging.info('Starting problem solve.')
    prob.solve(solver=cp.ECOS, **self.problem_kwargs)
    logging.info('Problem status is %s', prob.status)
    val = obj.value
    if val is not None:
      val += objective_bias
    solution = jax.tree_map(lambda x: jnp.array(x.value),
                            self.solver_variables)
    return val, solution, prob.status in (cp.settings.OPTIMAL,
                                          cp.settings.OPTIMAL_INACCURATE)


class CvxpyMIPSolver(CvxpySolver, relaxation.MIPSolver):  # pytype: disable=signature-mismatch  # maybe_create_solver_variable uses parameter contravariance
  """Holder class to represent problem being built."""

  def __init__(self):
    """Build and solve the MIP encoding using CVXPY."""
    super(CvxpyMIPSolver, self).__init__()

    if not MIP_SOLVERS:
      raise NotImplementedError('No MIP solver is available.')

    self.problem_kwargs.update(solver=MIP_SOLVERS[0])

  def _create_solver_bool_variable(
      self,
      bin_var: relaxation.BinaryVariable,
      index: int):
    """Create a new bound-constrained CVXPY variable based on a BinaryVariable.

    Args:
      bin_var: Variable generated by the MIP formulation.
      index: Index in the batch for which to build the variable.
    """
    shape = np.prod(bin_var.shape[1:])
    self.solver_variables[bin_var.name] = cp.Variable(shape, boolean=True)

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

    Encodes outvar =(>)(<) slope * invar + binslope * binvar + bias.

    Args:
      constraint: Constraint generated by the relaxation bound propagation.
      act_index: Index of the activation to encode (in the variables involved
        in constraint)
      binslope: Multiplicative coefficient for the binary variable.
      slope : Slope coefficients of the linear inequality.
      bias: Bias of the linear inequality
    """
    binvar = self.solver_variables[constraint.binvar.name]
    outvar = self.solver_variables[constraint.outvar.name]
    invar = self.solver_variables[constraint.invar.name]
    expression = (invar[act_index] * slope + binvar[act_index] * binslope
                  + bias - outvar[act_index])
    if constraint.sense == 0:
      self.constraints += [expression == 0]
    if constraint.sense == 1:
      self.constraints += [expression <= 0]
    if constraint.sense == -1:
      self.constraints += [expression >= 0]
