/* ***********************************************************************
//
//   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 <memory>
#include <utility>

#include "integrators.h"
#include "4DSTARTypes.h"
#include "polytropeOperator.h"
#include "config.h"
#include "meshIO.h"
#include "probe.h"
#include "quill/Logger.h"

namespace serif {
namespace polytrope {

/**
 * @brief Namespace for Lane-Emden equation related utility functions.
 *
 * Provides functions to compute coefficients and evaluate the series expansion
 * solution to the Lane-Emden equation, which describes the structure of a
 * spherically symmetric polytropic star.
 * The Lane-Emden equation is given by:
 * \f[
 * \frac{1}{\xi^2} \frac{d}{d\xi} \left( \xi^2 \frac{d\theta}{d\xi} \right) = -\theta^n
 * \f]
 * where \f$\xi\f$ is a dimensionless radius and \f$\theta\f$ is related to the density,
 * and \f$n\f$ is the polytropic index.
 */
namespace laneEmden {
    /**
     * @brief Computes the coefficient \f$a_k\f$ for the Lane-Emden series expansion.
     *
     * The series solution for \f$\theta(\xi)\f$ is given by \f$\theta(\xi) = \sum_{k=0}^{\infty} a_k \xi^k\f$.
     * The coefficients \f$a_k\f$ are determined by substituting the series into the Lane-Emden equation.
     * Specifically, \f$a_0 = 1\f$, \f$a_1 = 0\f$, and for \f$k \ge 2\f$,
     * \f$a_k = -\frac{c_{k-2,n}}{k(k+1)}\f$.
     *
     * @param k The index of the coefficient.
     * @param n The polytropic index.
     * @return The value of the coefficient \f$a_k\f$.
     * @see c(const int m, const double n)
     */
    double a (const int k, const double n);

    /**
     * @brief Computes the auxiliary coefficient \f$c_{m,n}\f$ used in determining \f$a_k\f$.
     *
     * The term \f$\theta^n\f$ in the Lane-Emden equation can also be expanded as a series
     * \f$\theta^n(\xi) = \sum_{m=0}^{\infty} c_{m,n} \xi^m\f$.
     * The coefficients \f$c_{m,n}\f$ are related to \f$a_k\f$ by:
     * \f$c_{0,n} = a_0^n\f$
     * \f$c_{m,n} = \frac{1}{m a_0} \sum_{j=1}^{m} (j n - m + j) a_j c_{m-j,n}\f$ for \f$m > 0\f$.
     *
     * @param m The index of the coefficient.
     * @param n The polytropic index.
     * @return The value of the coefficient \f$c_{m,n}\f$.
     * @see a(const int k, const double n)
     */
    double c(const int m, const double n);

    /**
     * @brief Computes the Lane-Emden function \f$\theta(\xi)\f$ using a series expansion.
     *
     * Evaluates the series \f$\theta(\xi) = \sum_{k=0}^{\text{order}-1} a_k \xi^k\f$ up to a specified order.
     * This provides an approximate solution to the Lane-Emden equation, particularly accurate for small \f$\xi\f$.
     *
     * @param xi The dimensionless radius \f$\xi\f$.
     * @param n The polytropic index.
     * @param order The number of terms to include in the series expansion.
     * @return The approximate value of \f$\theta(\xi)\f$.
     *
     * @example
     * @code
     * double xi = 0.5;
     * double n = 1.5;
     * int series_order = 10;
     * double theta_val = laneEmden::thetaSeriesExpansion(xi, n, series_order);
     * // theta_val will be an approximation of the Lane-Emden function at xi=0.5 for n=1.5
     * @endcode
     */
    double thetaSeriesExpansion(const double xi, const double n, const int order);
} // namespace laneEmden

/**
 * @brief Structure to manage the lifetime of MFEM solver objects.
 *
 * This structure ensures that the `mfem::GMRESSolver` outlives the `mfem::NewtonSolver`
 * that might use it. The `mfem::NewtonSolver` often takes a pointer or reference to a
 * linear solver, and if the linear solver is destroyed first, it can lead to dangling
 * pointers and crashes.
 *
 * @note The order of declaration of members is crucial: `solver` must be declared
 *       before `newton` to ensure it is constructed first and destroyed last.
 */
