#ifndef POLY_UTILS_OPERATOR_H
#define POLY_UTILS_OPERATOR_H

#include "mfem.hpp"
#include "4DSTARTypes.h"
#include <memory>

#include "probe.h"

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 SSE::MFEMArrayPair& theta_ess_tdofs, const SSE::MFEMArrayPair& phi_ess_tdofs);
    void SetEssentialTrueDofs(const SSE::MFEMArrayPairSet& ess_tdof_pair_set);

    SSE::MFEMArrayPairSet GetEssentialTrueDofs() const;

    bool isFinalized() const { return m_isFinalized; }

    void finalize();

    const mfem::Array<int>& GetBlockOffsets() const { return m_blockOffsets; }

    const mfem::BlockOperator &GetJacobianOperator() const;

    // Get the diagonals of -M, -Q, and D. I use J01, J10, and J11 for those here since the diagonals
    // are not just from the matrixes, but are the scaled versions in the jacobian specificially (-1*M & -1*Q)
    mfem::Vector GetJ00Diag() const;
    mfem::Vector GetJ01Diag() const;
    mfem::Vector GetJ10diag() const;
    mfem::Vector GetJ11diag() const;

private:
    Probe::LogManager& m_logManager = Probe::LogManager::getInstance();
    quill::Logger* m_logger = m_logManager.getLogger("log");
    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;

    SSE::MFEMArrayPair m_theta_ess_tdofs;
    SSE::MFEMArrayPair m_phi_ess_tdofs;

    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;

    // TODO I think these need to be calculated in the GetGradient every time since they will always change
    mutable std::unique_ptr<mfem::SparseMatrix> m_invSchurCompliment;
    mutable std::unique_ptr<mfem::SparseMatrix> m_invNonlinearJacobian;

    /*
     * The schur preconditioner has the form
     *
     * | df/drtheta^-1 0    |
     * | 0             S^-1 |
     *
     * Where S is the Schur compliment of the system
     *
     */

    // TODO: I have not combined these parts yet and they need to be combined
    mutable std::unique_ptr<mfem::SparseMatrix> m_schurPreconditioner;

    bool m_isFinalized = false;

private:
    void updateInverseNonlinearJacobian(const mfem::Operator &grad) const;
    void updateInverseSchurCompliment() const;
};


#endif // POLY_UTILS_OPERATOR_H