import alpaqa
import numpy as np

class MyProblem:
    def __init__(self):
        self.num_variables = 3 # Number of variables
        self.num_constraints = 2 # Number of constraints
    def eval_projecting_difference_constraints(self, z: np.ndarray, e: np.ndarray) -> None: ...
    def eval_projection_multipliers(self, y: np.ndarray, M: float) -> None: ...
    def eval_proximal_gradient_step(self, γ: float, x: np.ndarray, grad_ψ: np.ndarray, x_hat: np.ndarray, p: np.ndarray) -> float: ...
    def eval_inactive_indices_res_lna(self, γ: float, x: np.ndarray, grad_ψ: np.ndarray, J: np.ndarray) -> int: ...
    def eval_objective(self, x: np.ndarray) -> float: ...
    def eval_objective_gradient(self, x: np.ndarray, grad_fx: np.ndarray) -> None: ...
    def eval_constraints(self, x: np.ndarray, gx: np.ndarray) -> None: ...
    def eval_constraints_gradient_product(self, x: np.ndarray, y: np.ndarray, grad_gxy: np.ndarray) -> None: ...
    def eval_grad_gi(self, x: np.ndarray, i: int, grad_gi: np.ndarray) -> None: ...
    def eval_lagrangian_hessian_product(self, x: np.ndarray, y: np.ndarray, scale: float, v: np.ndarray, Hv: np.ndarray) -> None: ...
    def eval_augmented_lagrangian_hessian_product(self, x: np.ndarray, y: np.ndarray, Σ: np.ndarray, scale: float, v: np.ndarray, Hv: np.ndarray) -> None: ...
    def eval_objective_and_gradient(self, x: np.ndarray, grad_fx: np.ndarray) -> float: ...
    def eval_objective_and_constraints(self, x: np.ndarray, g: np.ndarray) -> float: ...
    def eval_objective_gradient_and_constraints_gradient_product(self, x: np.ndarray, y: np.ndarray, grad_f: np.ndarray, grad_gxy: np.ndarray) -> None: ...
    def eval_lagrangian_gradient(self, x: np.ndarray, y: np.ndarray, grad_L: np.ndarray, work_n: np.ndarray) -> None: ...
    def eval_augmented_lagrangian(self, x: np.ndarray, y: np.ndarray, Σ: np.ndarray, ŷ: np.ndarray) -> float: ...
    def eval_augmented_lagrangian_gradient(self, x: np.ndarray, y: np.ndarray, Σ: np.ndarray, grad_ψ: np.ndarray, work_n: np.ndarray, work_m: np.ndarray) -> None: ...
    def eval_augmented_lagrangian_and_gradient(self, x: np.ndarray, y: np.ndarray, Σ: np.ndarray, grad_ψ: np.ndarray, work_n: np.ndarray, work_m: np.ndarray) -> float: ...
    def get_variable_bounds(self) -> alpaqa.Box: ...
    def get_general_bounds(self) -> alpaqa.Box: ...
    def check(self): ...

def test_problem_obj():
    problem = alpaqa.Problem(MyProblem())

    assert problem.num_variables == 3
    assert problem.num_constraints == 2
    assert problem.provides_eval_inactive_indices_res_lna()
    assert not problem.provides_eval_constraints_jacobian()
    assert not problem.provides_get_constraints_jacobian_sparsity()
    assert problem.provides_eval_grad_gi()
    assert problem.provides_eval_lagrangian_hessian_product()
    assert not problem.provides_eval_lagrangian_hessian()
    assert not problem.provides_get_lagrangian_hessian_sparsity()
    assert problem.provides_eval_augmented_lagrangian_hessian_product()
    assert not problem.provides_eval_augmented_lagrangian_hessian()
    assert not problem.provides_get_augmented_lagrangian_hessian_sparsity()
    assert problem.provides_eval_objective_and_gradient()
    assert problem.provides_eval_objective_and_constraints()
    assert problem.provides_eval_objective_gradient_and_constraints_gradient_product()
    assert problem.provides_eval_lagrangian_gradient()
    assert problem.provides_eval_augmented_lagrangian()
    assert problem.provides_eval_augmented_lagrangian_gradient()
    assert problem.provides_eval_augmented_lagrangian_and_gradient()
    assert problem.provides_get_variable_bounds()
    assert problem.provides_get_general_bounds()
    assert problem.provides_check()
