#pragma once
#include "asgard_time_advance.hpp"

#ifdef ASGARD_USE_HIGHFIVE
#include "asgard_io.hpp"
#endif

/*!
 * \internal
 * \file asgard_discretization.hpp
 * \brief Defines the container class discretization_manager
 * \author The ASGarD Team
 * \ingroup asgard_discretization
 *
 * \endinternal
 */

namespace asgard
{

/*!
 * \ingroup asgard_discretization
 * \brief Wrapper around several aspects of the pde discretization
 *
 * Assumes ownership of the loaded PDE and builds the sparse grid and operators.
 * The current state is set to the initial conditions and time is set to 0
 * (if a restart file is provided the state and time are loaded form the file).
 *
 * Time integration can be performed with the advance_time() function.
 */
template<typename precision = default_precision>
class discretization_manager
{
public:
  //! allows the creation of a null manager, has to be reinitialized later
  discretization_manager() {
    #ifdef ASGARD_ENABLE_DOUBLE
    #ifdef ASGARD_ENABLE_FLOAT
    static_assert(is_double<precision> or is_float<precision>,
                  "invalid precision type, must use 'double' or 'float'");
    #else
    static_assert(is_double<precision>, "invalid precision type, must use 'double'");
    #endif
    #else
    static_assert(is_float<precision>, "invalid precision type, must use 'float'");
    #endif
  }
  //! take ownership of the pde object and discretize the pde
  discretization_manager(pde_scheme<precision> pde,
                         verbosity_level verbosity = verbosity_level::quiet);

  //! returns the degree of the discretization
  int degree() const { return hier.degree(); }

  //! returns the number of dimensions
  int num_dims() const { return grid.num_dims(); }
  //! returns the max level of the grid
  int max_level() const { return terms.max_level; }
  //! returns the user provided program options
  prog_opts const &options() const { return options_; }
  //! returns the discretization domain
  pde_domain<precision> const &domain() const { return domain_; }

  //! returns the current simulation time
  double time() const { return stepper.data.time(); }
  //! returns the stop time that, the end of the simulation
  double stop_time() const { return stepper.data.stop_time(); }
  //! returns the time step
  double dt() const { return stepper.data.dt(); }
  //! returns the number of remaining time-steps
  int64_t remaining_steps() const { return stepper.data.num_remain(); }
  //! returns the current time step
  int64_t current_step() const { return stepper.data.step(); }
  //! extend the number of steps
  void extend_steps(int64_t num_more) { stepper.data.extend_steps(num_more); }

  //! returns the non-separable initial conditions
  md_func<precision> const &initial_cond_md() const { return initial_md_; }
  //! returns the separable initial conditions
  std::vector<separable_func<precision>> const &initial_cond_sep() const { return initial_sep_; }

  //! set the time in the beginning of the simulation, time() must be zero to call this
  void set_time(precision t) {
    rassert(stepper.data.step() == 0,
            "cannot reset the current time after the simulation start");
    stepper.data.time() = t;
  }
  //! return the current state, in wavelet format, local to this mpi rank
  std::vector<precision> const &current_state() const { return state; }
  //! returns the size of the current state
  int64_t state_size() const { return static_cast<int64_t>(state.size()); }
  //! returns the number of degrees of freedom, will match state_size()
  int64_t num_dof() const {
    // developer purposes mostly, need to know the state inbetween computations
    // when the state vector has not been updated yet due to GPU/MPI considerations
    return grid.num_indexes() * hier.block_size();
  }

  //! return a snapshot of the current solution (in MPI context, only rank 0 gets a valid snapshot)
  reconstruct_solution get_snapshot() const
  {
    #ifdef ASGARD_USE_MPI
    if (not is_leader())
      return reconstruct_solution();
    #endif

    return get_local_snapshot();
  }
  //! return a snapshot of the current solution across all mpi ranks
  reconstruct_solution get_snapshot_mpi() const
  {
    #ifdef ASGARD_USE_MPI
    if (terms.resources.num_ranks() > 1)
      terms.resources.bcast(state);
    #endif

    return get_local_snapshot();
  }

  //! check if the terms have poisson dependence
  bool has_poisson() const { return poisson; }
  //! check if the terms have moment dependence
  bool has_moments() const { return terms.moms; }

