#ifndef POLY_UTILS_OPERATOR_H
#define POLY_UTILS_OPERATOR_H

#include "mfem.hpp"
#include <memory>

class PolytropeOperator : public mfem::Operator {
public:
    PolytropeOperator(
        std::unique_ptr<mfem::MixedBilinearForm> M,
        std::unique_ptr<mfem::MixedBilinearForm> Q,
        std::unique_ptr<mfem::BilinearForm> D,
        std::unique_ptr<mfem::NonlinearForm> f,
        const mfem::Array<int> &blockOffsets);
    ~PolytropeOperator() override = default;

    void Mult(const mfem::Vector &x, mfem::Vector &y) const override;
    mfem::Operator& GetGradient(const mfem::Vector &x) const override;

    void SetEssentialTrueDofs(const mfem::Array<int> &theta_ess_tofs,
                            const mfem::Array<int> &phi_ess_tofs);

private:
    std::unique_ptr<mfem::MixedBilinearForm> m_M;
    std::unique_ptr<mfem::MixedBilinearForm> m_Q;
    std::unique_ptr<mfem::BilinearForm> m_D;
    std::unique_ptr<mfem::NonlinearForm> m_f;

    const mfem::Array<int> &m_blockOffsets;

    mfem::Array<int> m_theta_ess_tofs;
    mfem::Array<int> m_phi_ess_tofs;

    std::unique_ptr<mfem::SparseMatrix> m_Mmat;
    std::unique_ptr<mfem::SparseMatrix> m_Qmat;
    std::unique_ptr<mfem::SparseMatrix> m_Dmat;


    std::unique_ptr<mfem::ScaledOperator> m_negM_op;
    std::unique_ptr<mfem::ScaledOperator> m_negQ_op;
    mutable std::unique_ptr<mfem::BlockOperator> m_jacobian;
};


#endif // POLY_UTILS_OPERATOR_H