# coding=utf-8
# Copyright 2020 The jax_verify Authors.
#
# 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.

# Lint as: python3
"""Routines to solve the subproblem during verification."""
from typing import Dict, List, Tuple, Optional
from absl import logging
import cvxpy as cp

from jax_verify.src import relaxation
import numpy as np

_EPS = 1e-10
CvxpyConstraint = cp.constraints.constraint.Constraint
Tensor = np.ndarray


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

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

  def create_solver_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) > _EPS:
            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: int,
                         objective: Tensor,
                         objective_bias: float,
                         time_limit: Optional[int] = None
                         ) -> Tuple[float, 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.
      status: Status of the optimization function.
    """
    if time_limit:
      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):
      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()
    logging.info('Problem status is %s', prob.status)
    val = obj.value
    if val is not None:
      val += objective_bias
    return val, (prob.status == cp.settings.OPTIMAL)