  //! computes the right-hand-side of the ode
  void ode_rhs(group_id gid, precision time, std::vector<precision> const &current,
               std::vector<precision> &R) const
  {
    ode_rhs_base(gid, time, current, R);
  }
  //! computes the right-hand-side of the ode
  void ode_rhs(precision time, std::vector<precision> const &current,
               std::vector<precision> &R) const
  {
    ode_rhs_base(group_id::all(), time, current, R);
  }
  //! takes an Euler-like step, next = current + ode-rhs(current), but terms and sources can be scaled separately
  void ode_euler(group_id gid, precision time, std::vector<precision> const &current,
                 terms_scale term_scal, sources_scale source_scal,
                 std::vector<precision> &next) const
  {
    ode_euler_base(gid, time, current, term_scal, source_scal, next);
  }
  //! takes an Euler-like step, next = current + scale * ode-rhs(current)
  void ode_euler(group_id gid, precision time, std::vector<precision> const &current,
                 precision scale, std::vector<precision> &next) const
  {
    ode_euler_base(gid, time, current, terms_scale{scale}, sources_scale{scale}, next);
  }
  //! takes an Euler-like step, next = current + ode-rhs(current), but terms and sources can be scaled separately
  void ode_euler(precision time, std::vector<precision> const &current,
                 terms_scale term_scal, sources_scale source_scal,
                 std::vector<precision> &next) const
  {
    ode_euler_base(group_id::all(), time, current, term_scal, source_scal, next);
  }
  //! takes an Euler-like step, next = current + scale * ode-rhs(current)
  void ode_euler(precision time, std::vector<precision> const &current,
                 precision scale, std::vector<precision> &next) const
  {
    ode_euler_base(group_id::all(), time, current, terms_scale{scale}, sources_scale{scale}, next);
  }

  //! computes the ode right-hand-side sources by projecting them onto the basis and setting them in src
  void set_ode_rhs_sources(precision time, std::vector<precision> &src) const {
    ode_rhs_sources<data_mode::replace>(group_id::all(), time, 1, src);
  }
  //! computes the ode right-hand-side sources by projecting them onto the basis and setting them in src
  void set_ode_rhs_sources(precision time, precision alpha, std::vector<precision> &src) const {
    ode_rhs_sources<data_mode::scal_rep>(group_id::all(), time, alpha, src);
  }
  //! computes the ode right-hand-side sources by projecting them onto the basis and adding them to src
  void add_ode_rhs_sources(precision time, std::vector<precision> &src) const {
    ode_rhs_sources<data_mode::increment>(group_id::all(), time, 1, src);
  }
  //! computes the ode right-hand-side sources by projecting them onto the basis and adding them to src
  void add_ode_rhs_sources(precision time, precision alpha, std::vector<precision> &src) const {
    ode_rhs_sources<data_mode::scal_inc>(group_id::all(), time, alpha, src);
  }

  //! computes the ode right-hand-side sources by projecting them onto the basis and setting them in src
  void set_ode_rhs_sources_group(group_id gid, precision time, std::vector<precision> &src) const {
    ode_rhs_sources<data_mode::replace>(gid, time, 1, src);
  }
  //! computes the ode right-hand-side sources by projecting them onto the basis and adding them to src
  void add_ode_rhs_sources_group(group_id gid, precision time, std::vector<precision> &src) const {
    ode_rhs_sources<data_mode::increment>(gid, time, 1, src);
  }
  //! computes the ode right-hand-side sources by projecting them onto the basis and adding them to src
  void add_ode_rhs_sources_group(group_id gid, precision time, precision alpha, std::vector<precision> &src) const {
    ode_rhs_sources<data_mode::scal_inc>(gid, time, alpha, src);
  }

