#include "mfem.hpp"
#include <string>
#include <iostream>
#include <cmath>
#include <numbers>

#include "polyMFEMUtils.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)*stdDev_,3))) {}

  double GaussianCoefficient::Eval(mfem::ElementTransformation &T, const mfem::IntegrationPoint &ip) {
    mfem::Vector r;
    T.Transform(ip, r);
    double r2 = r * r;
    return norm_coeff * std::exp(-r2/(2*std::pow(stdDev, 2)));
  }


  AugmentedOperator::AugmentedOperator(mfem::NonlinearForm &nfl_, mfem::LinearForm &C_, int lambdaDofOffset_)
    : 
    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_), 
    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;

    for (int i = 0; i < lambdaDofOffset; i++) {
      y[i] = F[i];
    }
    mfem::GridFunction u_gf(C.FESpace());

    DEPRECATION_WARNING_OFF
      u_gf.SetData(u);
    DEPRECATION_WARNING_ON

    y[lambdaDofOffset] = 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 constraint_tdofs;
    C.Assemble();
    lambda_C = C.GetData();
    lambda_C *= -lambda;  // Multiply by -λ (this is now the term −λ ∫ vη(r)dΩ)

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

  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->Finalize();

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

  AugmentedOperator::~AugmentedOperator() {
    delete lastJacobian;
  }
} // namespace polyMFEMUtils