#include "bout/build_defines.hxx"

#if BOUT_HAS_PETSC

#include "snes.hxx"

#include <bout/boutcomm.hxx>
#include <bout/boutexception.hxx>
#include <bout/globals.hxx>
#include <bout/output.hxx>
#include <bout/output_bout_types.hxx>
#include <bout/petsc_interface.hxx>
#include <bout/utils.hxx>

#include <algorithm>
#include <cmath>
#include <cstddef>
#include <set>
#include <vector>

#include "petscerror.h"
#include "petscmat.h"
#include "petscpc.h"
#include "petscsnes.h"
#include "petscsys.h"
#include "petscsystypes.h"
#include "petscvec.h"

class ColoringStencil {
private:
  bool static isInSquare(int const i, int const j, int const n_square) {
    return std::abs(i) <= n_square && std::abs(j) <= n_square;
  }
  bool static isInCross(int const i, int const j, int const n_cross) {
    if (i == 0) {
      return std::abs(j) <= n_cross;
    }
    if (j == 0) {
      return std::abs(i) <= n_cross;
    }
    return false;
  }
  bool static isInTaxi(int const i, int const j, int const n_taxi) {
    return std::abs(i) + std::abs(j) <= n_taxi;
  }

public:
  auto static getOffsets(int n_square, int n_taxi, int n_cross) {
    ASSERT2(n_square >= 0 && n_cross >= 0 && n_taxi >= 0
            && n_square + n_cross + n_taxi > 0);
    auto inside = [&](int i, int j) {
      return isInSquare(i, j, n_square) || isInTaxi(i, j, n_taxi)
             || isInCross(i, j, n_cross);
    };
    std::vector<std::pair<int, int>> xy_offsets;
    auto loop_bound = std::max({n_square, n_taxi, n_cross});
    for (int i = -loop_bound; i <= loop_bound; ++i) {
      for (int j = -loop_bound; j <= loop_bound; ++j) {
        if (inside(i, j)) {
          xy_offsets.emplace_back(i, j);
        }
      }
    }
    return xy_offsets;
  }
};

namespace {
/*
 * PETSc callback function, which evaluates the nonlinear
 * function to be solved by SNES.
 *
 * This function assumes the context void pointer is a pointer
 * to an SNESSolver object.
 */
PetscErrorCode FormFunction(SNES UNUSED(snes), Vec x, Vec f, void* ctx) {
  return static_cast<SNESSolver*>(ctx)->snes_function(x, f, false);
}

/*!
 * PETSc callback function for forming Jacobian
 *
 * This function can be a linearised form of FormFunction
 */
PetscErrorCode FormFunctionForDifferencing(void* ctx, Vec x, Vec f) {
  return static_cast<SNESSolver*>(ctx)->snes_function(x, f, true);
}

/*!
 * SNES callback for forming Jacobian with coloring
 *
 * This can be a linearised and simplified form of FormFunction
 */
PetscErrorCode FormFunctionForColoring(void* UNUSED(snes), Vec x, Vec f, void* ctx) {
  return static_cast<SNESSolver*>(ctx)->snes_function(x, f, true);
}

PetscErrorCode snesPCapply(PC pc, Vec x, Vec y) {
  // Get the context
  SNESSolver* s;
  PetscCall(PCShellGetContext(pc, reinterpret_cast<void**>(&s)));

  PetscFunctionReturn(s->precon(x, y));
}
} // namespace

