#include "mfem.hpp"
#include <string>
#include <iostream>
#include <cmath>
#include <numbers>
#include <csignal>
#include <fstream>
#include <array>
#include <vector>

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

#include "warning_control.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);
    }
  }

  ConstraintIntegrator::ConstraintIntegrator(mfem::Coefficient &eta_) : eta(eta_) {}

  void ConstraintIntegrator::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 shape(dof);
    const int intOrder = 2 * el.GetOrder() + 3;

    const mfem::IntegrationRule &ir = mfem::IntRules.Get(el.GetGeomType(), intOrder);

    for (int i = 0; i < ir.GetNPoints(); i++) {
      const mfem::IntegrationPoint &ip = ir.IntPoint(i);
      Trans.SetIntPoint(&ip);
      el.CalcShape(ip, shape);

      double eta_val = eta.Eval(Trans, ip);

      double weight = ip.weight * Trans.Weight();
      elvect.Add(eta_val * weight, shape);
    }
  }

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

  GaussianCoefficient::GaussianCoefficient(double stdDev_)
   : stdDev(stdDev_),
     norm_coeff(1.0/std::pow(std::sqrt(2*std::numbers::pi*std::pow(stdDev_, 2)), 3.0/2.0)) {}

  double GaussianCoefficient::Eval(mfem::ElementTransformation &T, const mfem::IntegrationPoint &ip) {
    mfem::Vector r;
    T.Transform(ip, r);
    double rnorm = std::sqrt(r * r);
    // TODO: return to this to resolve why the Guassian does not always have peek at g(0) = 1 without the factor of 3.0145 (manually calibrated).
    // Open a file (append if already exists) to write the Gaussian values
    return norm_coeff * std::exp(-std::pow(rnorm, 2)/(2*std::pow(stdDev, 2)));
  }


  AugmentedOperator::AugmentedOperator(mfem::NonlinearForm &nfl_, mfem::LinearForm &C_, int lambdaDofOffset_, double C_val_)
    : 
    mfem::Operator( // Call the base class constructor
      nfl_.FESpace()->GetTrueVSize()+1, // Sets the height attribute (+1 for the lambda component)
      nfl_.FESpace()->GetTrueVSize()+1 // Sets the width attribute (+1 for the lambda component)
    ),
    nfl(nfl_),
    C(C_), 
    C_val(C_val_),
    lambdaDofOffset(lambdaDofOffset_),
    lastJacobian(nullptr) {}

  void AugmentedOperator::Mult(const mfem::Vector &x, mfem::Vector &y) const {
    // Select the lambda component of the input vector and seperate it from the θ component
    mfem::Vector u(x.GetData(), lambdaDofOffset);
    double lambda = x[lambdaDofOffset];

    // Compute the residual of the nonlinear form (F(u) - λC)
    mfem::Vector F(lambdaDofOffset);
    nfl.Mult(u, F); // This now includes the -λ∫vη(r) dΩ term

    // Compute the transpose of C for the off diagonal terms of the augmented operator
    y.SetSize(height);
    y = 0.0;

    mfem::GridFunction u_gf(C.FESpace());
    mfem::Vector C_u(1);

    DEPRECATION_WARNING_OFF
      u_gf.SetData(u);
    DEPRECATION_WARNING_ON
    C_u[0] = C.operator()(u_gf);

    // add -lambda * C to the residual
    mfem::Vector lambda_C(lambdaDofOffset);
    mfem::GridFunction constraint_gf(C.FESpace());
    constraint_gf = 0.0;

    mfem::Vector CTmp(lambdaDofOffset);
    CTmp = C.GetData();
    lambda_C = CTmp;
    lambda_C *= -lambda;  // Multiply by -λ (this is now the term −λ ∫ vη(r)dΩ)

    for (int i = 0; i < lambdaDofOffset; i++) {
      y[i] = F[i] + lambda_C[i];
    }

    // Get Global Debug Options for Poly
    std::string defaultKeyset = config.get<std::string>("Poly:Debug:Global:GLVis:Keyset", "");
    bool defaultView = config.get<bool>("Poly:Debug:Global:GLVis:View", false);
    bool defaultExit = config.get<bool>("Poly:Debug:Global:GLVis:Exit", false);


    if (config.get<bool>("Poly:Debug:GLVis:C_gf_View:View", defaultView)) {
      Probe::glVisView(CTmp, *C.FESpace(), "CTmp", config.get<std::string>("Poly:Debug:C_gf_View:Keyset", defaultKeyset));
      if (config.get<bool>("Poly:Debug:GLVis:C_gf_View:Exit", defaultExit)) {
        std::raise(SIGINT);
      }
    }

    if (config.get<bool>("Poly:Debug:GLVis:F_gf_View:View", defaultView)) {
      Probe::glVisView(lambda_C, *nfl.FESpace(), "lambda_C", config.get<std::string>("Poly:Debug:F_gf_View:Keyset", defaultKeyset));
      if (config.get<bool>("Poly:Debug:GLVis:F_gf_View:Exit", defaultExit)) {
        std::raise(SIGINT);
      }
    }

    if (config.get<bool>("Poly:Debug:GLVis:M_gf_View:View", defaultView)) {
      Probe::glVisView(y, *nfl.FESpace(), "y", config.get<std::string>("Poly:Debug:M_gf_View:Keyset", defaultKeyset));
      if (config.get<bool>("Poly:Debug:GLVis:M_gf_View:Exit", defaultExit)) {
        std::raise(SIGINT);
      }
    }

    // Compute the constraint residual (C(u))
    y[lambdaDofOffset] = C_u[0] - C_val; // Enforce the constraint C(u) = C_val
  }

  mfem::Operator &AugmentedOperator::GetGradient(const mfem::Vector &x) const {
    // Select the lambda component of the input vector and seperate it from the θ component
    mfem::Vector u(x.GetData(), lambdaDofOffset);


    // Fill in the blocks of the augmented Jacobian matrix
    // Top-Left: Jacobian of the nonlinear form
    mfem::SparseMatrix *Jnfl_sparse = dynamic_cast<mfem::SparseMatrix*>(&nfl.GetGradient(u));
    if (!Jnfl_sparse) {
      MFEM_ABORT("GetGradient did not return a SparseMatrix");
    }

    mfem::SparseMatrix *J_aug = new mfem::SparseMatrix(height, width);
    // Copy the original Jacobian into the augmented Jacobian
    for (int i = 0; i < Jnfl_sparse->Height(); i++) {
      const int *J_cols = Jnfl_sparse->GetRowColumns(i); 
      const double *J_vals = Jnfl_sparse->GetRowEntries(i);

      for (int jj = 0; jj < Jnfl_sparse->RowSize(i); jj++) {
        int j = J_cols[jj];
        double val = J_vals[jj];
        J_aug->Set(i, j, val);
      }
    }

    // Bottom-left C (the constraint matrix)
    mfem::Vector CVec(lambdaDofOffset);
    mfem::GridFunction tempCGrid(C.FESpace());
    C.Assemble();
    CVec = C.GetData();
    for (int i = 0; i < CVec.Size(); i++) {
      J_aug->Set(lambdaDofOffset, i, CVec[i]);
    }

    // Top-right -Cᵀ (the negative transpose of the constraint matrix)
    for (int i = 0; i < CVec.Size(); i++) {
      J_aug->Set(i, lambdaDofOffset, -CVec[i]);
    }

    J_aug->Set(lambdaDofOffset, lambdaDofOffset, 0.0); // The bottom right corner is zero

    J_aug->Finalize();

    delete lastJacobian;
    lastJacobian = J_aug;
    return *lastJacobian;
  }

  AugmentedOperator::~AugmentedOperator() {
    delete lastJacobian;
  }

  double calculateGaussianIntegral(mfem::Mesh &mesh, polyMFEMUtils::GaussianCoefficient &gaussianCoeff) {
    // Use a discontinuous L2 finite element space (order 0) for integrating the Gaussian.
    // We use L2 because the Gaussian is not necessarily continuous across element boundaries
    // if the Gaussian is narrow relative to the element size.
    mfem::L2_FECollection feCollection(0, mesh.SpaceDimension());
    mfem::FiniteElementSpace feSpace(&mesh, &feCollection);

    mfem::LinearForm gaussianIntegral(&feSpace);
    gaussianIntegral.AddDomainIntegrator(new mfem::DomainLFIntegrator(gaussianCoeff));
    gaussianIntegral.Assemble();

    // Create a GridFunction with a constant value of 1.0 on the L2 space.
    mfem::GridFunction one_gf(&feSpace);
    one_gf = 1.0;

    // Evaluate the linear form on the constant GridFunction.  This gives the integral.
    return gaussianIntegral(one_gf);

  }

  ZeroSlopeNewtonSolver::ZeroSlopeNewtonSolver(double alpha_, std::vector<double> zeroSlopeCoordinate_)
    : alpha(alpha_), zeroSlopeCoordinate(zeroSlopeCoordinate_) {}

  ZeroSlopeNewtonSolver::~ZeroSlopeNewtonSolver() {}

  void ZeroSlopeNewtonSolver::SetOperator(const mfem::Operator &op) {
    LOG_INFO(logger, "Setting operator for zero slope constraint...");
    mfem::NewtonSolver::SetOperator(op); // Call the base class method
    LOG_INFO(logger, "Setting operator for zero slope constraint...done");
    LOG_INFO(logger, "Building location of zero slope constraint...");

    mfem::NonlinearForm *nlf = dynamic_cast<mfem::NonlinearForm*>(const_cast<mfem::Operator*>(&op));
    if (!nlf) {
      LOG_ERROR(logger, "ZeroSlopeNewtonSolver::SetOperator: input operator is not a NonlinearForm");
      MFEM_ABORT("ZeroSlopeNewtonSolver::SetOperator: input operator is not a NonlinearForm");
    }

    mfem::FiniteElementSpace *fes = nlf->FESpace();

    if (!fes) {
      LOG_ERROR(logger, "ZeroSlopeNewtonSolver::SetOperator: input operator does not have a finite element space");
      MFEM_ABORT("ZeroSlopeNewtonSolver::SetOperator: input operator does not have a finite element space");
    }
    u_gf = std::make_unique<mfem::GridFunction>(fes);

    mfem::Mesh *mesh = fes->GetMesh();

    if (!mesh) {
      LOG_ERROR(logger, "ZeroSlopeNewtonSolver::SetOperator: input operator does not have a mesh");
      MFEM_ABORT("ZeroSlopeNewtonSolver::SetOperator: input operator does not have a mesh");
    }

    if (mesh->SpaceDimension() != static_cast<int>(zeroSlopeCoordinate.size())) {
      LOG_ERROR(logger, "ZeroSlopeNewtonSolver::SetOperator: input operator mesh dimension does not match the zero slope coordinate dimension");
      MFEM_ABORT("ZeroSlopeNewtonSolver::SetOperator: input operator mesh dimension does not match the zero slope coordinate dimension");
    }

    mfem::DenseMatrix zeroSlopeCoordinateMatrix(mesh->SpaceDimension(), 1);
    for (int dimID = 0; dimID < mesh->SpaceDimension(); dimID++) {
      zeroSlopeCoordinateMatrix(dimID, 0) = zeroSlopeCoordinate[dimID];
    }

    mfem::Array<int> elementsIDs;
    mfem::Array<mfem::IntegrationPoint> ips;
    mesh->FindPoints(zeroSlopeCoordinateMatrix, elementsIDs, ips);

    zeroSlopeElemID = elementsIDs[0];
    zeroSlopeIP = ips[0];

    LOG_INFO(logger, "Getting element dofs for zero slope constraint...");
    fes->GetElementDofs(zeroSlopeElemID, zeroSlopeDofs);
    LOG_INFO(logger, "Getting element dofs for zero slope constraint...done");
    LOG_INFO(logger, "Building location of zero slope constraint...done");
  }
  
  // void ZeroSlopeNewtonSolver::ProcessNewState(const mfem::Vector &x) const {
  //   LOG_INFO(logger, "Processing new state for zero slope constraint...");
  //   if (zeroSlopeElemID < 0) {
  //     LOG_ERROR(logger, "ZeroSlopeNewtonSolver::ProcessNewState: zero slope element ID is not set");
  //     MFEM_ABORT("ZeroSlopeNewtonSolver::ProcessNewState: zero slope element ID is not set");
  //   }

  //   mfem::NonlinearForm *nlf = dynamic_cast<mfem::NonlinearForm*>(const_cast<mfem::Operator*>(oper));
  //   if (!nlf) {
  //     LOG_ERROR(logger, "ZeroSlopeNewtonSolver::ProcessNewState: input operator is not a NonlinearForm");
  //     MFEM_ABORT("ZeroSlopeNewtonSolver::ProcessNewState: input operator is not a NonlinearForm");
  //   }

  //   mfem::FiniteElementSpace *fes = nlf->FESpace();
  //   if (!fes) {
  //     LOG_ERROR(logger, "ZeroSlopeNewtonSolver::ProcessNewState: input operator does not have a finite element space");
  //     MFEM_ABORT("ZeroSlopeNewtonSolver::ProcessNewState: input operator does not have a finite element space");
  //   }

  //   mfem::Mesh *mesh = fes->GetMesh();
  //   if (!mesh) {
  //     LOG_ERROR(logger, "ZeroSlopeNewtonSolver::ProcessNewState: input operator does not have a mesh");
  //     MFEM_ABORT("ZeroSlopeNewtonSolver::ProcessNewState: input operator does not have a mesh");
  //   }

  //   mfem::ElementTransformation *T = mesh->GetElementTransformation(zeroSlopeElemID);
  //   if (!T) {
  //     LOG_ERROR(logger, "ZeroSlopeNewtonSolver::ProcessNewState: element transformation is not found");
  //     MFEM_ABORT("ZeroSlopeNewtonSolver::ProcessNewState: element transformation is not found");
  //   }

  //   mfem::Vector grad_u(3);

  //   mfem::GridFunction u_gf(fes);
  //   DEPRECATION_WARNING_OFF
  //     u_gf.SetData(x.GetData());
  //   DEPRECATION_WARNING_ON

  //   T->SetIntPoint(&zeroSlopeIP);
  //   u_gf.GetGradient(*T, grad_u);

  //   int dof; 
  //   LOG_DEBUG(logger, "Adjusting the residual to enforce the zero slope constraint by {:0.4E}...", -alpha*grad_u[0]);
  //   double rNorm = r.Norml2();
  //   LOG_INFO(logger, "||r_B|| = {:0.4E}", rNorm);
  //   for (int i = 0; i < zeroSlopeDofs.Size(); i++) {
  //     dof = zeroSlopeDofs[i];
  //     r[dof] -= alpha * grad_u[0];
  //     r[dof] -= alpha * grad_u[1];
  //     r[dof] -= alpha * grad_u[2];
  //   }
  //   rNorm = r.Norml2();
  //   LOG_INFO(logger, "||r_A|| = {:0.4E}", rNorm);
  //   // This still is not working; however, I think I am close. I also need to modify the jacobain.
  // }

  void ZeroSlopeNewtonSolver::Mult(const mfem::Vector &b, mfem::Vector &x) const {
    using namespace mfem;
    using namespace std;

    MFEM_VERIFY(oper != NULL, "the Operator is not set (use SetOperator).");
    MFEM_VERIFY(prec != NULL, "the Solver is not set (use SetSolver).");
 
    int it;
    real_t norm0, norm, norm_goal;
    const bool have_b = (b.Size() == Height());
 
    if (!iterative_mode)
    {
       x = 0.0;
    }
 
    ProcessNewState(x);
 
    oper->Mult(x, r);
    if (have_b)
    {
       r -= b;
    }
    // ComputeConstrainedResidual(x, r);
 
    norm0 = norm = initial_norm = Norm(r);
    if (print_options.first_and_last && !print_options.iterations)
    {
       mfem::out << "Zero slope newton iteration " << setw(2) << 0
                 << " : ||r|| = " << norm << "...\n";
    }
    norm_goal = std::max(rel_tol*norm, abs_tol);
 
    prec->iterative_mode = false;
 
    // x_{i+1} = x_i - [DF(x_i)]^{-1} [F(x_i)-b]
    for (it = 0; true; it++)
    {
      MFEM_VERIFY(IsFinite(norm), "norm = " << norm);
      if (print_options.iterations)
      {
        mfem::out << "Zero slope newton iteration " << setw(2) << it
                  << " : ||r|| = " << norm;
        if (it > 0)
        {
            mfem::out << ", ||r||/||r_0|| = " << norm/norm0;
        }
        mfem::out << '\n';
      }
      Monitor(it, norm, r, x);

      if (norm <= norm_goal)
      {
        converged = true;
        break;
      }

      if (it >= max_iter)
      {
        converged = false;
        break;
      }

      grad = dynamic_cast<mfem::SparseMatrix*>(&oper->GetGradient(x));
      if (!grad)
      {
        LOG_ERROR(logger, "ZeroSlopeNewtonSolver::Mult: Operator does not return a SparseMatrix");
        MFEM_ABORT("ZeroSlopeNewtonSolver::Mult: Operator does not return a SparseMatrix");
      }
      ComputeConstrainedGradient(x);
      prec->SetOperator(*grad);

      if (lin_rtol_type)
      {
        AdaptiveLinRtolPreSolve(x, it, norm);
      }

      prec->Mult(r, c); // c = [DF(x_i)]^{-1} [F(x_i)-b]

      if (lin_rtol_type)
      {
        AdaptiveLinRtolPostSolve(c, r, it, norm);
      }

      const real_t c_scale = ComputeScalingFactor(x, b);
      if (c_scale == 0.0)
      {
        converged = false;
        break;
      }
      add(x, -c_scale, c, x);

      ProcessNewState(x);

      oper->Mult(x, r);
      if (have_b)
      {
        r -= b;
      }
      // ComputeConstrainedResidual(x, r);
      norm = Norm(r);
    }
 
    final_iter = it;
    final_norm = norm;
 
    if (print_options.summary || (!converged && print_options.warnings) ||
        print_options.first_and_last)
    {
       mfem::out << "Newton: Number of iterations: " << final_iter << '\n'
                 << "   ||r|| = " << final_norm
                 << ",  ||r||/||r_0|| = " << final_norm/norm0 << '\n';
    }
    if (!converged && (print_options.summary || print_options.warnings))
    {
       mfem::out << "Newton: No convergence!\n";
    }
  }

  void ZeroSlopeNewtonSolver::ComputeConstrainedResidual(const mfem::Vector &x, mfem::Vector &residual) const {
    mfem::NonlinearForm *nlf = dynamic_cast<mfem::NonlinearForm*>(const_cast<mfem::Operator*>(oper));
    if (!nlf) {
      LOG_ERROR(logger, "ZeroSlopeNewtonSolver::ProcessNewState: input operator is not a NonlinearForm");
      MFEM_ABORT("ZeroSlopeNewtonSolver::ProcessNewState: input operator is not a NonlinearForm");
    }

    mfem::FiniteElementSpace *fes = nlf->FESpace();
    if (!fes) {
      LOG_ERROR(logger, "ZeroSlopeNewtonSolver::ProcessNewState: input operator does not have a finite element space");
      MFEM_ABORT("ZeroSlopeNewtonSolver::ProcessNewState: input operator does not have a finite element space");
    }

    mfem::Mesh *mesh = fes->GetMesh();
    if (!mesh) {
      LOG_ERROR(logger, "ZeroSlopeNewtonSolver::ProcessNewState: input operator does not have a mesh");
      MFEM_ABORT("ZeroSlopeNewtonSolver::ProcessNewState: input operator does not have a mesh");
    }

    mfem::ElementTransformation *T = mesh->GetElementTransformation(zeroSlopeElemID);
    if (!T) {
      LOG_ERROR(logger, "ZeroSlopeNewtonSolver::ProcessNewState: element transformation is not found");
      MFEM_ABORT("ZeroSlopeNewtonSolver::ProcessNewState: element transformation is not found");
    }

    DEPRECATION_WARNING_OFF
      u_gf->SetData(x.GetData());
    DEPRECATION_WARNING_ON

    T->SetIntPoint(&zeroSlopeIP);
    mfem::Vector grad_u(3); // TODO make this a unique pointer so it can be dimensionally adaptive
    u_gf->GetGradient(*T, grad_u);

    for (int i = 0; i < zeroSlopeDofs.Size(); i++) {
      int dof = zeroSlopeDofs[i];
      residual[dof] -= alpha * grad_u[0];
      residual[dof] -= alpha * grad_u[1];
      residual[dof] -= alpha * grad_u[2];
    }

  }

  void ZeroSlopeNewtonSolver::ComputeConstrainedGradient(const mfem::Vector &x) const {
    mfem::NonlinearForm *nlf = dynamic_cast<mfem::NonlinearForm*>(const_cast<mfem::Operator*>(oper));
    if (!nlf) {
      LOG_ERROR(logger, "ZeroSlopeNewtonSolver::ProcessNewState: input operator is not a NonlinearForm");
      MFEM_ABORT("ZeroSlopeNewtonSolver::ProcessNewState: input operator is not a NonlinearForm");
    }

    mfem::FiniteElementSpace *fes = nlf->FESpace();
    if (!fes) {
      LOG_ERROR(logger, "ZeroSlopeNewtonSolver::ProcessNewState: input operator does not have a finite element space");
      MFEM_ABORT("ZeroSlopeNewtonSolver::ProcessNewState: input operator does not have a finite element space");
    }

    mfem::Mesh *mesh = fes->GetMesh();
    if (!mesh) {
      LOG_ERROR(logger, "ZeroSlopeNewtonSolver::ProcessNewState: input operator does not have a mesh");
      MFEM_ABORT("ZeroSlopeNewtonSolver::ProcessNewState: input operator does not have a mesh");
    }

    mfem::ElementTransformation *T = mesh->GetElementTransformation(zeroSlopeElemID);
    if (!T) {
      LOG_ERROR(logger, "ZeroSlopeNewtonSolver::ProcessNewState: element transformation is not found");
      MFEM_ABORT("ZeroSlopeNewtonSolver::ProcessNewState: element transformation is not found");
    }

    const mfem::FiniteElement* fe = fes->GetFE(zeroSlopeElemID); // Get FE *once*.
    mfem::DenseMatrix dshape;    // For shape function derivatives.
    dshape.SetSize(fe->GetDof(), mesh->Dimension());
    T->SetIntPoint(&zeroSlopeIP);
    fe->CalcDShape(zeroSlopeIP, dshape);
    if (!grad) {
      LOG_ERROR(logger, "ZeroSlopeNewtonSolver::ComputeConstrainedGradient: Grad is not set");
      MFEM_ABORT("ZeroSlopeNewtonSolver::ComputeConstrainedGradient: Grad is not set");
    }
    // --- Modify Jacobian ---
    LOG_INFO(logger, "Adjusting the Jacobian to enforce the zero slope constraint...");
    for (int i = 0; i < zeroSlopeDofs.Size(); i++) {
        for (int j = 0; j < zeroSlopeDofs.Size(); j++) {
            grad->Add(zeroSlopeDofs[i], zeroSlopeDofs[j], alpha * dshape(j, 0));
            grad->Add(zeroSlopeDofs[i], zeroSlopeDofs[j], alpha * dshape(j, 1));
            grad->Add(zeroSlopeDofs[i], zeroSlopeDofs[j], alpha * dshape(j, 2));
        }
    }
    LOG_INFO(logger, "Adjusting the Jacobian to enforce the zero slope constraint...done");
  }

} // namespace polyMFEMUtils