/* ***********************************************************************
//
//   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 <string>


// static std::ofstream debugOut("gradient.csv", std::ios::trunc);

namespace polyMFEMUtils {
  NonlinearPowerIntegrator::NonlinearPowerIntegrator(const double n) :
    m_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);
        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);
        }
        const double u_safe = std::max(u_val, 0.0);
        const double u_nl = std::pow(u_safe, m_polytropicIndex);

        const double x2_u_nl = 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);
        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 gradPhys(3);
        mfem::Vector physCoord(3);

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

    
          el.CalcShape(ip, shape);

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

          // Calculate the Jacobian
          const double u_safe = std::max(u_val, 0.0);
          const double d_u_nl = m_polytropicIndex * std::pow(u_safe, m_polytropicIndex - 1);
          const 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;
            }
          }

    //       // --- Debug Code to write out gradient ---
    //       Trans.Transform(ip,physCoord);
    //       el.CalcDShape(ip, dshape);
    //
    //       mfem::CalcInverse(Trans.Jacobian(), invJ);
    //
    //       mfem::DenseMatrix dshapePhys;
    //       dshapePhys.SetSize(dof, physCoord.Size());
    //       mfem::Mult(dshape, invJ, dshapePhys);
    //
    //       gradPhys = 0.0;
    //       for (int j = 0;  j < dof; j++) {
    //         for (int d = 0; d < gradPhys.Size(); d++) {
    //           gradPhys(d) += elfun(j)*dshapePhys(j, d);
    //         }
    //       }
    //
    //       debugOut
    //         << physCoord(0) << ", " << physCoord(1) << ", " << physCoord(2)
    //         << ", " << gradPhys(0) << ", " << gradPhys(1) << ", " << gradPhys(2) << '\n';
    }
    // debugOut.flush();
    }
} // namespace polyMFEMUtils