PetscErrorCode SNESSolver::FDJinitialise() {
  if (use_coloring) {
    // Use matrix coloring.
    // This greatly reduces the number of times the rhs() function
    // needs to be evaluated when calculating the Jacobian.

    // Use global mesh for now
    Mesh* mesh = bout::globals::mesh;

    //////////////////////////////////////////////////
    // Get the local indices by starting at 0
    Field3D index = globalIndex(0);

    //////////////////////////////////////////////////
    // Pre-allocate PETSc storage

    output_progress.write("Setting Jacobian matrix sizes\n");

    const int n2d = f2d.size();
    const int n3d = f3d.size();

    // Set size of Matrix on each processor to nlocal x nlocal
    MatCreate(BoutComm::get(), &Jfd);
    MatSetOption(Jfd, MAT_KEEP_NONZERO_PATTERN, PETSC_TRUE);
    MatSetSizes(Jfd, nlocal, nlocal, PETSC_DETERMINE, PETSC_DETERMINE);
    MatSetFromOptions(Jfd);
    // Determine which row/columns of the matrix are locally owned
    int Istart, Iend;
    MatGetOwnershipRange(Jfd, &Istart, &Iend);
    // Convert local into global indices
    // Note: Not in the boundary cells, to keep -1 values
    for (const auto& i : mesh->getRegion3D("RGN_NOBNDRY")) {
      index[i] += Istart;
    }
    // Now communicate to fill guard cells
    mesh->communicate(index);

    // Non-zero elements on this processor
    std::vector<PetscInt> d_nnz;
    std::vector<PetscInt> o_nnz;
    auto n_square = (*options)["stencil:square"]
                        .doc("Extent of stencil (square)")
                        .withDefault<int>(0);
    auto n_cross =
        (*options)["stencil:cross"].doc("Extent of stencil (cross)").withDefault<int>(0);
    // Set n_taxi 2 if nothing else is set
    auto n_taxi = (*options)["stencil:taxi"]
                      .doc("Extent of stencil (taxi-cab norm)")
                      .withDefault<int>((n_square == 0 && n_cross == 0) ? 2 : 0);

    auto const xy_offsets = ColoringStencil::getOffsets(n_square, n_taxi, n_cross);
    {
      // This is ugly but can't think of a better and robust way to
      // count the non-zeros for some arbitrary stencil
      // effectively the same loop as the one that sets the non-zeros below
      std::vector<std::set<int>> d_nnz_map2d(nlocal);
      std::vector<std::set<int>> o_nnz_map2d(nlocal);
      std::vector<std::set<int>> d_nnz_map3d(nlocal);
      std::vector<std::set<int>> o_nnz_map3d(nlocal);
      // Loop over every element in 2D to count the *unique* non-zeros
      for (int x = mesh->xstart; x <= mesh->xend; x++) {
        for (int y = mesh->ystart; y <= mesh->yend; y++) {

          const int ind0 = ROUND(index(x, y, 0)) - Istart;

          // 2D fields
          for (int i = 0; i < n2d; i++) {
            const PetscInt row = ind0 + i;
            // Loop through each point in the stencil
            for (const auto& [x_off, y_off] : xy_offsets) {
              const int xi = x + x_off;
              const int yi = y + y_off;
              if ((xi < 0) || (yi < 0) || (xi >= mesh->LocalNx)
                  || (yi >= mesh->LocalNy)) {
                continue;
              }

              const int ind2 = ROUND(index(xi, yi, 0));
              if (ind2 < 0) {
                continue; // A boundary point
              }

              // Depends on all variables on this cell
              for (int j = 0; j < n2d; j++) {
                const PetscInt col = ind2 + j;
                if (col >= Istart && col < Iend) {
                  d_nnz_map2d[row].insert(col);
                } else {
                  o_nnz_map2d[row].insert(col);
                }
              }
            }
          }
          // 3D fields
          for (int z = mesh->zstart; z <= mesh->zend; z++) {
            const int ind = ROUND(index(x, y, z)) - Istart;

            for (int i = 0; i < n3d; i++) {
              PetscInt row = ind + i;
              if (z == 0) {
                row += n2d;
              }

              // Depends on 2D fields
              for (int j = 0; j < n2d; j++) {
                const PetscInt col = ind0 + j;
                if (col >= Istart && col < Iend) {
                  d_nnz_map2d[row].insert(col);
                } else {
                  o_nnz_map2d[row].insert(col);
                }
              }

              // Star pattern
              for (const auto& [x_off, y_off] : xy_offsets) {
                const int xi = x + x_off;
                const int yi = y + y_off;

                if ((xi < 0) || (yi < 0) || (xi >= mesh->LocalNx)
                    || (yi >= mesh->LocalNy)) {
                  continue;
                }

                int ind2 = ROUND(index(xi, yi, 0));
                if (ind2 < 0) {
                  continue; // Boundary point
                }

                if (z == 0) {
                  ind2 += n2d;
                }

                // 3D fields on this cell
                for (int j = 0; j < n3d; j++) {
                  const PetscInt col = ind2 + j;
                  if (col >= Istart && col < Iend) {
                    d_nnz_map3d[row].insert(col);
                  } else {
                    o_nnz_map3d[row].insert(col);
                  }
                }
              }
            }
          }
        }
      }

      d_nnz.reserve(nlocal);
      d_nnz.reserve(nlocal);

      for (int i = 0; i < nlocal; ++i) {
        // Assume all elements in the z direction are potentially coupled
        d_nnz.emplace_back((d_nnz_map3d[i].size() * mesh->LocalNz)
                           + d_nnz_map2d[i].size());
        o_nnz.emplace_back((o_nnz_map3d[i].size() * mesh->LocalNz)
                           + o_nnz_map2d[i].size());
      }
    }

    output_progress.write("Pre-allocating Jacobian\n");
    // Pre-allocate
    MatMPIAIJSetPreallocation(Jfd, 0, d_nnz.data(), 0, o_nnz.data());
    MatSeqAIJSetPreallocation(Jfd, 0, d_nnz.data());
    MatSetUp(Jfd);
    MatSetOption(Jfd, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_TRUE);

    //////////////////////////////////////////////////
    // Mark non-zero entries

    output_progress.write("Marking non-zero Jacobian entries\n");
    const PetscScalar val = 1.0;
    for (int x = mesh->xstart; x <= mesh->xend; x++) {
      for (int y = mesh->ystart; y <= mesh->yend; y++) {

        const int ind0 = ROUND(index(x, y, 0));

        // 2D fields
        for (int i = 0; i < n2d; i++) {
          const PetscInt row = ind0 + i;

          // Loop through each point in the stencil
          for (const auto& [x_off, y_off] : xy_offsets) {
            const int xi = x + x_off;
            const int yi = y + y_off;
            if ((xi < 0) || (yi < 0) || (xi >= mesh->LocalNx) || (yi >= mesh->LocalNy)) {
              continue;
            }

            const int ind2 = ROUND(index(xi, yi, 0));
            if (ind2 < 0) {
              continue; // A boundary point
            }

            // Depends on all variables on this cell
            for (int j = 0; j < n2d; j++) {
              const PetscInt col = ind2 + j;
              PetscCall(MatSetValues(Jfd, 1, &row, 1, &col, &val, INSERT_VALUES));
            }
          }
        }
        // 3D fields
        for (int z = mesh->zstart; z <= mesh->zend; z++) {
          int ind = ROUND(index(x, y, z));

          for (int i = 0; i < n3d; i++) {
            PetscInt row = ind + i;
            if (z == 0) {
              row += n2d;
            }

            // Depends on 2D fields
            for (int j = 0; j < n2d; j++) {
              const PetscInt col = ind0 + j;
              PetscCall(MatSetValues(Jfd, 1, &row, 1, &col, &val, INSERT_VALUES));
            }

            // Star pattern
            for (const auto& [x_off, y_off] : xy_offsets) {
              int xi = x + x_off;
              int yi = y + y_off;

              if ((xi < 0) || (yi < 0) || (xi >= mesh->LocalNx)
                  || (yi >= mesh->LocalNy)) {
                continue;
              }
              for (int zi = mesh->zstart; zi <= mesh->zend; ++zi) {
                int ind2 = ROUND(index(xi, yi, zi));
                if (ind2 < 0) {
                  continue; // Boundary point
                }

                if (z == 0) {
                  ind2 += n2d;
                }

                // 3D fields on this cell
                for (int j = 0; j < n3d; j++) {
                  const PetscInt col = ind2 + j;
                  PetscErrorCode ierr =
                      MatSetValues(Jfd, 1, &row, 1, &col, &val, INSERT_VALUES);

                  if (ierr != PETSC_SUCCESS) {
                    output.write("ERROR: {} {} : ({}, {}) -> ({}, {}) : {} -> {}\n", row,
                                 col, x, y, xi, yi, ind2, ind2 + n3d - 1);
                  }
                  CHKERRQ(ierr);
                }
              }
            }
          }
        }
      }
    }

    // Finished marking non-zero entries

    output_progress.write("Assembling Jacobian matrix\n");

    // Assemble Matrix
    MatAssemblyBegin(Jfd, MAT_FINAL_ASSEMBLY);
    MatAssemblyEnd(Jfd, MAT_FINAL_ASSEMBLY);

    {
      // Test if the matrix is symmetric
      // Values are 0 or 1 so tolerance (1e-5) shouldn't matter
      PetscBool symmetric;
      PetscCall(MatIsSymmetric(Jfd, 1e-5, &symmetric));
      if (!static_cast<bool>(symmetric)) {
        output_warn.write("Jacobian pattern is not symmetric\n");
      }
    }

    // The above can miss entries around the X-point branch cut:
    // The diagonal terms are complicated because moving in X then Y
    // is different from moving in Y then X at the X-point.
    // Making sure the colouring matrix is symmetric does not
    // necessarily give the correct stencil but may help.
    if ((*options)["force_symmetric_coloring"]
            .doc("Modifies coloring matrix to force it to be symmetric")
            .withDefault<bool>(false)) {
      Mat Jfd_T;
      MatCreateTranspose(Jfd, &Jfd_T);
      MatAXPY(Jfd, 1, Jfd_T, DIFFERENT_NONZERO_PATTERN);
    }

    output_progress.write("Creating Jacobian coloring\n");
    updateColoring();

    if (prune_jacobian) {
      // Will remove small elements from the Jacobian.
      // Save a copy to recover from over-pruning
      PetscCall(MatDuplicate(Jfd, MAT_SHARE_NONZERO_PATTERN, &Jfd_original));
    }
  } else {
    // Brute force calculation
    // There is usually no reason to use this, except as a check of
    // the coloring calculation.

    MatCreateAIJ(
        BoutComm::get(), nlocal, nlocal,  // Local sizes
        PETSC_DETERMINE, PETSC_DETERMINE, // Global sizes
        3, // Number of nonzero entries in diagonal portion of local submatrix
        nullptr,
        0, // Number of nonzeros per row in off-diagonal portion of local submatrix
        nullptr, &Jfd);

    if (matrix_free_operator) {
      SNESSetJacobian(snes, Jmf, Jfd, SNESComputeJacobianDefault, this);
    } else {
      SNESSetJacobian(snes, Jfd, Jfd, SNESComputeJacobianDefault, this);
    }

    MatSetOption(Jfd, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_FALSE);
  }
  return PETSC_SUCCESS;
}

PetscErrorCode SNESSolver::FDJpruneJacobian() {
#if PETSC_VERSION_GE(3, 20, 0)

  // Remove small elements from the Jacobian and recompute the coloring
  // Only do this if there are a significant number of small elements.
  int small_elements = 0;
  int total_elements = 0;

  // Get index of rows owned by this processor
  int rstart, rend;
  MatGetOwnershipRange(Jfd, &rstart, &rend);

  PetscInt ncols;
  const PetscScalar* vals;
  for (int row = rstart; row < rend; row++) {
    MatGetRow(Jfd, row, &ncols, nullptr, &vals);
    for (int col = 0; col < ncols; col++) {
      if (std::abs(vals[col]) < prune_abstol) {
        ++small_elements;
      }
      ++total_elements;
    }
    MatRestoreRow(Jfd, row, &ncols, nullptr, &vals);
  }

  if (small_elements > prune_fraction * total_elements) {
    if (diagnose) {
      output.write("\nPruning Jacobian elements: {} / {}\n", small_elements,
                   total_elements);
    }

    // Prune Jacobian, keeping diagonal elements
    PetscCall(MatFilter(Jfd, prune_abstol, PETSC_TRUE, PETSC_TRUE));

    // Update the coloring from Jfd matrix
    updateColoring();

    // Mark the Jacobian as pruned. This is so that it is only restored if pruned.
    jacobian_pruned = true;
  }
#endif // PETSC_VERSION_GE(3,20,0)
  return PETSC_SUCCESS;
}

PetscErrorCode SNESSolver::FDJrestoreFromPruning() {
  // Restore pruned non-zero elements
  PetscCall(MatCopy(Jfd_original, Jfd, DIFFERENT_NONZERO_PATTERN));
  // The non-zero pattern has changed, so update coloring
  updateColoring();
  jacobian_pruned = false; // Reset flag. Will be set after pruning.
  return PETSC_SUCCESS;
}

