/* ***********************************************************************
//
//   Copyright (C) 2025 -- The 4D-STAR Collaboration
//   File Author: Emily Boudreaux
//   Last Modified: February 12, 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 <vector>

#include "polyMFEMUtils.h"
#include "probe.h"

namespace polyMFEMUtils {
  NonlinearPowerIntegrator::NonlinearPowerIntegrator(
  mfem::Coefficient &coeff,
  double n) : coeff_(coeff), polytropicIndex(n) { 

  }

  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);
    int dof = el.GetDof();
    elvect.SetSize(dof);
    elvect = 0.0;

    mfem::Vector shape(dof);

    for (int iqp = 0; iqp < ir->GetNPoints(); iqp++) {
      mfem::IntegrationPoint ip = ir->IntPoint(iqp);
      Trans.SetIntPoint(&ip);
      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_safe = std::max(u_val, 0.0);
      double u_nl = std::pow(u_safe, polytropicIndex);

      double coeff_val = coeff_.Eval(Trans, ip);
      double x2_u_nl = coeff_val * u_nl;

      for (int i = 0; i < dof; i++){
        elvect(i) += shape(i) * x2_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);
    int dof = el.GetDof();
    elmat.SetSize(dof);
    elmat = 0.0;
    mfem::Vector shape(dof);

    for (int iqp = 0; iqp < ir->GetNPoints(); iqp++) {
      const mfem::IntegrationPoint &ip = ir->IntPoint(iqp);
      Trans.SetIntPoint(&ip);
      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 coeff_val = coeff_.Eval(Trans, ip);
      

      // Calculate the Jacobian
      double u_safe = std::max(u_val, 0.0);
      double d_u_nl = coeff_val * polytropicIndex * std::pow(u_safe, polytropicIndex - 1);
      double x2_d_u_nl = d_u_nl;

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

    }
  }

  BilinearIntegratorWrapper::BilinearIntegratorWrapper(
  mfem::BilinearFormIntegrator *integratorInput
  ) : integrator(integratorInput) { }

  BilinearIntegratorWrapper::~BilinearIntegratorWrapper() {
    delete integrator;
  }

  void BilinearIntegratorWrapper::AssembleElementVector(
  const mfem::FiniteElement &el,
  mfem::ElementTransformation &Trans,
  const mfem::Vector &elfun,
  mfem::Vector &elvect) {
    int dof = el.GetDof();
    mfem::DenseMatrix elMat(dof);
    integrator->AssembleElementMatrix(el, Trans, elMat);
    elvect.SetSize(dof);
    elvect = 0.0;
    for (int i = 0; i < dof; i++)
    {
      double sum = 0.0;
      for (int j = 0; j < dof; j++)
      {
          sum += elMat(i, j) * elfun(j);
      }
      elvect(i) = sum;
    }
  }

  void BilinearIntegratorWrapper::AssembleElementGrad(const mfem::FiniteElement &el,
  mfem::ElementTransformation &Trans,
  const mfem::Vector &elfun,
  mfem::DenseMatrix &elmat) {
    int dof = el.GetDof();
    elmat.SetSize(dof, dof);
    elmat = 0.0;
    integrator->AssembleElementMatrix(el, Trans, elmat);
  }

  CompositeNonlinearIntegrator::CompositeNonlinearIntegrator() { }


  CompositeNonlinearIntegrator::~CompositeNonlinearIntegrator() { }

  void CompositeNonlinearIntegrator::add_integrator(mfem::NonlinearFormIntegrator *integrator) {
    integrators.push_back(integrator);
  }

  void CompositeNonlinearIntegrator::AssembleElementVector(
  const mfem::FiniteElement &el,
  mfem::ElementTransformation &Trans,
  const mfem::Vector &elfun,
  mfem::Vector &elvect) {
    int dof = el.GetDof();
    elvect.SetSize(dof);
    elvect = 0.0;
    mfem::Vector temp(dof);

    for (size_t i = 0; i < integrators.size(); i++) {
      temp= 0.0;
      integrators[i]->AssembleElementVector(el, Trans, elfun, temp);
      elvect.Add(1.0, temp);
    }
  }

  void CompositeNonlinearIntegrator::AssembleElementGrad(
  const mfem::FiniteElement &el,
  mfem::ElementTransformation &Trans,
  const mfem::Vector &elfun,
  mfem::DenseMatrix &elmat) {
    int dof = el.GetDof();
    elmat.SetSize(dof, dof);
    elmat = 0.0;
    mfem::DenseMatrix temp(dof);
    temp.SetSize(dof, dof);
    for (size_t i = 0; i < integrators.size(); i++) {
      temp = 0.0;
      integrators[i] -> AssembleElementGrad(el, Trans, elfun, temp);
      elmat.Add(1.0, temp);
    }
  }

} // namespace polyMFEMUtils