/* ***********************************************************************
//
//   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"

namespace serif::polytrope {

/**
 * @brief Represents the Schur complement operator used in the solution process.
 *
 * This class computes S = D - Q * GradInv * M.
 */
class SchurCompliment final : public mfem::Operator {
public:
    /**
     * @brief Constructs a SchurCompliment operator.
     * @param QOp The Q matrix operator.
     * @param DOp The D matrix operator.
     * @param MOp The M matrix operator.
     * @param GradInvOp The inverse of the gradient operator.
     */
    SchurCompliment(
        const mfem::SparseMatrix &QOp,
        const mfem::SparseMatrix &DOp,
        const mfem::SparseMatrix &MOp,
        const mfem::Solver &GradInvOp
    );

    /**
     * @brief Constructs a SchurCompliment operator without the inverse gradient initially.
     * The inverse gradient must be set later using updateInverseNonlinearJacobian.
     * @param QOp The Q matrix operator.
     * @param DOp The D matrix operator.
     * @param MOp The M matrix operator.
     */
    SchurCompliment(
        const mfem::SparseMatrix &QOp,
        const mfem::SparseMatrix &DOp,
        const mfem::SparseMatrix &MOp
    );

    /**
     * @brief Destructor.
     */
    ~SchurCompliment() override = default;

    /**
     * @brief Applies the Schur complement operator: y = (D - Q * GradInv * M) * x.
     * @param x The input vector.
     * @param y The output vector.
     */
    void Mult(const mfem::Vector &x, mfem::Vector &y) const override;

    /**
     * @brief Sets all operators for the Schur complement.
     * @param QOp The Q matrix operator.
     * @param DOp The D matrix operator.
     * @param MOp The M matrix operator.
     * @param GradInvOp The inverse of the gradient operator.
     */
    void SetOperator(const mfem::SparseMatrix &QOp, const mfem::SparseMatrix &DOp, const mfem::SparseMatrix &MOp, const mfem::Solver &GradInvOp);

    /**
     * @brief Updates the inverse of the nonlinear Jacobian (GradInv).
     * @param gradInv The new inverse gradient solver.
     */
    void updateInverseNonlinearJacobian(const mfem::Solver &gradInv);

private:
    /**
     * @brief Updates the constant matrix terms (Q, D, M).
     * @param QOp The Q matrix operator.
     * @param DOp The D matrix operator.
     * @param MOp The M matrix operator.
     */
    void updateConstantTerms(const mfem::SparseMatrix &QOp, const mfem::SparseMatrix &DOp, const mfem::SparseMatrix &MOp);

private:
    // Pointers to external operators (not owned by this class)
    const mfem::SparseMatrix* m_QOp = nullptr;         ///< Pointer to the Q matrix operator.
    const mfem::SparseMatrix* m_DOp = nullptr;         ///< Pointer to the D matrix operator.
    const mfem::SparseMatrix* m_MOp = nullptr;         ///< Pointer to the M matrix operator.
    const mfem::Solver* m_GradInvOp = nullptr;         ///< Pointer to the inverse of the gradient operator.

    // Dimensions
    int m_nPhi = 0;                                    ///< Dimension related to the phi variable (typically number of rows/cols of D).
    int m_nTheta = 0;                                  ///< Dimension related to the theta variable (typically number of rows of M).

    // Owned resources
    mutable std::unique_ptr<mfem::SparseMatrix> m_matrixForm; ///< Optional: if a matrix representation is ever explicitly formed.
};

/**
 * @brief Provides an approximate inverse of the SchurCompliment operator using GMRES.
 */
class GMRESInverter final : public mfem::Operator {
public:
    /**
     * @brief Constructs a GMRESInverter.
     * @param op The SchurCompliment operator to invert.
     */
    explicit GMRESInverter(const SchurCompliment& op);

    /**
     * @brief Destructor.
     */
    ~GMRESInverter() override = default;

    /**
     * @brief Applies the GMRES solver to approximate op^-1 * x.
     * @param x The input vector.
     * @param y The output vector (approximation of op^-1 * x).
     */
    void Mult(const mfem::Vector &x, mfem::Vector &y) const override;

private:
    const SchurCompliment& m_op; ///< Reference to the SchurCompliment operator to be inverted.
    mfem::GMRESSolver m_solver;   ///< GMRES solver instance.
};