SNESSolver::SNESSolver(Options* opts)
    : Solver(opts),
      output_trigger(
          (*options)["output_trigger"]
              .doc("Decides when to save outputs. fixed_time_interval, residual_ratio")
              .withDefault(BoutSnesOutput::fixed_time_interval)),
      output_residual_ratio(
          (*options)["output_residual_ratio"]
              .doc("Trigger an output when residual falls by this ratio")
              .withDefault(0.5)),
      timestep(
          (*options)["timestep"].doc("Initial backward Euler timestep").withDefault(1.0)),
      dt_min_reset((*options)["dt_min_reset"]
                       .doc("If dt falls below this, reset to starting dt")
                       .withDefault(1e-6)),
      max_timestep((*options)["max_timestep"].doc("Maximum timestep").withDefault(1e37)),
      equation_form(
          (*options)["equation_form"]
              .doc("Form of equation to solve: rearranged_backward_euler (default);"
                   " pseudo_transient; backward_euler; direct_newton")
              .withDefault(BoutSnesEquationForm::rearranged_backward_euler)),
      snes_type((*options)["snes_type"]
                    .doc("PETSc nonlinear solver method to use")
                    .withDefault("anderson")),
      atol((*options)["atol"].doc("Absolute tolerance in SNES solve").withDefault(1e-16)),
      rtol((*options)["rtol"].doc("Relative tolerance in SNES solve").withDefault(1e-10)),
      stol((*options)["stol"]
               .doc("Convergence tolerance in terms of the norm of the change in "
                    "the solution between steps")
               .withDefault(1e-8)),
      maxf((*options)["maxf"]
               .doc("Maximum number of function evaluations per SNES solve")
               .withDefault(10000)),
      maxits((*options)["max_nonlinear_iterations"]
                 .doc("Maximum number of nonlinear iterations per SNES solve")
                 .withDefault(50)),
      lower_its((*options)["lower_its"]
                    .doc("Iterations below which the next timestep is increased")
                    .withDefault(static_cast<int>(maxits * 0.5))),
      upper_its((*options)["upper_its"]
                    .doc("Iterations above which the next timestep is reduced")
                    .withDefault(static_cast<int>(maxits * 0.8))),
      max_snes_failures((*options)["max_snes_failures"]
                            .doc("Abort after this number of consecutive failures")
                            .withDefault(10)),
      timestep_factor_on_failure((*options)["timestep_factor_on_failure"]
                                     .doc("Multiply timestep on convergence failure")
                                     .withDefault(0.5)),
      timestep_factor_on_upper_its(
          (*options)["timestep_factor_on_upper_its"]
              .doc("Multiply timestep if iterations exceed upper_its")
              .withDefault(0.9)),
      timestep_factor_on_lower_its(
          (*options)["timestep_factor_on_lower_its"]
              .doc("Multiply timestep if iterations are below lower_its")
              .withDefault(1.4)),
      pseudo_strategy((*options)["pseudo_strategy"]
                          .doc("PTC strategy to use when setting timesteps")
                          .withDefault(BoutPTCStrategy::inverse_residual)),
      pseudo_alpha((*options)["pseudo_alpha"]
                       .doc("Sets timestep using dt = alpha / residual")
                       .withDefault(100. * atol * timestep)),
      pseudo_growth_factor((*options)["pseudo_growth_factor"]
                               .doc("PTC growth factor on success")
                               .withDefault(1.1)),
      pseudo_reduction_factor((*options)["pseudo_reduction_factor"]
                                  .doc("PTC reduction factor on failure")
                                  .withDefault(0.5)),
      pseudo_max_ratio((*options)["pseudo_max_ratio"]
                           .doc("PTC maximum timestep ratio between neighbors")
                           .withDefault(2.)),
      timestep_control((*options)["timestep_control"]
                           .doc("Timestep control method")
                           .withDefault(BoutSnesTimestep::pid_nonlinear_its)),
      timestep_factor((*options)["timestep_factor"]
                          .doc("When timestep_control=residual_ratio, multiply timestep "
                               "by this factor each step")
                          .withDefault(1.1)),
      target_its((*options)["target_its"].doc("Target snes iterations").withDefault(7)),
      kP((*options)["kP"].doc("Proportional PID parameter").withDefault(0.7)),
      kI((*options)["kI"].doc("Integral PID parameter").withDefault(0.3)),
      kD((*options)["kD"].doc("Derivative PID parameter").withDefault(0.2)),
      pid_consider_failures(
          (*options)["pid_consider_failures"]
              .doc("Reduce timestep increases if recent solves have failed")
              .withDefault<bool>(false)),
      last_failure_weight((*options)["last_failure_weight"]
                              .doc("Weighting of last timestep in recent failure rate")
                              .withDefault(0.1)),
      diagnose(
          (*options)["diagnose"].doc("Print additional diagnostics").withDefault(false)),
      diagnose_failures((*options)["diagnose_failures"]
                            .doc("Print more diagnostics when SNES fails")
                            .withDefault<bool>(false)),
      predictor((*options)["predictor"].doc("Use linear predictor?").withDefault(true)),
      use_precon((*options)["use_precon"]
                     .doc("Use user-supplied preconditioner?")
                     .withDefault<bool>(false)),
      ksp_type((*options)["ksp_type"]
                   .doc("Linear solver type. By default let PETSc decide (gmres)")
                   .withDefault("default")),
      kspsetinitialguessnonzero((*options)["kspsetinitialguessnonzero"]
                                    .doc("Set the initial guess to be non-zero")
                                    .withDefault<bool>(false)),
      maxl((*options)["maxl"].doc("Maximum number of linear iterations").withDefault(20)),
      pc_type(
          (*options)["pc_type"]
              .doc("Preconditioner type. By default lets PETSc decide (ilu or bjacobi)")
              .withDefault("default")),
      pc_hypre_type((*options)["pc_hypre_type"]
                        .doc("hypre preconditioner type: euclid, pilut, parasails, "
                             "boomeramg, ams, ads")
                        .withDefault("pilut")),
      line_search_type((*options)["line_search_type"]
                           .doc("Line search type: basic, bt, l2, cp, nleqerr")
                           .withDefault("default")),
      matrix_free((*options)["matrix_free"]
                      .doc("Use matrix free Jacobian?")
                      .withDefault<bool>(false)),
      matrix_free_operator((*options)["matrix_free_operator"]
                               .doc("Use matrix free Jacobian-vector operator?")
                               .withDefault<bool>(false)),
      lag_jacobian((*options)["lag_jacobian"]
                       .doc("Re-use the Jacobian this number of SNES iterations")
                       .withDefault(50)),
      jacobian_persists(
          (*options)["jacobian_persists"]
              .doc("Re-use Jacobian and preconditioner across nonlinear solves")
              .withDefault<bool>(false)),
      use_coloring((*options)["use_coloring"]
                       .doc("Use matrix coloring to calculate Jacobian?")
                       .withDefault<bool>(true)),
      jacobian_recalculated(false),
      prune_jacobian((*options)["prune_jacobian"]
                         .doc("Remove small elements in the Jacobian?")
                         .withDefault<bool>(false)),
      prune_abstol((*options)["prune_abstol"]
                       .doc("Prune values with absolute values smaller than this")
                       .withDefault<BoutReal>(1e-16)),
      prune_fraction((*options)["prune_fraction"]
                         .doc("Prune if fraction of small elements is larger than this")
                         .withDefault<BoutReal>(0.2)),
      scale_rhs((*options)["scale_rhs"]
                    .doc("Scale time derivatives (Jacobian row scaling)?")
                    .withDefault<bool>(false)),
      scale_vars((*options)["scale_vars"]
                     .doc("Scale variables (Jacobian column scaling)?")
                     .withDefault<bool>(false)),
      asinh_vars((*options)["asinh_vars"]
                     .doc("Apply asinh() to all variables?")
                     .withDefault<bool>(false)) {}

