#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# SPDX-License-Identifier: LGPL-3.0-or-later
# Copyright 2023 Inria

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

**Warm-start:** this solver interface does not support warm starting ❄️
"""

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 ..exceptions import ProblemError
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 :
        This argument is not used by Clarabel.
    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.* [Caron2022]_ for a
    primer on solver tolerances and residuals.
    """
    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("clarabel", 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:
        warnings.warn(
            "QP is unconstrained: "
            "solving with SciPy's LSQR rather than clarabel"
        )
        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)
    try:
        solver = clarabel.DefaultSolver(
            P, q, A_stack, b_stack, cones, settings
        )
    except BaseException as exn:
        # The one we want to catch is a pyo3_runtime.PanicException
        # But see https://github.com/PyO3/pyo3/issues/2880
        raise ProblemError("Solver failed to build problem") from exn

    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 :
        This argument is not used by Clarabel.
    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
