//
// Created by xtkuang on 2025/7/7.
//

#ifndef QP_SOLVER_H
#define QP_SOLVER_H

#pragma once

#include <exception>
#include <memory>
#include <optional>

#include <eigen3/Eigen/Core>
#include <OsqpEigen/OsqpEigen.h>

namespace cmvr::math {

class QPSolverException : public std::exception {
public:
    static constexpr unsigned int kStatusOffset = 100;

    explicit QPSolverException(int error_code);

    const char* what() const noexcept override;

    int code() const noexcept;

    static std::string GenerateMessage(int code);

private:
    int error_code_;
    std::string message_;
};

class QPSolverImpl {
public:
    QPSolverImpl() = default;
    ~QPSolverImpl() = default;

    void SetupImpl(int n_var, int n_const, double time_limit);
    void InitFunctionImpl();
    void AddCostFunctionImpl(const Eigen::MatrixXd &A, const Eigen::VectorXd &b);
    void SetCostFunctionImpl(const Eigen::MatrixXd &A, const Eigen::VectorXd &b);
    void SetConstraintsFunctionImpl(const Eigen::MatrixXd &A, const Eigen::VectorXd &lb, const Eigen::VectorXd &ub);
    void SetPrimalVariableImpl(const Eigen::VectorXd &pv);
    void ResetIsFirstImpl();
    Eigen::VectorXd SolveImpl();
    Eigen::MatrixXd GetACostImpl() const;
    Eigen::VectorXd GetBCostImpl() const;
    Eigen::MatrixXd GetAConstImpl() const;
    Eigen::VectorXd GetLowerBoundImpl() const;
    Eigen::VectorXd GetUpperBoundImpl() const;

private:
    OsqpEigen::Solver solver_;
    int n_var_{};
    int n_const_{};
    int err_code_{};

    int is_first_{};

    int n_hessian_element_{};

    Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> A_cost_;
    Eigen::Matrix<double, Eigen::Dynamic, 1> b_cost_;

    Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> A_const_;
    Eigen::Matrix<double, Eigen::Dynamic, 1> lb_;
    Eigen::Matrix<double, Eigen::Dynamic, 1> ub_;

    Eigen::SparseMatrix<double> hessian_;
    Eigen::Matrix<double, Eigen::Dynamic, 1> gradient_;

    Eigen::SparseMatrix<double> linearMatrix_;
    Eigen::Matrix<double, Eigen::Dynamic, 1> lowerBound_;
    Eigen::Matrix<double, Eigen::Dynamic, 1> upperBound_;

    Eigen::Matrix<double, Eigen::Dynamic, 1> primal_variable_for_warmstart_;
};


class QPSolver {
public:
    QPSolver();

    ~QPSolver();

    void Setup(int n_var, int n_const, double time_limit = 2e-3);

    void InitFunction();

    void AddCostFunction(const Eigen::MatrixXd& A, const Eigen::VectorXd& b);

    void SetCostFunction(const Eigen::MatrixXd& A, const Eigen::VectorXd& b);

    void SetConstraintsFunction(const Eigen::MatrixXd& A, const Eigen::VectorXd& lb, const Eigen::VectorXd& ub);

    void SetPrimalVariable(const Eigen::VectorXd& pv);

    void ResetIsFirst();

    /**
     * Solve the QP problem
     * @return Solution
     * @throw QPSolverException
     */
    Eigen::VectorXd Solve();

    Eigen::MatrixXd GetACost() const;  // NOLINT

    Eigen::VectorXd GetBCost() const;  // NOLINT

    Eigen::MatrixXd GetAConst() const;  // NOLINT

    Eigen::VectorXd GetLowerBound() const;  // NOLINT

    Eigen::VectorXd GetUpperBound() const;  // NOLINT

private:
    std::unique_ptr<QPSolverImpl> impl_;
};

}

#endif //QP_SOLVER_H