  #ifdef ASGARD_USE_GPU
  //! same as the CPU version, arrays have size num_dof() and sit on GPU-device 0
  void ode_rhs_gpu(group_id gid, precision time, precision const current[], precision R[]) const
  {
    ode_rhs_base_gpu(gid, time, current, R);
  }
  //! same as the CPU version, arrays have size num_dof() and sit on GPU-device 0
  void ode_rhs_gpu(precision time, precision const current[], precision R[]) const
  {
    ode_rhs_base_gpu(group_id::all(), time, current, R);
  }
  //! same as the CPU version, arrays have size num_dof() and sit on GPU-device 0
  void ode_euler_gpu(group_id gid, precision time, precision const current[],
                     terms_scale term_scal, sources_scale source_scal, precision next[]) const
  {
    ode_euler_base_gpu(gid, time, current, term_scal, source_scal, next);
  }
  //! same as the CPU version, arrays have size num_dof() and sit on GPU-device 0
  void ode_euler_gpu(group_id gid, precision time, precision const current[],
                     precision scale, precision next[]) const
  {
    ode_euler_base_gpu(gid, time, current, terms_scale{scale}, sources_scale{scale}, next);
  }
  //! same as the CPU version, arrays have size num_dof() and sit on GPU-device 0
  void ode_euler_gpu(precision time, precision const current[],
                     terms_scale term_scal, sources_scale source_scal, precision next[]) const
  {
    ode_euler_base_gpu(group_id::all(), time, current, term_scal, source_scal, next);
  }
  //! same as the CPU version, arrays have size num_dof() and sit on GPU-device 0
  void ode_euler_gpu(precision time, precision const current[],
                     precision scale, precision next[]) const
  {
    ode_euler_base_gpu(group_id::all(), time, current, terms_scale{scale}, sources_scale{scale}, next);
  }
  //! same as the CPU version, arrays have size num_dof() and sit on GPU-device 0
  void set_ode_rhs_sources_group_gpu(group_id gid, precision time, precision src[]) const {
    ode_rhs_sources_gpu<data_mode::replace>(gid, time, 1, src);
  }
  //! same as the CPU version, arrays have size num_dof() and sit on GPU-device 0
  void add_ode_rhs_sources_group_gpu(group_id gid, precision time, precision src[]) const {
    ode_rhs_sources_gpu<data_mode::increment>(gid, time, 1, src);
  }
  //! same as the CPU version, arrays have size num_dof() and sit on GPU-device 0
  void add_ode_rhs_sources_group_gpu(group_id gid, precision time, precision alpha, precision src[]) const {
    ode_rhs_sources_gpu<data_mode::scal_inc>(gid, time, alpha, src);
  }
  #endif

  //! computes the l-2 norm, taking the mass matrix into account
  precision normL2(std::vector<precision> const &x) const {
    rassert(x.size() == state.size(), "the vector size must match the state_size()");
    return terms.normL2(grid, conn, x);
  }

  //! applies all terms, does not recompute moments
  void terms_apply(precision alpha, std::vector<precision> const &x, precision beta,
                   std::vector<precision> &y) const
  {
    #ifdef ASGARD_USE_FLOPCOUNTER
    int64_t const flops = terms.flop_count(group_id::all(), grid, conn);
    tools::time_event performance_("terms_apply_all kronmult", flops);
    #else
    tools::time_event performance_("terms_apply_all kronmult");
    #endif
    terms.apply(group_id::all(), grid, conn, alpha, x, beta, y);
  }
  //! applies all terms, non-owning array signature
  void terms_apply(precision alpha, precision const x[], precision beta,
                   precision y[]) const
  {
    #ifdef ASGARD_USE_FLOPCOUNTER
    int64_t const flops = terms.flop_count(group_id::all(), grid, conn);
    tools::time_event performance_("terms_apply_all kronmult", flops);
    #else
    tools::time_event performance_("terms_apply_all kronmult");
    #endif
    terms.apply(group_id::all(), grid, conn, alpha, x, beta, y);
  }
  //! applies terms for the given group, does not recompute moments
  void terms_apply(group_id gid, precision alpha, std::vector<precision> const &x, precision beta,
                   std::vector<precision> &y) const
  {
    #ifdef ASGARD_USE_FLOPCOUNTER
    int64_t const flops = terms.flop_count(gid, grid, conn);
    tools::time_event performance_("terms_apply kronmult", flops);
    #else
    tools::time_event performance_("terms_apply kronmult");
    #endif
    terms.apply(gid, grid, conn, alpha, x, beta, y);
  }
  //! applies all terms, non-owning array signature
  void terms_apply(group_id gid, precision alpha, precision const x[], precision beta,
                   precision y[]) const
  {
    #ifdef ASGARD_USE_FLOPCOUNTER
    int64_t const flops = terms.flop_count(gid, grid, conn);
    tools::time_event performance_("terms_apply kronmult", flops);
    #else
    tools::time_event performance_("terms_apply kronmult");
    #endif
    terms.apply(gid, grid, conn, alpha, x, beta, y);
  }
  #ifdef ASGARD_USE_GPU
  //! applies all terms, non-owning array signature
  void terms_apply_gpu(precision alpha, precision const x[], precision beta,
                       precision y[]) const
  {
    terms_apply_gpu(group_id::all(), alpha, x, beta, y);
  }
  //! applies all terms, non-owning array signature
  void terms_apply_gpu(group_id gid, precision alpha, precision const x[], precision beta,
                       precision y[]) const
  {
    #ifdef ASGARD_USE_FLOPCOUNTER
    int64_t const flops = terms.flop_count(gid, grid, conn);
    tools::time_event performance_("terms_apply kronmult", flops);
    #else
    tools::time_event performance_("terms_apply kronmult");
    #endif
    terms.apply_gpu(gid, grid, conn, alpha, x, beta, y);
  }
  #endif