int SNESSolver::init() {
  Solver::init();
  output << "\n\tSNES steady state solver\n";

  // Calculate number of variables
  nlocal = getLocalN();

  // Get total problem size
  int ntmp;
  if (bout::globals::mpi->MPI_Allreduce(&nlocal, &ntmp, 1, MPI_INT, MPI_SUM,
                                        BoutComm::get())
      != 0) {
    throw BoutException("MPI_Allreduce failed!");
  }
  neq = ntmp;

  output_info.write("\t3d fields = {:d}, 2d fields = {:d} neq={:d}, local_N={:d}\n",
                    n3Dvars(), n2Dvars(), neq, nlocal);

  // Initialise PETSc components

  // Vectors
  output_info.write("Creating vector\n");
  PetscCall(VecCreate(BoutComm::get(), &snes_x));
  PetscCall(VecSetSizes(snes_x, nlocal, PETSC_DECIDE));
  PetscCall(VecSetFromOptions(snes_x));

  PetscCall(VecDuplicate(snes_x, &snes_f));
  PetscCall(VecDuplicate(snes_x, &x0));

  if ((equation_form == BoutSnesEquationForm::rearranged_backward_euler)
      || (equation_form == BoutSnesEquationForm::pseudo_transient)) {
    // Need an intermediate vector for rearranged Backward Euler or Pseudo-Transient Continuation
    PetscCall(VecDuplicate(snes_x, &delta_x));
  }

  if (predictor) {
    // Storage for previous solution
    PetscCall(VecDuplicate(snes_x, &x1));
  }

  if (scale_rhs) {
    // Storage for rhs factors, one per evolving variable
    PetscCall(VecDuplicate(snes_x, &rhs_scaling_factors));
    // Set all factors to 1 to start with
    PetscCall(VecSet(rhs_scaling_factors, 1.0));
    // Array to store inverse Jacobian row norms
    PetscCall(VecDuplicate(snes_x, &jac_row_inv_norms));
  }

  if (scale_vars) {
    // Storage for var factors, one per evolving variable
    PetscCall(VecDuplicate(snes_x, &var_scaling_factors));
    // Set all factors to 1 to start with
    PetscCall(VecSet(var_scaling_factors, 1.0));
    // Storage for scaled 'x' state vectors
    PetscCall(VecDuplicate(snes_x, &scaled_x));
  } else if (asinh_vars) {
    PetscCall(VecDuplicate(snes_x, &scaled_x));
  }

  if (equation_form == BoutSnesEquationForm::pseudo_transient) {
    PetscCall(initPseudoTimestepping());
  }
  // Per-cell residuals
  local_residual = 0.0;
  local_residual_2d = 0.0;
  global_residual = 0.0;

  // Nonlinear solver interface (SNES)
  output_info.write("Create SNES\n");
  SNESCreate(BoutComm::get(), &snes);

  // Set the callback function
  SNESSetFunction(snes, snes_f, FormFunction, this);

  SNESSetType(snes, snes_type.c_str());

  // Line search
  if (line_search_type != "default") {
    SNESLineSearch linesearch;
    SNESGetLineSearch(snes, &linesearch);
    SNESLineSearchSetType(linesearch, line_search_type.c_str());
  }

  // Set up the Jacobian
  if (matrix_free or matrix_free_operator) {
    /*
      PETSc SNES matrix free Jacobian, using a different
      operator for differencing.

      See PETSc examples
      http://www.mcs.anl.gov/petsc/petsc-current/src/snes/examples/tests/ex7.c.html
      and this thread:
      http://lists.mcs.anl.gov/pipermail/petsc-users/2014-January/020075.html

     */
    MatCreateSNESMF(snes, &Jmf);

    // Set a function to be called for differencing
    // This can be a linearised form of the SNES function
    MatMFFDSetFunction(Jmf, FormFunctionForDifferencing, this);
  }

  if (matrix_free) {
    // Use matrix free for both operator and preconditioner
    // Calculate Jacobian matrix free using FormFunctionForDifferencing
    SNESSetJacobian(snes, Jmf, Jmf, MatMFFDComputeJacobian, this);

  } else {
    // Calculate the Jacobian using finite differences.  The finite
    // difference Jacobian (Jfd) may be used for both operator and
    // preconditioner or, if matrix_free_operator, in only the
    // preconditioner.

    // Create a vector to store interpolated output solution
    // Used so that the timestep does not have to be adjusted,
    // because that would require updating the preconditioner.
    PetscCall(VecDuplicate(snes_x, &output_x));

    // Initialize the Finite Difference Jacobian
    PetscCall(FDJinitialise());

    // Re-use Jacobian
    // Note: If the 'Amat' Jacobian is matrix free, SNESComputeJacobian
    //       always updates its reference 'u' vector every nonlinear iteration
    SNESSetLagJacobian(snes, lag_jacobian);
    nl_its_prev = target_its;
    nl_its_prev2 = target_its;
    PetscCall(
        SNESSetLagJacobianPersists(snes, static_cast<PetscBool>(jacobian_persists)));
    PetscCall(SNESSetLagPreconditionerPersists(
        snes, static_cast<PetscBool>(jacobian_persists)));
    PetscCall(SNESSetLagPreconditioner(snes, 1)); // Rebuild when Jacobian is rebuilt
  }

  // Set tolerances
  SNESSetTolerances(snes, atol, rtol, stol, maxits, maxf);

  // Force SNES to take at least one nonlinear iteration.
  // This may prevent the solver from getting stuck in false steady state conditions
#if PETSC_VERSION_GE(3, 8, 0)
  SNESSetForceIteration(snes, PETSC_TRUE);
#endif

  // Get KSP context from SNES
  KSP ksp;
  SNESGetKSP(snes, &ksp);

  if (ksp_type != "default") {
    KSPSetType(ksp, ksp_type.c_str());
  }

  if (kspsetinitialguessnonzero) {
    // Set the initial guess to be non-zero
    KSPSetInitialGuessNonzero(ksp, PETSC_TRUE);
  }

  KSPSetTolerances(ksp,
                   PETSC_DEFAULT, // rtol
                   PETSC_DEFAULT, // abstol
                   PETSC_DEFAULT, // dtol (divergence tolerance)
                   maxl);         // Maximum number of iterations

  // Get PC context from KSP
  PC pc;
  KSPGetPC(ksp, &pc);

  if (use_precon && hasPreconditioner()) {
    output_info.write("\tUsing user-supplied preconditioner\n");

    // Set a Shell (matrix-free) preconditioner type
    PCSetType(pc, PCSHELL);

    // Specify the preconditioner function
    PCShellSetApply(pc, snesPCapply);
    // Context used to supply object pointer
    PCShellSetContext(pc, this);
  } else if (matrix_free) {
    // Can't use preconditioner because no Jacobian matrix available
    PCSetType(pc, PCNONE);
  } else {
    // Set PC type from input
    if (pc_type != "default") {
      PCSetType(pc, pc_type.c_str());

      if (pc_type == "hypre") {
#if PETSC_HAVE_HYPRE
        // Set the type of hypre preconditioner
        PCHYPRESetType(pc, pc_hypre_type.c_str());
#else
        throw BoutException("PETSc was not configured with Hypre.");
#endif
      }
    }
  }

  // Get runtime options
  lib.setOptionsFromInputFile(snes);

  {
    // Some reporting
    PCType pctype;
    PCGetType(pc, &pctype);
    KSPType ksptype;
    KSPGetType(ksp, &ksptype);
    SNESType snestype;
    SNESGetType(snes, &snestype);
    output_info.write("SNES Type : {}\n", snestype);
    if (ksptype != nullptr) {
      output_info.write("KSP Type  : {}\n", ksptype);
    }
    if (pctype != nullptr) {
      output_info.write("PC Type   : {}\n", pctype);
    }
  }

  return 0;
}