struct solverBundle {
    mfem::GMRESSolver solver; ///< The linear solver (e.g., GMRES). Must be declared first.
    mfem::NewtonSolver newton; ///< The nonlinear solver (e.g., Newton). Must be declared second.
};

/**
 * @brief Structure to hold the various bilinear and nonlinear forms for the polytrope problem.
 *
 * This structure bundles the MFEM forms that represent the discretized weak formulation
 * of the mixed variable polytropic equations.
 * The system being solved is typically:
 * \f[
 * \begin{aligned}
 *   \boldsymbol{\phi} + \nabla \theta &= \mathbf{0} \\
 *   \nabla \cdot \boldsymbol{\phi} - \theta^n &= 0
 * \end{aligned}
 * \f]
 * After integration by parts and discretization, these lead to matrix equations involving
 * the forms stored here.
 */
struct formBundle {
    /**
     * @brief Mixed bilinear form \f$ M(\psi^\theta, N^\phi) = \int_\Omega \nabla\psi^\theta \cdot N^\phi \,dV \f$.
     * This form arises from the term \f$\nabla \theta\f$ in the first equation, tested with a vector test function \f$N^\phi\f$.
     * It couples the scalar field \f$\theta\f$ (potential) with the vector field \f$\boldsymbol{\phi}\f$ (flux).
     */
    std::unique_ptr<mfem::MixedBilinearForm> M;

    /**
     * @brief Mixed bilinear form \f$ Q(\psi^\phi, N^\theta) = \int_\Omega \psi^\phi \cdot \nabla N^\theta \,dV \f$.
     * This form arises from the term \f$\boldsymbol{\phi}\f$ in the first equation, tested with a scalar test function \f$N^\theta\f$,
     * after integration by parts of the \f$\nabla \theta\f$ term.
     * It can also be seen as related to the transpose of a discrete gradient operator.
     */
    std::unique_ptr<mfem::MixedBilinearForm> Q;

    /**
     * @brief Bilinear form \f$ D(\psi^\phi, N^\phi) = \int_\Omega \psi^\phi \cdot N^\phi \,dV \f$.
     * This is a mass matrix for the vector field \f$\boldsymbol{\phi}\f$. It arises from the \f$\boldsymbol{\phi}\f$ term
     * in the first equation when tested with a vector test function \f$N^\phi\f$.
     */
    std::unique_ptr<mfem::BilinearForm> D;

    /**
     * @brief Bilinear form \f$ S(\psi^\theta, N^\theta) = \int_\Omega \nabla\psi^\theta \cdot \nabla N^\theta \,dV \f$.
     * This is a stiffness matrix (Laplacian) for the scalar field \f$\theta\f$. It is used for stabilization terms
     * or alternative formulations.
     */
    std::unique_ptr<mfem::BilinearForm> S;

    /**
     * @brief Nonlinear form \f$ f(\theta)(\psi^\theta) = \int_\Omega \psi^\theta \cdot \theta^n \,dV \f$.
     * This form arises from the nonlinear source term \f$\theta^n\f$ in the second equation, tested with a scalar
     * test function \f$\psi^\theta\f$.
     */
    std::unique_ptr<mfem::NonlinearForm> f;
};