  #ifdef ASGARD_USE_MPI
  //! initiate iterative loop on MPI for the given group and workspace
  void mpi_iteration_apply(group_id gid, std::vector<precision> &work) const {
    mpi_iteration_apply_base(gid, work);
  }
  //! initiate iterative loop on MPI for the all groups and given workspace
  void mpi_iteration_apply(std::vector<precision> &work) const {
    mpi_iteration_apply_base(group_id::all(), work);
  }
  //! stop the currently working iteration
  void mpi_iteration_stop() const;
  //! performs apply operation on the leader, assuming the non-leader ranks are running mpi_iteration_apply()
  void mpi_leader_apply(precision alpha, precision const x[], precision beta,
                        precision y[]) const
  {
    mpi_leader_apply_base(group_id::all(), alpha, x, beta, y);
  }
  //! performs apply operation on the leader, assuming the non-leader ranks are running mpi_iteration_apply()
  void mpi_leader_apply(group_id gid, precision alpha, precision const x[],
                        precision beta, precision y[]) const
  {
    mpi_leader_apply_base(gid, alpha, x, beta, y);
  }
  #ifdef ASGARD_USE_GPU
  //! same as the CPU version but the input array is on the GPU-device
  void mpi_iteration_apply_gpu(group_id gid, precision work[]) const {
    mpi_iteration_apply_base_gpu(gid, work);
  }
  //! initiate iterative loop on MPI for the all groups and given workspace
  void mpi_iteration_apply_gpu(precision work[]) const {
    mpi_iteration_apply_base_gpu(group_id::all(), work);
  }
  //! stop the currently working iteration
  void mpi_iteration_stop_gpu() const;
  //! performs apply operation on the leader, assuming the non-leader ranks are running mpi_iteration_apply()
  void mpi_leader_apply_gpu(precision alpha, precision const x[], precision beta,
                            precision y[]) const
  {
    mpi_leader_apply_base_gpu(group_id::all(), alpha, x, beta, y);
  }
  //! performs apply operation on the leader, assuming the non-leader ranks are running mpi_iteration_apply()
  void mpi_leader_apply_gpu(group_id gid, precision alpha, precision const x[],
                            precision beta, precision y[]) const
  {
    mpi_leader_apply_base_gpu(gid, alpha, x, beta, y);
  }
  #endif
  #else
  void mpi_iteration_apply(group_id, std::vector<precision> &) const {}
  void mpi_iteration_apply(std::vector<precision> &) const {}
  void mpi_iteration_stop() const {}
  void mpi_leader_apply(precision alpha, precision const x[], precision beta,
                        precision y[]) const
  {
    terms_apply(group_id::all(), alpha, x, beta, y);
  }
  void mpi_leader_apply(group_id gid, precision alpha, precision const x[],
                        precision beta, precision y[]) const
  {
    terms_apply(gid, alpha, x, beta, y);
  }
  #ifdef ASGARD_USE_GPU
  void mpi_iteration_apply_gpu(group_id, precision[]) const {}
  void mpi_iteration_apply_gpu(precision[]) const {}
  void mpi_iteration_stop_gpu() const {}
  void mpi_leader_apply_gpu(precision alpha, precision const x[], precision beta,
                            precision y[]) const
  {
    terms_apply_gpu(group_id::all(), alpha, x, beta, y);
  }
  void mpi_leader_apply_gpu(group_id gid, precision alpha, precision const x[],
                            precision beta, precision y[]) const
  {
    terms_apply_gpu(gid, alpha, x, beta, y);
  }
  #endif
  #endif

