/* ***********************************************************************
//
//   Copyright (C) 2025 -- The 4D-STAR Collaboration
//   File Author: Emily Boudreaux
//   Last Modified: February 14, 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
//
// *********************************************************************** */
#include "mfem.hpp"
#include <string>


void write_solution_to_csv(const mfem::GridFunction &u, const mfem::Mesh &mesh, const std::string &filename);

/**
 * @brief A class for nonlinear power integrator.
 */
class NonlinearPowerIntegrator: public mfem::NonlinearFormIntegrator {
private: 
  mfem::FunctionCoefficient coeff_;
  double polytropicIndex;
public:
  /**
   * @brief Constructor for NonlinearPowerIntegrator.
   * 
   * @param coeff The function coefficient.
   * @param n The polytropic index.
   */
  NonlinearPowerIntegrator(mfem::FunctionCoefficient &coeff, double n);

  /**
   * @brief Assembles the element vector.
   * 
   * @param el The finite element.
   * @param Trans The element transformation.
   * @param elfun The element function.
   * @param elvect The element vector to be assembled.
   */
  virtual void AssembleElementVector(const mfem::FiniteElement &el, mfem::ElementTransformation &Trans, const mfem::Vector &elfun, mfem::Vector &elvect) override;

  /**
   * @brief Assembles the element gradient.
   * 
   * @param el The finite element.
   * @param Trans The element transformation.
   * @param elfun The element function.
   * @param elmat The element matrix to be assembled.
   */
  virtual void AssembleElementGrad (const mfem::FiniteElement &el, mfem::ElementTransformation &Trans, const mfem::Vector &elfun, mfem::DenseMatrix &elmat) override;
};

/**
 * @brief A wrapper class for bilinear integrator.
 */
class BilinearIntegratorWrapper : public mfem::NonlinearFormIntegrator
{
private:
   mfem::BilinearFormIntegrator *integrator;
public:
   /**
    * @brief Constructor for BilinearIntegratorWrapper.
    * 
    * @param integratorInput The bilinear form integrator input.
    */
   BilinearIntegratorWrapper(mfem::BilinearFormIntegrator *integratorInput);

   /**
    * @brief Destructor for BilinearIntegratorWrapper.
    */
   virtual ~BilinearIntegratorWrapper();

   /**
    * @brief Assembles the element vector.
    * 
    * @param el The finite element.
    * @param Trans The element transformation.
    * @param elfun The element function.
    * @param elvect The element vector to be assembled.
    */
   virtual void AssembleElementVector(const mfem::FiniteElement &el, mfem::ElementTransformation &Trans, const mfem::Vector &elfun, mfem::Vector &elvect) override;

   /**
    * @brief Assembles the element gradient.
    * 
    * @param el The finite element.
    * @param Trans The element transformation.
    * @param elfun The element function.
    * @param elmat The element matrix to be assembled.
    */
   virtual void AssembleElementGrad(const mfem::FiniteElement &el, mfem::ElementTransformation &Trans, const mfem::Vector &elfun, mfem::DenseMatrix &elmat) override;
};

/**
 * @brief A class for composite nonlinear integrator.
 */
class CompositeNonlinearIntegrator: public mfem::NonlinearFormIntegrator {
  private:
    std::vector<mfem::NonlinearFormIntegrator*> integrators;
  public:
    /**
     * @brief Constructor for CompositeNonlinearIntegrator.
     */
    CompositeNonlinearIntegrator();

    /**
     * @brief Destructor for CompositeNonlinearIntegrator.
     */
    virtual ~CompositeNonlinearIntegrator();

    /**
     * @brief Adds an integrator to the composite integrator.
     * 
     * @param integrator The nonlinear form integrator to add.
     */
    void add_integrator(mfem::NonlinearFormIntegrator *integrator);

    /**
     * @brief Assembles the element vector.
     * 
     * @param el The finite element.
     * @param Trans The element transformation.
     * @param elfun The element function.
     * @param elvect The element vector to be assembled.
     */
    virtual void AssembleElementVector(const mfem::FiniteElement &el, mfem::ElementTransformation &Trans, const mfem::Vector &elfun, mfem::Vector &elvect) override;

    /**
     * @brief Assembles the element gradient.
     * 
     * @param el The finite element.
     * @param Trans The element transformation.
     * @param elfun The element function.
     * @param elmat The element matrix to be assembled.
     */
    virtual void AssembleElementGrad(const mfem::FiniteElement &el, mfem::ElementTransformation &Trans, const mfem::Vector &elfun, mfem::DenseMatrix &elmat) override;
};