/**
 * @brief Represents the coupled nonlinear operator for the polytropic system.
 *
 * This operator defines the system F(X) = 0, where X = [θ, φ]^T,
 * and F(X) = [ f(θ) + Mφ ]
 *            [ Dφ  - Qθ ].
 * It handles essential boundary conditions and the construction of the Jacobian.
 */
class PolytropeOperator final : public mfem::Operator {
public:
    /**
     * @brief Constructs the PolytropeOperator.
     * @param M The M bilinear form (coupling θ and φ).
     * @param Q The Q bilinear form (coupling φ and θ).
     * @param D The D bilinear form (acting on φ).
     * @param f The nonlinear form f(θ).
     * @param blockOffsets Offsets defining the blocks for θ and φ variables.
     */
    PolytropeOperator(
        std::unique_ptr<mfem::MixedBilinearForm> M,
        std::unique_ptr<mfem::MixedBilinearForm> Q,
        std::unique_ptr<mfem::BilinearForm> D,
        std::unique_ptr<mfem::BilinearForm> S,
        std::unique_ptr<mfem::NonlinearForm> f,
        const mfem::Array<int> &blockOffsets);

    /**
     * @brief Destructor.
     */
    ~PolytropeOperator() override = default;

    /**
     * @brief Applies the PolytropeOperator: y = F(x).
     * This computes the residual of the nonlinear system.
     * @param xFree The vector of free (non-essential) degrees of freedom.
     * @param yFree The resulting residual vector corresponding to free DOFs.
     */
    void Mult(const mfem::Vector &xFree, mfem::Vector &yFree) const override;

    /**
     * @brief Computes the Jacobian of the PolytropeOperator at a given state xFree.
     * The Jacobian is of the form:
     * J = [ G  M ]
     *     [ -Q D ]
     * where G is the gradient of f(θ).
     * @param xFree The vector of free DOFs at which to evaluate the gradient.
     * @return A reference to the Jacobian operator.
     */
    mfem::Operator& GetGradient(const mfem::Vector &xFree) const override;

    /**
     * @brief Finalizes the operator setup.
     * This must be called after setting essential true DOFs and before using Mult or GetGradient.
     * It constructs reduced matrices, block offsets, and populates free DOFs.
     * @param initTheta Initial guess for θ, used to evaluate the nonlinear form if necessary during setup.
     */
    void finalize(const mfem::Vector &initTheta);

    /**
     * @brief Checks if the operator has been finalized.
     * @return True if finalize() has been successfully called, false otherwise.
     */
    bool isFinalized() const { return m_isFinalized; }

    /**
     * @brief Sets the essential true degrees of freedom for both θ and φ variables.
     * Marks the operator as not finalized.
     * @param theta_ess_tdofs Pair of arrays: (indices of essential DOFs for θ, values at these DOFs).
     * @param phi_ess_tdofs Pair of arrays: (indices of essential DOFs for φ, values at these DOFs).
     */
    void set_essential_true_dofs(const serif::types::MFEMArrayPair& theta_ess_tdofs, const serif::types::MFEMArrayPair& phi_ess_tdofs);

    /**
     * @brief Sets the essential true degrees of freedom using a pair of pairs.
     * Marks the operator as not finalized.
     * @param ess_tdof_pair_set A pair containing the essential TDOF pairs for theta and phi.
     */
    void set_essential_true_dofs(const serif::types::MFEMArrayPairSet& ess_tdof_pair_set);


    /**
     * @brief Reconstructs the full state vector (including essential DOFs) from a reduced state vector (free DOFs).
     * @param reducedState The vector containing only the free degrees of freedom.
     * @return Constant reference to the internal full state vector, updated with the reducedState.
     */
    [[nodiscard]] const mfem::Vector &reconstruct_full_state_vector(const mfem::Vector &reducedState) const;

    /**
     * @breif Reconstruct the full state vector (including essential DOFs) from a reduced state vector (free DOFs) as well as the block offsets.
     * @param reducedState The vector containing only the free degrees of freedom.
     * @return Constant reference to the internal full state vector, updated with the reducedState as a block vector.
     */
    [[nodiscard]] const mfem::BlockVector reconstruct_full_block_state_vector(const mfem::Vector &reducedState) const;

    /**
     * @brief Populates the internal array of free (non-essential) degree of freedom indices.
     * This is called during finalize().
     */
    void populate_free_dof_array();