/**
 * @brief Solves the Lane-Emden equation for a polytropic star using a mixed finite element method.
 *
 * This class sets up and solves the system of equations describing a polytropic
 * stellar model. The Lane-Emden equation, in its second-order form, is:
 * \f[
 * \frac{1}{\xi^2} \frac{d}{d\xi} \left( \xi^2 \frac{d\theta}{d\xi} \right) = -\theta^n
 * \f]
 * where \f$\xi\f$ is a dimensionless radius, \f$\theta\f$ is a dimensionless temperature/potential
 * related to density (\f$\rho = \rho_c \theta^n\f$), and \f$n\f$ is the polytropic index.
 *
 * This solver uses a mixed formulation by introducing \f$\boldsymbol{\phi} = -\nabla \theta\f$:
 * \f[
 * \begin{aligned}
 *   \boldsymbol{\phi} + \nabla \theta &= \mathbf{0} \\
 *   \nabla \cdot \boldsymbol{\phi} - \theta^n &= 0
 * \end{aligned}
 * \f]
 * These equations are discretized using H1 finite elements for \f$\theta\f$ and
 * RT (Raviart-Thomas) finite elements for \f$\boldsymbol{\phi}\f$. The resulting nonlinear
 * system is solved using a Newton method, with the `PolytropeOperator` class
 * defining the system residual and Jacobian.
 *
 * Boundary conditions typically include \f$\theta(\xi_c) = 1\f$ at the center (\f$\xi_c \approx 0\f$),
 * \f$\theta(\xi_1) = 0\f$ at the surface (\f$\xi_1\f$ is the first root of \f$\theta\f$), and
 * \f$\boldsymbol{\phi} \cdot \mathbf{n} = -\frac{d\theta}{d\xi}\f$ at the surface, which is related to the
 * surface gravity. At the center, \f$\boldsymbol{\phi}(\xi_c) = \mathbf{0}\f$ due to symmetry.
 */
class PolySolver final{
public: // Public methods
    /**
     * @brief Constructs a PolySolver instance.
     *
     * Initializes the solver for a given polytropic index \f$n\f$ and finite element order.
     * This constructor internally calls `prepareMesh` to load and scale the mesh.
     *
     * @param n The polytropic index. Must be \f$0 \le n < 5\f$.
     * @param order The polynomial order of the H1 finite elements for \f$\theta\f$.
     *              The RT elements for \f$\boldsymbol{\phi}\f$ will be of order `order-1`.
     * @throw std::runtime_error If \f$n \ge 5\f$ or \f$n < 0\f$ (via `prepareMesh`).
     *
     * @example
     * @code
     * try {
     *     PolySolver solver(1.5, 2); // Polytropic index n=1.5, FE order 2 for theta
     *     solver.solve();
     * } catch (const std::exception& e) {
     *     std::cerr << "Error: " << e.what() << std::endl;
     * }
     * @endcode
     */
    PolySolver(const double n, const double order);

    /**
     * @brief Destructor for PolySolver.
     *
     * Defaulted, handles cleanup of owned resources like `std::unique_ptr` members.
     */
    ~PolySolver();

    /**
     * @brief Solves the polytropic system.
     *
     * This method orchestrates the solution process:
     * 1. Sets the initial guess for \f$\theta\f$ and \f$\boldsymbol{\phi}\f$.
     * 2. Applies essential boundary conditions to the `PolytropeOperator`.
     * 3. Finalizes the `PolytropeOperator` with the initial state.
     * 4. Sets up the Newton solver (including the GMRES linear solver for Jacobian systems).
     * 5. Runs the Newton iteration to find the solution for the free DOFs.
     * 6. Reconstructs the full solution vector from the free DOFs.
     * 7. Saves and/or views the solution based on configuration.
     *
     * @throw std::runtime_error or `mfem::ErrorException` if the solver fails to converge or encounters numerical issues.
     */
    void solve() const;

    /**
     * @brief Gets the polytropic index \f$n\f$.
     * @return The polytropic index.
     */
    double getN() const { return m_polytropicIndex; }

    /**
     * @brief Gets the finite element order used for the \f$\theta\f$ variable.
     * @return The polynomial order of the H1 finite elements.
     */
    double getOrder() const { return m_feOrder; }

    /**
     * @brief Gets a reference to the computational mesh.
     * @return A reference to the `mfem::Mesh` object.
     */
    mfem::Mesh& getMesh() const { return m_mesh; }

    /**
     * @brief Gets a reference to the solution grid function for \f$\theta\f$.
     * @return A reference to the `mfem::GridFunction` storing the \f$\theta\f$ solution.
     * @note The solution is populated after `solve()` has been successfully called.
     */
    mfem::GridFunction& getTheta() const { return *m_theta; }

