// Copyright (c) 2017, Lawrence Livermore National Security, LLC. Produced at
// the Lawrence Livermore National Laboratory. LLNL-CODE-734707. All Rights
// reserved. See files LICENSE and NOTICE for details.
//
// This file is part of CEED, a collection of benchmarks, miniapps, software
// libraries and APIs for efficient high-order finite element and spectral
// element discretizations for exascale applications. For more information and
// source code availability see http://github.com/ceed.
//
// The CEED research is supported by the Exascale Computing Project 17-SC-20-SC,
// a collaborative effort of two U.S. Department of Energy organizations (Office
// of Science and the National Nuclear Security Administration) responsible for
// the planning and preparation of a capable exascale ecosystem, including
// software, applications, hardware, advanced system engineering and early
// testbed platforms, in support of the nation's exascale computing imperative.
//
//                     __                __
//                    / /   ____  ____  / /_  ____  _____
//                   / /   / __ `/ __ `/ __ \/ __ \/ ___/
//                  / /___/ /_/ / /_/ / / / / /_/ (__  )
//                 /_____/\__,_/\__, /_/ /_/\____/____/
//                             /____/
//
//             High-order Lagrangian Hydrodynamics Miniapp
//
// Laghos(LAGrangian High-Order Solver) is a miniapp that solves the
// time-dependent Euler equation of compressible gas dynamics in a moving
// Lagrangian frame using unstructured high-order finite element spatial
// discretization and explicit high-order time-stepping. Laghos is based on the
// numerical algorithm described in the following article:
//
//    V. Dobrev, Tz. Kolev and R. Rieben, "High-order curvilinear finite element
//    methods for Lagrangian hydrodynamics", SIAM Journal on Scientific
//    Computing, (34) 2012, pp. B606–B641, https://doi.org/10.1137/120864672.
//
//             *** THIS IS AN AUTOMATIC MESH REFINEMENT DEMO ***
//
// Sample runs:
//    mpirun -np 8 laghos -p 1 -m ../data/square01_quad.mesh -rs 4 -tf 0.8 -amr
//    mpirun -np 8 laghos -p 1 -m ../data/square01_quad.mesh -rs 4 -tf 0.8 -ok 3 -ot 2 -amr
//    mpirun -np 8 laghos -p 1 -m ../data/cube01_hex.mesh -rs 3 -tf 0.6 -amr
//    mpirun -np 8 laghos -p 1 -m ../data/cube01_hex.mesh -rs 4 -tf 0.6 -rt 1e-3 -amr
//
// Test problems:
//    p = 1  --> Sedov blast.


#include "laghos_solver.hpp"
#include <memory>
#include <iostream>
#include <fstream>

using namespace std;
using namespace mfem;
using namespace mfem::hydrodynamics;

// Choice for the problem setup.
int problem;

void display_banner(ostream & os);

void AMRUpdate(BlockVector &S, BlockVector &S_tmp,
               Array<int> &true_offset,
               ParGridFunction &x_gf,
               ParGridFunction &v_gf,
               ParGridFunction &e_gf);

void GetZeroBCDofs(ParMesh *pmesh, ParFiniteElementSpace *pspace,
                   int bdr_attr_max, Array<int> &ess_tdofs);

void FindElementsWithVertex(const Mesh* mesh, const Vertex &vert,
                            const double size, Array<int> &elements);

void GetPerElementMinMax(const GridFunction &gf,
                         Vector &elem_min, Vector &elem_max,
                         int int_order = -1);


