/* ***********************************************************************
//
//   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 "operator.h"
#include "config.h"
#include "probe.h"
#include "quill/Logger.h"


namespace laneEmden {
    double a (const int k, const double n);
    double c(const int m, const double n);
    double thetaSeriesExpansion(const double xi, const double n, const int order);
}

// Struct to persist lifetime of the linear and nonlinear solvers
struct solverBundle {
    mfem::GMRESSolver solver; // Must be first so it lives longer than the newton solver
    mfem::NewtonSolver newton; // Must be second so that when it is destroyed the solver is still alive preventing a double delete
};

struct formBundle {
    std::unique_ptr<mfem::MixedBilinearForm> M; //<-- M ∫∇ψ^θ·N^φ dV
    std::unique_ptr<mfem::MixedBilinearForm> Q; //<-- Q ∫ψ^φ·∇N^θ dV
    std::unique_ptr<mfem::BilinearForm> D; //<-- D ∫ψ^φ·N^φ dV
    std::unique_ptr<mfem::NonlinearForm> f; //<-- f(θ) ∫ψ^θ·θ^n dV
};

class PolySolver final{
public: // Public methods
    PolySolver(const double n, const double order);
    ~PolySolver();

    void solve() const;

    double getN() const { return m_polytropicIndex; }
    double getOrder() const { return m_feOrder; }
    mfem::Mesh* getMesh() const { return m_mesh.get(); }
    mfem::GridFunction& getSolution() const { return *m_theta; }

private: // Private Attributes
    Config& m_config = Config::getInstance();
    Probe::LogManager& m_logManager = Probe::LogManager::getInstance();
    quill::Logger* m_logger = m_logManager.getLogger("log");
    double m_polytropicIndex, m_feOrder;
    
    std::unique_ptr<mfem::Mesh> m_mesh;
    std::unique_ptr<mfem::H1_FECollection> m_fecH1;
    std::unique_ptr<mfem::RT_FECollection> m_fecRT;

    std::unique_ptr<mfem::FiniteElementSpace> m_feTheta;
    std::unique_ptr<mfem::FiniteElementSpace> m_fePhi;

    std::unique_ptr<mfem::GridFunction> m_theta;
    std::unique_ptr<mfem::GridFunction> m_phi;

    std::unique_ptr<PolytropeOperator> m_polytropOperator;

    std::unique_ptr<mfem::OperatorJacobiSmoother> m_prec;

    std::unique_ptr<mfem::VectorConstantCoefficient> m_negationCoeff;


private: // Private methods
    void assembleBlockSystem();

    /**
     * @breif 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 [0, thetaDofs, thetaDofs + phiDofs].
     * 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 thetaDofs + phiDofs. The theta dofs will span
     * from blockOffsets[0] -> blockOffsets[1] and the phiDofs will span from blockOffsets[1] -> 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 referees to the number of true dofs. That is the dofs which
     * still are present in the system after eliminating boundary conditions. This is the wrong size to use if we are
     * trying to account for the true size of the system. VSize on the other hand refers to the total number of dofs.
     *
     * @return blockOffsets The offsets for the blocks in the operator
     */
    mfem::Array<int> computeBlockOffsets() const;

    /**
     * @breif Build the individual forms for the block operator (M, Q, D, and f)
     *
     * @param blockOffsets The offsets for the blocks in the operator
     * @return forms The forms for the block operator
     *
     * @note These forms are build exactly how they are defined in the derivation. This means that Mform -> M not -M and Qform -> Q not -Q.
     *
     * @details
     * Computes the block offsets
     * \f$\{0,\;|\theta|,\;|\theta|+|\phi|\}\f$, then builds and finalizes
     * the MixedBilinearForms \c Mform, \c Qform and the BilinearForm \c Dform,
     * plus the NonlinearForm \c fform.  Finally, these are handed off to
     * \c PolytropeOperator along with the offsets.
     *
     * The discretized weak form is
     * \f[
     *   R(X)
     *   = \begin{pmatrix}
     *       f(\theta) - M\,\theta \\[6pt]
     *       D\,\theta   - Q\,\phi
     *     \end{pmatrix}
     *   = \mathbf{0},
     * \f]
     * with
     * \f[
     *   M = \int \nabla\psi^\theta \;\cdot\; N^\phi \,dV,\quad
     *   D = \int \psi^\phi \;\cdot\; N^\phi \,dV,
     *   \quad
     *   Q = \int \psi^\phi \;\cdot\; \nabla N^\theta \,dV,
     *   \quad
     *   f(\theta) = \int \psi^\theta \;\cdot\; \theta^n \,dV.
     * \f]
     *
     * @note  MFEM’s MixedVectorWeakDivergenceIntegrator implements
     *        \f$ -\nabla\!\cdot \f$ so we supply a –1 coefficient to make
     *        `Mform` represent the +M from the derivation.  The single negation
     *        in `PolytropeOperator` then restores the final block sign.
     *

     *
     * @pre  \c m_feTheta and \c m_fePhi must be valid, populated FiniteElementSpaces.
     * @post \c m_polytropOperator is constructed with assembled forms and offsets.
     *
     */
    std::unique_ptr<formBundle> buildIndividualForms(const mfem::Array<int>& blockOffsets);

    /**
     * @brief Assemble and finalize the form (Must be a form that can be assembled and finalized)
     *
     * @param f form which is to be assembled and finalized
     *
     * @pre f is a valid form that can be assembled and finalized (Such as Bilinear or MixedBilinearForm)
     * @post f is assembled and finalized
     */
    static void assembleAndFinalizeForm(auto &f);
    SSE::MFEMArrayPairSet getEssentialTrueDof() const;
    std::pair<mfem::Array<int>, mfem::Array<int>> findCenterElement() const;
    void setInitialGuess() const;
    void saveAndViewSolution(const mfem::BlockVector& state_vector) const;
    solverBundle setupNewtonSolver() const;
    void setOperatorEssentialTrueDofs() const;

    void LoadSolverUserParams(double &newtonRelTol, double &newtonAbsTol, int &newtonMaxIter, int &newtonPrintLevel,
                              double &gmresRelTol, double &gmresAbsTol, int &gmresMaxIter, int &gmresPrintLevel) const;

    static void GetDofCoordinates(const mfem::FiniteElementSpace &fes, const std::string& filename);

};