#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# SPDX-License-Identifier: LGPL-3.0-or-later
# Copyright 2016-2022 Stéphane Caron and the qpsolvers contributors

"""Solver interface for `SIP`_.

.. _SIP: https://github.com/joaospinto/sip_python

SIP is a general NLP solver based. It is based on the barrier augmented
Lagrangian method, which combines the interior point and augmented Lagrangian
methods. If you are using SIP in a scientific work, consider citing the
corresponding GitHub repository (or paper, if one has been released).

**Warm-start:** this solver interface supports warm starting 🔥
"""

import warnings
from typing import Optional, Union

import numpy as np
import scipy.sparse as spa
import sip_python as sip

from ..conversions import linear_from_box_inequalities, split_dual_linear_box
from ..exceptions import ProblemError
from ..problem import Problem
from ..solution import Solution

_cvxopt_available = True
try:
    from cvxopt import amd, spmatrix  # noqa: F401
except ImportError:
    _cvxopt_available = False


def sip_solve_problem(
    problem: Problem,
    initvals: Optional[np.ndarray] = None,
    verbose: bool = False,
    allow_non_psd_P: bool = False,
    **kwargs,
) -> Solution:
    """Solve a quadratic program using SIP.

    Parameters
    ----------
    problem :
        Quadratic program to solve.
    initvals :
        Warm-start guess vector for the primal solution.
    verbose :
        Set to `True` to print out extra information.

    Returns
    -------
    :
        Solution to the QP returned by the solver.

    Notes
    -----
    All other keyword arguments are forwarded as options to SIP. For
    instance, you can call ``sip_solve_qp(P, q, G, h, eps_abs=1e-6)``.
    For a quick overview, the solver accepts the following settings:

    .. list-table::
       :widths: 30 70
       :header-rows: 1

       * - Name
         - Effect
       * - max_iterations
         - The maximum number of iterations the solver can do.
       * - max_ls_iterations
         - The maximum cumulative number of line search iterations.
       * - min_iterations_for_convergence
         - The least number of iterations until we can declare convergence.
       * - num_iterative_refinement_steps
         - The number of iterative refinement steps.
       * - max_kkt_violation
         - The maximum allowed violation of the KKT system.
       * - max_merit_slope
         - The maximum allowed merit function slope.
       * - initial_regularization
         - The initial x-regularizatino to be applied on the LHS.
       * - regularization_decay_factor
         - The multiplicative decay of the x-regularization coefficient.
       * - tau
         - A parameter of the fraction-to-the-boundary rule.
       * - start_ls_with_alpha_s_max
         - Determines whether we start with alpha=alpha_s_max or alpha=1.
       * - initial_mu
         - The initial barrier function coefficient.
       * - mu_update_factor
         - Determines how much mu decreases per iteration.
       * - mu_min
         - The minimum barrier coefficient.
       * - initial_penalty_parameter
         - The initial penalty parameter of the Augmented Lagrangian.
       * - min_acceptable_constraint_violation_ratio
         - Least acceptable constraint violation ratio to not increase eta.
       * - penalty_parameter_increase_factor
         - By what factor to increase eta.
       * - penalty_parameter_decrease_factor
         - By what factor to decrease eta.
       * - max_penalty_parameter
         - The maximum allowed penalty parameter in the AL merit function.
       * - armijo_factor
         - Determines when we accept a line search step.
       * - line_search_factor
         - Determines how much to backtrack at each line search iteration.
       * - line_search_min_step_size
         - Determines when we declare a line search failure.
       * - min_merit_slope_to_skip_line_search
         - Min merit slope to skip the line search.
       * - dual_armijo_factor
         - Fraction of the primal merit decrease to allow on the dual update.
       * - min_allowed_merit_increase
         - The minimum allowed merit function increase in the dual update.
       * - enable_elastics
         - Whether to enable the usage of elastic variables.
       * - elastic_var_cost_coeff
         - Elastic variables cost coefficient.
       * - enable_line_search_failures
         - Halts the optimization process if a good step is not found.
       * - print_logs
         - Determines whether we should print the solver logs.
       * - print_line_search_logs
         - Determines whether we should print the line search logs.
       * - print_search_direction_logs
         - Whether we should print the search direction computation logs.
       * - print_derivative_check_logs
         - Whether to print derivative check logs when something looks off.
       * - only_check_search_direction_slope
         - Only derivative-check the search direction.
       * - assert_checks_pass
         - Handle checks with assert calls.

    This list may not be exhaustive.
    Check the `Settings` struct in the `solver code
    <https://github.com/joaospinto/sip/blob/main/sip/types.hpp>`__ for details.
    """
    P, q, G, h, A, b, lb, ub = problem.unpack()
    if lb is not None or ub is not None:
        G, h = linear_from_box_inequalities(
            G, h, lb, ub, use_sparse=problem.has_sparse
        )
    n: int = q.shape[0]

    # SIP does not support A, b, G, and h to be None.
    G = G if G is not None else spa.csr_matrix(np.zeros((0, n)))
    A = A if A is not None else spa.csr_matrix(np.zeros((0, n)))
    h = np.zeros((0,)) if h is None else h
    b = np.zeros((0,)) if b is None else b

    # Remove any infs from h.
    G[np.isinf(h), :] = 0.0
    h[np.isinf(h)] = 1.0

    if not isinstance(P, spa.csr_matrix):
        P = spa.csc_matrix(P)
        if verbose:
            warnings.warn("Converted P to a csc_matrix.")
    if not isinstance(G, spa.csr_matrix):
        G = spa.csr_matrix(G)
        if verbose:
            warnings.warn("Converted G to a csr_matrix.")
    if not isinstance(A, spa.csr_matrix):
        A = spa.csr_matrix(A)
        if verbose:
            warnings.warn("Converted A to a csr_matrix.")

    P.eliminate_zeros()
    G.eliminate_zeros()
    A.eliminate_zeros()

    P_T = spa.csc_matrix(P.T)
    if (
        (P.indices != P_T.indices).any()
        or (P.indptr != P_T.indptr).any()
        or (P.data != P_T.data).any()
    ):
        raise ProblemError("P should be symmetric.")

    if G is None and h is not None:
        raise ProblemError(
            "Inconsistent inequalities: G is not set but h is set"
        )
    if G is not None and h is None:
        raise ProblemError("Inconsistent inequalities: G is set but h is None")
    if A is None and b is not None:
        raise ProblemError(
            "Inconsistent inequalities: A is not set but b is set"
        )
    if A is not None and b is None:
        raise ProblemError("Inconsistent inequalities: A is set but b is None")

    k = None
    if allow_non_psd_P:
        eigenvalues, _eigenvectors = spa.linalg.eigsh(P, k=1, which="SM")
        k = -min(eigenvalues[0], 0.0) + 1e-3
    else:
        k = 1e-6

    # hess_L = P + k * spa.eye(n);
    # the code below avoids potential index cancellations.
    hess_L = spa.coo_matrix(P)
    upp_hess_L_rows = np.concatenate([hess_L.row, np.arange(n)])
    upp_hess_L_cols = np.concatenate([hess_L.col, np.arange(n)])
    upp_hess_L_data = np.concatenate([hess_L.data, k * np.ones(n)])
    hess_L = spa.coo_matrix(
        (upp_hess_L_data, (upp_hess_L_rows, upp_hess_L_cols)), shape=P.shape
    )
    hess_L.sum_duplicates()
    upp_hess_L = spa.triu(hess_L.tocsc())

    qs = sip.QDLDLSettings()
    qs.permute_kkt_system = True
    qs.kkt_pinv = sip.get_kkt_perm_inv(
        P=hess_L,
        A=A,
        G=G,
    )

    pd = sip.ProblemDimensions()
    pd.x_dim = n
    pd.s_dim = h.shape[0]
    pd.y_dim = b.shape[0]
    pd.upper_hessian_lagrangian_nnz = upp_hess_L.nnz
    pd.jacobian_c_nnz = A.nnz
    pd.jacobian_g_nnz = G.nnz
    pd.kkt_nnz, pd.kkt_L_nnz = sip.get_kkt_and_L_nnzs(
        P=hess_L,
        A=A,
        G=G,
        perm_inv=qs.kkt_pinv,
    )
    pd.is_jacobian_c_transposed = True
    pd.is_jacobian_g_transposed = True

    vars = sip.Variables(pd)

    if initvals is not None:
        vars.x[:] = initvals
    else:
        vars.x[:] = 0.0

    vars.s[:] = 1.0
    vars.y[:] = 0.0
    vars.e[:] = 0.0
    vars.z[:] = 1.0

    ss = sip.Settings()
    ss.max_iterations = 100
    ss.max_ls_iterations = 1000
    ss.max_kkt_violation = 1e-8
    ss.max_merit_slope = 1e-16
    ss.penalty_parameter_increase_factor = 2.0
    ss.mu_update_factor = 0.5
    ss.mu_min = 1e-16
    ss.max_penalty_parameter = 1e16
    ss.assert_checks_pass = True

    ss.print_logs = verbose
    ss.print_line_search_logs = verbose
    ss.print_search_direction_logs = verbose
    ss.print_derivative_check_logs = False

    for key, value in kwargs.items():
        try:
            setattr(ss, key, value)
        except AttributeError:
            if verbose:
                warnings.warn(
                    f"Received an undefined solver setting {key}\
                    with value {value}"
                )

    def mc(mci: sip.ModelCallbackInput) -> sip.ModelCallbackOutput:
        mco = sip.ModelCallbackOutput()

        Px = P.T @ mci.x

        mco.f = 0.5 * np.dot(Px, mci.x) + np.dot(q, mci.x)
        mco.c = A @ mci.x - b
        mco.g = G @ mci.x - h

        mco.gradient_f = Px + q
        mco.jacobian_c = A
        mco.jacobian_g = G
        mco.upper_hessian_lagrangian = upp_hess_L

        return mco

    solver = sip.Solver(ss, qs, pd, mc)

    output = solver.solve(vars)

    solution = Solution(problem)
    solution.extras = {"sip_output": output, "sip_vars": vars}
    solution.found = output.exit_status == sip.Status.SOLVED
    solution.obj = 0.5 * np.dot(P.T @ vars.x, vars.x) + np.dot(q, vars.x)
    solution.x = vars.x
    solution.y = vars.y
    if h is not None and vars.z is not None:
        z_sip = np.array(vars.z)
        z, z_box = split_dual_linear_box(z_sip, lb, ub)
        solution.z = z
        solution.z_box = z_box
    return solution


