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

#include <emscripten/emscripten.h>

#include <cmath>
#include <iomanip>
#include <sstream>
#include <stdexcept>
#include <string>

namespace
{
constexpr double pi = 3.141592653589793238462643383279502884;

double exact_solution(const mfem::Vector &point)
{
   return std::sin(pi * point[0]) * std::sin(pi * point[1]);
}

double forcing_function(const mfem::Vector &point)
{
   return 2.0 * pi * pi * exact_solution(point);
}

std::string json_escape(const std::string &input)
{
   std::ostringstream output;
   for (const unsigned char character : input)
   {
      switch (character)
      {
         case '"': output << "\\\""; break;
         case '\\': output << "\\\\"; break;
         case '\b': output << "\\b"; break;
         case '\f': output << "\\f"; break;
         case '\n': output << "\\n"; break;
         case '\r': output << "\\r"; break;
         case '\t': output << "\\t"; break;
         default:
            if (character < 0x20)
            {
               output << "\\u" << std::hex << std::setw(4)
                      << std::setfill('0') << static_cast<int>(character)
                      << std::dec;
            }
            else
            {
               output << character;
            }
      }
   }
   return output.str();
}

const char *error_response(const std::string &message)
{
   static std::string response;
   response = "{\"ok\":false,\"error\":\"" + json_escape(message) + "\"}";
   return response.c_str();
}
} // namespace

extern "C" EMSCRIPTEN_KEEPALIVE const char *mfem_demo_capabilities()
{
   static const std::string capabilities = []() {
      std::ostringstream output;
      output << "{\"mfemVersion\":\"" << json_escape(MFEM_VERSION_STRING)
             << "\",\"mpi\":" << (MESON_MFEM_HAS_MPI ? "true" : "false")
             << ",\"hypre\":" << (MESON_MFEM_HAS_HYPRE ? "true" : "false")
             << ",\"zlib\":" << (MESON_MFEM_HAS_ZLIB ? "true" : "false")
             << ",\"ceed\":" << (MESON_MFEM_HAS_CEED ? "true" : "false")
             << ",\"fms\":" << (MESON_MFEM_HAS_FMS ? "true" : "false")
             << ",\"wasm\":true}";
      return output.str();
   }();
   return capabilities.c_str();
}

extern "C" EMSCRIPTEN_KEEPALIVE const char *mfem_demo_solve(int cells, int order)
{
   static std::string response;
   try
   {
      if (cells < 2 || cells > 64)
      {
         throw std::invalid_argument("cells must be between 2 and 64");
      }
      if (order < 1 || order > 4)
      {
         throw std::invalid_argument("polynomial order must be between 1 and 4");
      }

      static mfem::Device device("cpu");
      (void)device;
      mfem::StopWatch timer;
      timer.Start();

      mfem::Mesh mesh = mfem::Mesh::MakeCartesian2D(
         cells, cells, mfem::Element::QUADRILATERAL, true, 1.0, 1.0);
      mfem::H1_FECollection elements(order, mesh.Dimension());
      mfem::FiniteElementSpace 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::FunctionCoefficient exact(exact_solution);
      mfem::FunctionCoefficient forcing(forcing_function);
      mfem::LinearForm rhs(&space);
      rhs.AddDomainIntegrator(new mfem::DomainLFIntegrator(forcing));
      rhs.Assemble();

      mfem::GridFunction solution(&space);
      solution = 0.0;
      solution.ProjectBdrCoefficient(exact, essential_boundary);

      mfem::BilinearForm diffusion(&space);
      mfem::ConstantCoefficient one(1.0);
      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::GSSmoother smoother(static_cast<mfem::SparseMatrix &>(*matrix));
      mfem::CGSolver solver;
      solver.SetOperator(*matrix);
      solver.SetPreconditioner(smoother);
      solver.SetRelTol(1e-10);
      solver.SetAbsTol(0.0);
      solver.SetMaxIter(400);
      solver.SetPrintLevel(0);
      solver.Mult(linear_rhs, linear_solution);
      diffusion.RecoverFEMSolution(linear_solution, rhs, solution);

      timer.Stop();
      const double l2_error = solution.ComputeL2Error(exact);
      if (!solver.GetConverged() || !std::isfinite(l2_error))
      {
         throw std::runtime_error("MFEM conjugate-gradient solve did not converge");
      }

      std::ostringstream output;
      output << std::setprecision(10)
             << "{\"ok\":true"
             << ",\"mfemVersion\":\"" << json_escape(MFEM_VERSION_STRING) << '"'
             << ",\"cells\":" << cells
             << ",\"elements\":" << mesh.GetNE()
             << ",\"order\":" << order
             << ",\"trueDofs\":" << space.GetTrueVSize()
             << ",\"iterations\":" << solver.GetNumIterations()
             << ",\"finalNorm\":" << solver.GetFinalNorm()
             << ",\"l2Error\":" << l2_error
             << ",\"elapsedSeconds\":" << timer.RealTime()
             << '}';
      response = output.str();
      return response.c_str();
   }
   catch (const std::exception &error)
   {
      return error_response(error.what());
   }
}

int main()
{
   return 0;
}