  //! write out snapshot data, same as checkpoint but can be invoked manually
  void save_snapshot(std::filesystem::path const &filename) const;
  //! calls save-snapshot for the final step, if requested with -outfile
  void save_final_snapshot() const;

  //! returns the title of the PDE
  std::string const &title() const { return options_.title; }
  //! returns the subtitle of the PDE
  std::string const &subtitle() const { return options_.subtitle; }
  //! returns true if the title contains the given sub-string
  bool title_contains(std::string const &substring) const {
    return (title().find(substring) != std::string::npos);
  }
  //! returns true if the subtitle contains the given sub-string
  bool subtitle_contains(std::string const &substring) const {
    return (subtitle().find(substring) != std::string::npos);
  }

  //! convenient check if we are using high verbosity level
  bool high_verbosity() const { return (verb == verbosity_level::high); }
  //! convenient check if we are using low verbosity level
  bool low_verbosity() const { return (verb == verbosity_level::low); }
  //! convenient check if we are using quiet verbosity level
  bool stop_verbosity() const { return (verb == verbosity_level::quiet); }
  //! resets the verbosity level
  void set_verbosity(verbosity_level v) const { verb = v; }

  /*!
   * \brief integrate in time for the given number of steps, -1 means until the end
   *
   * \param num_steps indicates the number of time-steps to perform.
   *
   * If \b num_steps is negative, integration proceeds until the end-time.
   *
   * If \b num_steps is more than the number of remaining steps, integration
   * will not exceed the final time and will stop as soon as end-time is reached.
   *
   * If solving for a steady state, \b num_steps has no effect.
   *
   * \b returns true on success or false if running with -safe-step and inf or nan was detected
   *
   * Example usage:
   * \code
   *   int64_t const stride = 10;
   *   while (disc.remaining_steps() > 0
   *          and disc.advance_time(stride)){
   *    // the while loop will try to advance for 10 time-steps
   *    // if safety is enabled and a step fails, the loop will exit
   *    // safe a snapshot with filename snapshot_10, snapshot_20 ...
   *     disc.save_snapshot("snapshot_" + std::to_string(disc.current_step()));
   *   }
   *   disc.save_final_snapshot();
   * \endcode
   *
   * \code
   *   auto success = disc.advance_time(); // integrate until the end
   *   if (not success) {
   *     int const step = disc.current_step();
   *     std::cerr << "at time step " << step << ", 'inf' and/or 'nan' detected "
   *                  "when computing step " << step + 1 << '\n';
   *   }
   * \endcode
   *
   * \code
   *   if (disc.advance_time())
   *     std::cout << "all good\n";
   *   else
   *     std::cout << "problem encountered\n";
   * \endcode
   *
   * \code
   *   disc.advance_time();
   *   disc.final_output();
   *   // if an error occurred, the final output will be saved for the time-step
   *   // right before inf/nan was encountered
   *   // the success can also be verified at a later time
   *   if (disc.remaining_steps() > 0)
   *     std::cerr << "the time stepping process encountered "
   *                  "'inf' and/or 'nan' and terminated early\n";
   * \endcode
   */
  bool advance_time(int64_t num_steps = -1) {
    return advance_in_time(*this, num_steps);
  }

  //! report time progress, ndof is the degrees-of-freedom, if different from current_state.size()
  void progress_report(std::ostream &os = std::cout, int64_t ndof = -1) const {
    if (stepper.is_steady_state())
    {
      os << "refinement iteration " << std::setw(10) << tools::split_style(stepper.data.step());
    }
    else
    {
      os << "time-step: " << std::setw(10) << tools::split_style(stepper.data.step()) << "  time: ";
      std::string s = std::to_string(stepper.data.time());

      if (s.size() < 7)
        os << std::setw(10) << s << std::string(7 - s.size(), ' ');
      else
        os << std::setw(10) << s;
    }
    os << "  grid size: " << std::setw(12) << tools::split_style(grid.num_indexes());
    if (ndof >= 0) {
      os << "  dof: " << std::setw(14) << tools::split_style(ndof);
    } else {
      os << "  dof: " << std::setw(14) << tools::split_style(state.size());
    }

    #ifdef ASGARD_USE_FLOPCOUNTER
    int64_t const flops = tools::timer.max_flops();
    if (flops > 0)
      os << "  maxGflops: " << std::to_string(1.E-9 * static_cast<double>(flops));
    #endif

    int64_t const num_appy = stepper.solver_iterations();
    if (num_appy > 0) { // using iterative solver
      os << "  av-iter: " << std::setw(14) << tools::split_style(num_appy / stepper.data.step())
         << '\n';
    } else {
      os << '\n';
    }
  }
  //! safe final result and print statistics, if verbosity allows it and output file is given
  void final_output() const {
    save_final_snapshot();
    if (not stop_verbosity()) {
      progress_report();
      if (asgard::tools::timer.enabled())
        std::cout << asgard::tools::timer.report() << '\n';
    }
  }