int SNESSolver::run() {
  // Set initial guess at the solution from variables
  {
    BoutReal* xdata = nullptr;
    PetscCall(VecGetArray(snes_x, &xdata));
    save_vars(xdata);

    if (asinh_vars) {
      // Evolving asinh(vars)
      PetscInt size;
      PetscCall(VecGetLocalSize(snes_x, &size));
      for (PetscInt i = 0; i != size; ++i) {
        xdata[i] = std::asinh(xdata[i] / asinh_scale);
      }
    }

    PetscCall(VecRestoreArray(snes_x, &xdata));
  }

  // Initialise residuals
  local_residual = 0.0;
  local_residual_2d = 0.0;
  global_residual = 0.0;
  PetscCall(updateResiduals(snes_x));
  if (diagnose) {
    output.write("\n    Residual: {}\n", global_residual);
  }

  BoutReal target = simtime;
  recent_failure_rate = 0.0;
  for (int s = 0; s < getNumberOutputSteps(); s++) {
    target += getOutputTimestep();

    bool looping = true;
    int snes_failures = 0; // Count SNES convergence failures
    int saved_jacobian_lag = 0;
    int loop_count = 0;

    BoutReal start_global_residual = global_residual;
    do {
      if ((output_trigger == BoutSnesOutput::fixed_time_interval && (simtime >= target))
          || (output_trigger == BoutSnesOutput::residual_ratio
              && (global_residual <= start_global_residual * output_residual_ratio))) {
        break; // Could happen if step over multiple outputs
      }

      if (scale_vars) {
        // Individual variable scaling
        // Note: If variables are rescaled then the Jacobian columns
        //       need to be scaled or recalculated

        if (loop_count % 100 == 0) {
          // Rescale state (snes_x) so that all quantities are around 1
          // If quantities are near zero then RTOL is used
          int istart, iend;
          VecGetOwnershipRange(snes_x, &istart, &iend);

          // Take ownership of snes_x and var_scaling_factors data
          PetscScalar* snes_x_data = nullptr;
          PetscCall(VecGetArray(snes_x, &snes_x_data));
          PetscScalar* x1_data;
          PetscCall(VecGetArray(x1, &x1_data));
          PetscScalar* var_scaling_factors_data;
          PetscCall(VecGetArray(var_scaling_factors, &var_scaling_factors_data));

          // Normalise each value in the state
          // Limit normalisation so scaling factor is never smaller than rtol
          for (int i = 0; i < iend - istart; ++i) {
            const PetscScalar norm =
                BOUTMAX(std::abs(snes_x_data[i]), rtol / var_scaling_factors_data[i]);
            snes_x_data[i] /= norm;
            x1_data[i] /= norm; // Update history for predictor
            var_scaling_factors_data[i] *= norm;
          }

          // Restore vector underlying data
          PetscCall(VecRestoreArray(var_scaling_factors, &var_scaling_factors_data));
          PetscCall(VecRestoreArray(x1, &x1_data));
          PetscCall(VecRestoreArray(snes_x, &snes_x_data));

          if (diagnose) {
            // Print maximum and minimum scaling factors
            PetscReal max_scale, min_scale;
            VecMax(var_scaling_factors, nullptr, &max_scale);
            VecMin(var_scaling_factors, nullptr, &min_scale);
            output.write("Var scaling: {} -> {}\n", min_scale, max_scale);
          }

          // Force recalculation of the Jacobian
          SNESGetLagJacobian(snes, &saved_jacobian_lag);
          SNESSetLagJacobian(snes, 1);
        }
      }
      ++loop_count;

      // Copy the state (snes_x) into initial values (x0)
      VecCopy(snes_x, x0);

      if (equation_form == BoutSnesEquationForm::pseudo_transient) {
        // Pseudo-Transient Continuation
        // Each evolving quantity may have its own timestep
        // Set timestep and dt scalars to the minimum of dt_vec

        PetscCall(VecMin(dt_vec, nullptr, &timestep));
        dt = timestep;

        if (output_trigger == BoutSnesOutput::fixed_time_interval
            && simtime + timestep >= target) {
          looping = false;
        }
      } else {
        // Timestepping

        if (timestep < dt_min_reset) {
          // Hit the minimum timestep, probably through repeated failures

          if (saved_jacobian_lag != 0) {
            // Already tried this and it didn't work
            throw BoutException("Solver failed after many attempts");
          }

          // Try resetting the preconditioner, turn off predictor, and use a large timestep
          SNESGetLagJacobian(snes, &saved_jacobian_lag);
          SNESSetLagJacobian(snes, 1);
          timestep = getOutputTimestep();
          predictor = false; // Predictor can cause problems in near steady-state.
        }

        // Set the timestep
        dt = timestep;
        if (output_trigger == BoutSnesOutput::fixed_time_interval
            && simtime + dt >= target) {
          // Note: When the timestep is changed the preconditioner needs to be updated
          // => Step over the output time and interpolate if not matrix free

          if (matrix_free) {
            // Ensure that the timestep goes to the next output time and then stops.
            // This avoids the need to interpolate
            dt = target - simtime;
          }
          looping = false;
        }

        if (predictor and (time1 > 0.0)) {
          // Use (time1, x1) and (simtime, x0) to make prediction
          // snes_x <- x0 + (dt / (simtime - time1)) * (x0 - x1)
          // snes_x <- -β * x1 + (1 + β) * snes_x
          const BoutReal beta = dt / (simtime - time1);
          VecAXPBY(snes_x, -beta, (1. + beta), x1);
        }

        SNESSetLagJacobian(snes, lag_jacobian);
      }

      // Run the solver
      PetscErrorCode ierr = SNESSolve(snes, nullptr, snes_x);

      // Find out if converged
      SNESConvergedReason reason;
      SNESGetConvergedReason(snes, &reason);

      // Get number of iterations
      int nl_its;
      SNESGetIterationNumber(snes, &nl_its);
      int lin_its;
      SNESGetLinearSolveIterations(snes, &lin_its);

      // Rolling average of recent failures
      recent_failure_rate *= 1. - last_failure_weight;

      if ((ierr != PETSC_SUCCESS) or (reason < 0)) {
        // Diverged or SNES failed

        recent_failure_rate += last_failure_weight;

        ++snes_failures;

        if (diagnose_failures or (snes_failures == max_snes_failures)) {
          // Print diagnostics to help identify source of the problem

          output.write("\n======== SNES failed =========\n");
          output.write("\nReturn code: {}, reason: {}\n", ierr, static_cast<int>(reason));
          for (const auto& f : f2d) {
            output.write(
                "{} : ({} -> {}), ddt: ({} -> {})\n", f.name,
                min(*f.var, true, "RGN_NOBNDRY"), max(*f.var, true, "RGN_NOBNDRY"),
                min(*f.F_var, true, "RGN_NOBNDRY"), max(*f.F_var, true, "RGN_NOBNDRY"));
          }
          for (const auto& f : f3d) {
            output.write(
                "{} : ({} -> {}), ddt: ({} -> {})\n", f.name,
                min(*f.var, true, "RGN_NOBNDRY"), max(*f.var, true, "RGN_NOBNDRY"),
                min(*f.F_var, true, "RGN_NOBNDRY"), max(*f.F_var, true, "RGN_NOBNDRY"));
          }
        }

        if (snes_failures == max_snes_failures) {
          output.write("Too many SNES failures ({}). Aborting.", snes_failures);
          return 1;
        }

        if (equation_form == BoutSnesEquationForm::pseudo_transient) {
          if (snes_failures == max_snes_failures - 1) {
            // Last chance. Set to uniform smallest timestep
            PetscCall(VecSet(dt_vec, dt_min_reset));

          } else if (snes_failures == 5) {
            // Set uniform timestep
            PetscCall(VecSet(dt_vec, timestep));
          } else {
            // Global scaling of timesteps
            // Note: A better strategy might be to reduce timesteps
            //       in problematic cells.
            PetscCall(VecScale(dt_vec, timestep_factor_on_failure));
          }
        } else {
          // Try a smaller timestep
          timestep *= timestep_factor_on_failure;
        }
        // Restore state
        VecCopy(x0, snes_x);

        // Recalculate the Jacobian
        if (jacobian_pruned and (snes_failures > 2) and (4 * lin_its > 3 * maxl)) {
          // Taking 3/4 of maximum linear iterations on average per linear step
          // May indicate a preconditioner problem.
          if (diagnose) {
            output.write("\nRestoring Jacobian\n");
          }
          PetscCall(FDJrestoreFromPruning());
        }

        if (saved_jacobian_lag == 0) {
          // This triggers a Jacobian recalculation
          SNESGetLagJacobian(snes, &saved_jacobian_lag);
          SNESSetLagJacobian(snes, 1);
        }

        // Check lock state
        PetscInt lock_state;
        VecLockGet(snes_x, &lock_state);
        if (lock_state > 0) {
          // Locked for read
          output_warn.write("WARNING: snes_x locked for reading\n");
        } else if (lock_state < 0) {
          // Locked for write
          output_warn.write("WARNING: snes_x locked for writing\n");
        }
        looping = true;
        continue; // Try again
      }

      if (saved_jacobian_lag != 0) {
        // Following successful step, reset Jacobian lag
        // to previous value
        SNESSetLagJacobian(snes, saved_jacobian_lag);
        saved_jacobian_lag = 0;
      }

      if (predictor) {
        // Save previous values: x1 <- x0
        VecCopy(x0, x1);
        time1 = simtime;
      }

      if (nl_its == 0) {
        // This can occur even with SNESSetForceIteration
        // Results in simulation state freezing and rapidly going
        // to the end

        if (scale_vars) {
          // scaled_x <- snes_x * var_scaling_factors
          PetscCall(VecPointwiseMult(scaled_x, snes_x, var_scaling_factors));

          const BoutReal* xdata = nullptr;
          PetscCall(VecGetArrayRead(scaled_x, &xdata));
          load_vars(const_cast<BoutReal*>(xdata));
          PetscCall(VecRestoreArrayRead(scaled_x, &xdata));
        } else {
          const BoutReal* xdata = nullptr;
          PetscCall(VecGetArrayRead(snes_x, &xdata));
          load_vars(const_cast<BoutReal*>(xdata));
          PetscCall(VecRestoreArrayRead(snes_x, &xdata));
        }

        try {
          run_rhs(simtime);
        } catch (BoutException& e) {
          output_error.write("ERROR: BoutException thrown: {}\n", e.what());
          // Abort simulation. There is no way to recover and
          // synchronise unless all processors throw exceptions
          // together
          BoutComm::abort(1);
        }

        // Copy derivatives back
        {
          BoutReal* fdata = nullptr;
          PetscCall(VecGetArray(snes_f, &fdata));
          save_derivs(fdata);
          PetscCall(VecRestoreArray(snes_f, &fdata));
        }

        // Forward Euler
        VecAXPY(snes_x, dt, snes_f);
      }

      simtime += dt;

      // Update local and global residuals
      PetscCall(updateResiduals(snes_x));

      if (diagnose) {
        // Gather and print diagnostic information

        output.print("\r"); // Carriage return for printing to screen
        output.write("Time: {}, timestep: {}, nl iter: {}, lin iter: {}, reason: {}",
                     simtime, timestep, nl_its, lin_its, static_cast<int>(reason));
        if (snes_failures > 0) {
          output.write(", SNES failures: {}", snes_failures);
        }
        output.write("\n");
        // Print residual on separate line, so post-processing isn't affected
        output.write("    Residual: {}\n", global_residual);
      }

      // MatFilter and MatEliminateZeros(Mat, bool) require PETSc >= 3.20
      if (jacobian_recalculated and prune_jacobian) {
        jacobian_recalculated = false; // Reset flag

        FDJpruneJacobian();
      }

      if (equation_form == BoutSnesEquationForm::pseudo_transient) {
        // Adjust pseudo_alpha to globally scale timesteps
        pseudo_alpha = updateGlobalTimestep(pseudo_alpha, nl_its, recent_failure_rate,
                                            max_timestep * atol * 100);

        // Adjust local timesteps
        PetscCall(updatePseudoTimestepping());

      } else {
        // Adjust timestep
        timestep =
            updateGlobalTimestep(timestep, nl_its, recent_failure_rate, max_timestep);
      }

      if (static_cast<BoutReal>(lin_its) / nl_its > 0.5 * maxl) {
        // Recompute Jacobian if number of linear iterations is too high
        if (saved_jacobian_lag == 0) {
          SNESGetLagJacobian(snes, &saved_jacobian_lag);
          SNESSetLagJacobian(snes, 1);
        }
      }

      snes_failures = 0;
    } while (looping);

    BoutReal output_time = simtime;
    if (output_trigger == BoutSnesOutput::fixed_time_interval && !matrix_free) {
      ASSERT2(simtime >= target);
      ASSERT2(simtime - dt <= target);
      // Stepped over output timestep => Interpolate
      // snes_x is the solution at t = simtime
      // x0 is the solution at t = simtime - dt
      // Calculate output_x at t = target
      VecCopy(snes_x, output_x);

      // Note: If simtime = target then alpha = 0
      //       and output_x = snes_x
      BoutReal alpha = (simtime - target) / dt;

      // output_x <- alpha * x0 + (1 - alpha) * output_x
      VecAXPBY(output_x, alpha, 1. - alpha, x0);
      output_time = target;

    } else {
      // Timestep was adjusted to hit target output time
      output_x = snes_x;
    }

    // Put the result into variables
    if (scale_vars) {
      // scaled_x <- output_x * var_scaling_factors
      PetscCall(VecPointwiseMult(scaled_x, output_x, var_scaling_factors));
    } else if (asinh_vars) {
      PetscCall(VecCopy(output_x, scaled_x));
    } else {
      scaled_x = output_x;
    }

    if (asinh_vars) {
      PetscInt size;
      PetscCall(VecGetLocalSize(scaled_x, &size));

      BoutReal* scaled_data = nullptr;
      PetscCall(VecGetArray(scaled_x, &scaled_data));
      for (PetscInt i = 0; i != size; ++i) {
        scaled_data[i] = asinh_scale * std::sinh(scaled_data[i]);
      }
      PetscCall(VecRestoreArray(scaled_x, &scaled_data));
    }

    const BoutReal* xdata = nullptr;
    PetscCall(VecGetArrayRead(scaled_x, &xdata));
    load_vars(const_cast<BoutReal*>(xdata));
    PetscCall(VecRestoreArrayRead(scaled_x, &xdata));

    try {
      run_rhs(output_time); // Run RHS to calculate auxilliary variables
    } catch (BoutException& e) {
      output_error.write("ERROR: BoutException thrown: {}\n", e.what());
      // Abort simulation. There is no way to recover unless
      // all processors throw an exception at the same point.
      BoutComm::abort(1);
    }

    if (call_monitors(output_time, s, getNumberOutputSteps()) != 0) {
      break; // User signalled to quit
    }
  }

  return 0;
}