    /**
     * @brief Gets a reference to the solution grid function for \f$\phi\f$.
     * @return A reference to the `mfem::GridFunction` storing the \f$\phi\f$ solution.
     * @note The solution is populated after `solve()` has been successfully called.
     */
    mfem::GridFunction& getPhi() const { return *m_phi; }

private: // Private Attributes
    // --- Configuration and Logging ---
    serif::config::Config& m_config;                                 ///< Reference to the global configuration manager instance.
    serif::probe::LogManager& m_logManager;                  ///< Reference to the global log manager instance.
    quill::Logger* m_logger;                          ///< Pointer to the specific logger instance for this class.

    // --- Physical and Discretization Parameters ---
    double m_polytropicIndex;                         ///< The polytropic index \f$n\f$.
    double m_feOrder;                                 ///< The polynomial order for H1 elements (\f$\theta\f$). RT elements for \f$\boldsymbol{\phi}\f$ are of order `m_feOrder - 1`.

    // --- MFEM Core Objects ---
    mfem::Mesh& m_mesh;                               ///< Reference to the computational mesh (owned by ResourceManager).
    std::unique_ptr<mfem::H1_FECollection> m_fecH1;   ///< Finite Element Collection for \f$\theta\f$ (H1 elements).
    std::unique_ptr<mfem::RT_FECollection> m_fecRT;   ///< Finite Element Collection for \f$\boldsymbol{\phi}\f$ (Raviart-Thomas elements).

    std::unique_ptr<mfem::FiniteElementSpace> m_feTheta; ///< Finite Element Space for \f$\theta\f$.
    std::unique_ptr<mfem::FiniteElementSpace> m_fePhi;   ///< Finite Element Space for \f$\boldsymbol{\phi}\f$.

    // --- Solution Vectors ---
    std::unique_ptr<mfem::GridFunction> m_theta;      ///< Grid function for the scalar potential \f$\theta\f$.
    std::unique_ptr<mfem::GridFunction> m_phi;        ///< Grid function for the vector flux \f$\boldsymbol{\phi}\f$.

    // --- Operator and Solver Components ---
    std::unique_ptr<PolytropeOperator> m_polytropOperator; ///< The main nonlinear operator for the mixed system.
    std::unique_ptr<mfem::OperatorJacobiSmoother> m_prec;  ///< Preconditioner (currently seems unused in `polySolver.cpp`).

private: // Private methods
    /**
     * @brief Private constructor that takes an existing mesh.
     *
     * Initializes FE collections, spaces, grid functions, and assembles the block system.
     * This is called by the public constructor after `prepareMesh`.
     *
     * @param mesh A reference to an initialized `mfem::Mesh`.
     * @param n The polytropic index.
     * @param order The polynomial order for H1 finite elements.
     */
    PolySolver(mfem::Mesh& mesh, double n, double order);

    /**
     * @brief Prepares the mesh for the simulation.
     *
     * Loads a generic sphere mesh from the `ResourceManager` and scales it
     * to the dimensionless radius \f$\xi_1(n)\f$ (the first zero of the Lane-Emden function
     * for the given polytropic index \f$n\f$).
     *
     * @param n The polytropic index.
     * @return A reference to the prepared `mfem::Mesh` (owned by `ResourceManager`).
     * @throw std::runtime_error If \f$n \ge 5.0\f$ or \f$n < 0.0\f$, as \f$\xi_1(n)\f$ is typically
     *         undefined or infinite outside this range for physical polytropes.
     * @throw std::runtime_error If the mesh resource "mesh:polySphere" cannot be found or loaded.
     *
     * @details The scaling factor \f$\xi_1(n)\f$ is obtained from `polycoeff::x1(n)`.
     * The mesh is expected to be a unit sphere initially.
     */
    static mfem::Mesh& prepareMesh(double n);

    /**
     * @brief Assembles the block system operator `m_polytropOperator`.
     *
     * This involves:
     * 1. Computing block offsets for \f$\theta\f$ and \f$\boldsymbol{\phi}\f$ variables.
     * 2. Building individual bilinear and nonlinear forms (M, Q, D, S, f) using `buildIndividualForms`.
     * 3. Constructing the `PolytropeOperator` with these forms and offsets.
     */
    void assembleBlockSystem();