  //! projects and sum-of-separable functions and md_func onto the current basis
  void project_function(std::vector<separable_func<precision>> const &sep,
                        md_func<precision> const &fmd, std::vector<precision> &out) const;

  //! projects and sum-of-separable functions and md_func onto the current basis
  std::vector<precision> project_function(
      std::vector<separable_func<precision>> const &sep = {},
      md_func<precision> const &fmd = nullptr) const
  {
    if (sep.empty() and not fmd)
      return std::vector<precision>(state.size());

    std::vector<precision> result;
    project_function(sep, fmd, result);
    return result;
  }
  //! projects a single separable function and md_func onto the current basis
  std::vector<precision> project_function(
      separable_func<precision> const &sep,
      md_func<precision> const &fmd = nullptr) const
  {
    std::vector<precision> result;
    project_function({sep, }, fmd, result);
    return result;
  }
  //! computes a specific moment for the current state
  std::vector<precision> get_moment(moment_id id) const;
  //! computes a specific moment for the current state
  std::vector<precision> get_moment_level(moment_id id) const;
  //! computes and returns the electric field for the current state
  std::vector<precision> get_electric() const;

  //! allows an auxiliary field to be saved for post-processing
  void add_aux_field(aux_field_entry<precision> f) {
    aux_fields.emplace_back(std::move(f));
    if (aux_fields.back().grid.empty()) // if grid provided
      aux_fields.back().grid = grid.get_cells(); // assume the current grid
    if (aux_fields.back().num_dimensions == -1) // default num-dims is the current
      aux_fields.back().num_dimensions = grid.num_dims();
    rassert(aux_fields.back().data.size()
            == static_cast<size_t>(hier.block_size()
                                   * (aux_fields.back().grid.size() / num_dims())),
            "incompatible data size and number of cells");
  }
  //! return reference to the saved fields
  std::vector<aux_field_entry<precision>> const &
  get_aux_fields() const { return aux_fields; }
  //! deletes the current list of auxiliary fields
  void clear_aux_fields() { aux_fields.clear(); }

  //! recomputes the moments with the current state, if groupid is negative all groups will be computed
  void compute_moments(group_id gid = group_id::all()) const {
    compute_moments(gid, state);
  }
  //! recomputes the moments given the state of interest and this term group
  void compute_moments(std::vector<precision> const &f) const {
    compute_moments(group_id::all(), f);
  }
  //! recomputes the moments given the state of interest and this term group
  void compute_moments(group_id gid, std::vector<precision> const &f) const {
    rassert(terms.moms, "no moments set for this PDE");
    compute_moments_(gid, f);
  }
  #ifdef ASGARD_USE_GPU
  //! recomputes the moments with the current state, if groupid is negative all groups will be computed
  void compute_moments_gpu(group_id gid = group_id::all()) const {
    gpu::vector<precision> gpu_state = state;
    compute_moments_gpu(gid, gpu_state.data());
  }
  //! recomputes the moments given the state of interest and this term group
  void compute_moments_gpu(precision const f[]) const {
    compute_moments_gpu(group_id::all(), f);
  }
  //! recomputes the moments given the state of interest and this term group
  void compute_moments_gpu(group_id gid, precision const f[]) const {
    rassert(terms.moms, "no moments set for this PDE");
    compute_moments_gpu_(gid, f);
  }
  #endif

#ifndef __ASGARD_DOXYGEN_SKIP_INTERNAL
  //! returns a ref to the sparse grid
  sparse_grid const &get_grid() const { return grid; }
  //! returns the current grid generation
  int grid_generation() const { return grid.generation(); }
  //! synchronizes the grid across MPI ranks and GPU devices
  void grid_sync() {
    #ifdef ASGARD_USE_MPI
    grid.mpi_sync(terms.resources, grid_synced_gen_);
    grid_synced_gen_ = grid.generation();
    #endif
    #ifdef ASGARD_USE_GPU
    grid.gpu_sync();
    #endif
  }
  //! returns the term manager
  term_manager<precision> const &get_terms() const { return terms; }
  //! returns the compute resources meta structure
  resource_set const &get_resources() const { return terms.resources; }