int main(int argc, char *argv[])
{
   // Initialize MPI.
   MPI_Session mpi(argc, argv);
   int myid = mpi.WorldRank();

   // Print the banner.
   if (mpi.Root()) { display_banner(cout); }

   // Parse command-line options.
   const char *mesh_file = "data/square01_quad.mesh";
   int rs_levels = 0;
   int rp_levels = 0;
   int order_v = 2;
   int order_e = 1;
   int ode_solver_type = 4;
   double t_final = 0.5;
   double cfl = 0.5;
   double cg_tol = 1e-8;
   int cg_max_iter = 300;
   int max_tsteps = -1;
   bool p_assembly = true;
   bool visualization = false;
   int vis_steps = 5;
   bool visit = false;
   bool gfprint = false;
   const char *basename = "results/Laghos";
   int partition_type = 111;
   bool amr = false;
   double ref_threshold = 2e-4;
   double deref_threshold = 0.75;
   const int nc_limit = 1;
   const double blast_energy = 0.25;
   const double blast_position[] = {0.0, 0.0, 0.0};
   const double blast_amr_size = 1e-10;

   OptionsParser args(argc, argv);
   args.AddOption(&mesh_file, "-m", "--mesh",
                  "Mesh file to use.");
   args.AddOption(&rs_levels, "-rs", "--refine-serial",
                  "Number of times to refine the mesh uniformly in serial.");
   args.AddOption(&rp_levels, "-rp", "--refine-parallel",
                  "Number of times to refine the mesh uniformly in parallel.");

   args.AddOption(&problem, "-p", "--problem", "Problem setup to use.");
   args.AddOption(&order_v, "-ok", "--order-kinematic",
                  "Order (degree) of the kinematic finite element space.");
   args.AddOption(&order_e, "-ot", "--order-thermo",
                  "Order (degree) of the thermodynamic finite element space.");
   args.AddOption(&ode_solver_type, "-s", "--ode-solver",
                  "ODE solver: 1 - Forward Euler,\n\t"
                  "            2 - RK2 SSP, 3 - RK3 SSP, 4 - RK4, 6 - RK6.");
   args.AddOption(&t_final, "-tf", "--t-final",
                  "Final time; start time is 0.");

   args.AddOption(&cfl, "-cfl", "--cfl", "CFL-condition number.");
   args.AddOption(&cg_tol, "-cgt", "--cg-tol",
                  "Relative CG tolerance (velocity linear solve).");
   args.AddOption(&cg_max_iter, "-cgm", "--cg-max-steps",
                  "Maximum number of CG iterations (velocity linear solve).");
   args.AddOption(&max_tsteps, "-ms", "--max-steps",
                  "Maximum number of steps (negative means no restriction).");
   args.AddOption(&p_assembly, "-pa", "--partial-assembly", "-fa",
                  "--full-assembly",
                  "Activate 1D tensor-based assembly (partial assembly).");

   args.AddOption(&amr, "-amr", "--enable-amr", "-no-amr", "--disable-amr",
                  "Experimental adaptive mesh refinement (problem 1 only).");
   args.AddOption(&ref_threshold, "-rt", "--ref-threshold",
                  "AMR refinement threshold.");
   args.AddOption(&deref_threshold, "-dt", "--deref-threshold",
                  "AMR derefinement threshold (0 = no derefinement).");

   args.AddOption(&visualization, "-vis", "--visualization", "-no-vis",
                  "--no-visualization",
                  "Enable or disable GLVis visualization.");
   args.AddOption(&vis_steps, "-vs", "--visualization-steps",
                  "Visualize every n-th timestep.");
   args.AddOption(&visit, "-visit", "--visit", "-no-visit", "--no-visit",
                  "Enable or disable VisIt visualization.");
   args.AddOption(&gfprint, "-print", "--print", "-no-print", "--no-print",
                  "Enable or disable result output (files in mfem format).");
   args.AddOption(&basename, "-k", "--outputfilename",
                  "Name of the visit dump files");

   args.AddOption(&partition_type, "-pt", "--partition",
                  "Customized x/y/z Cartesian MPI partitioning of the serial mesh.\n\t"
                  "Here x,y,z are relative task ratios in each direction.\n\t"
                  "Example: with 48 mpi tasks and -pt 321, one would get a Cartesian\n\t"
                  "partition of the serial mesh by (6,4,2) MPI tasks in (x,y,z).\n\t"
                  "NOTE: the serially refined mesh must have the appropriate number\n\t"
                  "of zones in each direction, e.g., the number of zones in direction x\n\t"
                  "must be divisible by the number of MPI tasks in direction x.\n\t"
                  "Available options: 11, 21, 111, 211, 221, 311, 321, 322, 432.");
   args.Parse();
   if (!args.Good())
   {
      if (mpi.Root()) { args.PrintUsage(cout); }
      return 1;
   }

   if (amr && problem != 1)
   {
      if (mpi.Root()) { cout << "AMR only supported for problem 1." << endl; }
      return 0;
   }

   if (mpi.Root()) { args.PrintOptions(cout); }

   // Read the serial mesh from the given mesh file on all processors.
   // Refine the mesh in serial to increase the resolution.
   Mesh *mesh = new Mesh(mesh_file, 1, 1);
   const int dim = mesh->Dimension();
   if (!amr)
   {
      for (int lev = 0; lev < rs_levels; lev++)
      {
         mesh->UniformRefinement();
      }
   }
   else
   {
      // Initial refinement for AMR demo on Sedov
      mesh->EnsureNCMesh();
      for (int lev = 0; lev < rs_levels; lev++)
      {
         mesh->RefineAtVertex(Vertex(blast_position[0], blast_position[1],
                                     blast_position[2]),
                              blast_amr_size);
      }
   }

   if (p_assembly && dim == 1)
   {
      p_assembly = false;
      if (mpi.Root())
      {
         cout << "Laghos does not support PA in 1D. Switching to FA." << endl;
      }
   }
   if (p_assembly && amr)
   {
      p_assembly = false;
      if (mpi.Root())
      {
         cout << "Laghos currently does not support PA for AMR. "
              "Switching to FA." << endl;
      }
   }

   // Parallel partitioning of the mesh.
   ParMesh *pmesh = NULL;
   const int num_tasks = mpi.WorldSize(); int unit;
   int *nxyz = new int[3];
   switch (partition_type)
   {
      case 11:
      case 111:
         unit = floor(pow(num_tasks, 1.0 / dim) + 1e-2);
         for (int d = 0; d < dim; d++) { nxyz[d] = unit; }
         if (dim == 2) { nxyz[2] = 0; }
         break;
      case 21: // 2D
         unit = floor(pow(num_tasks / 2, 1.0 / 2) + 1e-2);
         nxyz[0] = 2 * unit; nxyz[1] = unit; nxyz[2] = 0;
         break;
      case 211: // 3D.
         unit = floor(pow(num_tasks / 2, 1.0 / 3) + 1e-2);
         nxyz[0] = 2 * unit; nxyz[1] = unit; nxyz[2] = unit;
         break;
      case 221: // 3D.
         unit = floor(pow(num_tasks / 4, 1.0 / 3) + 1e-2);
         nxyz[0] = 2 * unit; nxyz[1] = 2 * unit; nxyz[2] = unit;
         break;
      case 311: // 3D.
         unit = floor(pow(num_tasks / 3, 1.0 / 3) + 1e-2);
         nxyz[0] = 3 * unit; nxyz[1] = unit; nxyz[2] = unit;
         break;
      case 321: // 3D.
         unit = floor(pow(num_tasks / 6, 1.0 / 3) + 1e-2);
         nxyz[0] = 3 * unit; nxyz[1] = 2 * unit; nxyz[2] = unit;
         break;
      case 322: // 3D.
         unit = floor(pow(2 * num_tasks / 3, 1.0 / 3) + 1e-2);
         nxyz[0] = 3 * unit / 2; nxyz[1] = unit; nxyz[2] = unit;
         break;
      case 432: // 3D.
         unit = floor(pow(num_tasks / 3, 1.0 / 3) + 1e-2);
         nxyz[0] = 2 * unit; nxyz[1] = 3 * unit / 2; nxyz[2] = unit;
         break;
      default:
         if (myid == 0)
         {
            cout << "Unknown partition type: " << partition_type << '\n';
         }
         delete mesh;
         MPI_Finalize();
         return 3;
   }
   int product = 1;
   for (int d = 0; d < dim; d++) { product *= nxyz[d]; }
   if (myid == 0)
   {
      cout << nxyz[0] << " " << nxyz[1] << " " << nxyz[2] << " "
           << product << " " << num_tasks << endl;
   }
   if (product == num_tasks)
   {
      int *partitioning = mesh->CartesianPartitioning(nxyz);
      pmesh = new ParMesh(MPI_COMM_WORLD, *mesh, partitioning);
      delete partitioning;
   }
   else
   {
      if (myid == 0)
      {
         cout << "Non-Cartesian partitioning through METIS will be used.\n";
#ifndef MFEM_USE_METIS
         cout << "MFEM was built without METIS. "
              << "Adjust the number of tasks to use a Cartesian split." << endl;
#endif
      }
#ifndef MFEM_USE_METIS
      return 1;
#endif
      pmesh = new ParMesh(MPI_COMM_WORLD, *mesh);
   }
   delete [] nxyz;
   delete mesh;

   // Refine the mesh further in parallel to increase the resolution.
   for (int lev = 0; lev < rp_levels; lev++)
   {
      pmesh->UniformRefinement();
   }

   int nzones = pmesh->GetNE(), nzones_min, nzones_max;
   MPI_Reduce(&nzones, &nzones_min, 1, MPI_INT, MPI_MIN, 0, pmesh->GetComm());
   MPI_Reduce(&nzones, &nzones_max, 1, MPI_INT, MPI_MAX, 0, pmesh->GetComm());
   if (myid == 0)
   { cout << "Zones min/max: " << nzones_min << " " << nzones_max << endl; }

   int amr_max_level = rs_levels + rp_levels;

   // Define the parallel finite element spaces. We use:
   // - H1 (Gauss-Lobatto, continuous) for position and velocity.
   // - L2 (Bernstein, discontinuous) for specific internal energy.
   L2_FECollection L2FEC(order_e, dim/*, BasisType::Positive*/);
   H1_FECollection H1FEC(order_v, dim);
   ParFiniteElementSpace L2FESpace(pmesh, &L2FEC);
   ParFiniteElementSpace H1FESpace(pmesh, &H1FEC, pmesh->Dimension());

   // Boundary conditions: all tests use v.n = 0 on the boundary, and we assume
   // that the boundaries are straight.
   Array<int> ess_tdofs;
   int bdr_attr_max = pmesh->bdr_attributes.Max();
   GetZeroBCDofs(pmesh, &H1FESpace, bdr_attr_max, ess_tdofs);

   // Define the explicit ODE solver used for time integration.
   ODESolver *ode_solver = NULL;
   switch (ode_solver_type)
   {
      case 1: ode_solver = new ForwardEulerSolver; break;
      case 2: ode_solver = new RK2Solver(0.5); break;
      case 3: ode_solver = new RK3SSPSolver; break;
      case 4: ode_solver = new RK4Solver; break;
      case 6: ode_solver = new RK6Solver; break;
      default:
         if (myid == 0)
         {
            cout << "Unknown ODE solver type: " << ode_solver_type << '\n';
         }
         delete pmesh;
         MPI_Finalize();
         return 3;
   }

   HYPRE_Int glob_size_l2 = L2FESpace.GlobalTrueVSize();
   HYPRE_Int glob_size_h1 = H1FESpace.GlobalTrueVSize();

   if (mpi.Root())
   {
      cout << "Number of kinematic (position, velocity) dofs: "
           << glob_size_h1 << endl;
      cout << "Number of specific internal energy dofs: "
           << glob_size_l2 << endl;
   }

   int Vsize_l2 = L2FESpace.GetVSize();
   int Vsize_h1 = H1FESpace.GetVSize();

   // The monolithic BlockVector stores unknown fields as:
   // - 0 -> position
   // - 1 -> velocity
   // - 2 -> specific internal energy

   Array<int> true_offset(4);
   true_offset[0] = 0;
   true_offset[1] = true_offset[0] + Vsize_h1;
   true_offset[2] = true_offset[1] + Vsize_h1;
   true_offset[3] = true_offset[2] + Vsize_l2;
   BlockVector S(true_offset);

   // Define GridFunction objects for the position, velocity and specific
   // internal energy.  There is no function for the density, as we can always
   // compute the density values given the current mesh position, using the
   // property of pointwise mass conservation.
   ParGridFunction x_gf, v_gf, e_gf;
   x_gf.MakeRef(&H1FESpace, S, true_offset[0]);
   v_gf.MakeRef(&H1FESpace, S, true_offset[1]);
   e_gf.MakeRef(&L2FESpace, S, true_offset[2]);

   // Initialize x_gf using the starting mesh coordinates. This also links the
   // mesh positions to the values in x_gf.
   pmesh->SetNodalGridFunction(&x_gf);

   // Initialize the velocity.
   VectorFunctionCoefficient v_coeff(pmesh->Dimension(), v0);
   v_gf.ProjectCoefficient(v_coeff);

   // Initialize density and specific internal energy values. We interpolate in
   // a non-positive basis to get the correct values at the dofs.  Then we do an
   // L2 projection to the positive basis in which we actually compute. The goal
   // is to get a high-order representation of the initial condition. Note that
   // this density is a temporary function and it will not be updated during the
   // time evolution.
   ParGridFunction rho0_gf(&L2FESpace);
   {
      FunctionCoefficient rho_coeff(hydrodynamics::rho0);
      L2_FECollection l2_fec(order_e, pmesh->Dimension());
      ParFiniteElementSpace l2_fes(pmesh, &l2_fec);

      ParGridFunction l2_rho(&l2_fes), l2_e(&l2_fes);
      l2_rho.ProjectCoefficient(rho_coeff);
      rho0_gf.ProjectGridFunction(l2_rho);

      if (problem == 1)
      {
         // For the Sedov test, we use a delta function at the origin.
         DeltaCoefficient e_coeff(blast_position[0], blast_position[1],
                                  blast_position[2], blast_energy);
         l2_e.ProjectCoefficient(e_coeff);
      }
      else
      {
         FunctionCoefficient e_coeff(e0);
         l2_e.ProjectCoefficient(e_coeff);
      }
      e_gf.ProjectGridFunction(l2_e);
   }

   // Space-dependent ideal gas coefficient over the Lagrangian mesh.
   Coefficient *material_pcf = new FunctionCoefficient(hydrodynamics::gamma);

   // Additional details, depending on the problem.
   int source = 0; bool visc = false;
   switch (problem)
   {
      case 0: if (pmesh->Dimension() == 2) { source = 1; }
         visc = false; break;
      case 1: visc = true; break;
      case 2: visc = true; break;
      case 3: visc = true; break;
      default: MFEM_ABORT("Wrong problem specification!");
   }

   LagrangianHydroOperator oper(S.Size(), H1FESpace, L2FESpace,
                                ess_tdofs, rho0_gf, source, cfl, material_pcf,
                                visc, p_assembly, cg_tol, cg_max_iter);

   if (amr)
   {
      // set a base for h0, this will be further divided in UpdateQuadratureData
      // TODO: for AMR, the treatment of h0 needs more work
      double elem_size = 0.5; // coarse element size (TODO calculate)
      double h0 = elem_size / order_v;
      oper.SetH0(h0);
   }

   socketstream vis_rho, vis_v, vis_e;
   char vishost[] = "localhost";
   int  visport   = 19916;

   ParGridFunction rho_gf;
   if (visualization || visit)
   {
      oper.ComputeDensity(rho_gf);
   }

   if (visualization)
   {
      // Make sure all MPI ranks have sent their 'v' solution before initiating
      // another set of GLVis connections (one from each rank):
      MPI_Barrier(pmesh->GetComm());

      int Wx = 0, Wy = 0; // window position
      const int Ww = 500, Wh = 500; // window size
      int offx = Ww+10; // window offsets

      VisualizeField(vis_rho, vishost, visport, rho_gf,
                     "Density", Wx, Wy, Ww, Wh);
      Wx += offx;
      VisualizeField(vis_v, vishost, visport, v_gf,
                     "Velocity", Wx, Wy, Ww, Wh);
      Wx += offx;
      VisualizeField(vis_e, vishost, visport, e_gf,
                     "Specific Internal Energy", Wx, Wy, Ww, Wh);
   }

   // Save data for VisIt visualization.
   VisItDataCollection visit_dc(basename, pmesh);
   if (visit)
   {
      visit_dc.RegisterField("Density",  &rho_gf);
      visit_dc.RegisterField("Velocity", &v_gf);
      visit_dc.RegisterField("Specific Internal Energy", &e_gf);
      visit_dc.SetCycle(0);
      visit_dc.SetTime(0.0);
      visit_dc.Save();
   }

   // Perform time-integration (looping over the time iterations, ti, with a
   // time-step dt). The object oper is of type LagrangianHydroOperator that
   // defines the Mult() method that used by the time integrators.
   ode_solver->Init(oper);
   oper.ResetTimeStepEstimate();
   double t = 0.0, dt = oper.GetTimeStepEstimate(S), t_old;
   bool last_step = false;
   int steps = 0;
   BlockVector S_old(S);
   for (int ti = 1; !last_step; ti++)
   {
      if (t + dt >= t_final)
      {
         dt = t_final - t;
         last_step = true;
      }
      if (steps == max_tsteps) { last_step = true; }

      S_old = S;
      t_old = t;
      oper.ResetTimeStepEstimate();

      // S is the vector of dofs, t is the current time, and dt is the time step
      // to advance.
      ode_solver->Step(S, t, dt);
      steps++;

      // Adaptive time step control.
      const double dt_est = oper.GetTimeStepEstimate(S);
      if (dt_est < dt)
      {
         // Repeat (solve again) with a decreased time step - decrease of the
         // time estimate suggests appearance of oscillations.
         dt *= 0.85;
         if (dt < numeric_limits<double>::epsilon())
         { MFEM_ABORT("The time step crashed!"); }
         t = t_old;
         S = S_old;
         oper.ResetQuadratureData();
         if (mpi.Root()) { cout << "Repeating step " << ti << endl; }
         ti--; continue;
      }
      else if (dt_est > 1.25 * dt) { dt *= 1.02; }

      // Make sure that the mesh corresponds to the new solution state.
      pmesh->NewNodes(x_gf, false);

      if (last_step || (ti % vis_steps) == 0)
      {
         double loc_norm = e_gf * e_gf, tot_norm;
         MPI_Allreduce(&loc_norm, &tot_norm, 1, MPI_DOUBLE, MPI_SUM,
                       pmesh->GetComm());
         if (mpi.Root())
         {
            cout << fixed;
            cout << "step " << setw(5) << ti
                 << ",\tt = " << setw(5) << setprecision(4) << t
                 << ",\tdt = " << setw(5) << setprecision(6) << dt
                 << ",\t|e| = " << setprecision(10)
                 << sqrt(tot_norm) << endl;
         }

         // Make sure all ranks have sent their 'v' solution before initiating
         // another set of GLVis connections (one from each rank):
         MPI_Barrier(pmesh->GetComm());

         if (visualization || visit || gfprint)
         {
            oper.ComputeDensity(rho_gf);
         }
         if (visualization)
         {
            int Wx = 0, Wy = 0; // window position
            int Ww = 500, Wh = 500; // window size
            int offx = Ww+10; // window offsets

            VisualizeField(vis_rho, vishost, visport, rho_gf,
                           "Density", Wx, Wy, Ww, Wh);
            Wx += offx;
            VisualizeField(vis_v, vishost, visport,
                           v_gf, "Velocity", Wx, Wy, Ww, Wh);
            Wx += offx;
            VisualizeField(vis_e, vishost, visport, e_gf,
                           "Specific Internal Energy", Wx, Wy, Ww,Wh);
            Wx += offx;
         }

         if (visit)
         {
            visit_dc.SetCycle(ti);
            visit_dc.SetTime(t);
            visit_dc.Save();
         }

         if (gfprint)
         {
            ostringstream mesh_name, rho_name, v_name, e_name;
            mesh_name << basename << "_" << ti
                      << "_mesh." << setfill('0') << setw(6) << myid;
            rho_name  << basename << "_" << ti
                      << "_rho." << setfill('0') << setw(6) << myid;
            v_name << basename << "_" << ti
                   << "_v." << setfill('0') << setw(6) << myid;
            e_name << basename << "_" << ti
                   << "_e." << setfill('0') << setw(6) << myid;

            ofstream mesh_ofs(mesh_name.str().c_str());
            mesh_ofs.precision(8);
            pmesh->Print(mesh_ofs);
            mesh_ofs.close();

            ofstream rho_ofs(rho_name.str().c_str());
            rho_ofs.precision(8);
            rho_gf.Save(rho_ofs);
            rho_ofs.close();

            ofstream v_ofs(v_name.str().c_str());
            v_ofs.precision(8);
            v_gf.Save(v_ofs);
            v_ofs.close();

            ofstream e_ofs(e_name.str().c_str());
            e_ofs.precision(8);
            e_gf.Save(e_ofs);
            e_ofs.close();
         }
      }

      if (amr)
      {
         Vector &error_est = oper.GetZoneMaxVisc();

         Vector v_max, v_min;
         GetPerElementMinMax(v_gf, v_min, v_max);

         bool mesh_changed = false;

         // make a list of elements to refine
         Array<int> refs;
         for (int i = 0; i < pmesh->GetNE(); i++)
         {
            if (error_est(i) > ref_threshold
                && pmesh->pncmesh->GetElementDepth(i) < amr_max_level
                && (v_min(i) < 1e-3 || ti < 50) // only refine the still area
               )
            {
               refs.Append(i);
            }
         }

         int nref = pmesh->ReduceInt(refs.Size());
         if (nref)
         {
            pmesh->GeneralRefinement(refs, 1, nc_limit);
            mesh_changed = true;

            if (myid == 0)
            {
               cout << "Refined " << nref << " elements." << endl;
            }
         }
         else if (deref_threshold)
         {
            oper.ComputeDensity(rho_gf);

            Vector rho_max, rho_min;
            GetPerElementMinMax(rho_gf, rho_min, rho_max);

            // simple derefinement based on zone maximum rho in post-shock region
            double rho_max_max = rho_max.Size() ? rho_max.Max() : 0.0;
            double threshold, loc_threshold = deref_threshold * rho_max_max;
            MPI_Allreduce(&loc_threshold, &threshold, 1, MPI_DOUBLE, MPI_MAX,
                          pmesh->GetComm());

            // make sure the blast point is never derefined
            Array<int> elements;
			FindElementsWithVertex(pmesh, Vertex(blast_position[0],
                                                 blast_position[1],
                                                 blast_position[2]),
                                   blast_amr_size, elements);
			for (int i = 0; i < elements.Size(); i++)
            {
               int index = elements[i];
			   if (index >= 0) { rho_max(index) = 1e10; }
            }

            // also, only derefine where the mesh is in motion, i.e. after the shock
            for (int i = 0; i < pmesh->GetNE(); i++)
            {
               if (v_min(i) < 0.1) { rho_max(i) = 1e10; }
            }

            const int op = 2; // maximum value of fine elements
            mesh_changed = pmesh->DerefineByError(rho_max, threshold,
                                                  nc_limit, op);
            if (mesh_changed && myid == 0)
            {
               cout << "Derefined, threshold = " << threshold << endl;
            }
         }

         if (mesh_changed)
         {
            // update state and operator
            AMRUpdate(S, S_old, true_offset, x_gf, v_gf, e_gf);
            oper.AMRUpdate(S, true);

            pmesh->Rebalance();

            // update state and operator
            AMRUpdate(S, S_old, true_offset, x_gf, v_gf, e_gf);
            oper.AMRUpdate(S, false);

            GetZeroBCDofs(pmesh, &H1FESpace, bdr_attr_max, ess_tdofs);

            ode_solver->Init(oper);

            //H1FESpace.PrintPartitionStats();
         }
      }
   }

   switch (ode_solver_type)
   {
      case 2: steps *= 2; break;
      case 3: steps *= 3; break;
      case 4: steps *= 4; break;
      case 6: steps *= 6;
   }
   oper.PrintTimingData(mpi.Root(), steps);

   if (visualization)
   {
      vis_v.close();
      vis_e.close();
   }

   // Free the used memory.
   delete ode_solver;
   delete pmesh;
   delete material_pcf;

   return 0;
}


