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

#include <cmath>
#include <cstdlib>
#include <functional>
#include <iomanip>
#include <iostream>
#include <memory>
#include <sstream>
#include <string>

#if !MESON_MFEM_HAS_MPI || !MESON_MFEM_HAS_HYPRE
#error "The backend benchmark requires MFEM with MPI and Hypre"
#endif

namespace
{

constexpr const char *json_prefix = "MFEM_BENCHMARK_JSON ";

double AnalyticField(const mfem::Vector &point)
{
   constexpr double pi = 3.141592653589793238462643383279502884;
   return std::sin(pi * point[0]) * std::sin(pi * point[1]) *
          std::sin(pi * point[2]);
}

void SynchronizeDevice(bool uses_gpu)
{
   if (uses_gpu) { MFEM_DEVICE_SYNC; }
}

double MaxTimedRegion(MPI_Comm communicator, bool uses_gpu,
                      const std::function<void()> &operation)
{
   SynchronizeDevice(uses_gpu);
   MPI_Barrier(communicator);
   const double begin = MPI_Wtime();
   operation();
   SynchronizeDevice(uses_gpu);
   const double local_seconds = MPI_Wtime() - begin;
   double maximum_seconds = 0.0;
   MPI_Allreduce(&local_seconds, &maximum_seconds, 1, MPI_DOUBLE, MPI_MAX,
                 communicator);
   return maximum_seconds;
}

std::string JsonEscape(const std::string &value)
{
   std::ostringstream escaped;
   for (const char character : value)
   {
      switch (character)
      {
         case '\\': escaped << "\\\\"; break;
         case '"': escaped << "\\\""; break;
         case '\n': escaped << "\\n"; break;
         case '\r': escaped << "\\r"; break;
         case '\t': escaped << "\\t"; break;
         default: escaped << character; break;
      }
   }
   return escaped.str();
}

int PositiveEnvironmentInteger(const char *name, int fallback)
{
   const char *value = std::getenv(name);
   if (value == nullptr) { return fallback; }
   char *end = nullptr;
   const long parsed = std::strtol(value, &end, 10);
   return end != value && *end == '\0' && parsed > 0
             ? static_cast<int>(parsed)
             : fallback;
}

}  // namespace