  //! return the hierarchy_manipulator
  hierarchy_manipulator<precision> const &get_hier() const { return hier; }
  //! return the connection patterns
  connection_patterns const &get_conn() const { return conn; }

  //! recomputes the Poisson term for the given group
  void compute_poisson(group_id gid = group_id::all()) const {
    if (not poisson or not terms.has_poisson(gid))
      return;

    // currently we only support 1d in position space, so the solver is trivial
    // the cost is so low, that everyone can do it even if it is repeated work
    // when we get to multi-d Poisson problems, the leader will be needed
    // to help the communication process
    poisson.solve_periodic(terms.moms.get_cached_level(poisson.moment0(), hier),
                           terms.moms.edit_poisson_level());
  }
  //! (testing/debugging) copy ns to the current state, e.g., force an initial condition
  void set_current_state(std::vector<precision> const &ns) {
    rassert(ns.size() == state.size(), "cannot set state with different size");
    state = ns;
  }
  //! get the current moment manager, allows detailed access to loaded moments
  moment_manager<precision> const &get_moment_manager() const {
    return terms.moms;
  }

  //! (debugging) prints the term-matrices
  void print_mats() const;

  #ifdef ASGARD_USE_MPI
  //! returns true if this is mpi rank 0, always true when MPI is not enabled
  bool is_leader() const { return terms.resources.is_leader(); }
  #else
  static constexpr bool is_leader() { return true; }
  #endif

  #ifdef ASGARD_USE_MPI
  //! returns persistent vector for mpi operations
  std::vector<precision> &get_mpiwork() const { return terms.mpiwork; }
  //! sync the state across the mpi communicator
  void sync_mpi_state() const {
    if (terms.resources.num_ranks() > 1)
      terms.resources.bcast(state);
  }
  //! sync the state across the mpi communicator, then return
  std::vector<precision> const &current_state_mpi() const {
    sync_mpi_state();
    return state;
  }
  #else
  void sync_mpi_state() const {}
  std::vector<precision> const &current_state_mpi() const { return state; }
  #endif

  //! report memory usage by different componets
  void report_memusage(std::ostream &os = std::cout) const;

  // performs integration in time
  friend bool advance_in_time<precision>(
      discretization_manager<precision> &disc, int64_t num_steps);
  // this is the I/O manager
  friend class h5manager<precision>;
  // handles the time-integration meta-data
  friend struct time_advance_manager<precision>;
#endif // __ASGARD_DOXYGEN_SKIP_INTERNAL

protected:
#ifndef __ASGARD_DOXYGEN_SKIP_INTERNAL
  //! sets the initial conditions, performs adaptivity in the process
  void set_initial_condition();