    /// --- Getters for key internal state and operators ---
    /**
     * @brief Gets the Jacobian operator.
     * Asserts that the operator is finalized and the Jacobian has been computed.
     * @return Constant reference to the block Jacobian operator.
     */
    const mfem::BlockOperator &get_jacobian_operator() const;

    /**
     * @brief Gets the block diagonal preconditioner for the Schur complement system.
     * Asserts that the operator is finalized and the preconditioner has been computed.
     * @return Reference to the block diagonal preconditioner.
     */
    mfem::BlockDiagonalPreconditioner &get_preconditioner() const;


    /// --- Getters for information on internal state ---
    /**
     * @brief Gets the full state vector, including essential DOFs.
     * @return Constant reference to the internal state vector.
     */
    const mfem::Array<int>& get_free_dofs() const { return m_freeDofs; } ///< Getter for the free DOFs array.

    /**
     * @brief Gets the size of the reduced system (number of free DOFs).
     * Asserts that the operator is finalized.
     * @return The total number of free degrees of freedom.
     */
    int get_reduced_system_size() const;

    /**
     * @brief Gets the currently set essential true degrees of freedom.
     * @return A pair containing the essential TDOF pairs for theta and phi.
     */
    serif::types::MFEMArrayPairSet get_essential_true_dofs() const;

    /**
     * @brief Gets the block offsets for the full (unreduced) system.
     * @return Constant reference to the array of block offsets.
     */
    const mfem::Array<int>& get_block_offsets() const { return m_blockOffsets; }

    /**
     * @brief Gets the block offsets for the reduced system (after eliminating essential DOFs).
     * @return Constant reference to the array of reduced block offsets.
     */
    const mfem::Array<int>& get_reduced_block_offsets() const {return m_reducedBlockOffsets; }

private:
    // --- Logging ---
    serif::probe::LogManager& m_logManager = serif::probe::LogManager::getInstance(); ///< Reference to the global log manager.
    quill::Logger* m_logger = m_logManager.getLogger("log");            ///< Pointer to the specific logger instance.

    // --- Input Bilinear/Nonlinear Forms (owned) ---
    std::unique_ptr<mfem::MixedBilinearForm> m_M; ///< Bilinear form M, coupling θ and φ.
    std::unique_ptr<mfem::MixedBilinearForm> m_Q; ///< Bilinear form Q, coupling φ and θ.
    std::unique_ptr<mfem::BilinearForm> m_D;      ///< Bilinear form D, acting on φ.
    std::unique_ptr<mfem::BilinearForm> m_S;      ///< Bilinear form S, used for least squares stabilization.
    std::unique_ptr<mfem::NonlinearForm> m_f;     ///< Nonlinear form f, acting on θ.

    // --- Full Matrix Representations (owned, derived from forms) ---
    std::unique_ptr<mfem::SparseMatrix> m_Mmat;   ///< Sparse matrix representation of M.
    std::unique_ptr<mfem::SparseMatrix> m_Qmat;   ///< Sparse matrix representation of Q.
    std::unique_ptr<mfem::SparseMatrix> m_Dmat;   ///< Sparse matrix representation of D.
    std::unique_ptr<mfem::SparseMatrix> m_Smat;

    // --- Reduced Matrix Representations (owned, after eliminating essential DOFs) ---
    std::unique_ptr<mfem::SparseMatrix> m_MReduced; ///< Reduced M matrix (free DOFs only).
    std::unique_ptr<mfem::SparseMatrix> m_QReduced; ///< Reduced Q matrix (free DOFs only).
    std::unique_ptr<mfem::SparseMatrix> m_DReduced; ///< Reduced D matrix (free DOFs only).
    std::unique_ptr<mfem::SparseMatrix> m_SReduced; ///< Reduced S matrix (free DOFs only, used for least squares stabilization).
    mutable std::unique_ptr<mfem::SparseMatrix> m_gradReduced; ///< Reduced gradient operator (G) for the nonlinear part f(θ).

