/* ***********************************************************************
//
//   Copyright (C) 2025 -- The 4D-STAR Collaboration
//   File Author: Emily Boudreaux
//   Last Modified: April 21, 2025
//
//   4DSSE is free software; you can use it and/or modify
//   it under the terms and restrictions the GNU General Library Public
//   License version 3 (GPLv3) as published by the Free Software Foundation.
//
//   4DSSE is distributed in the hope that it will be useful,
//   but WITHOUT ANY WARRANTY; without even the implied warranty of
//   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
//   See the GNU Library General Public License for more details.
//
//   You should have received a copy of the GNU Library General Public License
//   along with this software; if not, write to the Free Software
//   Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
//
// *********************************************************************** */
#pragma once

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

#include "probe.h"

class SchurCompliment final : public mfem::Operator {
public:
    SchurCompliment(const mfem::SparseMatrix &QOp, const mfem::SparseMatrix &DOp, const mfem::SparseMatrix &MOp, const mfem::Solver &GradInvOp);
    SchurCompliment(const mfem::SparseMatrix &QOp, const mfem::SparseMatrix &DOp, const mfem::SparseMatrix &MOp);
    ~SchurCompliment() override = default;
    void Mult(const mfem::Vector &x, mfem::Vector &y) const override;
    void SetOperator(const mfem::SparseMatrix &QOp, const mfem::SparseMatrix &DOp, const mfem::SparseMatrix &MOp, const mfem::Solver &
                     GradInvOp);
    void updateInverseNonlinearJacobian(const mfem::Solver &gradInv);

private:
    void updateConstantTerms(const mfem::SparseMatrix &QOp, const mfem::SparseMatrix &DOp, const mfem::SparseMatrix &MOp);

private:

    // Note that these are not owned by this class
    const mfem::SparseMatrix* m_QOp = nullptr;
    const mfem::SparseMatrix* m_DOp = nullptr;
    const mfem::SparseMatrix* m_MOp = nullptr;
    const mfem::Solver* m_GradInvOp = nullptr;
    int m_nPhi = 0;
    int m_nTheta = 0;

    mutable std::unique_ptr<mfem::SparseMatrix> m_matrixForm;
};

class GMRESInverter final : public mfem::Operator {
public:
    explicit GMRESInverter(const SchurCompliment& op);
    ~GMRESInverter() override = default;
    void Mult(const mfem::Vector &x, mfem::Vector &y) const override;

private:
    const SchurCompliment& m_op;
    mfem::GMRESSolver m_solver;
};

class PolytropeOperator final : 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,
        const double index);
    ~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::Vector &initTheta);

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

    const mfem::BlockOperator &GetJacobianOperator() const;

    mfem::BlockDiagonalPreconditioner &GetPreconditioner() 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;

    // These are used to store the matrix representations
    // for the bi-linear forms. This might seem counterintuitive
    // for a matrix-free approach. However, these will be computed
    // regardless due to MFEM's implementation of these operators.
    // Further since these do not change it is not a performance issue.

    // We need to store these separately because the jacobian and preconditioner
    // must be computed from the boundary aware operators (which will be these
    // matrices) whereas the residuals must be computed from the raw, physical,
    // operators
    std::unique_ptr<mfem::SparseMatrix> m_Mmat;
    std::unique_ptr<mfem::SparseMatrix> m_Qmat;
    std::unique_ptr<mfem::SparseMatrix> m_Dmat;

    const mfem::Array<int> m_blockOffsets;

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

    std::unique_ptr<mfem::ScaledOperator> m_negM_mat;
    std::unique_ptr<mfem::ScaledOperator> m_negQ_mat;
    mutable std::unique_ptr<mfem::BlockOperator> m_jacobian;

    mutable std::unique_ptr<SchurCompliment> m_schurCompliment;
    mutable std::unique_ptr<GMRESInverter> m_invSchurCompliment;
    mutable std::unique_ptr<mfem::Solver> m_invNonlinearJacobian;

    mutable std::unique_ptr<mfem::BlockDiagonalPreconditioner> m_schurPreconditioner;

    bool m_isFinalized = false;

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