BoutReal SNESSolver::updateGlobalTimestep(BoutReal timestep, int nl_its,
                                          BoutReal recent_failure_rate, BoutReal max_dt) {
  // Note: The preconditioner depends on the timestep,
  // so if it is not recalculated the it will be less
  // effective.

  switch (timestep_control) {
  case BoutSnesTimestep::pid_nonlinear_its:
    // Changing the timestep using a PID controller.
    return pid(timestep, nl_its, max_dt);

  case BoutSnesTimestep::threshold_nonlinear_its:
    // Consider changing the timestep, based on thresholds in NL iterations
    if ((nl_its <= lower_its) && (timestep < max_timestep)
        && (recent_failure_rate < 0.5)) {
      // Increase timestep slightly
      timestep *= timestep_factor_on_lower_its;
      return std::min(timestep, max_timestep);
    }

    if (nl_its >= upper_its) {
      // Reduce timestep slightly
      return timestep * timestep_factor_on_upper_its;
    }
    return timestep; // No change

  case BoutSnesTimestep::residual_ratio:
    // Use ratio of previous and current global residual
    // Intended to be the same as https://petsc.org/release/manualpages/TS/TSPSEUDO/
    // (Note the PETSc manual has the expression for dt_n upside down)

    return std::min({timestep_factor * timestep * global_residual_prev / global_residual,
                     max_timestep});

  case BoutSnesTimestep::fixed:
    break;
  }
  return timestep; // No change
}

PetscErrorCode SNESSolver::updateResiduals(Vec x) {
  // Push residuals to previous residuals
  // Copying so that data is not shared
  local_residual_prev = copy(local_residual);
  local_residual_2d_prev = copy(local_residual_2d);
  global_residual_prev = global_residual;

  // Call RHS function to get time derivatives
  PetscCall(rhs_function(x, snes_f, false));

  // Reading the residual vectors
  const BoutReal* current_residual = nullptr;
  PetscCall(VecGetArrayRead(snes_f, &current_residual));

  // Note: The ordering of quantities in the PETSc vectors
  // depends on the Solver::loop_vars function
  Mesh* mesh = bout::globals::mesh;
  int idx = 0; // Index into PETSc Vecs

  // Boundary cells
  for (const auto& i2d : mesh->getRegion2D("RGN_BNDRY")) {
    // Field2D quantities evolved together
    BoutReal residual = 0.0;
    int count = 0;

    for (const auto& f : f2d) {
      if (!f.evolve_bndry) {
        continue; // Not evolving boundary => Skip
      }
      residual += SQ(current_residual[idx++]);
      ++count;
    }
    if (count > 0) {
      local_residual_2d[i2d] = sqrt(residual / count);
    }

    // Field3D quantities evolved together within a cell
    for (int jz = mesh->zstart; jz <= mesh->zend; jz++) {
      count = 0;
      residual = 0.0;
      for (const auto& f : f3d) {
        if (!f.evolve_bndry) {
          continue; // Not evolving boundary => Skip
        }
        residual += SQ(current_residual[idx++]);
        ++count;
      }
      if (count > 0) {
        auto i3d = mesh->ind2Dto3D(i2d, jz);
        local_residual[i3d] = sqrt(residual / count);
      }
    }
  }

  // Bulk of domain.
  // These loops don't check the boundary flags
  for (const auto& i2d : mesh->getRegion2D("RGN_NOBNDRY")) {
    // Field2D quantities evolved together
    if (!f2d.empty()) {
      BoutReal residual = 0.0;
      for (std::size_t i = 0; i != f2d.size(); ++i) {
        residual += SQ(current_residual[idx++]);
      }
      local_residual_2d[i2d] = sqrt(residual / static_cast<BoutReal>(f2d.size()));
    }

    // Field3D quantities evolved together within a cell
    if (!f3d.empty()) {
      for (int jz = mesh->zstart; jz <= mesh->zend; jz++) {
        auto i3d = mesh->ind2Dto3D(i2d, jz);

        BoutReal residual = 0.0;
        for (std::size_t i = 0; i != f3d.size(); ++i) {
          residual += SQ(current_residual[idx++]);
        }
        local_residual[i3d] = sqrt(residual / static_cast<BoutReal>(f3d.size()));
      }
    }
  }

  // Restore Vec data arrays
  PetscCall(VecRestoreArrayRead(snes_f, &current_residual));

  // Global residual metric (RMS)
  global_residual = std::sqrt(mean(SQ(local_residual), true));

  return PETSC_SUCCESS;
}

PetscErrorCode SNESSolver::initPseudoTimestepping() {
  // Storage for per-variable timestep
  PetscCall(VecDuplicate(snes_x, &dt_vec));
  // Starting timestep
  PetscCall(VecSet(dt_vec, timestep));

  // Diagnostic outputs
  pseudo_timestep = timestep;

  return PETSC_SUCCESS;
}

