#include <meson_mfem_template/config.hpp>
#include <mfem.hpp>

#include <cmath>
#include <iostream>

#if !MESON_MFEM_HAS_MPI || !MESON_MFEM_HAS_HYPRE
#error "This example requires the MPI/Hypre capability"
#endif

int main(int argc, char **argv)
{
   mfem::Mpi::Init(argc, argv);
   mfem::Hypre::Init();
   const int rank = mfem::Mpi::WorldRank();
   const char *device_name = argc > 1 ? argv[1] : "cpu";
   mfem::Device device(device_name);

   mfem::Mesh serial_mesh = mfem::Mesh::MakeCartesian2D(
      8, 8, mfem::Element::QUADRILATERAL, true, 1.0, 1.0);
   mfem::ParMesh mesh(MPI_COMM_WORLD, serial_mesh);
   serial_mesh.Clear();

   mfem::H1_FECollection elements(2, mesh.Dimension());
   mfem::ParFiniteElementSpace space(&mesh, &elements);
   mfem::Array<int> essential_boundary(mesh.bdr_attributes.Max());
   essential_boundary = 1;
   mfem::Array<int> essential_dofs;
   space.GetEssentialTrueDofs(essential_boundary, essential_dofs);

   mfem::ConstantCoefficient one(1.0);
   mfem::ParLinearForm rhs(&space);
   rhs.AddDomainIntegrator(new mfem::DomainLFIntegrator(one));
   rhs.Assemble();

   mfem::ParGridFunction solution(&space);
   solution = 0.0;
   mfem::ParBilinearForm diffusion(&space);
   diffusion.AddDomainIntegrator(new mfem::DiffusionIntegrator(one));
   diffusion.Assemble();

   mfem::OperatorPtr matrix;
   mfem::Vector linear_rhs;
   mfem::Vector linear_solution;
   diffusion.FormLinearSystem(
      essential_dofs, solution, rhs, matrix, linear_solution, linear_rhs);

   mfem::HypreBoomerAMG boomer_amg;
   boomer_amg.SetPrintLevel(0);
   mfem::CGSolver solver(MPI_COMM_WORLD);
   solver.SetPreconditioner(boomer_amg);
   solver.SetOperator(*matrix);
   solver.SetRelTol(1e-10);
   solver.SetAbsTol(0.0);
   solver.SetMaxIter(300);
   solver.SetPrintLevel(0);
   solver.Mult(linear_rhs, linear_solution);
   diffusion.RecoverFEMSolution(linear_solution, rhs, solution);

#if MESON_MFEM_HAS_GSLIB
   mfem::FindPointsGSLIB find_points(MPI_COMM_WORLD);
#endif
#if MESON_MFEM_HAS_SUNDIALS
   mfem::CVODESolver cvode(MPI_COMM_WORLD, CV_BDF);
#endif

   const double local_norm = solution.Norml2();
   double global_norm = 0.0;
   MPI_Allreduce(&local_norm, &global_norm, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);
   if (rank == 0)
   {
      std::cout << "MFEM parallel Poisson: ranks=" << mfem::Mpi::WorldSize()
                << " global_dofs=" << space.GlobalTrueVSize()
                << " norm_sum=" << global_norm
                << " device=" << device_name << '\n';
   }
   return solver.GetConverged() && std::isfinite(global_norm) && global_norm > 0.0
             ? 0
             : 2;
}