def sip_solve_qp(
    P: Union[np.ndarray, spa.csc_matrix],
    q: np.ndarray,
    G: Optional[Union[np.ndarray, spa.csc_matrix]] = None,
    h: Optional[np.ndarray] = None,
    A: Optional[Union[np.ndarray, spa.csc_matrix]] = None,
    b: Optional[np.ndarray] = None,
    lb: Optional[np.ndarray] = None,
    ub: Optional[np.ndarray] = None,
    initvals: Optional[np.ndarray] = None,
    allow_non_psd_P: bool = False,
    verbose: bool = False,
    **kwargs,
) -> Optional[np.ndarray]:
    r"""Solve a quadratic program using SIP.

    The quadratic program is defined as:

    .. math::

        \begin{split}\begin{array}{ll}
        \underset{\mbox{minimize}}{x} &
            \frac{1}{2} x^T P x + q^T x \\
        \mbox{subject to}
            & G x \leq h                \\
            & A x = b                   \\
            & lb \leq x \leq ub
        \end{array}\end{split}

    It is solved using `SIP
    <https://github.com/joaospinto/sip>`__.

    Parameters
    ----------
    P :
        Positive semidefinite cost matrix.
    q :
        Cost vector.
    G :
        Linear inequality constraint matrix.
    h :
        Linear inequality constraint vector.
    A :
        Linear equality constraint matrix.
    b :
        Linear equality constraint vector.
    lb :
        Lower bound constraint vector.
    ub :
        Upper bound constraint vector.
    verbose :
        Set to `True` to print out extra information.
    initvals :
        Warm-start guess vector for the primal solution.

    Returns
    -------
    :
        Primal solution to the QP, if found, otherwise ``None``.
    """
    problem = Problem(P, q, G, h, A, b, lb, ub)
    solution = sip_solve_problem(problem, initvals, verbose, **kwargs)
    return solution.x if solution.found else None
