#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Copyright (C) 2023 Inria
#
# This file is part of qpsolvers.
#
# qpsolvers is free software: you can redistribute it and/or modify it under
# the terms of the GNU Lesser General Public License as published by the Free
# Software Foundation, either version 3 of the License, or (at your option) any
# later version.
#
# qpsolvers is distributed in the hope that it will be useful, but WITHOUT ANY
# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
# A PARTICULAR PURPOSE.  See the GNU Lesser General Public License for more
# details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with qpsolvers. If not, see <http://www.gnu.org/licenses/>.

"""Solver interface for `Clarabel.rs`_.

.. _Clarabel.rs: https://github.com/oxfordcontrol/Clarabel.rs

Clarabel.rs is a Rust implementation of an interior point numerical solver for
convex optimization problems using a novel homogeneous embedding. A paper
describing the Clarabel solver algorithm and implementation will be forthcoming
soon (retrieved: 2023-02-06). Until then, the authors ask that you cite its
documentation if you have found Clarabel.rs useful in your work.
"""

import warnings
from typing import Optional, Union

import clarabel
import numpy as np
import scipy.sparse as spa

from ..conversions import (
    ensure_sparse_matrices,
    linear_from_box_inequalities,
    split_dual_linear_box,
)
from ..problem import Problem
from ..solution import Solution
from ..solve_unconstrained import solve_unconstrained


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

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

    Returns
    -------
    :
        Solution to the QP, if found, otherwise ``None``.

    Notes
    -----
    Keyword arguments are forwarded as options to Clarabel.rs. For instance, we
    can call ``clarabel_solve_qp(P, q, G, h, u, tol_feas=1e-6)``. Clarabel
    options include the following:

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

       * - Name
         - Description
       * - ``max_iter``
         - Maximum number of iterations.
       * - ``time_limit``
         - Time limit for solve run in seconds (can be fractional).
       * - ``tol_gap_abs``
         - absolute duality-gap tolerance
       * - ``tol_gap_rel``
         - relative duality-gap tolerance
       * - ``tol_feas``
         - feasibility check tolerance (primal and dual)

    Check out the `API reference
    <https://oxfordcontrol.github.io/ClarabelDocs/stable/api_settings/#Clarabel.Settings>`_
    for details.

    Lower values for absolute or relative tolerances yield more precise
    solutions at the cost of computation time. See *e.g.* [tolerances]_ for an
    overview of solver tolerances.
    """
    if initvals is not None and verbose:
        warnings.warn("Clarabel: warm-start values are ignored")
    P, q, G, h, A, b, lb, ub = problem.unpack()
    P, G, A = ensure_sparse_matrices(P, G, A)
    if lb is not None or ub is not None:
        G, h = linear_from_box_inequalities(G, h, lb, ub, use_sparse=True)

    cones = []
    A_list = []
    b_list = []
    if A is not None and b is not None:
        A_list.append(A)
        b_list.append(b)
        cones.append(clarabel.ZeroConeT(b.shape[0]))
    if G is not None and h is not None:
        A_list.append(G)
        b_list.append(h)
        cones.append(clarabel.NonnegativeConeT(h.shape[0]))
    if not A_list:
        return solve_unconstrained(problem)

    settings = clarabel.DefaultSettings()
    settings.verbose = verbose
    for key, value in kwargs.items():
        setattr(settings, key, value)

    A_stack = spa.vstack(A_list, format="csc")
    b_stack = np.concatenate(b_list)
    solver = clarabel.DefaultSolver(P, q, A_stack, b_stack, cones, settings)
    result = solver.solve()

    solution = Solution(problem)
    solution.obj = result.obj_val
    solution.extras = {
        "s": result.s,
        "status": result.status,
        "solve_time": result.solve_time,
    }

    solution.found = result.status == clarabel.SolverStatus.Solved
    if not solution.found:
        warnings.warn(f"Clarabel.rs terminated with status {result.status}")

    solution.x = np.array(result.x)
    meq = A.shape[0] if A is not None else 0
    solution.y = result.z[:meq] if meq > 0 else np.empty((0,))
    if G is not None:
        z, z_box = split_dual_linear_box(np.array(result.z[meq:]), lb, ub)
        solution.z = z
        solution.z_box = z_box
    else:  # G is None
        solution.z = np.empty((0,))
        solution.z_box = np.empty((0,))
    return solution


def clarabel_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,
    verbose: bool = False,
    **kwargs,
) -> Optional[np.ndarray]:
    r"""Solve a quadratic program using Clarabel.rs.

    The quadratic program is defined as:

    .. math::

        \begin{split}\begin{array}{ll}
            \underset{x}{\mbox{minimize}} &
                \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 `Clarabel.rs`_.

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

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