int main(int argc, char **argv)
{
   mfem::Mpi::Init(argc, argv);
   mfem::Hypre::Init();
   MPI_Comm communicator = MPI_COMM_WORLD;
   const int rank = mfem::Mpi::WorldRank();
   const int ranks = mfem::Mpi::WorldSize();
   const double process_begin = MPI_Wtime();

   const char *device_name = "cpu";
   int mesh_n = 32;
   int order = 3;
   int applications = 50;
   int warmup_applications = 5;
   int maximum_applications = 1000000;
   int maximum_iterations = 1000;
   int trial = 0;
   double relative_tolerance = 1.0e-6;
   double minimum_apply_seconds = 0.0;
   bool run_solve = false;

   mfem::OptionsParser options(argc, argv);
   options.AddOption(&device_name, "-d", "--device",
                     "MFEM device string (cpu, omp, ceed-cpu, cuda, ...). ");
   options.AddOption(&mesh_n, "-n", "--mesh-n",
                     "Elements per dimension in the fixed global mesh.");
   options.AddOption(&order, "-o", "--order", "H1 polynomial order.");
   options.AddOption(&applications, "-a", "--applications",
                     "Number of timed distributed operator applications.");
   options.AddOption(&warmup_applications, "-w", "--warmup-applications",
                     "Untimed operator applications before measurement.");
   options.AddOption(&maximum_applications, "-ma", "--max-applications",
                     "Safety cap for automatically calibrated applications.");
   options.AddOption(&minimum_apply_seconds, "-mt", "--minimum-apply-seconds",
                     "Minimum final timing window; zero keeps the requested count.");
   options.AddOption(&maximum_iterations, "-m", "--max-iterations",
                     "Maximum iterations for the identically configured CG solve.");
   options.AddOption(&relative_tolerance, "-r", "--relative-tolerance",
                     "Relative tolerance for the identically configured CG solve.");
   options.AddOption(&run_solve, "-s", "--solve", "-no-s", "--no-solve",
                     "Run the optional Jacobi-preconditioned CG validation solve.");
   options.AddOption(&trial, "-t", "--trial",
                     "Trial identifier copied into the JSON record.");
   options.Parse();
   if (!options.Good())
   {
      if (rank == 0) { options.PrintUsage(std::cerr); }
      return 2;
   }
   if (mesh_n < 2 || order < 1 || applications < 1 ||
       warmup_applications < 0 || maximum_applications < applications ||
       maximum_iterations < 1 || relative_tolerance <= 0.0 ||
       minimum_apply_seconds < 0.0)
   {
      if (rank == 0)
      {
         std::cerr << "All sizes/counts must be positive (warmups may be zero).\n";
      }
      return 2;
   }

   std::unique_ptr<mfem::Device> device;
   MPI_Barrier(communicator);
   const double device_begin = MPI_Wtime();
   device = std::make_unique<mfem::Device>(device_name);
   const bool uses_gpu = mfem::Device::Allows(
      mfem::Backend::CUDA_MASK | mfem::Backend::HIP_MASK);
   SynchronizeDevice(uses_gpu);
   const double local_device_seconds = MPI_Wtime() - device_begin;
   double device_seconds = 0.0;
   MPI_Allreduce(&local_device_seconds, &device_seconds, 1, MPI_DOUBLE, MPI_MAX,
                 communicator);

   std::unique_ptr<mfem::ParMesh> mesh;
   const double mesh_seconds = MaxTimedRegion(communicator, uses_gpu, [&]() {
      mfem::Mesh serial_mesh = mfem::Mesh::MakeCartesian3D(
         mesh_n, mesh_n, mesh_n, mfem::Element::HEXAHEDRON);
      mfem::Array<int> partitioning(serial_mesh.GetNE());
      for (int element = 0; element < serial_mesh.GetNE(); ++element)
      {
         partitioning[element] = static_cast<long long>(element) * ranks /
                                 serial_mesh.GetNE();
      }
      mesh = std::make_unique<mfem::ParMesh>(communicator, serial_mesh,
                                             partitioning.GetData());
   });

   std::unique_ptr<mfem::H1_FECollection> elements;
   std::unique_ptr<mfem::ParFiniteElementSpace> space;
   const double space_seconds = MaxTimedRegion(communicator, uses_gpu, [&]() {
      elements = std::make_unique<mfem::H1_FECollection>(
         order, mesh->Dimension(), mfem::BasisType::GaussLobatto);
      space = std::make_unique<mfem::ParFiniteElementSpace>(mesh.get(),
                                                            elements.get());
   });

   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);
   auto right_hand_side = std::make_unique<mfem::ParLinearForm>(space.get());
   auto solution = std::make_unique<mfem::ParGridFunction>(space.get());
   auto diffusion = std::make_unique<mfem::ParBilinearForm>(space.get());
   const double assembly_seconds = MaxTimedRegion(communicator, uses_gpu, [&]() {
      right_hand_side->AddDomainIntegrator(new mfem::DomainLFIntegrator(one));
      right_hand_side->Assemble();
      *solution = 0.0;
      diffusion->SetAssemblyLevel(mfem::AssemblyLevel::PARTIAL);
      diffusion->AddDomainIntegrator(new mfem::DiffusionIntegrator(one));
      diffusion->Assemble();
   });

   mfem::OperatorPtr system_operator;
   mfem::Vector linear_right_hand_side;
   mfem::Vector linear_solution;
   const double form_seconds = MaxTimedRegion(communicator, uses_gpu, [&]() {
      diffusion->FormLinearSystem(essential_dofs, *solution, *right_hand_side,
                                  system_operator, linear_solution,
                                  linear_right_hand_side);
   });

   std::unique_ptr<mfem::OperatorJacobiSmoother> preconditioner;
   double preconditioner_seconds = 0.0;
   if (run_solve)
   {
      preconditioner_seconds = MaxTimedRegion(
         communicator, uses_gpu, [&]() {
         preconditioner = std::make_unique<mfem::OperatorJacobiSmoother>(
            *diffusion, essential_dofs);
      });
   }

   mfem::FunctionCoefficient analytic(AnalyticField);
   mfem::ParGridFunction probe_field(space.get());
   probe_field.ProjectCoefficient(analytic);
   mfem::Vector probe_input;
   probe_field.GetTrueDofs(probe_input);
   mfem::Vector probe_output(system_operator->Height());
   probe_input.UseDevice(true);
   probe_output.UseDevice(true);

   for (int application = 0; application < warmup_applications; ++application)
   {
      system_operator->Mult(probe_input, probe_output);
   }
   SynchronizeDevice(uses_gpu);

   const int requested_applications = applications;
   int timed_applications = requested_applications;
   if (minimum_apply_seconds > 0.0)
   {
      while (true)
      {
         const double calibration_seconds = MaxTimedRegion(
            communicator, uses_gpu, [&]() {
               for (int application = 0; application < timed_applications;
                    ++application)
               {
                  system_operator->Mult(probe_input, probe_output);
               }
            });
         if (calibration_seconds >= minimum_apply_seconds ||
             timed_applications == maximum_applications)
         {
            break;
         }
         timed_applications = timed_applications > maximum_applications / 2
            ? maximum_applications
            : timed_applications * 2;
      }
   }
   const double apply_seconds = MaxTimedRegion(
      communicator, uses_gpu, [&]() {
         for (int application = 0; application < timed_applications;
              ++application)
         {
            system_operator->Mult(probe_input, probe_output);
         }
      });

   const double probe_norm = std::sqrt(
      mfem::InnerProduct(communicator, probe_output, probe_output));
   double solve_seconds = 0.0;
   int cg_iterations = 0;
   bool cg_converged = false;
   double cg_final_norm = 0.0;
   double relative_residual = 0.0;
   double solution_norm = 0.0;
   if (run_solve)
   {
      mfem::CGSolver solver(communicator);
      solver.SetPreconditioner(*preconditioner);
      solver.SetOperator(*system_operator);
      solver.SetRelTol(relative_tolerance);
      solver.SetAbsTol(0.0);
      solver.SetPrintLevel(-1);
      solver.iterative_mode = false;

      linear_solution = 0.0;
      solver.SetMaxIter(2);
      solver.Mult(linear_right_hand_side, linear_solution);
      SynchronizeDevice(uses_gpu);

      linear_solution = 0.0;
      solver.SetMaxIter(maximum_iterations);
      solve_seconds = MaxTimedRegion(communicator, uses_gpu, [&]() {
         solver.Mult(linear_right_hand_side, linear_solution);
      });
      cg_iterations = solver.GetNumIterations();
      cg_converged = solver.GetConverged();
      cg_final_norm = solver.GetFinalNorm();

      mfem::Vector applied_solution(system_operator->Height());
      mfem::Vector residual(linear_right_hand_side);
      system_operator->Mult(linear_solution, applied_solution);
      residual -= applied_solution;
      SynchronizeDevice(uses_gpu);

      const double right_hand_side_norm_squared = mfem::InnerProduct(
         communicator, linear_right_hand_side, linear_right_hand_side);
      const double residual_norm_squared = mfem::InnerProduct(
         communicator, residual, residual);
      relative_residual = right_hand_side_norm_squared > 0.0
         ? std::sqrt(residual_norm_squared / right_hand_side_norm_squared)
         : std::sqrt(residual_norm_squared);
      solution_norm = std::sqrt(
         mfem::InnerProduct(communicator, linear_solution, linear_solution));
   }

   const auto global_dofs = space->GlobalTrueVSize();
   const auto global_elements = mesh->GetGlobalNE();
   const double mdof_per_second =
      1.0e-6 * static_cast<double>(global_dofs) * timed_applications /
      apply_seconds;
   const int omp_threads = PositiveEnvironmentInteger("OMP_NUM_THREADS", 1);
   const bool solve_valid = !run_solve ||
      (cg_converged && std::isfinite(relative_residual) &&
       relative_residual <= 10.0 * relative_tolerance &&
       std::isfinite(solution_norm) && solution_norm > 0.0);
   const bool valid = solve_valid && std::isfinite(probe_norm) &&
                      probe_norm > 0.0;

   SynchronizeDevice(uses_gpu);
   MPI_Barrier(communicator);
   const double local_total_seconds = MPI_Wtime() - process_begin;
   double total_seconds = 0.0;
   MPI_Allreduce(&local_total_seconds, &total_seconds, 1, MPI_DOUBLE, MPI_MAX,
                 communicator);

   if (rank == 0)
   {
      std::cout << std::setprecision(17) << json_prefix
                << '{'
                << "\"schema_version\":1,"
                << "\"trial\":" << trial << ','
                << "\"device\":\"" << JsonEscape(device_name) << "\","
                << "\"ranks\":" << ranks << ','
                << "\"omp_threads\":" << omp_threads << ','
                << "\"dimension\":3,"
                << "\"mesh_n\":" << mesh_n << ','
                << "\"order\":" << order << ','
                << "\"global_elements\":" << global_elements << ','
                << "\"global_true_dofs\":" << global_dofs << ','
                << "\"assembly\":\"partial\","
                << "\"warmup_applications\":" << warmup_applications << ','
                << "\"requested_applications\":" << requested_applications << ','
                << "\"applications\":" << timed_applications << ','
                << "\"minimum_apply_seconds\":" << minimum_apply_seconds << ','
                << "\"device_seconds\":" << device_seconds << ','
                << "\"mesh_seconds\":" << mesh_seconds << ','
                << "\"space_seconds\":" << space_seconds << ','
                << "\"assembly_seconds\":" << assembly_seconds << ','
                << "\"form_seconds\":" << form_seconds << ','
                << "\"preconditioner_seconds\":" << preconditioner_seconds << ','
                << "\"apply_seconds\":" << apply_seconds << ','
                << "\"apply_mdof_per_second\":" << mdof_per_second << ','
                << "\"solve_ran\":" << (run_solve ? "true" : "false") << ','
                << "\"solve_seconds\":" << solve_seconds << ','
                << "\"cg_iterations\":" << cg_iterations << ','
                << "\"cg_converged\":"
                << (cg_converged ? "true" : "false") << ','
                << "\"cg_final_norm\":" << cg_final_norm << ','
                << "\"verified_relative_residual\":" << relative_residual << ','
                << "\"probe_norm\":" << probe_norm << ','
                << "\"solution_norm\":" << solution_norm << ','
                << "\"total_seconds\":" << total_seconds << ','
                << "\"valid\":" << (valid ? "true" : "false")
                << "}\n";
   }

   return valid ? 0 : 3;
}