    // --- Scaled Reduced Matrix Representations (owned, after eliminating essential DOFs and scaling by stabilization coefficients) ---
    std::unique_ptr<mfem::SparseMatrix> m_MScaledReduced; ///< Scaled M matrix (free DOFs only, scaled by stabilization coefficient).
    std::unique_ptr<mfem::SparseMatrix> m_QScaledReduced; ///< Scaled Q matrix (free DOFs only, scaled by stabilization coefficient).
    std::unique_ptr<mfem::SparseMatrix> m_DScaledReduced; ///< Scaled D matrix (free DOFs only, scaled by stabilization coefficient).
    std::unique_ptr<mfem::SparseMatrix> m_ScaledSReduced; ///< Scaled S matrix (free DOFs only, scaled by stabilization coefficient).


    // --- Stabilization Coefficient --- (Perhapses this should be a user parameter...) // TODO: Sort out why this is negative (sign error?)
    static constexpr double m_stabilizationCoefficient = -2.0;       ///< Stabilization coefficient for the system, used to more tightly couple ∇θ and φ.
    static constexpr double m_IncrementedStabilizationCoefficient = 1.0 + m_stabilizationCoefficient; ///< 1 + Stabilization coefficient for the system, used to more tightly couple ∇θ and φ.

    // --- State Vectors and DOF Management ---
    mutable mfem::Vector m_state;                 ///< Full state vector [θ, φ]^T, including essential DOFs.
    mfem::Array<int> m_freeDofs;                  ///< Array of indices for free (non-essential) DOFs.

    // --- Block Offsets ---
    const mfem::Array<int> m_blockOffsets;        ///< Block offsets for the full system [0, size(θ), size(θ)+size(φ)].
    mfem::Array<int> m_reducedBlockOffsets;       ///< Block offsets for the reduced system (free DOFs).

    // --- Essential Boundary Conditions ---
    serif::types::MFEMArrayPair m_theta_ess_tdofs;         ///< Essential true DOFs for θ (indices and values).
    serif::types::MFEMArrayPair m_phi_ess_tdofs;           ///< Essential true DOFs for φ (indices and values).

    // --- Jacobian and Preconditioner Components (owned, mutable) ---
    std::unique_ptr<mfem::ScaledOperator> m_negQ_mat; ///< Scaled operator for -Q_reduced.
    mutable std::unique_ptr<mfem::BlockOperator> m_jacobian; ///< Jacobian operator J = [G M; -Q D]_reduced.
    mutable std::unique_ptr<SchurCompliment> m_schurCompliment; ///< Schur complement S = D_reduced - Q_reduced * G_inv_reduced * M_reduced.
    mutable std::unique_ptr<GMRESInverter> m_invSchurCompliment; ///< Approximate inverse of the Schur complement.
    mutable std::unique_ptr<mfem::Solver> m_invNonlinearJacobian; ///< Solver for the inverse of the G block (gradient of f(θ)_reduced).
    mutable std::unique_ptr<mfem::BlockDiagonalPreconditioner> m_schurPreconditioner; ///< Block diagonal preconditioner for the system.

    // --- State Flags ---
    bool m_isFinalized = false;                   ///< Flag indicating if finalize() has been called.

private:
    /**
     * @brief Constructs the sparse matrix representations of M, Q, and D, and their reduced forms.
     * Called during finalize().
     */
    void construct_matrix_representations();

    /**
     * @brief Constructs the block offsets for the reduced system.
     * Called during finalize().
     */
    void construct_reduced_block_offsets();

    /**
     * @brief Constructs the constant terms of the Jacobian operator (M, -Q, D).
     * The (0,0) block (gradient of f) is set in GetGradient.
     * Called during finalize().
     */
    void construct_jacobian_constant_terms();

    /**
     * @brief Scatters the values of essential boundary conditions into the full state vector.
     * Called during finalize().
     */
    void scatter_boundary_conditions();

    /**
     * @brief Updates the solver for the inverse of the nonlinear Jacobian block (G_00).
     * @param grad The gradient operator (G_00) of the nonlinear part f(θ).
     */
    void update_inverse_nonlinear_jacobian(const mfem::Operator &grad) const;

    /**
     * @brief Updates the inverse Schur complement operator and its components.
     * This is typically called after the nonlinear Jacobian part has been updated.
     */
    void update_inverse_schur_compliment() const;

    /**
     * @brief Updates the preconditioner components.
     * This involves updating the inverse nonlinear Jacobian and then the inverse Schur complement.
     * @param grad The gradient operator (G_00) of the nonlinear part f(θ).
     */
    void update_preconditioner(const mfem::Operator &grad) const;
};

} // namespace serif::polytrope