// Copyright (c) 2025 Charlie Vanaret
// Licensed under the MIT license. See LICENSE file in the project directory for details.

#ifndef UNO_PYTHONMODEL_H
#define UNO_PYTHONMODEL_H

#include <vector>
#include "../unopy.hpp"
#include "linear_algebra/SparseVector.hpp"
#include "linear_algebra/Vector.hpp"
#include "model/Model.hpp"
#include "symbolic/CollectionAdapter.hpp"
#include "tools/NumberModelEvaluations.hpp"

namespace uno {
   class PythonModel: public Model {
   public:
      explicit PythonModel(const PythonUserModel& user_model);
      ~PythonModel() override = default;

      [[nodiscard]] ProblemType get_problem_type() const override;

      // availability of linear operators
      [[nodiscard]] bool has_jacobian_operator() const override;
      [[nodiscard]] bool has_jacobian_transposed_operator() const override;
      [[nodiscard]] bool has_hessian_operator() const override;
      [[nodiscard]] bool has_hessian_matrix() const override;

      // function evaluations
      [[nodiscard]] double evaluate_objective(const Vector<double>& x) const override;
      void evaluate_constraints(const Vector<double>& x, Vector<double>& constraints) const override;

      // dense objective gradient
      void evaluate_objective_gradient(const Vector<double>& x, Vector<double>& gradient) const override;

      // sparsity patterns of Jacobian and Hessian
      void compute_jacobian_sparsity(uno_int * row_indices, uno_int * column_indices, uno_int solver_indexing,
         MatrixOrder matrix_format) const override;
      void compute_hessian_sparsity(uno_int *row_indices, uno_int *column_indices, uno_int solver_indexing) const override;

      // numerical evaluations of Jacobian and Hessian
      void evaluate_jacobian(const Vector<double>& x, double* jacobian_values) const override;
      void evaluate_lagrangian_hessian(const Vector<double>& x, double objective_multiplier, const Vector<double>& multipliers,
         double* hessian_values) const override;

      // linear operators for Jacobian-, Jacobian^T-, and Hessian-vector products
      void compute_jacobian_vector_product(const double* x, const double* vector, double* result) const override;
      void compute_jacobian_transposed_vector_product(const double* x, const double* vector, double* result) const override;
      void compute_hessian_vector_product(const double* x, const double* vector, double objective_multiplier,
         const Vector<double>& multipliers, double* result) const override;

      // purely functions
      [[nodiscard]] double variable_lower_bound(size_t variable_index) const override;
      [[nodiscard]] double variable_upper_bound(size_t variable_index) const override;
      [[nodiscard]] const SparseVector<size_t>& get_slacks() const override;
      [[nodiscard]] const Vector<size_t>& get_fixed_variables() const override;

      [[nodiscard]] double constraint_lower_bound(size_t constraint_index) const override;
      [[nodiscard]] double constraint_upper_bound(size_t constraint_index) const override;
      [[nodiscard]] const Collection<size_t>& get_equality_constraints() const override;
      [[nodiscard]] const Collection<size_t>& get_inequality_constraints() const override;
      [[nodiscard]] const Collection<size_t>& get_linear_constraints() const override;
      [[nodiscard]] const Collection<size_t>& get_nonlinear_constraints() const override;

      void initial_primal_point(Vector<double>& x) const override;
      void initial_dual_point(Vector<double>& multipliers) const override;
      void postprocess_solution(Iterate& iterate) const override;

      [[nodiscard]] size_t number_jacobian_nonzeros() const override;
      [[nodiscard]] size_t number_hessian_nonzeros() const override;

      [[nodiscard]] size_t number_model_objective_evaluations() const override;
      [[nodiscard]] size_t number_model_constraints_evaluations() const override;
      [[nodiscard]] size_t number_model_objective_gradient_evaluations() const override;
      [[nodiscard]] size_t number_model_jacobian_evaluations() const override;
      [[nodiscard]] size_t number_model_hessian_evaluations() const override;
      void reset_number_evaluations() const override;

   protected:
      const PythonUserModel& user_model;
      mutable NumberModelEvaluations number_model_evaluations{};
      const SparseVector<size_t> slacks{};
      Vector<size_t> fixed_variables{};
      const ForwardRange linear_constraints{0};
      const ForwardRange nonlinear_constraints;
      std::vector<size_t> equality_constraints;
      CollectionAdapter<std::vector<size_t>> equality_constraints_collection;
      std::vector<size_t> inequality_constraints;
      CollectionAdapter<std::vector<size_t>> inequality_constraints_collection;
   };
} // namespace

#endif // UNO_PYTHONMODEL_H