void GetZeroBCDofs(ParMesh *pmesh, ParFiniteElementSpace *pspace,
                   int bdr_attr_max, Array<int> &ess_tdofs)
{
   ess_tdofs.SetSize(0);
   Array<int> ess_bdr(bdr_attr_max), tdofs1d;
   for (int d = 0; d < pmesh->Dimension(); d++)
   {
      // Attributes 1/2/3 correspond to fixed-x/y/z boundaries, i.e., we must
      // enforce v_x/y/z = 0 for the velocity components.
      ess_bdr = 0; ess_bdr[d] = 1;
      pspace->GetEssentialTrueDofs(ess_bdr, tdofs1d, d);
      ess_tdofs.Append(tdofs1d);
   }
}

void AMRUpdate(BlockVector &S, BlockVector &S_tmp,
               Array<int> &true_offset,
               ParGridFunction &x_gf,
               ParGridFunction &v_gf,
               ParGridFunction &e_gf)
{
   ParFiniteElementSpace* H1FESpace = x_gf.ParFESpace();
   ParFiniteElementSpace* L2FESpace = e_gf.ParFESpace();

   H1FESpace->Update();
   L2FESpace->Update();

   int Vsize_h1 = H1FESpace->GetVSize();
   int Vsize_l2 = L2FESpace->GetVSize();

   true_offset[0] = 0;
   true_offset[1] = true_offset[0] + Vsize_h1;
   true_offset[2] = true_offset[1] + Vsize_h1;
   true_offset[3] = true_offset[2] + Vsize_l2;

   S_tmp = S;
   S.Update(true_offset);

   const Operator* H1Update = H1FESpace->GetUpdateOperator();
   const Operator* L2Update = L2FESpace->GetUpdateOperator();

   H1Update->Mult(S_tmp.GetBlock(0), S.GetBlock(0));
   H1Update->Mult(S_tmp.GetBlock(1), S.GetBlock(1));
   L2Update->Mult(S_tmp.GetBlock(2), S.GetBlock(2));

   x_gf.MakeRef(H1FESpace, S, true_offset[0]);
   v_gf.MakeRef(H1FESpace, S, true_offset[1]);
   e_gf.MakeRef(L2FESpace, S, true_offset[2]);

   S_tmp.Update(true_offset);
}