PetscErrorCode SNESSolver::updatePseudoTimestepping() {
  // Use a per-cell timestep so that e.g density and pressure
  // evolve in a way consistent with the equation of state.

  // Modifying the dt_vec values
  BoutReal* dt_data = nullptr;
  PetscCall(VecGetArray(dt_vec, &dt_data));

  // Note: The ordering of quantities in the PETSc vectors
  // depends on the Solver::loop_vars function
  Mesh* mesh = bout::globals::mesh;
  int idx = 0; // Index into PETSc Vecs

  // Boundary cells
  for (const auto& i2d : mesh->getRegion2D("RGN_BNDRY")) {
    // Field2D quantities evolved together
    int count = 0;
    for (const auto& f : f2d) {
      if (!f.evolve_bndry) {
        continue; // Not evolving boundary => Skip
      }
      ++count;
    }
    if (count > 0) {
      // Adjust timestep for these quantities
      BoutReal new_timestep = updatePseudoTimestep(
          dt_data[idx], local_residual_2d_prev[i2d], local_residual_2d[i2d]);
      for (int i = 0; i != count; ++i) {
        dt_data[idx++] = new_timestep;
      }
    }

    // Field3D quantities evolved together within a cell
    for (int jz = mesh->zstart; jz <= mesh->zend; jz++) {
      count = 0;
      for (const auto& f : f3d) {
        if (!f.evolve_bndry) {
          continue; // Not evolving boundary => Skip
        }
        ++count;
      }
      if (count > 0) {
        auto i3d = mesh->ind2Dto3D(i2d, jz);
        const BoutReal new_timestep = updatePseudoTimestep(
            dt_data[idx], local_residual_prev[i3d], local_residual[i3d]);

        for (int i = 0; i != count; ++i) {
          dt_data[idx++] = new_timestep;
        }
      }
    }
  }

  // Bulk of domain.
  // These loops don't check the boundary flags
  for (const auto& i2d : mesh->getRegion2D("RGN_NOBNDRY")) {
    // Field2D quantities evolved together
    if (!f2d.empty()) {
      // Adjust timestep for these quantities
      const BoutReal new_timestep = updatePseudoTimestep(
          dt_data[idx], local_residual_2d_prev[i2d], local_residual_2d[i2d]);
      for (std::size_t i = 0; i != f2d.size(); ++i) {
        dt_data[idx++] = new_timestep;
      }
    }

    // Field3D quantities evolved together within a cell
    if (!f3d.empty()) {
      for (int jz = mesh->zstart; jz <= mesh->zend; jz++) {
        auto i3d = mesh->ind2Dto3D(i2d, jz);

        BoutReal new_timestep = updatePseudoTimestep(
            dt_data[idx], local_residual_prev[i3d], local_residual[i3d]);

        // Compare to neighbors
        BoutReal min_neighboring_dt = max_timestep;
        if (i3d.x() != 0) {
          min_neighboring_dt = std::min(min_neighboring_dt, pseudo_timestep[i3d.xm()]);
        }
        if (i3d.x() != mesh->LocalNx - 1) {
          min_neighboring_dt = std::min(min_neighboring_dt, pseudo_timestep[i3d.xp()]);
        }
        if (i3d.y() != 0) {
          min_neighboring_dt = std::min(min_neighboring_dt, pseudo_timestep[i3d.ym()]);
        }
        if (i3d.x() != mesh->LocalNy - 1) {
          min_neighboring_dt = std::min(min_neighboring_dt, pseudo_timestep[i3d.yp()]);
        }

        // Limit ratio of timestep between neighboring cells
        new_timestep = std::min(new_timestep, pseudo_max_ratio * min_neighboring_dt);

        pseudo_timestep[i3d] = new_timestep;

        for (std::size_t i = 0; i != f3d.size(); ++i) {
          dt_data[idx++] = new_timestep;
        }
      }
    }
  }

  // Restore Vec data arrays
  PetscCall(VecRestoreArray(dt_vec, &dt_data));

  // Need timesteps on neighboring processors
  // to limit ratio
  mesh->communicate(pseudo_timestep);
  // Neumann boundary so timestep isn't pinned at boundaries
  pseudo_timestep.applyBoundary("neumann");

  return PETSC_SUCCESS;
}

/// Timestep inversely proportional to the residual, clipped to avoid
/// rapid changes in timestep.
BoutReal SNESSolver::updatePseudoTimestep_inverse_residual(BoutReal previous_timestep,
                                                           BoutReal current_residual) {
  return std::min(
      {std::max({pseudo_alpha / current_residual, previous_timestep / 1.5, dt_min_reset}),
       1.5 * previous_timestep, max_timestep});
}

// Strategy based on history of residuals
BoutReal SNESSolver::updatePseudoTimestep_history_based(BoutReal previous_timestep,
                                                        BoutReal previous_residual,
                                                        BoutReal current_residual) {
  const BoutReal converged_threshold = 10 * atol;
  const BoutReal transition_threshold = 100 * atol;

  if (current_residual < converged_threshold) {
    return max_timestep;
  }

  // Smoothly transition from SER updates to frozen timestep
  if (current_residual < transition_threshold) {
    const BoutReal reduction_ratio = current_residual / previous_residual;

    if (reduction_ratio < 0.8) {
      return std::min(0.5 * (pseudo_growth_factor + 1.) * previous_timestep,
                      max_timestep);
    } else if (reduction_ratio > 1.2) {
      return std::max(0.5 * (pseudo_reduction_factor + 1) * previous_timestep,
                      dt_min_reset);
    }
    // Leave unchanged if residual changes are small
    return previous_timestep;
  }

  if (current_residual <= previous_residual) {
    // Success
    return std::min(pseudo_growth_factor * previous_timestep, max_timestep);
  }
  // Failed
  // Consider rejecting the step
  return std::max(pseudo_reduction_factor * previous_timestep, dt_min_reset);
}

/// Switched Evolution Relaxation (SER)
BoutReal SNESSolver::updatePseudoTimestep(BoutReal previous_timestep,
                                          BoutReal previous_residual,
                                          BoutReal current_residual) {
  switch (pseudo_strategy) {
  case BoutPTCStrategy::inverse_residual:
    return updatePseudoTimestep_inverse_residual(previous_timestep, current_residual);

  case BoutPTCStrategy::history_based:
    return updatePseudoTimestep_history_based(previous_timestep, previous_residual,
                                              current_residual);

  case BoutPTCStrategy::hybrid:
    // A hybrid strategy may be most effective, in which the timestep is
    // inversely proportional to residual initially, or when residuals are large,
    // and then the method transitions to being history-based
    if (current_residual > 1000. * atol) {
      return updatePseudoTimestep_inverse_residual(previous_timestep, current_residual);
    }
    return updatePseudoTimestep_history_based(previous_timestep, previous_residual,
                                              current_residual);
  };
  throw BoutException("SNESSolver::updatePseudoTimestep invalid BoutPTCStrategy");
}

PetscErrorCode SNESSolver::rhs_function(Vec x, Vec f, bool linear) {
  // Get data from PETSc into BOUT++ fields
  if (scale_vars) {
    // scaled_x <- x * var_scaling_factors
    PetscCall(VecPointwiseMult(scaled_x, x, var_scaling_factors));
  } else if (asinh_vars) {
    PetscCall(VecCopy(x, scaled_x));
  } else {
    scaled_x = x;
  }

  if (asinh_vars) {
    PetscInt size;
    PetscCall(VecGetLocalSize(scaled_x, &size));

    BoutReal* scaled_data = nullptr;
    PetscCall(VecGetArray(scaled_x, &scaled_data));
    for (PetscInt i = 0; i != size; ++i) {
      scaled_data[i] = asinh_scale * std::sinh(scaled_data[i]);
    }
    PetscCall(VecRestoreArray(scaled_x, &scaled_data));
  }

  const BoutReal* xdata = nullptr;
  PetscCall(VecGetArrayRead(scaled_x, &xdata));
  // const_cast needed due to load_vars API. Not writing to xdata.
  load_vars(const_cast<BoutReal*>(xdata));
  PetscCall(VecRestoreArrayRead(scaled_x, &xdata));

  try {
    // Call RHS function
    run_rhs(simtime + dt, linear);
  } catch (BoutException& e) {
    // Simulation might fail, e.g. negative densities
    // if timestep too large
    output_warn.write("WARNING: BoutException thrown: {}\n", e.what());
    // Abort simulation. Unless all processors threw an exception
    // at the same point, there is no way to synchronise again.
    BoutComm::abort(2);
  }

  // Copy derivatives back
  BoutReal* fdata = nullptr;
  PetscCall(VecGetArray(f, &fdata));
  save_derivs(fdata);

  if (asinh_vars) {
    // Modify time-derivatives for asinh(var) using chain rule
    // Evolving u = asinh(var / scale)
    //
    // du/dt = dvar/dt * du/dvar
    //
    // du/var = 1 / sqrt(var^2 + scale^2)
    PetscInt size;
    PetscCall(VecGetLocalSize(f, &size));
    const BoutReal* scaled_data = nullptr;
    PetscCall(VecGetArrayRead(scaled_x, &scaled_data));
    for (PetscInt i = 0; i != size; ++i) {
      fdata[i] /= std::sqrt(SQ(scaled_data[i]) + SQ(asinh_scale));
    }
    PetscCall(VecRestoreArrayRead(scaled_x, &scaled_data));
  }

  PetscCall(VecRestoreArray(f, &fdata));
  return PETSC_SUCCESS;
}

// Result in f depends on equation_form
PetscErrorCode SNESSolver::snes_function(Vec x, Vec f, bool linear) {

  // Call the RHS function
  if (rhs_function(x, f, linear) != PETSC_SUCCESS) {
    // Tell SNES that the input was out of domain
    SNESSetFunctionDomainError(snes);
    // Note: Returning non-zero error here leaves vectors in locked state
    return 0;
  }

  switch (equation_form) {
  case BoutSnesEquationForm::rearranged_backward_euler: {
    // Rearranged Backward Euler
    // f = (x0 - x)/Δt + f
    // First calculate x - x0 to minimise floating point issues
    VecWAXPY(delta_x, -1.0, x0, x); // delta_x = x - x0
    VecAXPY(f, -1. / dt, delta_x);  // f <- f - delta_x / dt
    break;
  }
  case BoutSnesEquationForm::pseudo_transient: {
    // Pseudo-transient timestepping. Same as Rearranged Backward Euler
    // except that Δt is a vector
    // f = (x0 - x)/Δt + f
    VecWAXPY(delta_x, -1.0, x0, x);
    VecPointwiseDivide(delta_x, delta_x, dt_vec); // delta_x /= dt
    VecAXPY(f, -1., delta_x);                     // f <- f - delta_x
    break;
  }
  case BoutSnesEquationForm::backward_euler: {
    // Backward Euler
    // Set f = x - x0 - Δt*f
    VecAYPX(f, -dt, x);   // f <- x - Δt*f
    VecAXPY(f, -1.0, x0); // f <- f - x0
    break;
  }
  case BoutSnesEquationForm::direct_newton: {
    // Direct Newton solve -> don't modify f
    break;
  }
  default: {
    throw BoutException("Invalid choice of equation_form. Try 0-3");
  }
  };

  if (scale_rhs) {
    // f <- f * rhs_scaling_factors
    PetscCall(VecPointwiseMult(f, f, rhs_scaling_factors));
  }

  return PETSC_SUCCESS;
}

