/**************************************************************************
 * Perpendicular Laplacian inversion.
 *                           Using PETSc Solvers
 *
 **************************************************************************
 * Copyright 2013 - 2025 BOUT++ developers
 *
 * Contact: Ben Dudson, dudson2@llnl.gov
 *
 * This file is part of BOUT++.
 *
 * BOUT++ is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * BOUT++ is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * along with BOUT++.  If not, see <http://www.gnu.org/licenses/>.
 *
 **************************************************************************/

#include "bout/build_defines.hxx"

#if BOUT_HAS_PETSC

#include "petsc_laplace.hxx"

#include <bout/assert.hxx>
#include <bout/bout_types.hxx>
#include <bout/boutcomm.hxx>
#include <bout/boutexception.hxx>
#include <bout/globals.hxx>
#include <bout/mesh.hxx>
#include <bout/output.hxx>
#include <bout/petsclib.hxx>
#include <bout/sys/timer.hxx>
#include <bout/utils.hxx>

#include <vector>

#define KSP_RICHARDSON "richardson"
#define KSP_CHEBYSHEV "chebyshev"
#define KSP_CG "cg"
#define KSP_GMRES "gmres"
#define KSP_TCQMR "tcqmr"
#define KSP_BCGS "bcgs"
#define KSP_CGS "cgs"
#define KSP_TFQMR "tfqmr"
#define KSP_CR "cr"
#define KSP_LSQR "lsqr"
#define KSP_BICG "bicg"
#define KSP_PREONLY "preonly"