    /**
     * @brief Compute the block offsets for the operator. These are the offsets that define which dofs belong to which variable.
     *
     * @details
     * Create the block offsets. These define the start of each block in the combined vector.
     * Block offsets will be \f$[0, N_\theta, N_\theta + N_\phi]\f$, where \f$N_\theta\f$ is the number of
     * degrees of freedom for \f$\theta\f$ and \f$N_\phi\f$ is for \f$\boldsymbol{\phi}\f$.
     * The interpretation of this is that each block tells the operator where in the flattened (1D) vector
     * the degrees of freedom or coefficients for that free parameter start and end. I.e.
     * we know that in any flattened vector will have a size \f$N_\theta + N_\phi\f$. The \f$\theta\f$ dofs will span
     * from `blockOffsets[0]` to `blockOffsets[1]` and the \f$\boldsymbol{\phi}\f$ dofs will span from `blockOffsets[1]` to `blockOffsets[2]`.
     *
     * This is the same for matrices only in 2D (rows and columns).
     *
     * The key point here is that this is fundamentally an accounting structure, it is here to keep track of what
     * parts of vectors and matrices belong to which variable.
     *
     * Also note that we use `VSize` rather than `Size`. `Size` refers to the number of true dofs (after eliminating
     * boundary conditions). `VSize` refers to the total number of dofs before BC elimination, which is needed here.
     *
     * @return `mfem::Array<int>` The offsets for the blocks in the operator.
     *
     * @pre `m_feTheta` and `m_fePhi` must be valid, populated `mfem::FiniteElementSpace` objects.
     */
    mfem::Array<int> computeBlockOffsets() const;

    /**
     * @brief Build the individual forms for the block operator (M, Q, D, S, and f).
     *
     * @param blockOffsets The offsets for the blocks in the operator, computed by `computeBlockOffsets`.
     * @return A `std::unique_ptr<formBundle>` containing the assembled forms.
     *
     * @note These forms are built exactly how they are defined in the derivation.
     *       For example, `Mform` corresponds to \f$+M\f$, not \f$-M\f$.
     *       The `PolytropeOperator` handles any necessary sign changes for the final system assembly.
     *
     * @details
     * This method initializes and assembles the discrete forms corresponding to the weak formulation
     * of the mixed polytropic system:
     * \f[
     *   M = \int_\Omega \nabla\psi^\theta \cdot N^\phi \,dV \quad (\text{from } \nabla\theta \text{ term})
     * \f]
     * \f[
     *   Q = \int_\Omega \psi^\phi \cdot \nabla N^\theta \,dV \quad (\text{from } \boldsymbol{\phi} \text{ term, related to } -M^T)
     * \f]
     * \f[
     *   D = \int_\Omega \psi^\phi \cdot N^\phi \,dV \quad (\text{mass matrix for } \boldsymbol{\phi})
     * \f]
     * \f[
     *   S = \int_\Omega \nabla\psi^\theta \cdot \nabla N^\theta \,dV \quad (\text{stiffness matrix for } \theta \text{, for stabilization})
     * \f]
     * \f[
     *   f(\theta)(\psi^\theta) = \int_\Omega \psi^\theta \cdot \theta^n \,dV \quad (\text{from nonlinear term } \theta^n)
     * \f]
     * The `PolytropeOperator` will then use these to form a system like:
     * \f[
     *   R(X) = \begin{pmatrix}
     *       \text{nonlin_op}(\theta) + M\,\boldsymbol{\phi} \\
     *       D\,\boldsymbol{\phi}   - Q\,\theta
     *     \end{pmatrix}
     *   = \mathbf{0}
     * \f]
     * (The exact structure depends on the `PolytropeOperator`'s internal assembly, which might involve stabilization terms using S).
     *
     * @pre  `m_feTheta` and `m_fePhi` must be valid, populated `mfem::FiniteElementSpace` objects.
     *       `m_polytropicIndex` must be set.
     * @post The returned `formBundle` contains unique pointers to assembled (and finalized where appropriate) MFEM forms.
     */
    std::unique_ptr<formBundle> buildIndividualForms(const mfem::Array<int>& blockOffsets);