void FindElementsWithVertex(const Mesh* mesh, const Vertex &vert,
                            const double size, Array<int> &elements)
{
   Array<int> v;

   for (int i = 0; i < mesh->GetNE(); i++)
   {
      mesh->GetElementVertices(i, v);
      for (int j = 0; j < v.Size(); j++)
      {
         double dist = 0.0;
         for (int l = 0; l < mesh->SpaceDimension(); l++)
         {
            double d = vert(l) - mesh->GetVertex(v[j])[l];
            dist += d*d;
         }
         if (dist <= size*size) { elements.Append(i); break; }
      }
   }
}

void Pow(Vector &vec, double p)
{
   for (int i = 0; i < vec.Size(); i++)
   {
      vec(i) = std::pow(vec(i), p);
   }
}

void GetPerElementMinMax(const GridFunction &gf,
                         Vector &elem_min, Vector &elem_max,
                         int int_order)
{
   const FiniteElementSpace *space = gf.FESpace();
   int ne = space->GetNE();

   if (int_order < 0) { int_order = space->GetOrder(0) + 1; }

   elem_min.SetSize(ne);
   elem_max.SetSize(ne);

   Vector vals, tmp;
   for (int i = 0; i < ne; i++)
   {
      int geom = space->GetFE(i)->GetGeomType();
      const IntegrationRule &ir = IntRules.Get(geom, int_order);

      gf.GetValues(i, ir, vals);

      if (space->GetVDim() > 1)
      {
         Pow(vals, 2.0);
         for (int vd = 1; vd < space->GetVDim(); vd++)
         {
            gf.GetValues(i, ir, tmp, vd+1);
            Pow(tmp, 2.0);
            vals += tmp;
         }
         Pow(vals, 0.5);
      }

      elem_min(i) = vals.Min();
      elem_max(i) = vals.Max();
   }
}