namespace {
PetscErrorCode laplacePCapply(PC pc, Vec x, Vec y) {
  PetscFunctionBegin; // NOLINT

  LaplacePetsc* laplace = nullptr;
  CHKERRQ(PCShellGetContext(pc, reinterpret_cast<void**>(&laplace))); // NOLINT

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

LaplacePetsc::LaplacePetsc(Options* opt, const CELL_LOC loc, Mesh* mesh_in,
                           [[maybe_unused]] Solver* solver)
    : Laplacian(opt, loc, mesh_in), A(0.0), C1(1.0), C2(1.0), D(1.0), Ex(0.0), Ez(0.0),
      issetD(false), issetC(false), issetE(false), comm(localmesh->getXcomm()),
      opts(opt == nullptr ? &(Options::root()["laplace"]) : opt),
      // WARNING: only a few of these options actually make sense: see the
      // PETSc documentation to work out which they are (possibly
      // pbjacobi, sor might be useful choices?)
      ksptype((*opts)["ksptype"].doc("KSP solver type").withDefault(KSPGMRES)),
      pctype((*opts)["pctype"]
                 .doc("Preconditioner type. See the PETSc documentation for options")
                 .withDefault("none")),
      richardson_damping_factor((*opts)["richardson_damping_factor"].withDefault(1.0)),
      chebyshev_max((*opts)["chebyshev_max"].withDefault(100)),
      chebyshev_min((*opts)["chebyshev_min"].withDefault(0.01)),
      gmres_max_steps((*opts)["gmres_max_steps"].withDefault(30)),
      rtol((*opts)["rtol"].doc("Relative tolerance for KSP solver").withDefault(1e-5)),
      atol((*opts)["atol"].doc("Absolute tolerance for KSP solver").withDefault(1e-50)),
      dtol((*opts)["dtol"].doc("Divergence tolerance for KSP solver").withDefault(1e5)),
      maxits(
          (*opts)["maxits"].doc("Maximum number of KSP iterations").withDefault(100000)),
      direct((*opts)["direct"].doc("Use direct (LU) solver?").withDefault(false)),
      fourth_order(
          (*opts)["fourth_order"].doc("Use fourth order stencil").withDefault(false)),
      lib(opts) {

  A.setLocation(location);
  C1.setLocation(location);
  C2.setLocation(location);
  D.setLocation(location);
  Ex.setLocation(location);
  Ez.setLocation(location);

#if CHECK > 0
  // Checking flags are set to something which is not implemented
  checkFlags();

  if (localmesh->periodicX) {
    throw BoutException("LaplacePetsc does not work with periodicity in the x direction "
                        "(localmesh->PeriodicX == true). Change boundary conditions or "
                        "use serial-tri or cyclic solver instead");
  }
#endif

  const bool first_x = localmesh->firstX();
  const bool last_x = localmesh->lastX();

  // Need to determine local size to use based on prior parallelisation
  // Coefficient values are stored only on local processors.
  const auto local_nz = localmesh->zend - localmesh->zstart + 1;
  localN = (localmesh->xend - localmesh->xstart + 1) * local_nz;
  if (first_x) {
    // If on first processor add on width of boundary region
    localN += localmesh->xstart * local_nz;
  }
  if (last_x) {
    // If on last processor add on width of boundary region
    localN += localmesh->xstart * local_nz;
  }

  // Calculate 'size' (the total number of points in physical grid)
  size = localN;
  if (bout::globals::mpi->MPI_Allreduce(&localN, &size, 1, MPI_INT, MPI_SUM, comm)
      != MPI_SUCCESS) {
    throw BoutException("Error in MPI_Allreduce during LaplacePetsc initialisation");
  }

  // Calculate total (physical) grid dimensions
  meshz = local_nz;
  meshx = size / meshz;

  // Create PETSc type of vectors for the solution and the RHS vector
  VecCreate(comm, &xs);
  VecSetSizes(xs, localN, size);
  VecSetFromOptions(xs);
  VecDuplicate(xs, &bs);

  // Set size of (the PETSc) Matrix on each processor to localN x localN
  MatCreate(comm, &MatA);
  MatSetSizes(MatA, localN, localN, size, size);
  MatSetFromOptions(MatA);

  // Pre allocate memory. See MatMPIAIJSetPreallocation docs for more info
  // Number of non-zero elements in each row of matrix owned by this processor
  // ("diagonal submatrices")
  std::vector<PetscInt> d_nnz(localN, 0);
  // Number of non-zero elements in each row of matrix owned by other processors
  // ("offdiagonal submatrices")
  std::vector<PetscInt> o_nnz(localN, 0);

  if (fourth_order) {
    // first and last 2*localmesh-LocalNz entries are the edge x-values that
    // (may) have 'off-diagonal' components (i.e. on another processor)

    const int first_first_off = first_x ? 0 : 10;
    const int first_last_off = last_x ? 0 : 10;
    const int second_first_off = first_x ? 0 : 5;
    const int second_last_off = last_x ? 0 : 5;

    for (int i = 0; i <= local_nz; i++) {
      d_nnz[i] = 15;
      d_nnz[localN - 1 - i] = 15;
      o_nnz[i] = first_first_off;
      o_nnz[localN - 1 - i] = first_last_off;
    }
    for (int i = local_nz; i < 2 * local_nz; i++) {
      d_nnz[i] = 20;
      d_nnz[localN - 1 - i] = 20;
      o_nnz[i] = second_first_off;
      o_nnz[localN - 1 - i] = second_last_off;
    }

    for (int i = 2 * local_nz; i < localN - (2 * local_nz); i++) {
      d_nnz[i] = 25;
      d_nnz[localN - 1 - i] = 25;
      o_nnz[i] = 0;
      o_nnz[localN - 1 - i] = 0;
    }
  } else {
    // first and last local_nz entries are the edge x-values that
    // (may) have 'off-diagonal' components (i.e. on another processor)
    const int first_off = first_x ? 0 : 3;
    const int last_off = last_x ? 0 : 3;

    for (int i = 0; i <= local_nz; i++) {
      d_nnz[i] = 6;
      d_nnz[localN - 1 - i] = 6;
      o_nnz[i] = first_off;
      o_nnz[localN - 1 - i] = last_off;
    }

    for (int i = local_nz; i < localN - local_nz; i++) {
      d_nnz[i] = 9;
      d_nnz[localN - 1 - i] = 9;
      o_nnz[i] = 0;
      o_nnz[localN - 1 - i] = 0;
    }
  }

  if (first_x and last_x) {
    // Only one processor in X
    MatSeqAIJSetPreallocation(MatA, 0, d_nnz.data());
  } else {
    MatMPIAIJSetPreallocation(MatA, 0, d_nnz.data(), 0, o_nnz.data());
  }

  // Sets up the internal matrix data structures for the later use.
  MatSetUp(MatA);

  // Declare KSP Context (abstract PETSc object that manages all Krylov methods)
  KSPCreate(comm, &ksp);

  // Let "user" be a synonym for "shell"
  if (pctype == "user") {
    pctype = PCSHELL;
  }

  if (direct) {
    output.write("\nUsing LU decompostion for direct solution of system\n");
  }

  if (pctype == PCSHELL) {
    rightprec = (*opts)["rightprec"].doc("Right preconditioning?").withDefault(true);

    // Options for preconditioner are in a subsection
    pcsolve = Laplacian::create(opts->getSection("precon"));
  }
}

FieldPerp LaplacePetsc::solve(const FieldPerp& b) { return solve(b, b); }

/*!
 * Solves Ax=b for x given a b and an initial guess for x (x0)
 *
 * This function will:
 *      1. Set the matrix element of the matrix A, used to solve Ax=b
 *         (this includes setting the values for the bounary condition)
 *      2. Solve the matrix Ax = b
 *
 * \param[in] b     The RHS of the equation Ax=b.
 *                  This is an y-slice of the original field. The field wil be
 *                  flattened to an 1D array in order to write the equation on
 *                  the form Ax=b
 * \param[in] x0    The initial guess for the solver.
 *                  May also contain the  boundary condition if flag 32 - INVERT_SET is set
 *
 * \returns sol     The solution x of the problem Ax=b.
 */
FieldPerp LaplacePetsc::solve(const FieldPerp& b, const FieldPerp& x0) {

  ASSERT1(localmesh == b.getMesh() && localmesh == x0.getMesh());
  ASSERT1(b.getLocation() == location);
  ASSERT1(x0.getLocation() == location);

#if CHECK > 0
  checkFlags();
#endif

  int y = b.getIndex(); // Get the Y index
  sol.setIndex(y);      // Initialize the solution field.
  sol = 0.;

  // Determine which row/columns of the matrix are locally owned
  MatGetOwnershipRange(MatA, &Istart, &Iend);

  int i = Istart; // The row in the PETSc matrix
  {
    Timer timer("petscsetup");

    //     if ((fourth_order) && !(lastflag&INVERT_4TH_ORDER)) throw BoutException("Should not change INVERT_4TH_ORDER flag in LaplacePetsc: 2nd order and 4th order require different pre-allocation to optimize PETSc solver");

    /* Set Matrix Elements
   *
   * Loop over locally owned rows of matrix A
   * i labels NODE POINT from
   * bottom left = (0,0) = 0
   * to
   * top right = (meshx-1,meshz-1) = meshx*meshz-1
   *
   * i increments by 1 for an increase of 1 in Z
   * i increments by meshz for an increase of 1 in X.
   *
   * In other word the indexing is done in a row-major order, but starting at
   * bottom left rather than top left
   */
    // X=0 to localmesh->xstart-1 defines the boundary region of the domain.
    // Set the values for the inner boundary region
    if (localmesh->firstX()) {
      for (int x = 0; x < localmesh->xstart; x++) {
        for (int z = localmesh->zstart; z <= localmesh->zend; z++) {
          PetscScalar val; // Value of element to be set in the matrix
          // If Neumann Boundary Conditions are set.
          if (isInnerBoundaryFlagSet(INVERT_AC_GRAD)) {
            // Set values corresponding to nodes adjacent in x
            if (fourth_order) {
              // Fourth Order Accuracy on Boundary
              Element(i, x, z, 0, 0,
                      -25.0 / (12.0 * coords->dx(x, y, z)) / sqrt(coords->g_11(x, y, z)),
                      MatA);
              Element(i, x, z, 1, 0,
                      4.0 / coords->dx(x, y, z) / sqrt(coords->g_11(x, y, z)), MatA);
              Element(i, x, z, 2, 0,
                      -3.0 / coords->dx(x, y, z) / sqrt(coords->g_11(x, y, z)), MatA);
              Element(i, x, z, 3, 0,
                      4.0 / (3.0 * coords->dx(x, y, z)) / sqrt(coords->g_11(x, y, z)),
                      MatA);
              Element(i, x, z, 4, 0,
                      -1.0 / (4.0 * coords->dx(x, y, z)) / sqrt(coords->g_11(x, y, z)),
                      MatA);
            } else {
              // Second Order Accuracy on Boundary
              //   Element(i,x,z, 0, 0, -3.0 / (2.0*coords->dx(x,y)), MatA );
              //   Element(i,x,z, 1, 0,  2.0 / coords->dx(x,y), MatA );
              //   Element(i,x,z, 2, 0, -1.0 / (2.0*coords->dx(x,y)), MatA );
              //   Element(i,x,z, 3, 0, 0.0, MatA );  // Reset these elements to 0
              //   in case 4th order flag was used previously: not allowed now
              //   Element(i,x,z, 4, 0, 0.0, MatA );
              // Second Order Accuracy on Boundary, set half-way between grid points
              Element(i, x, z, 0, 0,
                      -1.0 / coords->dx(x, y, z) / sqrt(coords->g_11(x, y, z)), MatA);
              Element(i, x, z, 1, 0,
                      1.0 / coords->dx(x, y, z) / sqrt(coords->g_11(x, y, z)), MatA);
              Element(i, x, z, 2, 0, 0.0, MatA);
              //                      Element(i,x,z, 3, 0, 0.0, MatA );  // Reset
              //                      these elements to 0 in case 4th order flag was
              //                      used previously: not allowed now
              //                      Element(i,x,z, 4, 0, 0.0, MatA );
            }
          } else {
            if (fourth_order) {
              // Set Diagonal Values to 1
              Element(i, x, z, 0, 0, 1., MatA);

              // Set off diagonal elements to zero
              Element(i, x, z, 1, 0, 0.0, MatA);
              Element(i, x, z, 2, 0, 0.0, MatA);
              Element(i, x, z, 3, 0, 0.0, MatA);
              Element(i, x, z, 4, 0, 0.0, MatA);
            } else {
              Element(i, x, z, 0, 0, 0.5, MatA);
              Element(i, x, z, 1, 0, 0.5, MatA);
              Element(i, x, z, 2, 0, 0., MatA);
            }
          }

          val = 0; // Initialize val

          // Set Components of RHS
          // If the inner boundary value should be set by b or x0
          if (isInnerBoundaryFlagSet(INVERT_RHS)) {
            val = b[x][z];
          } else if (isInnerBoundaryFlagSet(INVERT_SET)) {
            val = x0[x][z];
          }

          // Set components of the RHS (the PETSc vector bs)
          // 1 element is being set in row i to val
          // INSERT_VALUES replaces existing entries with new values
          VecSetValues(bs, 1, &i, &val, INSERT_VALUES);

          // Set components of the and trial solution (the PETSc vector xs)
          // 1 element is being set in row i to val
          // INSERT_VALUES replaces existing entries with new values
          val = x0[x][z];
          VecSetValues(xs, 1, &i, &val, INSERT_VALUES);

          i++; // Increment row in Petsc matrix
        }
      }
    }

    // Set the values for the main domain
    for (int x = localmesh->xstart; x <= localmesh->xend; x++) {
      for (int z = localmesh->zstart; z <= localmesh->zend; z++) {
        // NOTE: Only A0 is the A from setCoefA ()
        BoutReal A0, A1, A2, A3, A4, A5;
        A0 = A(x, y, z);

        ASSERT3(std::isfinite(A0));

        // Set the matrix coefficients
        Coeffs(x, y, z, A1, A2, A3, A4, A5);

        BoutReal dx = coords->dx(x, y, z);
        BoutReal dx2 = SQ(dx);
        BoutReal dz = coords->dz(x, y, z);
        BoutReal dz2 = SQ(dz);
        BoutReal dxdz = dx * dz;

        ASSERT3(std::isfinite(A1));
        ASSERT3(std::isfinite(A2));
        ASSERT3(std::isfinite(A3));
        ASSERT3(std::isfinite(A4));
        ASSERT3(std::isfinite(A5));

        // Set Matrix Elements
        PetscScalar val = 0.;
        if (fourth_order) {
          // f(i,j) = f(x,z)
          val = A0 - (5.0 / 2.0) * ((A1 / dx2) + (A2 / dz2));
          Element(i, x, z, 0, 0, val, MatA);

          // f(i-2,j-2)
          val = A3 / (144.0 * dxdz);
          Element(i, x, z, -2, -2, val, MatA);

          // f(i-2,j-1)
          val = -1.0 * A3 / (18.0 * dxdz);
          Element(i, x, z, -2, -1, val, MatA);

          // f(i-2,j)
          val = (1.0 / 12.0) * ((-1.0 * A1 / dx2) + (A4 / dx));
          Element(i, x, z, -2, 0, val, MatA);

          // f(i-2,j+1)
          val = A3 / (18.0 * dxdz);
          Element(i, x, z, -2, 1, val, MatA);

          // f(i-2,j+2)
          val = -1.0 * A3 / (144.0 * dxdz);
          Element(i, x, z, -2, 2, val, MatA);

          // f(i-1,j-2)
          val = -1.0 * A3 / (18.0 * dxdz);
          Element(i, x, z, -1, -2, val, MatA);

          // f(i-1,j-1)
          val = 4.0 * A3 / (9.0 * dxdz);
          Element(i, x, z, -1, -1, val, MatA);

          // f(i-1,j)
          val = (4.0 * A1 / (3.0 * dx2)) - (2.0 * A4 / (3.0 * dx));
          Element(i, x, z, -1, 0, val, MatA);

          // f(i-1,j+1)
          val = -4.0 * A3 / (9.0 * dxdz);
          Element(i, x, z, -1, 1, val, MatA);

          // f(i-1,j+2)
          val = A3 / (18.0 * dxdz);
          Element(i, x, z, -1, 2, val, MatA);

          // f(i,j-2)
          val = (1.0 / 12.0) * ((-1.0 * A2 / dz2) + (A5 / dz));
          Element(i, x, z, 0, -2, val, MatA);

          // f(i,j-1)
          val = (4.0 * A2 / (3.0 * dz2)) - (2.0 * A5 / (3.0 * dz));
          Element(i, x, z, 0, -1, val, MatA);

          // f(i,j+1)
          val = (4.0 * A2 / (3.0 * dz2)) + (2.0 * A5 / (3.0 * dz));
          Element(i, x, z, 0, 1, val, MatA);

          // f(i,j+2)
          val = (-1.0 / 12.0) * ((A2 / dz2) + (A5 / dz));
          Element(i, x, z, 0, 2, val, MatA);

          // f(i+1,j-2)
          val = A3 / (18.0 * dxdz);
          Element(i, x, z, 1, -2, val, MatA);

          // f(i+1,j-1)
          val = -4.0 * A3 / (9.0 * dxdz);
          Element(i, x, z, 1, -1, val, MatA);

          // f(i+1,j)
          val = (4.0 * A1 / (3.0 * dx2)) + (2.0 * A4 / (3.0 * dx));
          Element(i, x, z, 1, 0, val, MatA);

          // f(i+1,j+1)
          val = 4.0 * A3 / (9.0 * dxdz);
          Element(i, x, z, 1, 1, val, MatA);

          // f(i+1,j+2)
          val = -1.0 * A3 / (18.0 * dxdz);
          Element(i, x, z, 1, 2, val, MatA);

          // f(i+2,j-2)
          val = -1.0 * A3 / (144.0 * dxdz);
          Element(i, x, z, 2, -2, val, MatA);

          // f(i+2,j-1)
          val = A3 / (18.0 * dxdz);
          Element(i, x, z, 2, -1, val, MatA);

          // f(i+2,j)
          val = (-1.0 / 12.0) * ((A1 / dx2) + (A4 / dx));
          Element(i, x, z, 2, 0, val, MatA);

          // f(i+2,j+1)
          val = -1.0 * A3 / (18.0 * dxdz);
          Element(i, x, z, 2, 1, val, MatA);

          // f(i+2,j+2)
          val = A3 / (144.0 * dxdz);
          Element(i, x, z, 2, 2, val, MatA);
        } else {
          // Second order

          // f(i,j) = f(x,z)
          val = A0 - 2.0 * ((A1 / dx2) + (A2 / dz2));
          Element(i, x, z, 0, 0, val, MatA);

          // f(i-1,j-1)
          val = A3 / (4.0 * dxdz);
          Element(i, x, z, -1, -1, val, MatA);

          // f(i-1,j)
          val = (A1 / dx2) - A4 / (2.0 * dx);
          Element(i, x, z, -1, 0, val, MatA);

          // f(i-1,j+1)
          val = -1.0 * A3 / (4.0 * dxdz);
          Element(i, x, z, -1, 1, val, MatA);

          // f(i,j-1)
          val = (A2 / dz2) - (A5 / (2.0 * dz));
          Element(i, x, z, 0, -1, val, MatA);

          // f(i,j+1)
          val = (A2 / dz2) + (A5 / (2.0 * dz));
          Element(i, x, z, 0, 1, val, MatA);

          // f(i+1,j-1)
          val = -1.0 * A3 / (4.0 * dxdz);
          Element(i, x, z, 1, -1, val, MatA);

          // f(i+1,j)
          val = (A1 / dx2) + (A4 / (2.0 * dx));
          Element(i, x, z, 1, 0, val, MatA);

          // f(i+1,j+1)
          val = A3 / (4.0 * dxdz);
          Element(i, x, z, 1, 1, val, MatA);
        }
        // Set Components of RHS Vector
        val = b[x][z];
        VecSetValues(bs, 1, &i, &val, INSERT_VALUES);

        // Set Components of Trial Solution Vector
        val = x0[x][z];
        VecSetValues(xs, 1, &i, &val, INSERT_VALUES);
        i++;
      }
    }

    // X=localmesh->xend+1 to localmesh->LocalNx-1 defines the upper boundary region of the domain.
    // Set the values for the outer boundary region
    if (localmesh->lastX()) {
      for (int x = localmesh->xend + 1; x < localmesh->LocalNx; x++) {
        for (int z = localmesh->zstart; z <= localmesh->zend; z++) {
          // Set Diagonal Values to 1
          PetscScalar val = 1;
          Element(i, x, z, 0, 0, val, MatA);

          // If Neumann Boundary Conditions are set.
          if (isOuterBoundaryFlagSet(INVERT_AC_GRAD)) {
            // Set values corresponding to nodes adjacent in x
            if (fourth_order) {
              // Fourth Order Accuracy on Boundary
              Element(i, x, z, 0, 0,
                      25.0 / (12.0 * coords->dx(x, y, z)) / sqrt(coords->g_11(x, y, z)),
                      MatA);
              Element(i, x, z, -1, 0,
                      -4.0 / coords->dx(x, y, z) / sqrt(coords->g_11(x, y, z)), MatA);
              Element(i, x, z, -2, 0,
                      3.0 / coords->dx(x, y, z) / sqrt(coords->g_11(x, y, z)), MatA);
              Element(i, x, z, -3, 0,
                      -4.0 / (3.0 * coords->dx(x, y, z)) / sqrt(coords->g_11(x, y, z)),
                      MatA);
              Element(i, x, z, -4, 0,
                      1.0 / (4.0 * coords->dx(x, y, z)) / sqrt(coords->g_11(x, y, z)),
                      MatA);
            } else {
              // // Second Order Accuracy on Boundary
              // Element(i,x,z,  0, 0,  3.0 / (2.0*coords->dx(x,y)), MatA );
              // Element(i,x,z, -1, 0, -2.0 / coords->dx(x,y), MatA );
              // Element(i,x,z, -2, 0,  1.0 / (2.0*coords->dx(x,y)), MatA );
              // Element(i,x,z, -3, 0,  0.0, MatA );  // Reset these elements to 0
              // in case 4th order flag was used previously: not allowed now
              // Element(i,x,z, -4, 0,  0.0, MatA );
              // Second Order Accuracy on Boundary, set half-way between grid
              // points
              Element(i, x, z, 0, 0,
                      1.0 / coords->dx(x, y, z) / sqrt(coords->g_11(x, y, z)), MatA);
              Element(i, x, z, -1, 0,
                      -1.0 / coords->dx(x, y, z) / sqrt(coords->g_11(x, y, z)), MatA);
              Element(i, x, z, -2, 0, 0.0, MatA);
              // Element(i,x,z, -3, 0,  0.0, MatA );  // Reset these elements to 0
              // in case 4th order flag was used previously: not allowed now
              // Element(i,x,z, -4, 0,  0.0, MatA );
            }
          } else {
            if (fourth_order) {
              // Set off diagonal elements to zero
              Element(i, x, z, -1, 0, 0.0, MatA);
              Element(i, x, z, -2, 0, 0.0, MatA);
              Element(i, x, z, -3, 0, 0.0, MatA);
              Element(i, x, z, -4, 0, 0.0, MatA);
            } else {
              Element(i, x, z, 0, 0, 0.5, MatA);
              Element(i, x, z, -1, 0, 0.5, MatA);
              Element(i, x, z, -2, 0, 0., MatA);
            }
          }

          // Set Components of RHS
          // If the inner boundary value should be set by b or x0
          val = 0;
          if (isOuterBoundaryFlagSet(INVERT_RHS)) {
            val = b[x][z];
          } else if (isOuterBoundaryFlagSet(INVERT_SET)) {
            val = x0[x][z];
          }

          // Set components of the RHS (the PETSc vector bs)
          // 1 element is being set in row i to val
          // INSERT_VALUES replaces existing entries with new values
          VecSetValues(bs, 1, &i, &val, INSERT_VALUES);

          // Set components of the and trial solution (the PETSc vector xs)
          // 1 element is being set in row i to val
          // INSERT_VALUES replaces existing entries with new values
          val = x0[x][z];
          VecSetValues(xs, 1, &i, &val, INSERT_VALUES);

          i++; // Increment row in Petsc matrix
        }
      }
    }

    if (i != Iend) {
      throw BoutException("Petsc index sanity check failed: i={} != Iend={}", i, Iend);
    }

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

    //   // Record which flags were used for this matrix
    //   lastflag = flags;

    // Assemble RHS Vector
    VecAssemblyBegin(bs);
    VecAssemblyEnd(bs);

    // Assemble Trial Solution Vector
    VecAssemblyBegin(xs);
    VecAssemblyEnd(xs);

    // Configure Linear Solver
#if PETSC_VERSION_GE(3, 5, 0)
    KSPSetOperators(ksp, MatA, MatA);
#else
    KSPSetOperators(ksp, MatA, MatA, DIFFERENT_NONZERO_PATTERN);
#endif
    PC pc; // The preconditioner option

    if (direct) { // If a direct solver has been chosen
      // Get the preconditioner
      KSPGetPC(ksp, &pc);
      // Set the preconditioner
      PCSetType(pc, PCLU);
      // Set the solver type
#if PETSC_VERSION_GE(3, 9, 0)
      PCFactorSetMatSolverType(pc, "mumps");
#else
      PCFactorSetMatSolverPackage(pc, "mumps");
#endif
    } else {                            // If a iterative solver has been chosen
      KSPSetType(ksp, ksptype.c_str()); // Set the type of the solver

      if (ksptype == KSPRICHARDSON) {
        KSPRichardsonSetScale(ksp, richardson_damping_factor);
      }
#ifdef KSPCHEBYSHEV
      else if (ksptype == KSPCHEBYSHEV) {
        KSPChebyshevSetEigenvalues(ksp, chebyshev_max, chebyshev_min);
      }
#endif
      else if (ksptype == KSPGMRES) {
        KSPGMRESSetRestart(ksp, gmres_max_steps);
      }

      // Set the relative and absolute tolerances
      KSPSetTolerances(ksp, rtol, atol, dtol, maxits);

      // If the initial guess is not set to zero
      if (!isGlobalFlagSet(INVERT_START_NEW)) {
        KSPSetInitialGuessNonzero(ksp, static_cast<PetscBool>(true));
      }

      // Get the preconditioner
      KSPGetPC(ksp, &pc);

      // Set the type of the preconditioner
      PCSetType(pc, pctype.c_str());

      // If pctype = user in BOUT.inp, it will be translated to PCSHELL upon
      // construction of the object
      if (pctype == PCSHELL) {
        // User-supplied preconditioner function
        PCShellSetApply(pc, laplacePCapply);
        PCShellSetContext(pc, this);
        if (rightprec) {
          KSPSetPCSide(ksp, PC_RIGHT); // Right preconditioning
        } else {
          KSPSetPCSide(ksp, PC_LEFT); // Left preconditioning
        }
        //ierr = PCShellSetApply(pc,laplacePCapply);CHKERRQ(ierr);
        //ierr = PCShellSetContext(pc,this);CHKERRQ(ierr);
        //ierr = KSPSetPCSide(ksp, PC_RIGHT);CHKERRQ(ierr);
      }

      lib.setOptionsFromInputFile(ksp);
    }
  }

  // Call the actual solver
  {
    Timer timer("petscsolve");
    KSPSolve(ksp, bs, xs); // Call the solver to solve the system
  }

  KSPConvergedReason reason;
  KSPGetConvergedReason(ksp, &reason);
  if (reason == -3) { // Too many iterations, might be fixed by taking smaller timestep
    throw BoutIterationFail("petsc_laplace: too many iterations");
  }
  if (reason <= 0) {
    throw BoutException(
        "petsc_laplace: inversion failed to converge. KSPConvergedReason: {} ({})",
        KSPConvergedReasons[reason], static_cast<int>(reason));
  }

  // Add data to FieldPerp Object
  i = Istart;
  // Set the inner boundary values
  if (localmesh->firstX()) {
    for (int x = 0; x < localmesh->xstart; x++) {
      for (int z = localmesh->zstart; z <= localmesh->zend; z++) {
        PetscScalar val = 0;
        VecGetValues(xs, 1, &i, &val);
        sol[x][z] = val;
        i++; // Increment row in Petsc matrix
      }
    }
  }

  // Set the main domain values
  for (int x = localmesh->xstart; x <= localmesh->xend; x++) {
    for (int z = localmesh->zstart; z <= localmesh->zend; z++) {
      PetscScalar val = 0;
      VecGetValues(xs, 1, &i, &val);
      sol[x][z] = val;
      i++; // Increment row in Petsc matrix
    }
  }

  // Set the outer boundary values
  if (localmesh->lastX()) {
    for (int x = localmesh->xend + 1; x < localmesh->LocalNx; x++) {
      for (int z = localmesh->zstart; z <= localmesh->zend; z++) {
        PetscScalar val = 0;
        VecGetValues(xs, 1, &i, &val);
        sol[x][z] = val;
        i++; // Increment row in Petsc matrix
      }
    }
  }

  if (i != Iend) {
    throw BoutException("Petsc index sanity check 2 failed");
  }

  checkData(sol);

  // Return the solution
  return sol;
}

/*!
 * Sets the elements of the matrix A, which is used to solve the problem Ax=b.
 *
 * \param[in]
 * i
 * The row of the PETSc matrix
 * \param[in] x         Local x index of the mesh
 * \param[in] z         Local z index of the mesh
 * \param[in] xshift    The shift in rows from the index x
 * \param[in] zshift    The shift in columns from the index z
 * \param[in] ele       Value of the element
 * \param[in] MatA      The matrix A used in the inversion
 *
 * \param[out] MatA     The matrix A used in the inversion
 */
void LaplacePetsc::Element(int i, int x, int z, int xshift, int zshift, PetscScalar ele,
                           Mat& MatA) {

  // Need to convert LOCAL x to GLOBAL x in order to correctly calculate
  // PETSC Matrix Index.
  int xoffset = Istart / meshz;
#if CHECK > 2
  const int rem = Istart % meshz;
  if (rem != 0) {
    throw BoutException("Petsc index sanity check 3 failed: Istart={} % meshz={} == {}",
                        Istart, meshz, rem);
  }
#endif

  // Calculate the row to be set
  int row_new = x + xshift; // should never be out of range.
  if (!localmesh->firstX()) {
    row_new += (xoffset - localmesh->xstart);
  }

  // Calculate the column to be set
  int col_new = z + zshift - localmesh->zstart;
  if (col_new < 0) {
    col_new += meshz;
  } else if (col_new > meshz - 1) {
    col_new -= meshz;
  }

  // Convert to global indices
  int index = (row_new * meshz) + col_new;

#if CHECK > 2
  if (!std::isfinite(ele)) {
    throw BoutException("Non-finite element at x={:d}, z={:d}, row={:d}, col={:d}\n", x,
                        z, i, index);
  }
#endif

  /* Inserts or adds a block of values into a matrix
   * Input:
   * MatA   - The matrix to set the values in
   * 1      - The number of rows to be set
   * &i     - The global index of the row
   * 1      - The number of columns to be set
   * &index - The global index of the column
   * &ele   - The vlaue to be set
   * INSERT_VALUES replaces existing entries with new values
   */
  MatSetValues(MatA, 1, &i, 1, &index, &ele, INSERT_VALUES);
}

/*!
 * Set the matrix components of A in Ax=b, solving
 * D*Laplace_perp(x) + (1/C1)Grad_perp(C2)*Grad_perp(x) + Ax = B
 *
 * \note "A" in the equation above is not added here.
 * For calculations of the coefficients, please refer to the user manual.
 *
 * \param[in] x The current x index
 * \param[in] y The current y index
 * \param[in] z The current y index
 * \param[in] coef1  Placeholder for convenient variable used to set matrix
 *                   (see manual for details)
 * \param[in] coef2  Convenient variable used to set matrix
 *                   (see manual for details)
 * \param[in] coef3  Placeholder for convenient variable used to set matrix
 *                   (see manual for details)
 * \param[in] coef4  Placeholder for convenient variable used to set matrix
 *                   (see manual for details)
 * \param[in] coef5  Placeholder for convenient variable used to set matrix
 *                   (see manual for details)
 *
 * \param[out] coef1    Convenient variable used to set matrix
 *                      (see manual for details)
 * \param[out] coef2    Convenient variable used to set matrix
 *                      (see manual for details)
 * \param[out] coef3    Convenient variable used to set matrix
 *                      (see manual for details)
 * \param[out] coef4    Convenient variable used to set matrix
 *                      (see manual for details)
 * \param[out] coef5    Convenient variable used to set matrix
 *                      (see manual for details)
 */
void LaplacePetsc::Coeffs(int x, int y, int z, BoutReal& coef1, BoutReal& coef2,
                          BoutReal& coef3, BoutReal& coef4, BoutReal& coef5) {

  coef1 = coords->g11(x, y, z);      // X 2nd derivative coefficient
  coef2 = coords->g33(x, y, z);      // Z 2nd derivative coefficient
  coef3 = 2. * coords->g13(x, y, z); // X-Z mixed derivative coefficient

  coef4 = 0.0;
  coef5 = 0.0;
  // If global flag all_terms are set (true by default)
  if (all_terms) {
    coef4 = coords->G1(x, y, z); // X 1st derivative
    coef5 = coords->G3(x, y, z); // Z 1st derivative

    ASSERT3(std::isfinite(coef4));
    ASSERT3(std::isfinite(coef5));
  }

  if (nonuniform) {
    // non-uniform mesh correction
    if ((x != 0) && (x != (localmesh->LocalNx - 1))) {
      coef4 -= 0.5
               * ((coords->dx(x + 1, y, z) - coords->dx(x - 1, y, z))
                  / SQ(coords->dx(x, y, z)))
               * coef1; // BOUT-06 term
    }
  }

  if (localmesh->IncIntShear) {
    // d2dz2 term
    coef2 += coords->g11(x, y, z) * coords->IntShiftTorsion(x, y, z)
             * coords->IntShiftTorsion(x, y, z);
    // Mixed derivative
    coef3 = 0.0; // This cancels out
  }

  if (issetD) {
    coef1 *= D(x, y, z);
    coef2 *= D(x, y, z);
    coef3 *= D(x, y, z);
    coef4 *= D(x, y, z);
    coef5 *= D(x, y, z);
  }

  // A second/fourth order derivative term
  if (issetC) {
    //   if( (x > 0) && (x < (localmesh->LocalNx-1)) ) //Valid if doing second order derivative, not if fourth: should only be called for xstart<=x<=xend anyway
    if ((x > 1) && (x < (localmesh->LocalNx - 2))) {
      int zp = z + 1; // z plus 1
      if (zp > meshz - 1) {
        zp -= meshz;
      }
      int zm = z - 1; // z minus 1
      if (zm < 0) {
        zm += meshz;
      }
      BoutReal ddx_C;
      BoutReal ddz_C;

      if (fourth_order) {
        int zpp = z + 2; // z plus 1 plus 1
        if (zpp > meshz - 1) {
          zpp -= meshz;
        }
        int zmm = z - 2; // z minus 1 minus 1
        if (zmm < 0) {
          zmm += meshz;
        }
        // Fourth order discretization of C in x
        ddx_C = (-C2(x + 2, y, z) + 8. * C2(x + 1, y, z) - 8. * C2(x - 1, y, z)
                 + C2(x - 2, y, z))
                / (12. * coords->dx(x, y, z) * (C1(x, y, z)));
        // Fourth order discretization of C in z
        ddz_C = (-C2(x, y, zpp) + 8. * C2(x, y, zp) - 8. * C2(x, y, zm) + C2(x, y, zmm))
                / (12. * coords->dz(x, y, z) * (C1(x, y, z)));
      } else {
        // Second order discretization of C in x
        ddx_C = (C2(x + 1, y, z) - C2(x - 1, y, z))
                / (2. * coords->dx(x, y, z) * (C1(x, y, z)));
        // Second order discretization of C in z
        ddz_C =
            (C2(x, y, zp) - C2(x, y, zm)) / (2. * coords->dz(x, y, z) * (C1(x, y, z)));
      }

      coef4 += coords->g11(x, y, z) * ddx_C + coords->g13(x, y, z) * ddz_C;
      coef5 += coords->g13(x, y, z) * ddx_C + coords->g33(x, y, z) * ddz_C;
    }
  }

  /* Ex and Ez
   * Additional 1st derivative terms to allow for solution field to be
   * components of a vector
   *
   * NB multiply by D or Grad_perp(C)/C as appropriate before passing to
   * setCoefEx()/setCoefEz() because (in principle) both are needed and we
   * don't know how to split them up here
   */
  if (issetE) {
    // These coefficients are 0 by default
    coef4 += Ex(x, y, z);
    coef5 += Ez(x, y, z);
  }
}

void LaplacePetsc::vecToField(Vec xs, FieldPerp& f) {

  ASSERT1(localmesh == f.getMesh());

  f.allocate();
  int i = Istart;
  if (localmesh->firstX()) {
    for (int x = 0; x < localmesh->xstart; x++) {
      for (int z = localmesh->zstart; z <= localmesh->zend; z++) {
        PetscScalar val;
        VecGetValues(xs, 1, &i, &val);
        f[x][z] = val;
        i++; // Increment row in Petsc matrix
      }
    }
  }

  for (int x = localmesh->xstart; x <= localmesh->xend; x++) {
    for (int z = localmesh->zstart; z <= localmesh->zend; z++) {
      PetscScalar val;
      VecGetValues(xs, 1, &i, &val);
      f[x][z] = val;
      i++; // Increment row in Petsc matrix
    }
  }

  if (localmesh->lastX()) {
    for (int x = localmesh->xend + 1; x < localmesh->LocalNx; x++) {
      for (int z = localmesh->zstart; z <= localmesh->zend; z++) {
        PetscScalar val;
        VecGetValues(xs, 1, &i, &val);
        f[x][z] = val;
        i++; // Increment row in Petsc matrix
      }
    }
  }
  ASSERT1(i == Iend);
}

void LaplacePetsc::fieldToVec(const FieldPerp& f, Vec bs) {
  ASSERT1(localmesh == f.getMesh());

  int i = Istart;
  if (localmesh->firstX()) {
    for (int x = 0; x < localmesh->xstart; x++) {
      for (int z = localmesh->zstart; z <= localmesh->zend; z++) {
        PetscScalar val = f[x][z];
        VecSetValues(bs, 1, &i, &val, INSERT_VALUES);
        i++; // Increment row in Petsc matrix
      }
    }
  }

  for (int x = localmesh->xstart; x <= localmesh->xend; x++) {
    for (int z = localmesh->zstart; z <= localmesh->zend; z++) {
      PetscScalar val = f[x][z];
      VecSetValues(bs, 1, &i, &val, INSERT_VALUES);
      i++; // Increment row in Petsc matrix
    }
  }

  if (localmesh->lastX()) {
    for (int x = localmesh->xend + 1; x < localmesh->LocalNx; x++) {
      for (int z = localmesh->zstart; z <= localmesh->zend; z++) {
        PetscScalar val = f[x][z];
        VecSetValues(bs, 1, &i, &val, INSERT_VALUES);
        i++; // Increment row in Petsc matrix
      }
    }
  }
  ASSERT1(i == Iend);

  VecAssemblyBegin(bs);
  VecAssemblyEnd(bs);
}

/// Preconditioner function
int LaplacePetsc::precon(Vec x, Vec y) {
  // Get field to be preconditioned
  FieldPerp xfield;
  vecToField(x, xfield);
  xfield.setIndex(sol.getIndex()); // y index stored in sol variable

  // Call the preconditioner solver
  FieldPerp yfield = pcsolve->solve(xfield);

  // Put result into y
  fieldToVec(yfield, y);
  return 0;
}

void LaplacePetsc::checkFlags() {
  if (isGlobalFlagSet(~implemented_flags)) {
    if (isGlobalFlagSet(INVERT_4TH_ORDER)) {
      output_error.write(
          "For PETSc based Laplacian inverter, use 'fourth_order=true' instead of "
          "setting INVERT_4TH_ORDER flag\n");
    }
    throw BoutException("Attempted to set Laplacian inversion flag that is not "
                        "implemented in petsc_laplace.cxx");
  }
  if (isInnerBoundaryFlagSet(~implemented_boundary_flags)) {
    throw BoutException("Attempted to set Laplacian inversion boundary flag that is not "
                        "implemented in petsc_laplace.cxx");
  }
  if (isOuterBoundaryFlagSet(~implemented_boundary_flags)) {
    throw BoutException("Attempted to set Laplacian inversion boundary flag that is not "
                        "implemented in petsc_laplace.cxx");
  }
}

#endif // BOUT_HAS_PETSC_3_3