    /**
     * @brief Assemble and finalize a given MFEM form.
     *
     * This template function calls `Assemble()` and `Finalize()` on the provided form.
     * `Assemble()` computes local element matrices and adds them to the global matrix.
     * `Finalize()` builds the sparsity pattern and allows the `SparseMatrix` representation to be extracted.
     *
     * @tparam FormType The type of the MFEM form (e.g., `mfem::BilinearForm`, `mfem::MixedBilinearForm`).
     * @param f A reference to the form to be assembled and finalized.
     *
     * @pre `f` must be a valid form object that supports `Assemble()` and `Finalize()` methods.
     * @post `f` is assembled and finalized, and its sparse matrix representation is available.
     *
     * @note Nonlinear forms like `mfem::NonlinearForm` are typically not "finalized" in this sense,
     *       as they don't directly produce a sparse matrix but are evaluated. This function is
     *       intended for linear forms.
     */
    static void assembleAndFinalizeForm(auto &f);

    /**
     * @brief Computes the essential true degrees of freedom (DOFs) and their values for boundary conditions.
     *
     * This method determines the DOFs that correspond to:
     * 1. \f$\theta = 1\f$ at the center of the star.
     * 2. \f$\theta = 0\f$ on the surface of the star (boundary attribute 1).
     * 3. \f$\boldsymbol{\phi} \cdot \mathbf{n} = \text{surface\_flux}\f$ on the surface, where the flux is derived from \f$\theta'(\xi_1)\f$.
     *    (This is applied to the normal component of \f$\boldsymbol{\phi}\f$).
     * 4. Potentially \f$\boldsymbol{\phi} = \mathbf{0}\f$ at the center (though the current implementation in `polySolver.cpp`
     *    seems to set components of \f$\boldsymbol{\phi}\f$ to zero at DOFs associated with the center element(s)).
     *
     * @return An `serif::types::MFEMArrayPairSet` containing two pairs:
     *         - The first pair is for \f$\theta\f$: (`mfem::Array<int>` of essential TDOF indices, `mfem::Array<double>` of corresponding values).
     *         - The second pair is for \f$\boldsymbol{\phi}\f$: (`mfem::Array<int>` of essential TDOF indices, `mfem::Array<double>` of corresponding values).
     *
     * @details "True DOFs" (tdof) in MFEM refer to the actual degrees of freedom in the global system,
     *          as opposed to local DOFs within an element. Essential boundary conditions fix the values of these DOFs.
     *          The center condition for \f$\theta\f$ is applied to DOFs identified by `findCenterElement()`.
     *          Surface conditions are applied to boundary attributes marked as essential (typically attribute 1).
     *          The surface flux for \f$\boldsymbol{\phi}\f$ is obtained from `polycoeff::thetaSurfaceFlux(m_polytropicIndex)`.
     *
     * @pre `m_mesh`, `m_feTheta`, `m_fePhi` must be initialized. `m_polytropicIndex` must be set.
     * @throw `mfem::ErrorException` (via `MFEM_ABORT`) if `findCenterElement()` fails to locate the center vertex.
     */
    serif::types::MFEMArrayPairSet getEssentialTrueDof() const;

    /**
     * @brief Finds the degrees of freedom (DOFs) associated with the geometric center (origin) of the mesh.
     *
     * This method identifies:
     * 1. The mesh vertex located at coordinates (0,0,0).
     * 2. The H1 DOFs (\f$\theta\f$) associated with this center vertex.
     * 3. The RT DOFs (\f$\boldsymbol{\phi}\f$) associated with mesh elements that share this center vertex.
     *
     * These DOFs are used to apply boundary conditions at the center of the polytrope,
     * such as \f$\theta(\xi_c)=1\f$ and \f$\boldsymbol{\phi}(\xi_c)=\mathbf{0}\f$.
     *
     * @return A `std::pair` of `mfem::Array<int>`:
     *         - `first`: Array of \f$\theta\f$ (H1) DOF indices at the center.
     *         - `second`: Array of \f$\boldsymbol{\phi}\f$ (RT) DOF indices associated with central elements.
     *
     * @throw `mfem::ErrorException` (via `MFEM_ABORT`) if no vertex is found at the origin (within tolerance).
     *
     * @details For RT elements, DOFs are typically associated with faces (or edges in 2D). This method collects
     *          DOFs from all elements connected to the center vertex. For H1 elements, DOFs can be directly
     *          associated with vertices.
     *
     * @pre `m_mesh`, `m_feTheta`, `m_fePhi` must be initialized.
     */
    std::pair<mfem::Array<int>, mfem::Array<int>> findCenterElement() const;