/*
 * Preconditioner function
 */
PetscErrorCode SNESSolver::precon(Vec x, Vec f) {
  if (!hasPreconditioner()) {
    // No user preconditioner
    throw BoutException("No user preconditioner");
  }

  // Get data from PETSc into BOUT++ fields
  Vec solution;
  PetscCall(SNESGetSolution(snes, &solution));
  BoutReal* soldata;
  PetscCall(VecGetArray(solution, &soldata));
  load_vars(soldata);
  PetscCall(VecRestoreArray(solution, &soldata));

  // Load vector to be inverted into ddt() variables
  const BoutReal* xdata;
  PetscCall(VecGetArrayRead(x, &xdata));
  load_derivs(const_cast<BoutReal*>(xdata)); // Note: load_derivs does not modify data
  PetscCall(VecRestoreArrayRead(x, &xdata));

  // Run the preconditioner
  runPreconditioner(simtime + dt, dt, 0.0);

  // Save the solution from F_vars
  BoutReal* fdata;
  PetscCall(VecGetArray(f, &fdata));
  save_derivs(fdata);
  PetscCall(VecRestoreArray(f, &fdata));

  return 0;
}

PetscErrorCode SNESSolver::scaleJacobian(Mat Jac_new) {
  jacobian_recalculated = true;

  if (!scale_rhs) {
    return PETSC_SUCCESS; // Not scaling the RHS values
  }

  // Get index of rows owned by this processor
  int rstart, rend;
  MatGetOwnershipRange(Jac_new, &rstart, &rend);

  // Check that the vector has the same ownership range
  int istart, iend;
  VecGetOwnershipRange(jac_row_inv_norms, &istart, &iend);
  if ((rstart != istart) or (rend != iend)) {
    throw BoutException("Ownership ranges different: [{}, {}) and [{}, {})\n", rstart,
                        rend, istart, iend);
  }

  // Calculate the norm of each row of the Jacobian
  PetscScalar* row_inv_norm_data;
  PetscCall(VecGetArray(jac_row_inv_norms, &row_inv_norm_data));

  PetscInt ncols;
  const PetscScalar* vals;
  for (int row = rstart; row < rend; ++row) {
    MatGetRow(Jac_new, row, &ncols, nullptr, &vals);

    // Calculate a norm of this row of the Jacobian
    PetscScalar norm = 0.0;
    for (int col = 0; col < ncols; col++) {
      PetscScalar absval = std::abs(vals[col]);
      if (absval > norm) {
        norm = absval;
      }
      // Can we identify small elements and remove them?
      // so we don't need to calculate them next time
    }

    // Store in the vector as 1 / norm
    row_inv_norm_data[row - rstart] = 1. / norm;

    MatRestoreRow(Jac_new, row, &ncols, nullptr, &vals);
  }

  PetscCall(VecRestoreArray(jac_row_inv_norms, &row_inv_norm_data));

  // Modify the RHS scaling: factor = factor / norm
  PetscCall(
      VecPointwiseMult(rhs_scaling_factors, rhs_scaling_factors, jac_row_inv_norms));

  if (diagnose) {
    // Print maximum and minimum scaling factors
    PetscReal max_scale, min_scale;
    VecMax(rhs_scaling_factors, nullptr, &max_scale);
    VecMin(rhs_scaling_factors, nullptr, &min_scale);
    output.write("RHS scaling: {} -> {}\n", min_scale, max_scale);
  }

  // Scale the Jacobian rows by multiplying on the left by 1/norm
  PetscCall(MatDiagonalScale(Jac_new, jac_row_inv_norms, nullptr));

  return PETSC_SUCCESS;
}

///
/// Input Parameters:
///   snes - nonlinear solver object
///   x1 - location at which to evaluate Jacobian
///   ctx - MatFDColoring context or NULL
///
/// Output Parameters:
///   Jac - Jacobian matrix (not altered in this routine)
///   Jac_new - newly computed Jacobian matrix to use with preconditioner (generally the same as
///   Jac)
static PetscErrorCode ComputeJacobianScaledColor(SNES snes, Vec x1, Mat Jac, Mat Jac_new,
                                                 void* ctx) {
  PetscErrorCode err = SNESComputeJacobianDefaultColor(snes, x1, Jac, Jac_new, ctx);
  CHKERRQ(err);

  if ((err != 0) or (ctx == nullptr)) {
    return err;
  }

  // Get the the SNESSolver pointer from the function call context
  SNESSolver* fctx = nullptr;
  err = MatFDColoringGetFunction(static_cast<MatFDColoring>(ctx), nullptr,
                                 reinterpret_cast<void**>(&fctx));
  CHKERRQ(err);

  // Call the SNESSolver function
  return fctx->scaleJacobian(Jac_new);
}

void SNESSolver::updateColoring() {
  // Re-calculate the coloring
  MatColoring coloring = NULL;
  MatColoringCreate(Jfd, &coloring);
  // MatColoringSetType(coloring, MATCOLORINGSL);  // Serial algorithm. Better for smale-to-medium size problems.
  MatColoringSetType(
      coloring, MATCOLORINGGREEDY); // Parallel algorith. Better for large parallel runs
  // MatColoringSetType(coloring, MATCOLORINGJP);  // This didn't work
  MatColoringSetFromOptions(coloring);

  // Calculate new index sets
  ISColoring iscoloring = NULL;
  MatColoringApply(coloring, &iscoloring);
  MatColoringDestroy(&coloring);

  // Replace the old coloring with the new one
  MatFDColoringDestroy(&fdcoloring);
  MatFDColoringCreate(Jfd, iscoloring, &fdcoloring);
  MatFDColoringSetFunction(fdcoloring,
                           bout::cast_MatFDColoringFn(FormFunctionForColoring), this);
  MatFDColoringSetFromOptions(fdcoloring);
  MatFDColoringSetUp(Jfd, iscoloring, fdcoloring);
  ISColoringDestroy(&iscoloring);

  // Replace the CTX pointer in SNES Jacobian
  if (matrix_free_operator) {
    // Use matrix-free calculation for operator, finite difference for preconditioner
    SNESSetJacobian(snes, Jmf, Jfd, ComputeJacobianScaledColor, fdcoloring);
  } else {
    SNESSetJacobian(snes, Jfd, Jfd, ComputeJacobianScaledColor, fdcoloring);
  }
}

BoutReal SNESSolver::pid(BoutReal timestep, int nl_its, BoutReal max_dt) {

  /* ---------- multiplicative PID factors ---------- */
  const BoutReal facP = std::pow(double(target_its) / double(nl_its), kP);
  const BoutReal facI = std::pow(double(nl_its_prev) / double(nl_its), kI);
  const BoutReal facD = std::pow(double(nl_its_prev) * double(nl_its_prev)
                                     / double(nl_its) / double(nl_its_prev2),
                                 kD);

  // clamp growth factor to avoid huge changes
  BoutReal fac = std::clamp(facP * facI * facD, 0.2, 5.0);

  if (pid_consider_failures && (fac > 1.0)) {
    // Reduce aggressiveness if recent steps have failed often
    fac = pow(fac, std::max(0.3, 1.0 - 2.0 * recent_failure_rate));
  }

  /* ---------- update timestep and history ---------- */
  const BoutReal dt_new = std::min(timestep * fac, max_dt);

  nl_its_prev2 = nl_its_prev;
  nl_its_prev = nl_its;

  return dt_new;
}

void SNESSolver::outputVars(Options& output_options, bool save_repeat) {
  // Call base class function
  Solver::outputVars(output_options, save_repeat);

  if (!save_repeat) {
    return; // Don't save diagnostics to restart files
  }

  output_options["snes_local_residual"].assignRepeat(local_residual, "t", save_repeat,
                                                     "SNESSolver");
  output_options["snes_global_residual"].assignRepeat(global_residual, "t", save_repeat,
                                                      "SNESSolver");

  if (equation_form == BoutSnesEquationForm::pseudo_transient) {
    output_options["snes_pseudo_alpha"].assignRepeat(pseudo_alpha, "t", save_repeat,
                                                     "SNESSolver");
    output_options["snes_pseudo_timestep"].assignRepeat(pseudo_timestep, "t", save_repeat,
                                                        "SNESSolver");
  }
}

#endif // BOUT_HAS_PETSC