  //! start from time 0 and nothing has been set
  void start_cold(pde_scheme<precision> &pde);
  //! restart from a file
  void restart_from_file(pde_scheme<precision> &pde);
  //! common operations for the two start methods
  void start_moments();
  //! recompute the moments, assuming moments are set, i.e., has_moments() is true
  void compute_moments_(group_id gid, std::vector<precision> const &f) const;
  //! recompute the moments locally, assuming moments are set, i.e., has_moments() is true
  void compute_moments_local(group_id gid, std::vector<precision> const &f) const;
  //! computes the right-hand-side of the ode, templated version
  void ode_rhs_base(group_id gid, precision time, std::vector<precision> const &current,
                    std::vector<precision> &R) const;
  //! computes next = current + scale_terms * F(x) + scale_src * sources(time)
  void ode_euler_base(group_id gid, precision time, std::vector<precision> const &current,
                      terms_scale term_scal, sources_scale source_scal,
                      std::vector<precision> &next) const;
  //! template version of ode right-hand-side sources
  template<data_mode mode>
  void ode_rhs_sources(group_id gid, precision time, precision alpha,
                       std::vector<precision> &src) const;
  #ifdef ASGARD_USE_GPU
  //! recompute the moments, assuming moments are set, i.e., has_moments() is true
  void compute_moments_gpu_(group_id gid, precision const f[]) const;
  //! recompute the moments locally, assuming moments are set, i.e., has_moments() is true
  void compute_moments_local_gpu(group_id gid, precision const f[]) const;
  //! same as ode_rhs_base() but the arrays are pre-allocated and on the GPU device 0
  void ode_rhs_base_gpu(group_id gid, precision time, precision const current[],
                        precision R[]) const;
  //! same as ode_euler_base() but the arrays are pre-allocated and on the GPU device 0
  void ode_euler_base_gpu(group_id gid, precision time, precision const current[],
                          terms_scale term_scal, sources_scale source_scal,
                          precision next[]) const;
  //! same as ode_rhs_sources() but the arrays are pre-allocated and on the GPU device 0
  template<data_mode mode>
  void ode_rhs_sources_gpu(group_id gid, precision time, precision alpha, precision src[]) const;
  #endif

  //! returns a snapshot of the state on the current MPI rank
  reconstruct_solution get_local_snapshot() const
  {
    reconstruct_solution shot(
        num_dims(), grid.num_indexes(), grid[0], degree(), state.data());

    std::array<double, max_num_dimensions> xmin, xmax;
    for (int d : iindexof(num_dims())) {
      xmin[d] = domain_.xleft(d);
      xmax[d] = domain_.xright(d);
    }

    shot.set_domain_bounds(xmin.data(), xmax.data());

    return shot;
  }
  //! refines the sparse grid using the given strategy and
  void refine(sparse_grid::strategy mode, std::vector<precision> const &f)
  {
    if (not is_leader())
      return;
    refinement.refine(conn, terms, f, mode, grid);
  }
  #ifdef ASGARD_USE_GPU
  //! refines the sparse grid using the given strategy and
  void refine(sparse_grid::strategy mode, gpu::vector<precision> const &f)
  {
    if (not is_leader())
      return;
    refinement.refine(conn, terms, f, mode, grid);
  }
  #endif

  #ifdef ASGARD_USE_MPI
  //! worker iteration apply
  void mpi_iteration_apply_base(group_id gid, std::vector<precision> &work) const;
  //! leader iteration apply
  void mpi_leader_apply_base(group_id gid, precision alpha, precision const x[],
                             precision beta, precision y[]) const;
  #ifdef ASGARD_USE_GPU
  void mpi_iteration_apply_base_gpu(group_id gid, precision work[]) const;
  void mpi_leader_apply_base_gpu(group_id gid, precision alpha, precision const x[],
                                 precision beta, precision y[]) const;
  #endif
  #endif
#endif // __ASGARD_DOXYGEN_SKIP_INTERNAL

private:
  // indicates the level of noise pushed to the cout
  mutable verbosity_level verb = verbosity_level::quiet;
  // indicates whether the time-stepper should perform sanity check accept/reject
  bool safe_step = false;
  // user provided options
  prog_opts options_;
  // initial conditions, non-separable
  md_func<precision> initial_md_;
  // initial conditions, separable
  std::vector<separable_func<precision>> initial_sep_;
  // pde-domain
  pde_domain<precision> domain_;

  sparse_grid grid;
  connection_patterns conn;
  hierarchy_manipulator<precision> hier;
  #ifdef ASGARD_USE_MPI
  int grid_synced_gen_ = -2;
  #endif

  refinement_manager<precision> refinement;

  // poisson solver data
  mutable solvers::poisson<precision> poisson;

  //! term manager holding coefficient matrices and kronmult meta-data
  mutable term_manager<precision> terms;
  //! time advance manager for the different methods
  time_advance_manager<precision> stepper;

  // constantly changing
  #ifdef ASGARD_USE_MPI
  // MPI processes may affect the state, e.g., sync state across the communicator
  mutable
  #endif
  std::vector<precision> state;

  #ifdef ASGARD_USE_GPU
  // remove when moments can be computed from the GPU directly
  mutable std::vector<precision> moments_workspace;
  #endif

  //! fields to store and save for plotting
  std::vector<aux_field_entry<precision>> aux_fields;
};

} // namespace asgard