    /**
     * @brief Sets the initial guess for the solution variables \f$\theta\f$ and \f$\boldsymbol{\phi}\f$.
     *
     * - For \f$\theta\f$:
     *   - The interior initial guess is based on the Lane-Emden series expansion:
     *     \f$\theta(\xi) = \sum a_k \xi^k\f$ (using `laneEmden::thetaSeriesExpansion`).
     *   - The boundary condition \f$\theta(\xi_1) = 0\f$ is projected onto the surface.
     * - For \f$\boldsymbol{\phi}\f$:
     *   - The initial guess is \f$\boldsymbol{\phi} = -\nabla \theta_{\text{init}}\f$, where \f$\theta_{\text{init}}\f$ is the initial guess for \f$\theta\f$.
     *   - The boundary condition for the normal component \f$\boldsymbol{\phi} \cdot \mathbf{n}\f$ at the surface is projected.
     *     This value is \f$\theta'(\xi_1)\f$, obtained from `polycoeff::thetaSurfaceFlux`.
     *   - \f$\boldsymbol{\phi}\f$ components corresponding to central DOFs (from `findCenterElement`) are set to 0.
     *
     * The method uses `mfem::GridFunction::ProjectCoefficient` and `ProjectBdrCoefficient` (or `ProjectBdrCoefficientNormal`)
     * for these projections.
     *
     * @note The projection of boundary conditions for \f$\boldsymbol{\phi}\f$ (RT elements) might not be exact due to
     *       H(div) continuity requirements. Inhomogeneous BC enforcement in `PolytropeOperator` handles this more robustly.
     *
     * @pre `m_theta`, `m_phi`, `m_mesh`, `m_feTheta`, `m_fePhi` must be initialized.
     *      `m_polytropicIndex` must be set.
     *      Configuration settings for viewing initial guess should be loaded if desired.
     * @post `m_theta` and `m_phi` grid functions are populated with the initial guess.
     */
    void setInitialGuess() const;

    /**
     * @brief Saves the 1D radial solution and optionally displays the 2D/3D solution using GLVis.
     *
     * Extracts the \f$\theta\f$ and \f$\boldsymbol{\phi}\f$ components from the converged `state_vector`.
     * - If configured (`Poly:Output:1D:Save`), it extracts a 1D solution along a specified ray
     *   (defined by co-latitude and longitude) using `Probe::getRaySolution` and saves it to a CSV file.
     * - If configured (`Poly:Output:View`), it displays \f$\theta\f$ and \f$\boldsymbol{\phi}\f$ using `Probe::glVisView`.
     *
     * @param state_vector The full solution vector (block vector containing \f$\theta\f$ and \f$\boldsymbol{\phi}\f$)
     *                     obtained from the Newton solver, typically after reconstruction by `PolytropeOperator`.
     *
     * @pre `m_polytropOperator`, `m_feTheta`, `m_fePhi` must be initialized.
     *      Configuration settings for output and viewing must be loaded.
     * @post Solution data is saved and/or visualized according to configuration.
     */
    void saveAndViewSolution(const mfem::BlockVector& state_vector) const;

    /**
     * @brief Sets up the Newton solver and its associated linear solver (GMRES).
     *
     * 1. Loads solver parameters (tolerances, max iterations, print levels) for both Newton
     *    and GMRES from the configuration file (via `LoadSolverUserParams`).
     * 2. Creates a `solverBundle` to manage the lifetimes of `mfem::GMRESSolver` and `mfem::NewtonSolver`.
     * 3. Configures the GMRES solver with its parameters.
     * 4. Configures the Newton solver with its parameters, sets the `PolytropeOperator` as the
     *    nonlinear system operator, and sets the configured GMRES solver as the linear solver
     *    for inverting Jacobians.
     *
     * @return A `solverBundle` struct containing the configured Newton and GMRES solvers.
     *         The ownership of the solvers within the bundle is managed by the bundle itself.
     *
     * @pre `m_polytropOperator` must be initialized and finalized.
     *      Configuration settings for solver parameters must be available.
     * @post A fully configured `solverBundle` is returned, ready for `newton.Mult()`.
     */
    solverBundle setupNewtonSolver() const;

