/* ***********************************************************************
//
//   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
//
// *********************************************************************** */
#include "mfem.hpp"
#include <cmath>

#include "integrators.h"
#include "config.h"
#include <string>

namespace serif::polytrope::polyMFEMUtils {
  NonlinearPowerIntegrator::NonlinearPowerIntegrator(const double n) :
    m_polytropicIndex(n),
    m_epsilon(serif::config::Config::getInstance().get<double>("Poly:Solver:Epsilon", 1.0e-8)) {

    if (m_polytropicIndex < 0.0) {
      throw std::invalid_argument("Polytropic index must be non-negative.");
    }
    if (m_polytropicIndex > 5.0) {
      throw std::invalid_argument("Polytropic index must be less than 5.0.");
    }
    if (m_epsilon <= 0.0) {
      throw std::invalid_argument("Epsilon (Poly:Solver:Epsilon) must be positive non-zero.");
    }
  }

  void NonlinearPowerIntegrator::AssembleElementVector(
    const mfem::FiniteElement &el,
    mfem::ElementTransformation &Trans,
    const mfem::Vector &elfun,
    mfem::Vector &elvect) {

    const mfem::IntegrationRule *ir = &mfem::IntRules.Get(el.GetGeomType(), 2 * el.GetOrder() + 3);
    const int dof = el.GetDof();
    elvect.SetSize(dof);
    elvect = 0.0;

    mfem::Vector shape(dof);
    mfem::Vector physCoord;
    for (int iqp = 0; iqp < ir->GetNPoints(); iqp++) {
      mfem::IntegrationPoint ip = ir->IntPoint(iqp);
      Trans.SetIntPoint(&ip);
      const double weight = ip.weight * Trans.Weight();

      el.CalcShape(ip, shape);

      double u_val = 0.0;
      for (int j = 0; j < dof; j++) {
        u_val += elfun(j) * shape(j);
      }

      double u_nl;
      Trans.Transform(ip, physCoord);
      const double r = physCoord.Norml2();
      if (r > m_regularizationRadius) {
        if (u_val < m_epsilon) {
          u_nl = fmod(u_val, m_polytropicIndex, m_epsilon);
        } else {
          u_nl = std::pow(u_val, m_polytropicIndex);
        }
      } else {
        u_nl = 1.0 - m_polytropicIndex * m_regularizationCoeff * std::pow(r, 2);
      }

      for (int i = 0; i < dof; i++){
        elvect(i) += shape(i) * u_nl * weight;
      }
    }
  }


  void NonlinearPowerIntegrator::AssembleElementGrad (
    const mfem::FiniteElement &el,
    mfem::ElementTransformation &Trans,
    const mfem::Vector &elfun,
    mfem::DenseMatrix &elmat) {

    const mfem::IntegrationRule *ir = &mfem::IntRules.Get(el.GetGeomType(), 2 * el.GetOrder() + 3);
    const int dof = el.GetDof();
    elmat.SetSize(dof);
    elmat = 0.0;
    mfem::Vector shape(dof);
    mfem::DenseMatrix dshape(dof, 3);
    mfem::DenseMatrix invJ(3, 3);
    mfem::Vector physCoord;

    for (int iqp = 0; iqp < ir->GetNPoints(); iqp++) {
      const mfem::IntegrationPoint &ip = ir->IntPoint(iqp);
      Trans.SetIntPoint(&ip);
      const double weight = ip.weight * Trans.Weight();
      Trans.Transform(ip, physCoord);
      double r = physCoord.Norml2();

      el.CalcShape(ip, shape);

      double u_val = 0.0;

      for (int j = 0; j < dof; j++) {
        u_val += elfun(j) * shape(j);
      }

      double d_u_nl;
      if (r > m_regularizationRadius) {
        if (u_val < m_epsilon) {
          d_u_nl = dfmod(m_epsilon, m_polytropicIndex);
        } else {
          d_u_nl = m_polytropicIndex * std::pow(u_val, m_polytropicIndex - 1.0);
        }
      } else {
        d_u_nl = 0.0;
      }

      for (int i = 0; i < dof; i++) {
        for (int j = 0; j < dof; j++) {
          elmat(i, j) += shape(i) * d_u_nl * shape(j) * weight;
        }
      }

    }
  }

}