namespace mfem
{

namespace hydrodynamics
{

double rho0(const Vector &x)
{
   switch (problem)
   {
      case 0: return 1.0;
      case 1: return 1.0;
      case 2: if (x(0) < 0.5) { return 1.0; }
         else { return 0.1; }
      case 3: if (x(0) > 1.0 && x(1) <= 1.5) { return 1.0; }
         else { return 0.125; }
      default: MFEM_ABORT("Bad number given for problem id!"); return 0.0;
   }
}

double gamma(const Vector &x)
{
   switch (problem)
   {
      case 0: return 5./3.;
      case 1: return 1.4;
      case 2: return 1.4;
      case 3: if (x(0) > 1.0 && x(1) <= 1.5) { return 1.4; }
         else { return 1.5; }
      default: MFEM_ABORT("Bad number given for problem id!"); return 0.0;
   }
}

void v0(const Vector &x, Vector &v)
{
   switch (problem)
   {
      case 0:
         v(0) =  sin(M_PI*x(0)) * cos(M_PI*x(1));
         v(1) = -cos(M_PI*x(0)) * sin(M_PI*x(1));
         if (x.Size() == 3)
         {
            v(0) *= cos(M_PI*x(2));
            v(1) *= cos(M_PI*x(2));
            v(2) = 0.0;
         }
         break;
      case 1: v = 0.0; break;
      case 2: v = 0.0; break;
      case 3: v = 0.0; break;
      default: MFEM_ABORT("Bad number given for problem id!");
   }
}

double e0(const Vector &x)
{
   switch (problem)
   {
      case 0:
      {
         const double denom = 2.0 / 3.0;  // (5/3 - 1) * density.
         double val;
         if (x.Size() == 2)
         {
            val = 1.0 + (cos(2*M_PI*x(0)) + cos(2*M_PI*x(1))) / 4.0;
         }
         else
         {
            val = 100.0 + ((cos(2*M_PI*x(2)) + 2) *
                           (cos(2*M_PI*x(0)) + cos(2*M_PI*x(1))) - 2) / 16.0;
         }
         return val/denom;
      }
      case 1: return 0.0; // This case is initialized in main().
      case 2: if (x(0) < 0.5) { return 1.0 / rho0(x) / (gamma(x) - 1.0); }
         else { return 0.1 / rho0(x) / (gamma(x) - 1.0); }
      case 3: if (x(0) > 1.0) { return 0.1 / rho0(x) / (gamma(x) - 1.0); }
         else { return 1.0 / rho0(x) / (gamma(x) - 1.0); }
      default: MFEM_ABORT("Bad number given for problem id!"); return 0.0;
   }
}

} // namespace hydrodynamics

} // namespace mfem

void display_banner(ostream & os)
{
   os << endl
      << "       __                __                 " << endl
      << "      / /   ____  ____  / /_  ____  _____   " << endl
      << "     / /   / __ `/ __ `/ __ \\/ __ \\/ ___/ " << endl
      << "    / /___/ /_/ / /_/ / / / / /_/ (__  )    " << endl
      << "   /_____/\\__,_/\\__, /_/ /_/\\____/____/  " << endl
      << "               /____/                       " << endl << endl;
}