    /**
     * @brief Sets the essential true DOFs on the `PolytropeOperator`.
     *
     * Calls `getEssentialTrueDof()` to compute the boundary condition DOFs and values,
     * and then passes them to `m_polytropOperator->set_essential_true_dofs()`.
     * This step is crucial for the `PolytropeOperator` to correctly handle
     * boundary conditions when forming reduced systems and applying residuals/Jacobians.
     *
     * @pre `m_polytropOperator` must be initialized.
     *      `m_mesh`, `m_feTheta`, `m_fePhi`, `m_polytropicIndex` must be set (for `getEssentialTrueDof`).
     * @post Essential boundary conditions are registered with the `m_polytropOperator`.
     *       The `PolytropeOperator` will likely be marked as not finalized after this call.
     */
    void setOperatorEssentialTrueDofs() const;

    /**
     * @brief Loads Newton and GMRES solver parameters from the configuration.
     *
     * Retrieves relative tolerance, absolute tolerance, maximum iterations, and print level
     * for both the Newton solver and the GMRES linear solver from the `Config` instance.
     * Default values are used if specific configuration keys are not found.
     *
     * Configuration keys are typically prefixed with `Poly:Solver:Newton:` and `Poly:Solver:GMRES:`.
     * Example keys: `Poly:Solver:Newton:RelTol`, `Poly:Solver:GMRES:MaxIter`.
     *
     * @param[out] newtonRelTol Relative tolerance for Newton solver.
     * @param[out] newtonAbsTol Absolute tolerance for Newton solver.
     * @param[out] newtonMaxIter Maximum iterations for Newton solver.
     * @param[out] newtonPrintLevel Print level for Newton solver.
     * @param[out] gmresRelTol Relative tolerance for GMRES solver.
     * @param[out] gmresAbsTol Absolute tolerance for GMRES solver.
     * @param[out] gmresMaxIter Maximum iterations for GMRES solver.
     * @param[out] gmresPrintLevel Print level for GMRES solver.
     *
     * @pre `m_config` must be a valid reference to a `Config` object.
     *      `m_logger` should be initialized if debug logging is enabled.
     * @post Output parameters are populated with values from configuration or defaults.
     *       Solver parameters are logged at DEBUG level.
     */
    void LoadSolverUserParams(double &newtonRelTol, double &newtonAbsTol, int &newtonMaxIter, int &newtonPrintLevel,
                              double &gmresRelTol, double &gmresAbsTol, int &gmresMaxIter, int &gmresPrintLevel) const;

    /**
     * @brief Utility function to get and save the coordinates of degrees of freedom for a finite element space.
     *
     * For each element in the mesh:
     * 1. Gets the DOFs associated with that element from the given `fes`.
     * 2. Gets the physical coordinates of the element's center (or a reference point).
     * 3. Writes the DOF indices, mesh radius, element center radius, and element center x,y,z coordinates
     *    to the specified output file in CSV format.
     *
     * This can be useful for debugging or analyzing the distribution of DOFs.
     *
     * @param fes The `mfem::FiniteElementSpace` for which to get DOF coordinates.
     * @param filename The name of the output CSV file.
     *
     * @pre `fes` must be a valid, initialized `mfem::FiniteElementSpace`.
     *      The mesh associated with `fes` must be valid.
     * @post A CSV file named `filename` is created/truncated and populated with DOF coordinate information.
     *
     * @note For DOFs shared by multiple elements, this function might list them multiple times,
     *       associated with each element they belong to. The output format lists all DOFs for an element
     *       on one line, pipe-separated, followed by coordinate data for that element.
     */
    static void GetDofCoordinates(const mfem::FiniteElementSpace &fes, const std::string& filename);

};

} // namespace polytrope
} // namespace serif