code-stellarator-optimization/include/sctl/boundary_quadrature.hpp

#ifndef _SCTL_BOUNDARY_QUADRATURE_HPP_
#define _SCTL_BOUNDARY_QUADRATURE_HPP_

#include SCTL_INCLUDE(tree.hpp)
#include SCTL_INCLUDE(tensor.hpp)
#include SCTL_INCLUDE(morton.hpp)
#include SCTL_INCLUDE(matrix.hpp)
#include SCTL_INCLUDE(vector.hpp)
#include SCTL_INCLUDE(common.hpp)
#include SCTL_INCLUDE(cheb_utils.hpp)
#include SCTL_INCLUDE(kernel_functions.hpp)

#include <biest.hpp>
#include <mutex>
#include <atomic>
#include <tuple>
#include <functional>
#include <Eigen/Core>
#include <LBFGS.h>

namespace SCTL_NAMESPACE {

template <class Real, Integer DIM, Integer ORDER> class Basis {
  public:
    using ValueType = Real;

    // class EvalOperator {
    //   public:
    // };
    using EvalOpType = Matrix<ValueType>;

    static constexpr Long Dim() {
      return DIM;
    }
    static constexpr Long Size() {
      return pow<DIM,Long>(ORDER);
    }
    static const Matrix<ValueType>& Nodes() {
      static Matrix<ValueType> nodes_(DIM,Size());
      auto nodes_1d = [](Integer i) {
        return 0.5 - 0.5 * cos<ValueType>((2*i+1) * const_pi<ValueType>() / (2*ORDER));
      };
      { // Set nodes_
        static std::mutex mutex;
        static std::atomic<Integer> first_time(true);
        if (first_time.load(std::memory_order_relaxed)) {
          std::lock_guard<std::mutex> guard(mutex);
          if (first_time.load(std::memory_order_relaxed)) {
            Integer N = 1;
            for (Integer d = 0; d < DIM; d++) {
              for (Integer j = 0; j < ORDER; j++) {
                for (Integer i = 0; i < N; i++) {
                  for (Integer k = 0; k < d; k++) {
                    nodes_[k][j*N+i] = nodes_[k][i];
                  }
                  nodes_[d][j*N+i] = nodes_1d(j);
                }
              }
              N *= ORDER;
            }

            std::atomic_thread_fence(std::memory_order_seq_cst);
            first_time.store(false);
          }
        }
      }
      return nodes_;
    }
    static const Vector<ValueType>& QuadWts() {
      static Vector<ValueType> wts(Size());
      { // Set nodes_
        static std::mutex mutex;
        static std::atomic<Integer> first_time(true);
        if (first_time.load(std::memory_order_relaxed)) {
          std::lock_guard<std::mutex> guard(mutex);
          if (first_time.load(std::memory_order_relaxed)) {
            StaticArray<ValueType,ORDER> wts_1d;
            { // Set wts_1d
              Vector<ValueType> x_(ORDER);
              ChebBasis<ValueType>::template Nodes<1>(ORDER, x_);

              Vector<ValueType> V_cheb(ORDER * ORDER);
              { // Set V_cheb
                Vector<ValueType> I(ORDER*ORDER);
                I = 0;
                for (Long i = 0; i < ORDER; i++) I[i*ORDER+i] = 1;
                ChebBasis<ValueType>::template Approx<1>(ORDER, I, V_cheb);
              }
              Matrix<ValueType> M(ORDER, ORDER, V_cheb.begin());

              Vector<ValueType> w_sample(ORDER);
              for (Integer i = 0; i < ORDER; i++) {
                w_sample[i] = (i % 2 ? 0 : -(ORDER/(ValueType)(i*i-1)));
              }
              for (Integer j = 0; j < ORDER; j++) {
                wts_1d[j] = 0;
                for (Integer i = 0; i < ORDER; i++) {
                  wts_1d[j] += M[j][i] * w_sample[i] / ORDER;
                }
              }
            }

            wts[0] = 1;
            Integer N = 1;
            for (Integer d = 0; d < DIM; d++) {
              for (Integer j = 1; j < ORDER; j++) {
                for (Integer i = 0; i < N; i++) {
                  wts[j*N+i] = wts[i] * wts_1d[j];
                }
              }
              for (Integer i = 0; i < N; i++) {
                wts[i] *= wts_1d[0];
              }
              N *= ORDER;
            }

            std::atomic_thread_fence(std::memory_order_seq_cst);
            first_time.store(false);
          }
        }
      }
      return wts;
    }

    static void Grad(Vector<Basis>& dX, const Vector<Basis>& X) {
      static Matrix<ValueType> GradOp[DIM];
      static std::mutex mutex;
      static std::atomic<Integer> first_time(true);
      if (first_time.load(std::memory_order_relaxed)) {
        std::lock_guard<std::mutex> guard(mutex);
        if (first_time.load(std::memory_order_relaxed)) {
          { // Set GradOp
            auto nodes = Basis<ValueType,1,ORDER>::Nodes();
            SCTL_ASSERT(nodes.Dim(1) == ORDER);
            Matrix<ValueType> M(ORDER, ORDER);
            for (Integer i = 0; i < ORDER; i++) { // Set M
              Real x = nodes[0][i];
              for (Integer j = 0; j < ORDER; j++) {
                M[j][i] = 0;
                for (Integer l = 0; l < ORDER; l++) {
                  if (l != j) {
                    Real M_ = 1;
                    for (Integer k = 0; k < ORDER; k++) {
                      if (k != j && k != l) M_ *= (x - nodes[0][k]);
                      if (k != j) M_ /= (nodes[0][j] - nodes[0][k]);
                    }
                    M[j][i] += M_;
                  }
                }
              }
            }
            for (Integer d = 0; d < DIM; d++) {
              GradOp[d].ReInit(Size(), Size());
              GradOp[d] = 0;
              Integer stride0 = pow<Integer>(ORDER, d);
              Integer repeat0 = pow<Integer>(ORDER, d);
              Integer stride1 = pow<Integer>(ORDER, d+1);
              Integer repeat1 = pow<Integer>(ORDER, DIM-d-1);
              for (Integer k1 = 0; k1 < repeat1; k1++) {
                for (Integer i = 0; i < ORDER; i++) {
                  for (Integer j = 0; j < ORDER; j++) {
                    for (Integer k0 = 0; k0 < repeat0; k0++) {
                      GradOp[d][k1*stride1 + i*stride0 + k0][k1*stride1 + j*stride0 + k0] = M[i][j];
                    }
                  }
                }
              }
            }
          }
          std::atomic_thread_fence(std::memory_order_seq_cst);
          first_time.store(false);
        }
      }

      if (dX.Dim() != X.Dim()*DIM) dX.ReInit(X.Dim()*DIM);
      for (Long i = 0; i < X.Dim(); i++) {
        const Matrix<ValueType> Vi(1, Size(), (Iterator<ValueType>)(ConstIterator<ValueType>)X[i].NodeValues_, false);
        for (Integer k = 0; k < DIM; k++) {
          Matrix<ValueType> Vo(1, Size(), dX[i*DIM+k].NodeValues_, false);
          Matrix<ValueType>::GEMM(Vo, Vi, GradOp[k]);
        }
      }
    }
    static EvalOpType SetupEval(const Matrix<ValueType>& X) {
      Long N = X.Dim(1);
      SCTL_ASSERT(X.Dim(0) == DIM);
      Matrix<ValueType> M(Size(), N);
      { // Set M
        auto nodes = Basis<ValueType,1,ORDER>::Nodes();
        Integer NN = Basis<ValueType,1,ORDER>::Size();
        Matrix<ValueType> M_(NN, DIM*N);
        for (Long i = 0; i < DIM*N; i++) {
          ValueType x = X[0][i];
          for (Integer j = 0; j < NN; j++) {
            ValueType y = 1;
            for (Integer k = 0; k < NN; k++) {
              y *= (j==k ? 1 : (nodes[0][k] - x) / (nodes[0][k] - nodes[0][j]));
            }
            M_[j][i] = y;
          }
        }
        if (DIM == 1) {
          SCTL_ASSERT(M.Dim(0) == M_.Dim(0));
          SCTL_ASSERT(M.Dim(1) == M_.Dim(1));
          M = M_;
        } else {
          Integer NNN = 1;
          M = 1;
          for (Integer d = 0; d < DIM; d++) {
            for (Integer k = 1; k < NN; k++) {
              for (Integer j = 0; j < NNN; j++) {
                for (Long i = 0; i < N; i++) {
                  M[k*NNN+j][i] = M[j][i] * M_[k][d*N+i];
                }
              }
            }
            { // k = 0
              for (Integer j = 0; j < NNN; j++) {
                for (Long i = 0; i < N; i++) {
                  M[j][i] *= M_[0][d*N+i];
                }
              }
            }
            NNN *= NN;
          }
        }
      }
      return M;
    }
    static void Eval(Matrix<ValueType>& Y, const Vector<Basis>& X, const EvalOpType& M) {
      Long N0 = X.Dim();
      Long N1 = M.Dim(1);
      SCTL_ASSERT(M.Dim(0) == Size());
      if (Y.Dim(0) != N0 || Y.Dim(1) != N1) Y.ReInit(N0, N1);
      for (Long i = 0; i < N0; i++) {
        const Matrix<ValueType> X_(1,Size(),(Iterator<ValueType>)(ConstIterator<ValueType>)X[i].NodeValues_,false);
        Matrix<ValueType> Y_(1,N1,Y[i],false);
        Matrix<ValueType>::GEMM(Y_,X_,M);
      }
    }

    Basis operator+(Basis X) const {
      for (Long i = 0; i < Size(); i++) X[i] = (*this)[i] + X[i];
      return X;
    }
    Basis operator-(Basis X) const {
      for (Long i = 0; i < Size(); i++) X[i] = (*this)[i] - X[i];
      return X;
    }
    Basis operator*(Basis X) const {
      for (Long i = 0; i < Size(); i++) X[i] = (*this)[i] * X[i];
      return X;
    }
    Basis operator*(Real a) const {
      Basis X = (*this);
      for (Long i = 0; i < Size(); i++) X[i] *= a;
      return X;
    }
    Basis operator+(Real a) const {
      Basis X = (*this);
      for (Long i = 0; i < Size(); i++) X[i] += a;
      return X;
    }
    Basis& operator+=(const Basis& X) {
      for (Long i = 0; i < Size(); i++) (*this)[i] += X[i];
      return *this;
    }
    Basis& operator-=(const Basis& X) {
      for (Long i = 0; i < Size(); i++) (*this)[i] -= X[i];
      return *this;
    }
    Basis& operator*=(const Basis& X) {
      for (Long i = 0; i < Size(); i++) (*this)[i] *= X[i];
      return *this;
    }
    Basis& operator*=(Real a) {
      for (Long i = 0; i < Size(); i++) (*this)[i] *= a;
      return *this;
    }
    Basis& operator+=(Real a) {
      for (Long i = 0; i < Size(); i++) (*this)[i] += a;
      return *this;
    }
    Basis& operator=(Real a) {
      for (Long i = 0; i < Size(); i++) (*this)[i] = a;
      return *this;
    }


    const ValueType& operator[](Long i) const {
      SCTL_ASSERT(i < Size());
      return NodeValues_[i];
    }
    ValueType& operator[](Long i) {
      SCTL_ASSERT(i < Size());
      return NodeValues_[i];
    }

  private:
    StaticArray<ValueType,Size()> NodeValues_;
};

template <Integer COORD_DIM, class Basis> class ElemList {
  public:

    using CoordBasis = Basis;
    using CoordType = typename CoordBasis::ValueType;

    static constexpr Integer CoordDim() {
      return COORD_DIM;
    }
    static constexpr Integer ElemDim() {
      return CoordBasis::Dim();
    }

    ElemList(Long Nelem = 0) {
      ReInit(Nelem);
    }
    void ReInit(Long Nelem = 0) {
      Nelem_ = Nelem;
      X_.ReInit(Nelem_ * COORD_DIM);
    }
    void ReInit(const Vector<CoordBasis>& X) {
      Nelem_ = X.Dim() / COORD_DIM;
      SCTL_ASSERT(X.Dim() == Nelem_ * COORD_DIM);
      X_ = X;
    }
    Long NElem() const {
      return Nelem_;
    }

    CoordBasis& operator()(Long elem, Integer dim) {
      SCTL_ASSERT(elem >= 0 && elem < Nelem_);
      SCTL_ASSERT(dim >= 0 && dim < COORD_DIM);
      return X_[elem*COORD_DIM+dim];
    }
    const CoordBasis& operator()(Long elem, Integer dim) const {
      if (!(elem >= 0 && elem < Nelem_)) exit(0);
      SCTL_ASSERT(elem >= 0 && elem < Nelem_);
      SCTL_ASSERT(dim >= 0 && dim < COORD_DIM);
      return X_[elem*COORD_DIM+dim];
    }
    const Vector<CoordBasis>& ElemVector() const {
      return X_;
    }

  private:
    static_assert(CoordBasis::Dim() <= CoordDim(), "Basis dimension can not be greater than COORD_DIM.");
    Vector<CoordBasis> X_;
    Long Nelem_;

    //mutable Vector<CoordBasis> dX_;
};

template <class Real> class Quadrature {

    static Real machine_epsilon() {
      Real eps=1;
      while(eps*(Real)0.5+(Real)1.0>1.0) eps*=0.5;
      return eps;
    }

    template <Integer DIM> static void DuffyQuad(Matrix<Real>& nodes, Vector<Real>& weights, const Vector<Real>& coord, Integer order, Real adapt = -1.0) {
      SCTL_ASSERT(coord.Dim() == DIM);
      static Real eps = machine_epsilon()*16;

      Matrix<Real> qx;
      Vector<Real> qw;
      { // Set qx, qw
        Vector<Real> qx0, qw0;
        ChebBasis<Real>::quad_rule(order, qx0, qw0);

        Integer N = pow<DIM,Integer>(order);
        qx.ReInit(DIM,N);
        qw.ReInit(N);
        qw[0] = 1;

        Integer N_ = 1;
        for (Integer d = 0; d < DIM; d++) {
          for (Integer j = 0; j < order; j++) {
            for (Integer i = 0; i < N_; i++) {
              for (Integer k = 0; k < d; k++) {
                qx[k][j*N_+i] = qx[k][i];
              }
              qx[d][j*N_+i] = qx0[j];
              qw[j*N_+i] = qw[i];
            }
          }
          for (Integer j = 0; j < order; j++) {
            for (Integer i = 0; i < N_; i++) {
              qw[j*N_+i] *= qw0[j];
            }
          }
          N_ *= order;
        }
      }

      Vector<Real> X;
      { // Set X
        StaticArray<Real,2*DIM+2> X_;
        X_[0] = 0;
        X_[1] = adapt;
        for (Integer i = 0; i < DIM; i++) {
          X_[2*i+2] = fabs<Real>(coord[i]);
          X_[2*i+3] = fabs<Real>(coord[i]-1);
        }
        std::sort((Iterator<Real>)X_, (Iterator<Real>)X_+2*DIM+2);

        X.PushBack(std::max<Real>(0, X_[2*DIM]-1));
        for (Integer i = 0; i < 2*DIM+2; i++) {
          if (X[X.Dim()-1] < X_[i]) {
            if (X.Dim())
            X.PushBack(X_[i]);
          }
        }

        /////////////////////////////////////////////////////////////////////////////////////////////////
        Vector<Real> r(1);
        r[0] = X[0];
        for (Integer i = 1; i < X.Dim(); i++) {
          while (r[r.Dim() - 1] > 0.0 && (order*0.5) * r[r.Dim() - 1] < X[i]) r.PushBack((order*0.5) * r[r.Dim() - 1]); // TODO
          r.PushBack(X[i]);
        }
        X = r;
        /////////////////////////////////////////////////////////////////////////////////////////////////
      }

      Vector<Real> nds, wts;
      for (Integer k = 0; k < X.Dim()-1; k++) {
        for (Integer dd = 0; dd < 2*DIM; dd++) {
          Integer d0 = (dd>>1);
          StaticArray<Real,2*DIM> range0, range1;
          { // Set range0, range1
            Integer d1 = (dd%2?1:-1);
            for (Integer d = 0; d < DIM; d++) {
              range0[d*2+0] = std::max<Real>(0,std::min<Real>(1,coord[d] - X[k]  ));
              range0[d*2+1] = std::max<Real>(0,std::min<Real>(1,coord[d] + X[k]  ));
              range1[d*2+0] = std::max<Real>(0,std::min<Real>(1,coord[d] - X[k+1]));
              range1[d*2+1] = std::max<Real>(0,std::min<Real>(1,coord[d] + X[k+1]));
            }
            range0[d0*2+0] = std::max<Real>(0,std::min<Real>(1,coord[d0] + d1*X[k+0]));
            range0[d0*2+1] = std::max<Real>(0,std::min<Real>(1,coord[d0] + d1*X[k+0]));
            range1[d0*2+0] = std::max<Real>(0,std::min<Real>(1,coord[d0] + d1*X[k+1]));
            range1[d0*2+1] = std::max<Real>(0,std::min<Real>(1,coord[d0] + d1*X[k+1]));
          }
          { // if volume(range0, range1) == 0 then continue
            Real v0 = 1, v1 = 1;
            for (Integer d = 0; d < DIM; d++) {
              if (d == d0) {
                v0 *= fabs<Real>(range0[d*2+0]-range1[d*2+0]);
                v1 *= fabs<Real>(range0[d*2+0]-range1[d*2+0]);
              } else {
                v0 *= range0[d*2+1]-range0[d*2+0];
                v1 *= range1[d*2+1]-range1[d*2+0];
              }
            }
            if (v0 < eps && v1 < eps) continue;
          }
          for (Integer i = 0; i < qx.Dim(1); i++) { // Set nds, wts
            Real w = qw[i];
            Real z = qx[d0][i];
            for (Integer d = 0; d < DIM; d++) {
              Real y = qx[d][i];
              nds.PushBack((range0[d*2+0]*(1-y) + range0[d*2+1]*y)*(1-z) + (range1[d*2+0]*(1-y) + range1[d*2+1]*y)*z);
              if (d == d0) {
                w *= abs(range1[d*2+0] - range0[d*2+0]);
              } else {
                w *= (range0[d*2+1] - range0[d*2+0])*(1-z) + (range1[d*2+1] - range1[d*2+0])*z;
              }
            }
            wts.PushBack(w);
          }
        }
      }
      nodes = Matrix<Real>(nds.Dim()/DIM,DIM,nds.begin()).Transpose();
      weights = wts;
    }

    template <Integer DIM> static void TensorProductGaussQuad(Matrix<Real>& nodes, Vector<Real>& weights, Integer order) {
      Vector<Real> coord(DIM);
      coord = 0;
      coord[0] = -10;
      DuffyQuad<DIM>(nodes, weights, coord, order);
    }



    template <class DensityBasis, class ElemList, class Kernel> static void SetupSingular(Matrix<Real>& M_singular, const Matrix<Real>& trg_nds, const ElemList& elem_lst, const Kernel& kernel, Integer order_singular = 10, Integer order_direct = 10, Real Rqbx = 0) {
      using CoordBasis = typename ElemList::CoordBasis;
      using CoordEvalOpType = typename CoordBasis::EvalOpType;
      using DensityEvalOpType = typename DensityBasis::EvalOpType;

      constexpr Integer CoordDim = ElemList::CoordDim();
      constexpr Integer ElemDim = ElemList::ElemDim();
      constexpr Integer KDIM0 = Kernel::SrcDim();
      constexpr Integer KDIM1 = Kernel::TrgDim();

      const Long Nelem = elem_lst.NElem();
      const Integer Ntrg = trg_nds.Dim(1);
      SCTL_ASSERT(trg_nds.Dim(0) == ElemDim);

      const Vector<CoordBasis>& X = elem_lst.ElemVector();
      Vector<CoordBasis> dX;
      CoordBasis::Grad(dX, X);

      Vector<Real> Xt, Xnt;
      { // Set Xt, Xnt
        auto Meval = CoordBasis::SetupEval(trg_nds);
        eval_basis(Xt, X, CoordDim, trg_nds.Dim(1), Meval);

        Xnt = Xt;
        Vector<Real> dX_;
        eval_basis(dX_, dX, 2*CoordDim, trg_nds.Dim(1), Meval);

        for (Long i = 0; i < Ntrg; i++) {
          for (Long j = 0; j < Nelem; j++) {
            auto Xn = Xnt.begin() + (j*Ntrg+i)*CoordDim;
            auto dX0 = dX_.begin() + (j*Ntrg+i)*2*CoordDim;

            StaticArray<Real,CoordDim> normal;
            normal[0] = dX0[2]*dX0[5] - dX0[4]*dX0[3];
            normal[1] = dX0[4]*dX0[1] - dX0[0]*dX0[5];
            normal[2] = dX0[0]*dX0[3] - dX0[2]*dX0[1];
            Real Xa = sqrt<Real>(normal[0]*normal[0]+normal[1]*normal[1]+normal[2]*normal[2]);
            Real invXa = 1/Xa;
            normal[0] *= invXa;
            normal[1] *= invXa;
            normal[2] *= invXa;

            Real sqrt_Xa = sqrt<Real>(Xa);
            Xn[0] = normal[0]*sqrt_Xa*Rqbx;
            Xn[1] = normal[1]*sqrt_Xa*Rqbx;
            Xn[2] = normal[2]*sqrt_Xa*Rqbx;
          }
        }
      }
      SCTL_ASSERT(Xt.Dim() == Nelem * Ntrg * CoordDim);

      auto& M = M_singular;
      M.ReInit(Nelem * KDIM0 * DensityBasis::Size(), KDIM1 * Ntrg);
      #pragma omp parallel for schedule(static)
      for (Long i = 0; i < Ntrg; i++) { // Set M (singular)
        Matrix<Real> quad_nds;
        Vector<Real> quad_wts;
        { // Set quad_nds, quad_wts
          StaticArray<Real,ElemDim> trg_node_;
          for (Integer k = 0; k < ElemDim; k++) {
            trg_node_[k] = trg_nds[k][i];
          }
          Vector<Real> trg_node(ElemDim, trg_node_, false);
          DuffyQuad<ElemDim>(quad_nds, quad_wts, trg_node, order_singular, fabs(Rqbx));
        }
        const CoordEvalOpType CoordEvalOp = CoordBasis::SetupEval(quad_nds);
        Integer Nnds = quad_wts.Dim();

        Vector<Real> X_, dX_, Xa_, Xn_;
        { // Set X_, dX_
          eval_basis(X_, X, CoordDim, Nnds, CoordEvalOp);
          eval_basis(dX_, dX, CoordDim * ElemDim, Nnds, CoordEvalOp);
        }
        if (CoordDim == 3 && ElemDim == 2) { // Compute Xa_, Xn_
          Long N = Nelem*Nnds;
          Xa_.ReInit(N);
          Xn_.ReInit(N*CoordDim);
          for (Long j = 0; j < N; j++) {
            StaticArray<Real,CoordDim> normal;
            normal[0] = dX_[j*6+2]*dX_[j*6+5] - dX_[j*6+4]*dX_[j*6+3];
            normal[1] = dX_[j*6+4]*dX_[j*6+1] - dX_[j*6+0]*dX_[j*6+5];
            normal[2] = dX_[j*6+0]*dX_[j*6+3] - dX_[j*6+2]*dX_[j*6+1];
            Xa_[j] = sqrt<Real>(normal[0]*normal[0]+normal[1]*normal[1]+normal[2]*normal[2]);
            Real invXa = 1/Xa_[j];
            Xn_[j*3+0] = normal[0] * invXa;
            Xn_[j*3+1] = normal[1] * invXa;
            Xn_[j*3+2] = normal[2] * invXa;
          }
        }

        DensityEvalOpType DensityEvalOp;
        if (std::is_same<CoordBasis,DensityBasis>::value) {
          DensityEvalOp = CoordEvalOp;
        } else {
          DensityEvalOp = DensityBasis::SetupEval(quad_nds);
        }

        for (Long j = 0; j < Nelem; j++) {
          Matrix<Real> M__(Nnds * KDIM0, KDIM1);
          if (Rqbx == 0) { // Set kernel matrix M__
            const Vector<Real> X0_(CoordDim, Xt.begin() + (j * Ntrg + i) * CoordDim, false);
            const Vector<Real> X__(Nnds * CoordDim, X_.begin() + j * Nnds * CoordDim, false);
            const Vector<Real> Xn__(Nnds * CoordDim, Xn_.begin() + j * Nnds * CoordDim, false);
            kernel.template KernelMatrix<Real>(M__, X0_, X__, Xn__);
          } else {
            Vector<Real> X0_(CoordDim);
            constexpr Integer qbx_order = 6;
            StaticArray<Matrix<Real>,qbx_order> M___;
            for (Integer k = 0; k < qbx_order; k++) { // Set kernel matrix M___
              for (Integer kk = 0; kk < CoordDim; kk++) X0_[kk] = Xt[(j * Ntrg + i) * CoordDim + kk] + (k+1) * Xnt[(j * Ntrg + i) * CoordDim + kk];
              const Vector<Real> X__(Nnds * CoordDim, X_.begin() + j * Nnds * CoordDim, false);
              const Vector<Real> Xn__(Nnds * CoordDim, Xn_.begin() + j * Nnds * CoordDim, false);
              kernel.template KernelMatrix<Real>(M___[k], X0_, X__, Xn__);
            }

            for (Long k = 0; k < Nnds * KDIM0 * KDIM1; k++) {
              M__[0][k] = 0;
              M__[0][k] +=   6*M___[0][0][k];
              M__[0][k] += -15*M___[1][0][k];
              M__[0][k] +=  20*M___[2][0][k];
              M__[0][k] += -15*M___[3][0][k];
              M__[0][k] +=   6*M___[4][0][k];
              M__[0][k] +=  -1*M___[5][0][k];
            }
          }

          for (Long k0 = 0; k0 < KDIM0; k0++) {
            for (Long k1 = 0; k1 < KDIM1; k1++) {
              for (Long l = 0; l < DensityBasis::Size(); l++) {
                Real M_lk = 0;
                for (Long n = 0; n < Nnds; n++) {
                  Real quad_wt = Xa_[j * Nnds + n] * quad_wts[n];
                  M_lk += DensityEvalOp[l][n] * quad_wt * M__[n*KDIM0+k0][k1];
                }
                M[(j * KDIM0 + k0) * DensityBasis::Size() + l][k1 * Ntrg + i] = M_lk;
              }
            }
          }
        }
      }
      { // Set M (subtract direct)
        Matrix<Real> quad_nds;
        Vector<Real> quad_wts;
        TensorProductGaussQuad<ElemDim>(quad_nds, quad_wts, order_direct);
        const CoordEvalOpType CoordEvalOp = CoordBasis::SetupEval(quad_nds);
        Integer Nnds = quad_wts.Dim();

        Vector<Real> X_, dX_, Xa_, Xn_;
        { // Set X_, dX_
          eval_basis(X_, X, CoordDim, Nnds, CoordEvalOp);
          eval_basis(dX_, dX, CoordDim * ElemDim, Nnds, CoordEvalOp);
        }
        if (CoordDim == 3 && ElemDim == 2) { // Compute Xa_, Xn_
          Long N = Nelem*Nnds;
          Xa_.ReInit(N);
          Xn_.ReInit(N*CoordDim);
          for (Long j = 0; j < N; j++) {
            StaticArray<Real,CoordDim> normal;
            normal[0] = dX_[j*6+2]*dX_[j*6+5] - dX_[j*6+4]*dX_[j*6+3];
            normal[1] = dX_[j*6+4]*dX_[j*6+1] - dX_[j*6+0]*dX_[j*6+5];
            normal[2] = dX_[j*6+0]*dX_[j*6+3] - dX_[j*6+2]*dX_[j*6+1];
            Xa_[j] = sqrt<Real>(normal[0]*normal[0]+normal[1]*normal[1]+normal[2]*normal[2]);
            Real invXa = 1/Xa_[j];
            Xn_[j*3+0] = normal[0] * invXa;
            Xn_[j*3+1] = normal[1] * invXa;
            Xn_[j*3+2] = normal[2] * invXa;
          }
        }

        DensityEvalOpType DensityEvalOp;
        if (std::is_same<CoordBasis,DensityBasis>::value) {
          DensityEvalOp = CoordEvalOp;
        } else {
          DensityEvalOp = DensityBasis::SetupEval(quad_nds);
        }

        #pragma omp parallel for schedule(static)
        for (Long i = 0; i < Ntrg; i++) { // Subtract direct contribution
          for (Long j = 0; j < Nelem; j++) {
            Matrix<Real> M__(Nnds * KDIM0, KDIM1);
            { // Set kernel matrix M__
              const Vector<Real> X0_(CoordDim, (Iterator<Real>)Xt.begin() + (j * Ntrg + i) * CoordDim, false);
              const Vector<Real> X__(Nnds * CoordDim, X_.begin() + j * Nnds * CoordDim, false);
              const Vector<Real> Xn__(Nnds * CoordDim, Xn_.begin() + j * Nnds * CoordDim, false);
              kernel.template KernelMatrix<Real>(M__, X0_, X__, Xn__);
            }
            for (Long k0 = 0; k0 < KDIM0; k0++) {
              for (Long k1 = 0; k1 < KDIM1; k1++) {
                for (Long l = 0; l < DensityBasis::Size(); l++) {
                  Real M_lk = 0;
                  for (Long n = 0; n < Nnds; n++) {
                    Real quad_wt = Xa_[j * Nnds + n] * quad_wts[n];
                    M_lk += DensityEvalOp[l][n] * quad_wt * M__[n*KDIM0+k0][k1];
                  }
                  M[(j * KDIM0 + k0) * DensityBasis::Size() + l][k1 * Ntrg + i] -= M_lk;
                }
              }
            }
          }
        }
      }
    }

    template <class DensityBasis> static void EvalSingular(Matrix<Real>& U, const Vector<DensityBasis>& density, const Matrix<Real>& M, Integer KDIM0_, Integer KDIM1_) {
      if (M.Dim(0) == 0 || M.Dim(1) == 0) {
        U.ReInit(0,0);
        return;
      }

      const Long Ntrg = M.Dim(1) / KDIM1_;
      SCTL_ASSERT(M.Dim(1) == KDIM1_ * Ntrg);

      const Long Nelem = M.Dim(0) / (KDIM0_ * DensityBasis::Size());
      SCTL_ASSERT(M.Dim(0) == Nelem * KDIM0_ * DensityBasis::Size());

      const Integer dof = density.Dim() / (Nelem * KDIM0_);
      SCTL_ASSERT(density.Dim() == Nelem * dof * KDIM0_);

      if (U.Dim(0) != Nelem * dof * KDIM1_ || U.Dim(1) != Ntrg) {
        U.ReInit(Nelem * dof * KDIM1_, Ntrg);
        U = 0;
      }
      for (Long j = 0; j < Nelem; j++) {
        const Matrix<Real> M_(KDIM0_ * DensityBasis::Size(), KDIM1_ * Ntrg, (Iterator<Real>)M[j * KDIM0_ * DensityBasis::Size()], false);
        Matrix<Real> U_(dof, KDIM1_ * Ntrg, U[j*dof*KDIM1_], false);
        Matrix<Real> F_(dof, KDIM0_ * DensityBasis::Size());
        for (Long i = 0; i < dof; i++) {
          for (Long k = 0; k < KDIM0_; k++) {
            for (Long l = 0; l < DensityBasis::Size(); l++) {
              F_[i][k * DensityBasis::Size() + l] = density[(j * dof + i) * KDIM0_ + k][l];
            }
          }
        }
        Matrix<Real>::GEMM(U_, F_, M_);
      }
    }



    template <Integer DIM> struct PointData {
      bool operator<(const PointData& p) const {
        return mid < p.mid;
      }

      Long rank;
      Long surf_rank;
      Morton<DIM> mid;
      StaticArray<Real,DIM> coord;
      Real radius2;
    };

    template <class T1, class T2> struct Pair {
      Pair() {}

      Pair(T1 x, T2 y) : first(x), second(y) {}

      bool operator<(const Pair& p) const {
        return (first < p.first) || (((first == p.first) && (second < p.second)));
      }

      T1 first;
      T2 second;
    };

    template <class ElemList> static void BuildNbrList(Vector<Pair<Long,Long>>& pair_lst, const Vector<Real>& Xt, const Vector<Long>& trg_surf, const ElemList& elem_lst, Real distance_factor, Real period_length, const Comm& comm) {
      using CoordBasis = typename ElemList::CoordBasis;
      constexpr Integer CoordDim = ElemList::CoordDim();
      constexpr Integer ElemDim = ElemList::ElemDim();
      using PtData = PointData<CoordDim>;
      const Integer rank = comm.Rank();

      Real R0 = 0;
      StaticArray<Real,CoordDim> X0;
      { // Find bounding box
        Long N = Xt.Dim() / CoordDim;
        SCTL_ASSERT(Xt.Dim() == N * CoordDim);
        SCTL_ASSERT(N);

        StaticArray<Real,CoordDim*2> Xloc;
        StaticArray<Real,CoordDim*2> Xglb;
        for (Integer k = 0; k < CoordDim; k++) {
          Xloc[0*CoordDim+k] = Xt[k];
          Xloc[1*CoordDim+k] = Xt[k];
        }
        for (Long i = 0; i < N; i++) {
          for (Integer k = 0; k < CoordDim; k++) {
            Xloc[0*CoordDim+k] = std::min<Real>(Xloc[0*CoordDim+k], Xt[i*CoordDim+k]);
            Xloc[1*CoordDim+k] = std::max<Real>(Xloc[1*CoordDim+k], Xt[i*CoordDim+k]);
          }
        }
        comm.Allreduce((ConstIterator<Real>)Xloc+0*CoordDim, (Iterator<Real>)Xglb+0*CoordDim, CoordDim, Comm::CommOp::MIN);
        comm.Allreduce((ConstIterator<Real>)Xloc+1*CoordDim, (Iterator<Real>)Xglb+1*CoordDim, CoordDim, Comm::CommOp::MAX);
        for (Integer k = 0; k < CoordDim; k++) {
          R0 = std::max(R0, Xglb[1*CoordDim+k]-Xglb[0*CoordDim+k]);
        }

        R0 = R0 * 2.0;
        for (Integer k = 0; k < CoordDim; k++) {
          X0[k] = Xglb[k] - R0*0.25;
        }
      }
      if (period_length > 0) {
        R0 = period_length;
      }

      Vector<PtData> PtSrc, PtTrg;
      Integer order_upsample = (Integer)(const_pi<Real>() / distance_factor + 0.5);
      { // Set PtSrc
        const Vector<CoordBasis>& X_elem_lst = elem_lst.ElemVector();
        Vector<CoordBasis> dX_elem_lst;
        CoordBasis::Grad(dX_elem_lst, X_elem_lst);

        Matrix<Real> nds;
        Vector<Real> wts;
        TensorProductGaussQuad<ElemDim>(nds, wts, order_upsample);
        const Long Nnds = nds.Dim(1);

        Vector<Real> X, dX;
        const auto CoordEvalOp = CoordBasis::SetupEval(nds);
        eval_basis(X, X_elem_lst, CoordDim, Nnds, CoordEvalOp);
        eval_basis(dX, dX_elem_lst, CoordDim * ElemDim, Nnds, CoordEvalOp);

        const Long N = X.Dim() / CoordDim;
        const Long Nelem = elem_lst.NElem();
        SCTL_ASSERT(X.Dim() == N * CoordDim);
        SCTL_ASSERT(N == Nelem * Nnds);

        Long rank_offset, surf_rank_offset;
        { // Set rank_offset, surf_rank_offset
          comm.Scan(Ptr2ConstItr<Long>(&N,1), Ptr2Itr<Long>(&rank_offset,1), 1, Comm::CommOp::SUM);
          comm.Scan(Ptr2ConstItr<Long>(&Nelem,1), Ptr2Itr<Long>(&surf_rank_offset,1), 1, Comm::CommOp::SUM);
          surf_rank_offset -= Nelem;
          rank_offset -= N;
        }

        PtSrc.ReInit(N);
        const Real R0inv = 1.0 / R0;
        for (Long i = 0; i < N; i++) { // Set coord
          for (Integer k = 0; k < CoordDim; k++) {
            PtSrc[i].coord[k] = (X[i*CoordDim+k] - X0[k]) * R0inv;
          }
        }
        if (period_length > 0) { // Wrap-around coord
          for (Long i = 0; i < N; i++) {
            auto& x = PtSrc[i].coord;
            for (Integer k = 0; k < CoordDim; k++) {
              x[k] -= (Long)(x[k]);
            }
          }
        }
        for (Long i = 0; i < N; i++) { // Set radius2, mid, rank
          Integer depth = 0;
          { // Set radius2, depth
            Real radius2 = 0;
            for (Integer k0 = 0; k0 < ElemDim; k0++) {
              Real R2 = 0;
              for (Integer k1 = 0; k1 < CoordDim; k1++) {
                Real dX_ = dX[(i*CoordDim+k1)*ElemDim+k0];
                R2 += dX_*dX_;
              }
              radius2 = std::max(radius2, R2);
            }
            radius2 *= R0inv*R0inv * distance_factor*distance_factor;
            PtSrc[i].radius2 = radius2;

            Long Rinv = (Long)(1.0/radius2);
            while (Rinv > 0) {
              Rinv = (Rinv>>2);
              depth++;
            }
          }
          PtSrc[i].mid = Morton<CoordDim>((Iterator<Real>)PtSrc[i].coord, std::min(Morton<CoordDim>::MaxDepth(),depth));
          PtSrc[i].rank = rank_offset + i;
        }
        for (Long i = 0 ; i < Nelem; i++) { // Set surf_rank
          for (Long j = 0; j < Nnds; j++) {
            PtSrc[i*Nnds+j].surf_rank = surf_rank_offset + i;
          }
        }

        Vector<PtData> PtSrcSorted;
        comm.HyperQuickSort(PtSrc, PtSrcSorted);
        PtSrc.Swap(PtSrcSorted);
      }
      { // Set PtTrg
        const Long N = Xt.Dim() / CoordDim;
        SCTL_ASSERT(Xt.Dim() == N * CoordDim);

        Long rank_offset;
        { // Set rank_offset
          comm.Scan(Ptr2ConstItr<Long>(&N,1), Ptr2Itr<Long>(&rank_offset,1), 1, Comm::CommOp::SUM);
          rank_offset -= N;
        }

        PtTrg.ReInit(N);
        const Real R0inv = 1.0 / R0;
        for (Long i = 0; i < N; i++) { // Set coord
          for (Integer k = 0; k < CoordDim; k++) {
            PtTrg[i].coord[k] = (Xt[i*CoordDim+k] - X0[k]) * R0inv;
          }
        }
        if (period_length > 0) { // Wrap-around coord
          for (Long i = 0; i < N; i++) {
            auto& x = PtTrg[i].coord;
            for (Integer k = 0; k < CoordDim; k++) {
              x[k] -= (Long)(x[k]);
            }
          }
        }
        for (Long i = 0; i < N; i++) { // Set radius2, mid, rank
          PtTrg[i].radius2 = 0;
          PtTrg[i].mid = Morton<CoordDim>((Iterator<Real>)PtTrg[i].coord);
          PtTrg[i].rank = rank_offset + i;
        }
        if (trg_surf.Dim()) { // Set surf_rank
          SCTL_ASSERT(trg_surf.Dim() == N);
          for (Long i = 0; i < N; i++) {
            PtTrg[i].surf_rank = trg_surf[i];
          }
        } else {
          for (Long i = 0; i < N; i++) {
            PtTrg[i].surf_rank = -1;
          }
        }

        Vector<PtData> PtTrgSorted;
        comm.HyperQuickSort(PtTrg, PtTrgSorted);
        PtTrg.Swap(PtTrgSorted);
      }

      Tree<CoordDim> tree(comm);
      { // Init tree
        Vector<Real> Xall(PtSrc.Dim()+PtTrg.Dim());
        { // Set Xall
          Xall.ReInit((PtSrc.Dim()+PtTrg.Dim())*CoordDim);
          Long Nsrc = PtSrc.Dim();
          Long Ntrg = PtTrg.Dim();
          for (Long i = 0; i < Nsrc; i++) {
            for (Integer k = 0; k < CoordDim; k++) {
              Xall[i*CoordDim+k] = PtSrc[i].coord[k];
            }
          }
          for (Long i = 0; i < Ntrg; i++) {
            for (Integer k = 0; k < CoordDim; k++) {
              Xall[(Nsrc+i)*CoordDim+k] = PtTrg[i].coord[k];
            }
          }
        }
        tree.UpdateRefinement(Xall, 1000, true, period_length>0);
      }
      { // Repartition PtSrc, PtTrg
        PtData splitter;
        splitter.mid = tree.GetPartitionMID()[rank];
        comm.PartitionS(PtSrc, splitter);
        comm.PartitionS(PtTrg, splitter);
      }
      { // Add tree data PtSrc
        const auto& node_mid = tree.GetNodeMID();
        const Long N = node_mid.Dim();
        SCTL_ASSERT(N);

        Vector<Long> dsp(N), cnt(N);
        for (Long i = 0; i < N; i++) {
          PtData m0;
          m0.mid = node_mid[i];
          dsp[i] = std::lower_bound(PtSrc.begin(), PtSrc.end(), m0) - PtSrc.begin();
        }
        for (Long i = 0; i < N-1; i++) {
          cnt[i] = dsp[i+1] - dsp[i];
        }
        cnt[N-1] = PtSrc.Dim() - dsp[N-1];
        tree.AddData("PtSrc", PtSrc, cnt);
      }
      tree.template Broadcast<PtData>("PtSrc");

      { // Build pair_lst
        Vector<Long> cnt;
        Vector<PtData> PtSrc;
        tree.GetData(PtSrc, cnt, "PtSrc");
        const auto& node_mid = tree.GetNodeMID();
        const auto& node_attr = tree.GetNodeAttr();

        Vector<Morton<CoordDim>> nbr_mid_tmp;
        for (Long i = 0; i < node_mid.Dim(); i++) {
          if (node_attr[i].Leaf && !node_attr[i].Ghost) {
            Vector<Morton<CoordDim>> child_mid;
            node_mid[i].Children(child_mid);
            for (const auto& trg_mid : child_mid) {
              Integer d0 = trg_mid.Depth();
              Vector<PtData> Src, Trg;
              { // Set Trg
                PtData m0, m1;
                m0.mid = trg_mid;
                m1.mid = trg_mid.Next();
                Long a = std::lower_bound(PtTrg.begin(), PtTrg.end(), m0) - PtTrg.begin();
                Long b = std::lower_bound(PtTrg.begin(), PtTrg.end(), m1) - PtTrg.begin();
                Trg.ReInit(b-a, PtTrg.begin()+a, false);
                if (!Trg.Dim()) continue;
              }

              Vector<std::set<Long>> near_elem(Trg.Dim());
              for (Integer d = 0; d <= d0; d++) {
                trg_mid.NbrList(nbr_mid_tmp, d, period_length>0);
                for (const auto& src_mid : nbr_mid_tmp) { // Set Src
                  PtData m0, m1;
                  m0.mid = src_mid;
                  m1.mid = (d==d0 ? src_mid.Next() : src_mid.Ancestor(d+1));
                  Long a = std::lower_bound(PtSrc.begin(), PtSrc.end(), m0) - PtSrc.begin();
                  Long b = std::lower_bound(PtSrc.begin(), PtSrc.end(), m1) - PtSrc.begin();
                  Src.ReInit(b-a, PtSrc.begin()+a, false);
                  if (!Src.Dim()) continue;

                  for (Long t = 0; t < Trg.Dim(); t++) { // set near_elem[t] <-- {s : dist(s,t) < radius(s)}
                    for (Long s = 0; s < Src.Dim(); s++) {
                      if (Trg[t].surf_rank != Src[s].surf_rank) {
                        Real R2 = 0;
                        for (Integer k = 0; k < CoordDim; k++) {
                          Real dx = (Src[s].coord[k] - Trg[t].coord[k]);
                          R2 += dx * dx;
                        }
                        if (R2 < Src[s].radius2) {
                          near_elem[t].insert(Src[s].surf_rank);
                        }
                      }
                    }
                  }
                }
              }

              for (Long t = 0; t < Trg.Dim(); t++) { // Set pair_lst
                for (Long elem_idx : near_elem[t]) {
                  pair_lst.PushBack(Pair<Long,Long>(elem_idx,Trg[t].rank));
                }
              }
            }
          }
        }
      }
      { // Sort and repartition pair_lst
        Vector<Pair<Long,Long>> pair_lst_sorted;
        comm.HyperQuickSort(pair_lst, pair_lst_sorted);

        Long surf_rank_offset;
        const Long Nelem = elem_lst.NElem();
        comm.Scan(Ptr2ConstItr<Long>(&Nelem,1), Ptr2Itr<Long>(&surf_rank_offset,1), 1, Comm::CommOp::SUM);
        surf_rank_offset -= Nelem;

        comm.PartitionS(pair_lst_sorted, Pair<Long,Long>(surf_rank_offset,0));
        pair_lst.Swap(pair_lst_sorted);
      }
    }

    template <class ElemList> static void BuildNbrListDeprecated(Vector<Pair<Long,Long>>& pair_lst, const Vector<Real>& Xt, const ElemList& elem_lst, const Matrix<Real>& surf_nds, Real distance_factor) {
      using CoordBasis = typename ElemList::CoordBasis;
      constexpr Integer CoordDim = ElemList::CoordDim();
      constexpr Integer ElemDim = ElemList::ElemDim();
      const Long Nelem = elem_lst.NElem();

      const Long Ntrg = Xt.Dim() / CoordDim;
      SCTL_ASSERT(Xt.Dim() == Ntrg * CoordDim);

      Long Nnds, Nsurf_nds;
      Vector<Real> X_surf, X, dX;
      Integer order_upsample = (Integer)(const_pi<Real>() / distance_factor + 0.5);
      { // Set X, dX
        const Vector<CoordBasis>& X_elem_lst = elem_lst.ElemVector();
        Vector<CoordBasis> dX_elem_lst;
        CoordBasis::Grad(dX_elem_lst, X_elem_lst);

        Matrix<Real> nds_upsample;
        Vector<Real> wts_upsample;
        TensorProductGaussQuad<ElemDim>(nds_upsample, wts_upsample, order_upsample);

        Nnds = nds_upsample.Dim(1);
        const auto CoordEvalOp = CoordBasis::SetupEval(nds_upsample);
        eval_basis(X, X_elem_lst, CoordDim, nds_upsample.Dim(1), CoordEvalOp);
        eval_basis(dX, dX_elem_lst, CoordDim * ElemDim, nds_upsample.Dim(1), CoordEvalOp);

        Nsurf_nds = surf_nds.Dim(1);
        const auto CoordEvalOp_surf = CoordBasis::SetupEval(surf_nds);
        eval_basis(X_surf, X_elem_lst, CoordDim, Nsurf_nds, CoordEvalOp_surf);
      }

      Real d2 = distance_factor * distance_factor;
      for (Long i = 0; i < Nelem; i++) {
        std::set<Long> near_pts;
        std::set<Long> self_pts;
        for (Long j = 0; j < Nnds; j++) {
          Real R2_max = 0;
          StaticArray<Real, CoordDim> X0;
          for (Integer k = 0; k < CoordDim; k++) {
            X0[k] = X[(i*Nnds+j)*CoordDim+k];
          }
          for (Integer k0 = 0; k0 < ElemDim; k0++) {
            Real R2 = 0;
            for (Integer k1 = 0; k1 < CoordDim; k1++) {
              Real dX_ = dX[((i*Nnds+j)*CoordDim+k1)*ElemDim+k0];
              R2 += dX_*dX_;
            }
            R2_max = std::max(R2_max, R2*d2);
          }

          for (Long k = 0; k < Ntrg; k++) {
            Real R2 = 0;
            for (Integer l = 0; l < CoordDim; l++) {
              Real dX = Xt[k*CoordDim+l]- X0[l];
              R2 += dX * dX;
            }
            if (R2 < R2_max) near_pts.insert(k);
          }
        }
        for (Long j = 0; j < Nsurf_nds; j++) {
          StaticArray<Real, CoordDim> X0;
          for (Integer k = 0; k < CoordDim; k++) {
            X0[k] = X_surf[(i*Nsurf_nds+j)*CoordDim+k];
          }
          for (Long k = 0; k < Ntrg; k++) {
            Real R2 = 0;
            for (Integer l = 0; l < CoordDim; l++) {
              Real dX = Xt[k*CoordDim+l]- X0[l];
              R2 += dX * dX;
            }
            if (R2 == 0) self_pts.insert(k);
          }
        }
        for (Long trg_idx : self_pts) {
          near_pts.erase(trg_idx);
        }
        for (Long trg_idx : near_pts) {
          pair_lst.PushBack(Pair<Long,Long>(i,trg_idx));
        }
      }
    }

    template <class DensityBasis, class ElemList, class Kernel> static void SetupNearSingular(Matrix<Real>& M_near_singular, Vector<Pair<Long,Long>>& pair_lst, const Vector<Real>& Xt_, const Vector<Long>& trg_surf, const ElemList& elem_lst, const Kernel& kernel, Integer order_singular, Integer order_direct, Real period_length, const Comm& comm) {
      static_assert(std::is_same<Real,typename DensityBasis::ValueType>::value);
      static_assert(std::is_same<Real,typename ElemList::CoordType>::value);
      static_assert(DensityBasis::Dim() == ElemList::ElemDim());
      using CoordBasis = typename ElemList::CoordBasis;
      using CoordEvalOpType = typename CoordBasis::EvalOpType;
      using DensityEvalOpType = typename DensityBasis::EvalOpType;

      constexpr Integer CoordDim = ElemList::CoordDim();
      constexpr Integer ElemDim = ElemList::ElemDim();
      constexpr Integer KDIM0 = Kernel::SrcDim();
      constexpr Integer KDIM1 = Kernel::TrgDim();
      const Long Nelem = elem_lst.NElem();

      BuildNbrList(pair_lst, Xt_, trg_surf, elem_lst, 2.5/order_direct, period_length, comm);
      const Long Ninterac = pair_lst.Dim();

      Vector<Real> Xt;
      { // Set Xt
        Integer rank = comm.Rank();
        Integer np = comm.Size();

        Vector<Long> splitter_ranks;
        { // Set splitter_ranks
          Vector<Long> cnt(np);
          const Long N = Xt_.Dim() / CoordDim;
          comm.Allgather(Ptr2ConstItr<Long>(&N,1), 1, cnt.begin(), 1);
          scan(splitter_ranks, cnt);
        }

        Vector<Long> scatter_index, recv_index, recv_cnt(np), recv_dsp(np);
        { // Set scatter_index, recv_index, recv_cnt, recv_dsp
          { // Set scatter_index, recv_index
            Vector<Pair<Long,Long>> scatter_pair(pair_lst.Dim());
            for (Long i = 0; i < pair_lst.Dim(); i++) {
              scatter_pair[i] = Pair<Long,Long>(pair_lst[i].second,i);
            }
            omp_par::merge_sort(scatter_pair.begin(), scatter_pair.end());

            recv_index.ReInit(scatter_pair.Dim());
            scatter_index.ReInit(scatter_pair.Dim());
            for (Long i = 0; i < scatter_index.Dim(); i++) {
              recv_index[i] = scatter_pair[i].first;
              scatter_index[i] = scatter_pair[i].second;
            }
          }
          for (Integer i = 0; i < np; i++) {
            recv_dsp[i] = std::lower_bound(recv_index.begin(), recv_index.end(), splitter_ranks[i]) - recv_index.begin();
          }
          for (Integer i = 0; i < np-1; i++) {
            recv_cnt[i] = recv_dsp[i+1] - recv_dsp[i];
          }
          recv_cnt[np-1] = recv_index.Dim() - recv_dsp[np-1];
        }

        Vector<Long> send_index, send_cnt(np), send_dsp(np);
        { // Set send_index, send_cnt, send_dsp
          comm.Alltoall(recv_cnt.begin(), 1, send_cnt.begin(), 1);
          scan(send_dsp, send_cnt);
          send_index.ReInit(send_cnt[np-1] + send_dsp[np-1]);
          comm.Alltoallv(recv_index.begin(), recv_cnt.begin(), recv_dsp.begin(), send_index.begin(), send_cnt.begin(), send_dsp.begin());
        }

        Vector<Real> Xt_send(send_index.Dim() * CoordDim);
        for (Long i = 0; i < send_index.Dim(); i++) { // Set Xt_send
          Long idx = send_index[i] - splitter_ranks[rank];
          for (Integer k = 0; k < CoordDim; k++) {
            Xt_send[i*CoordDim+k] = Xt_[idx*CoordDim+k];
          }
        }

        Vector<Real> Xt_recv(recv_index.Dim() * CoordDim);
        { // Set Xt_recv
          for (Long i = 0; i < np; i++) {
            send_cnt[i] *= CoordDim;
            send_dsp[i] *= CoordDim;
            recv_cnt[i] *= CoordDim;
            recv_dsp[i] *= CoordDim;
          }
          comm.Alltoallv(Xt_send.begin(), send_cnt.begin(), send_dsp.begin(), Xt_recv.begin(), recv_cnt.begin(), recv_dsp.begin());
        }

        Xt.ReInit(scatter_index.Dim() * CoordDim);
        for (Long i = 0; i < scatter_index.Dim(); i++) { // Set Xt
          Long idx = scatter_index[i];
          for (Integer k = 0; k < CoordDim; k++) {
            Xt[idx*CoordDim+k] = Xt_recv[i*CoordDim+k];
          }
        }
      }

      const Vector<CoordBasis>& X = elem_lst.ElemVector();
      Vector<CoordBasis> dX;
      CoordBasis::Grad(dX, X);

      Long elem_rank_offset;
      { // Set elem_rank_offset
        comm.Scan(Ptr2ConstItr<Long>(&Nelem,1), Ptr2Itr<Long>(&elem_rank_offset,1), 1, Comm::CommOp::SUM);
        elem_rank_offset -= Nelem;
      }

      auto& M = M_near_singular;
      M.ReInit(Ninterac * KDIM0 * DensityBasis::Size(), KDIM1);
      #pragma omp parallel for schedule(static)
      for (Long j = 0; j < Ninterac; j++) { // Set M (near-singular)
        const Long src_idx = pair_lst[j].first - elem_rank_offset;

        Real adapt = -1.0;
        Tensor<Real,true,ElemDim,1> u0;
        { // Set u0 (project target point to the surface patch in parameter space)
          ConstIterator<Real> Xt_ = Xt.begin() + j * CoordDim;
          const auto& nodes = CoordBasis::Nodes();

          Long min_idx = -1;
          Real min_R2 = 1e10;
          for (Long i = 0; i < CoordBasis::Size(); i++) {
            Real R2 = 0;
            for (Integer k = 0; k < CoordDim; k++) {
              Real dX = X[src_idx * CoordDim + k][i] - Xt_[k];
              R2 += dX * dX;
            }
            if (R2 < min_R2) {
              min_R2 = R2;
              min_idx = i;
            }
          }
          SCTL_ASSERT(min_idx >= 0);
          for (Integer k = 0; k < ElemDim; k++) {
            u0(k,0) = nodes[k][min_idx];
          }

          for (Integer i = 0; i < 2; i++) { // iterate
            Matrix<Real> X_, dX_;
            for (Integer k = 0; k < ElemDim; k++) {
              u0(k,0) = std::min<Real>(1.0, u0(k,0));
              u0(k,0) = std::max<Real>(0.0, u0(k,0));
            }
            const auto eval_op = CoordBasis::SetupEval(Matrix<Real>(ElemDim,1,u0.begin(),false));
            CoordBasis::Eval(X_, Vector<CoordBasis>(CoordDim,(Iterator<CoordBasis>)X.begin()+src_idx*CoordDim,false),eval_op);
            CoordBasis::Eval(dX_, Vector<CoordBasis>(CoordDim*ElemDim,dX.begin()+src_idx*CoordDim*ElemDim,false),eval_op);

            const Tensor<Real,false,CoordDim,1> x0((Iterator<Real>)Xt_);
            const Tensor<Real,false,CoordDim,1> x(X_.begin());
            const Tensor<Real,false,CoordDim,ElemDim> x_u(dX_.begin());
            auto inv = [](const Tensor<Real,true,2,2>& M) {
              Tensor<Real,true,2,2> Minv;
              Real det_inv = 1.0 / (M(0,0)*M(1,1) - M(1,0)*M(0,1));
              Minv(0,0) = M(1,1) * det_inv;
              Minv(0,1) =-M(0,1) * det_inv;
              Minv(1,0) =-M(1,0) * det_inv;
              Minv(1,1) = M(0,0) * det_inv;
              return Minv;
            };
            auto du = inv(x_u.RotateRight()*x_u) * x_u.RotateRight()*(x0-x);
            u0 = u0 + du;

            auto x_u_squared = x_u.RotateRight() * x_u;
            adapt = sqrt<Real>( ((x0-x).RotateRight()*(x0-x))(0,0) / std::max<Real>(x_u_squared(0,0),x_u_squared(1,1)) );
          }
        }

        Matrix<Real> quad_nds;
        Vector<Real> quad_wts;
        DuffyQuad<ElemDim>(quad_nds, quad_wts, Vector<Real>(ElemDim,u0.begin(),false), order_singular, adapt);
        const CoordEvalOpType CoordEvalOp = CoordBasis::SetupEval(quad_nds);
        Integer Nnds = quad_wts.Dim();

        Vector<Real> X_, dX_, Xa_, Xn_;
        { // Set X_, dX_
          const Vector<CoordBasis> X__(CoordDim, (Iterator<CoordBasis>)X.begin() + src_idx * CoordDim, false);
          const Vector<CoordBasis> dX__(CoordDim * ElemDim, (Iterator<CoordBasis>)dX.begin() + src_idx * CoordDim * ElemDim, false);
          eval_basis(X_, X__, CoordDim, Nnds, CoordEvalOp);
          eval_basis(dX_, dX__, CoordDim * ElemDim, Nnds, CoordEvalOp);
        }
        if (CoordDim == 3 && ElemDim == 2) { // Compute Xa_, Xn_
          Xa_.ReInit(Nnds);
          Xn_.ReInit(Nnds*CoordDim);
          for (Long j = 0; j < Nnds; j++) {
            StaticArray<Real,CoordDim> normal;
            normal[0] = dX_[j*6+2]*dX_[j*6+5] - dX_[j*6+4]*dX_[j*6+3];
            normal[1] = dX_[j*6+4]*dX_[j*6+1] - dX_[j*6+0]*dX_[j*6+5];
            normal[2] = dX_[j*6+0]*dX_[j*6+3] - dX_[j*6+2]*dX_[j*6+1];
            Xa_[j] = sqrt<Real>(normal[0]*normal[0]+normal[1]*normal[1]+normal[2]*normal[2]);
            Real invXa = 1/Xa_[j];
            Xn_[j*3+0] = normal[0] * invXa;
            Xn_[j*3+1] = normal[1] * invXa;
            Xn_[j*3+2] = normal[2] * invXa;
          }
        }

        DensityEvalOpType DensityEvalOp;
        if (std::is_same<CoordBasis,DensityBasis>::value) {
          DensityEvalOp = CoordEvalOp;
        } else {
          DensityEvalOp = DensityBasis::SetupEval(quad_nds);
        }

        Matrix<Real> M__(Nnds * KDIM0, KDIM1);
        { // Set kernel matrix M__
          const Vector<Real> X0_(CoordDim, (Iterator<Real>)Xt.begin() + j * CoordDim, false);
          kernel.template KernelMatrix<Real>(M__, X0_, X_, Xn_);
        }
        for (Long k0 = 0; k0 < KDIM0; k0++) {
          for (Long k1 = 0; k1 < KDIM1; k1++) {
            for (Long l = 0; l < DensityBasis::Size(); l++) {
              Real M_lk = 0;
              for (Long n = 0; n < Nnds; n++) {
                Real quad_wt = Xa_[n] * quad_wts[n];
                M_lk += DensityEvalOp[l][n] * quad_wt * M__[n*KDIM0+k0][k1];
              }
              M[(j * KDIM0 + k0) * DensityBasis::Size() + l][k1] = M_lk;
            }
          }
        }
      }
      { // Set M (subtract direct)
        Matrix<Real> quad_nds;
        Vector<Real> quad_wts;
        TensorProductGaussQuad<ElemDim>(quad_nds, quad_wts, order_direct);
        const CoordEvalOpType CoordEvalOp = CoordBasis::SetupEval(quad_nds);
        Integer Nnds = quad_wts.Dim();

        Vector<Real> X_, dX_, Xa_, Xn_;
        { // Set X_, dX_
          eval_basis(X_, X, CoordDim, Nnds, CoordEvalOp);
          eval_basis(dX_, dX, CoordDim * ElemDim, Nnds, CoordEvalOp);
        }
        if (CoordDim == 3 && ElemDim == 2) { // Compute Xa_, Xn_
          Long N = Nelem*Nnds;
          Xa_.ReInit(N);
          Xn_.ReInit(N*CoordDim);
          for (Long j = 0; j < N; j++) {
            StaticArray<Real,CoordDim> normal;
            normal[0] = dX_[j*6+2]*dX_[j*6+5] - dX_[j*6+4]*dX_[j*6+3];
            normal[1] = dX_[j*6+4]*dX_[j*6+1] - dX_[j*6+0]*dX_[j*6+5];
            normal[2] = dX_[j*6+0]*dX_[j*6+3] - dX_[j*6+2]*dX_[j*6+1];
            Xa_[j] = sqrt<Real>(normal[0]*normal[0]+normal[1]*normal[1]+normal[2]*normal[2]);
            Real invXa = 1/Xa_[j];
            Xn_[j*3+0] = normal[0] * invXa;
            Xn_[j*3+1] = normal[1] * invXa;
            Xn_[j*3+2] = normal[2] * invXa;
          }
        }

        DensityEvalOpType DensityEvalOp;
        if (std::is_same<CoordBasis,DensityBasis>::value) {
          DensityEvalOp = CoordEvalOp;
        } else {
          DensityEvalOp = DensityBasis::SetupEval(quad_nds);
        }

        #pragma omp parallel for schedule(static)
        for (Long j = 0; j < Ninterac; j++) { // Subtract direct contribution
          const Long src_idx = pair_lst[j].first - elem_rank_offset;

          Matrix<Real> M__(Nnds * KDIM0, KDIM1);
          { // Set kernel matrix M__
            const Vector<Real> X0_(CoordDim, (Iterator<Real>)Xt.begin() + j * CoordDim, false);
            Vector<Real> X__(Nnds * CoordDim, X_.begin() + src_idx * Nnds * CoordDim, false);
            Vector<Real> Xn__(Nnds * CoordDim, Xn_.begin() + src_idx * Nnds * CoordDim, false);
            kernel.template KernelMatrix<Real>(M__, X0_, X__, Xn__);
          }
          for (Long k0 = 0; k0 < KDIM0; k0++) {
            for (Long k1 = 0; k1 < KDIM1; k1++) {
              for (Long l = 0; l < DensityBasis::Size(); l++) {
                Real M_lk = 0;
                for (Long n = 0; n < Nnds; n++) {
                  Real quad_wt = Xa_[src_idx * Nnds + n] * quad_wts[n];
                  M_lk += DensityEvalOp[l][n] * quad_wt * M__[n*KDIM0+k0][k1];
                }
                M[(j * KDIM0 + k0) * DensityBasis::Size() + l][k1] -= M_lk;
              }
            }
          }
        }
      }
    }

    template <class DensityBasis> static void EvalNearSingular(Vector<Real>& U, const Vector<DensityBasis>& density, const Matrix<Real>& M, const Vector<Pair<Long,Long>>& pair_lst, Long Nelem_, Long Ntrg_, Integer KDIM0_, Integer KDIM1_, const Comm& comm) {
      const Long Ninterac = pair_lst.Dim();
      const Integer dof = density.Dim() / Nelem_ / KDIM0_;
      SCTL_ASSERT(density.Dim() == Nelem_ * dof * KDIM0_);

      Long elem_rank_offset;
      { // Set elem_rank_offset
        comm.Scan(Ptr2ConstItr<Long>(&Nelem_,1), Ptr2Itr<Long>(&elem_rank_offset,1), 1, Comm::CommOp::SUM);
        elem_rank_offset -= Nelem_;
      }

      Vector<Real> U_loc(Ninterac*dof*KDIM1_);
      for (Long j = 0; j < Ninterac; j++) {
        const Long src_idx = pair_lst[j].first - elem_rank_offset;
        const Matrix<Real> M_(KDIM0_ * DensityBasis::Size(), KDIM1_, (Iterator<Real>)M[j * KDIM0_ * DensityBasis::Size()], false);
        Matrix<Real> U_(dof, KDIM1_, U_loc.begin() + j*dof*KDIM1_, false);
        Matrix<Real> F_(dof, KDIM0_ * DensityBasis::Size());
        for (Long i = 0; i < dof; i++) {
          for (Long k = 0; k < KDIM0_; k++) {
            for (Long l = 0; l < DensityBasis::Size(); l++) {
              F_[i][k * DensityBasis::Size() + l] = density[(src_idx * dof + i) * KDIM0_ + k][l];
            }
          }
        }
        Matrix<Real>::GEMM(U_, F_, M_);
      }

      if (U.Dim() != Ntrg_ * dof * KDIM1_) {
        U.ReInit(Ntrg_ * dof * KDIM1_);
        U = 0;
      }
      { // Set U
        Integer rank = comm.Rank();
        Integer np = comm.Size();

        Vector<Long> splitter_ranks;
        { // Set splitter_ranks
          Vector<Long> cnt(np);
          comm.Allgather(Ptr2ConstItr<Long>(&Ntrg_,1), 1, cnt.begin(), 1);
          scan(splitter_ranks, cnt);
        }

        Vector<Long> scatter_index, send_index, send_cnt(np), send_dsp(np);
        { // Set scatter_index, send_index, send_cnt, send_dsp
          { // Set scatter_index, send_index
            Vector<Pair<Long,Long>> scatter_pair(pair_lst.Dim());
            for (Long i = 0; i < pair_lst.Dim(); i++) {
              scatter_pair[i] = Pair<Long,Long>(pair_lst[i].second,i);
            }
            omp_par::merge_sort(scatter_pair.begin(), scatter_pair.end());

            send_index.ReInit(scatter_pair.Dim());
            scatter_index.ReInit(scatter_pair.Dim());
            for (Long i = 0; i < scatter_index.Dim(); i++) {
              send_index[i] = scatter_pair[i].first;
              scatter_index[i] = scatter_pair[i].second;
            }
          }
          for (Integer i = 0; i < np; i++) {
            send_dsp[i] = std::lower_bound(send_index.begin(), send_index.end(), splitter_ranks[i]) - send_index.begin();
          }
          for (Integer i = 0; i < np-1; i++) {
            send_cnt[i] = send_dsp[i+1] - send_dsp[i];
          }
          send_cnt[np-1] = send_index.Dim() - send_dsp[np-1];
        }

        Vector<Long> recv_index, recv_cnt(np), recv_dsp(np);
        { // Set recv_index, recv_cnt, recv_dsp
          comm.Alltoall(send_cnt.begin(), 1, recv_cnt.begin(), 1);
          scan(recv_dsp, recv_cnt);
          recv_index.ReInit(recv_cnt[np-1] + recv_dsp[np-1]);
          comm.Alltoallv(send_index.begin(), send_cnt.begin(), send_dsp.begin(), recv_index.begin(), recv_cnt.begin(), recv_dsp.begin());
        }

        Vector<Real> U_send(scatter_index.Dim() * dof * KDIM1_);
        for (Long i = 0; i < scatter_index.Dim(); i++) {
          Long idx = scatter_index[i]*dof*KDIM1_;
          for (Long k = 0; k < dof * KDIM1_; k++) {
            U_send[i*dof*KDIM1_ + k] = U_loc[idx + k];
          }
        }

        Vector<Real> U_recv(recv_index.Dim() * dof * KDIM1_);
        { // Set U_recv
          for (Long i = 0; i < np; i++) {
            send_cnt[i] *= dof * KDIM1_;
            send_dsp[i] *= dof * KDIM1_;
            recv_cnt[i] *= dof * KDIM1_;
            recv_dsp[i] *= dof * KDIM1_;
          }
          comm.Alltoallv(U_send.begin(), send_cnt.begin(), send_dsp.begin(), U_recv.begin(), recv_cnt.begin(), recv_dsp.begin());
        }

        for (Long i = 0; i < recv_index.Dim(); i++) { // Set U
          Long idx = (recv_index[i] - splitter_ranks[rank]) * dof * KDIM1_;
          for (Integer k = 0; k < dof * KDIM1_; k++) {
            U[idx + k] += U_recv[i*dof*KDIM1_ + k];
          }
        }
      }
    }



    template <class ElemList, class DensityBasis, class Kernel> static void Direct(Vector<Real>& U, const Vector<Real>& Xt, const ElemList& elem_lst, const Vector<DensityBasis>& density, const Kernel& kernel, Integer order_direct, const Comm& comm) {
      using CoordBasis = typename ElemList::CoordBasis;
      using CoordEvalOpType = typename CoordBasis::EvalOpType;
      using DensityEvalOpType = typename DensityBasis::EvalOpType;

      constexpr Integer CoordDim = ElemList::CoordDim();
      constexpr Integer ElemDim = ElemList::ElemDim();
      constexpr Integer KDIM0 = Kernel::SrcDim();
      constexpr Integer KDIM1 = Kernel::TrgDim();
      const Long Nelem = elem_lst.NElem();
      const Integer dof = density.Dim() / Nelem / KDIM0;
      SCTL_ASSERT(density.Dim() == Nelem * dof * KDIM0);

      Matrix<Real> quad_nds;
      Vector<Real> quad_wts;
      TensorProductGaussQuad<ElemDim>(quad_nds, quad_wts, order_direct);
      const CoordEvalOpType CoordEvalOp = CoordBasis::SetupEval(quad_nds);
      Integer Nnds = quad_wts.Dim();

      const Vector<CoordBasis>& X = elem_lst.ElemVector();
      Vector<CoordBasis> dX;
      CoordBasis::Grad(dX, X);

      Vector<Real> X_, dX_, Xa_, Xn_;
      eval_basis(X_, X, CoordDim, Nnds, CoordEvalOp);
      eval_basis(dX_, dX, CoordDim*ElemDim, Nnds, CoordEvalOp);
      if (CoordDim == 3 && ElemDim == 2) { // Compute Xa_, Xn_
        Long N = Nelem*Nnds;
        Xa_.ReInit(N);
        Xn_.ReInit(N*CoordDim);
        for (Long j = 0; j < N; j++) {
          StaticArray<Real,CoordDim> normal;
          normal[0] = dX_[j*6+2]*dX_[j*6+5] - dX_[j*6+4]*dX_[j*6+3];
          normal[1] = dX_[j*6+4]*dX_[j*6+1] - dX_[j*6+0]*dX_[j*6+5];
          normal[2] = dX_[j*6+0]*dX_[j*6+3] - dX_[j*6+2]*dX_[j*6+1];
          Xa_[j] = sqrt<Real>(normal[0]*normal[0]+normal[1]*normal[1]+normal[2]*normal[2]);
          Real invXa = 1/Xa_[j];
          Xn_[j*3+0] = normal[0] * invXa;
          Xn_[j*3+1] = normal[1] * invXa;
          Xn_[j*3+2] = normal[2] * invXa;
        }
      }

      Vector<Real> Fa_;
      { // Set Fa_
        Vector<Real> F_;
        if (std::is_same<CoordBasis,DensityBasis>::value) {
          eval_basis(F_, density, dof * KDIM0, Nnds, CoordEvalOp);
        } else {
          const DensityEvalOpType EvalOp = DensityBasis::SetupEval(quad_nds);
          eval_basis(F_, density, dof * KDIM0, Nnds, EvalOp);
        }

        Fa_.ReInit(F_.Dim());
        const Integer DensityDOF = dof * KDIM0;
        SCTL_ASSERT(F_.Dim() == Nelem * Nnds * DensityDOF);
        for (Long j = 0; j < Nelem; j++) {
          for (Integer k = 0; k < Nnds; k++) {
            Long idx = j * Nnds + k;
            Real quad_wt = Xa_[idx] * quad_wts[k];
            for (Integer l = 0; l < DensityDOF; l++) {
              Fa_[idx * DensityDOF + l] = F_[idx * DensityDOF + l] * quad_wt;
            }
          }
        }
      }

      { // Evaluate potential
        const Long Ntrg = Xt.Dim() / CoordDim;
        SCTL_ASSERT(Xt.Dim() == Ntrg * CoordDim);
        if (U.Dim() != Ntrg * dof * KDIM1) {
          U.ReInit(Ntrg * dof * KDIM1);
          U = 0;
        }
        ParticleFMM<Real,CoordDim>::Eval(U, Xt, X_, Xn_, Fa_, kernel, comm);
      }
    }

  public:

    template <class DensityBasis, class ElemList, class Kernel> void Setup(const ElemList& elem_lst, const Vector<Real>& Xt, const Kernel& kernel, Integer order_singular, Integer order_direct, Real period_length, const Comm& comm) {
      Xt_.ReInit(0);
      M_singular.ReInit(0,0);
      M_near_singular.ReInit(0,0);
      pair_lst.ReInit(0);

      order_direct_ = order_direct;
      period_length_ = period_length;
      comm_ = comm;

      Profile::Tic("Setup", &comm_);
      static_assert(std::is_same<Real,typename DensityBasis::ValueType>::value);
      static_assert(std::is_same<Real,typename ElemList::CoordType>::value);
      static_assert(DensityBasis::Dim() == ElemList::ElemDim());

      Xt_ = Xt;
      M_singular.ReInit(0,0);

      Profile::Tic("SetupNearSingular", &comm_);
      SetupNearSingular<DensityBasis>(M_near_singular, pair_lst, Xt_, Vector<Long>(), elem_lst, kernel, order_singular, order_direct_, period_length_, comm_);
      Profile::Toc();

      Profile::Toc();
    }

    template <class DensityBasis, class PotentialBasis, class ElemList, class Kernel> void Setup(const ElemList& elem_lst, const Kernel& kernel, Integer order_singular, Integer order_direct, Real period_length, const Comm& comm, Real Rqbx = 0) {
      Xt_.ReInit(0);
      M_singular.ReInit(0,0);
      M_near_singular.ReInit(0,0);
      pair_lst.ReInit(0);

      order_direct_ = order_direct;
      period_length_ = period_length;
      comm_ = comm;

      Profile::Tic("Setup", &comm_);
      static_assert(std::is_same<Real,typename PotentialBasis::ValueType>::value);
      static_assert(std::is_same<Real,typename DensityBasis::ValueType>::value);
      static_assert(std::is_same<Real,typename ElemList::CoordType>::value);
      static_assert(PotentialBasis::Dim() == ElemList::ElemDim());
      static_assert(DensityBasis::Dim() == ElemList::ElemDim());

      Vector<Long> trg_surf;
      { // Set Xt_
        using CoordBasis = typename ElemList::CoordBasis;
        Matrix<Real> trg_nds = PotentialBasis::Nodes();
        auto Meval = CoordBasis::SetupEval(trg_nds);
        eval_basis(Xt_, elem_lst.ElemVector(), ElemList::CoordDim(), trg_nds.Dim(1), Meval);

        { // Set trg_surf
          const Long Nelem = elem_lst.NElem();
          const Long Nnds  = trg_nds.Dim(1);
          Long elem_offset;
          { // Set elem_offset
            comm.Scan(Ptr2ConstItr<Long>(&Nelem,1), Ptr2Itr<Long>(&elem_offset,1), 1, Comm::CommOp::SUM);
            elem_offset -= Nelem;
          }
          trg_surf.ReInit(elem_lst.NElem() * trg_nds.Dim(1));
          for (Long i = 0; i < Nelem; i++) {
            for (Long j = 0; j < Nnds; j++) {
              trg_surf[i*Nnds+j] = elem_offset + i;
            }
          }
        }
      }

      Profile::Tic("SetupSingular", &comm_);
      SetupSingular<DensityBasis>(M_singular, PotentialBasis::Nodes(), elem_lst, kernel, order_singular, order_direct_, Rqbx);
      Profile::Toc();

      Profile::Tic("SetupNearSingular", &comm_);
      SetupNearSingular<DensityBasis>(M_near_singular, pair_lst, Xt_, trg_surf, elem_lst, kernel, order_singular, order_direct_, period_length_, comm_);
      Profile::Toc();

      Profile::Toc();
    }

    template <class DensityBasis, class PotentialBasis, class ElemList, class Kernel> void Eval(Vector<PotentialBasis>& U, const ElemList& elements, const Vector<DensityBasis>& F, const Kernel& kernel) const {
      Profile::Tic("Eval", &comm_);
      Matrix<Real> U_singular;
      Vector<Real> U_direct, U_near_sing;

      Profile::Tic("EvalDirect", &comm_);
      Direct(U_direct, Xt_, elements, F, kernel, order_direct_, comm_);
      Profile::Toc();

      Profile::Tic("EvalSingular", &comm_);
      EvalSingular(U_singular, F, M_singular, kernel.SrcDim(), kernel.TrgDim());
      Profile::Toc();

      Profile::Tic("EvalNearSingular", &comm_);
      EvalNearSingular(U_near_sing, F, M_near_singular, pair_lst, elements.NElem(), Xt_.Dim() / ElemList::CoordDim(), kernel.SrcDim(), kernel.TrgDim(), comm_);
      SCTL_ASSERT(U_near_sing.Dim() == U_direct.Dim());
      Profile::Toc();

      const Long dof = U_direct.Dim() / (elements.NElem() * PotentialBasis::Size() * kernel.TrgDim());
      SCTL_ASSERT(U_direct   .Dim() == elements.NElem() * PotentialBasis::Size() * dof * kernel.TrgDim());
      SCTL_ASSERT(U_near_sing.Dim() == elements.NElem() * PotentialBasis::Size() * dof * kernel.TrgDim());

      if (U.Dim() != elements.NElem() * dof * kernel.TrgDim()) {
        U.ReInit(elements.NElem() * dof * kernel.TrgDim());
      }
      for (int i = 0; i < elements.NElem(); i++) {
        for (int j = 0; j < PotentialBasis::Size(); j++) {
          for (int k = 0; k < dof*kernel.TrgDim(); k++) {
            Real& U_ = U[i*dof*kernel.TrgDim()+k][j];
            U_ = 0;
            U_ += U_direct   [(i*PotentialBasis::Size()+j)*dof*kernel.TrgDim()+k];
            U_ += U_near_sing[(i*PotentialBasis::Size()+j)*dof*kernel.TrgDim()+k];
            U_ *= kernel.template ScaleFactor<Real>();
          }
        }
      }
      if (U_singular.Dim(1)) {
        SCTL_ASSERT(U_singular.Dim(0) == elements.NElem() * dof * kernel.TrgDim());
        SCTL_ASSERT(U_singular.Dim(1) == PotentialBasis::Size());
        for (int i = 0; i < elements.NElem(); i++) {
          for (int j = 0; j < PotentialBasis::Size(); j++) {
            for (int k = 0; k < dof*kernel.TrgDim(); k++) {
              U[i*dof*kernel.TrgDim()+k][j] += U_singular[i*dof*kernel.TrgDim()+k][j] * kernel.template ScaleFactor<Real>();
            }
          }
        }
      }
      Profile::Toc();
    }

    template <class DensityBasis, class ElemList, class Kernel> void Eval(Vector<Real>& U, const ElemList& elements, const Vector<DensityBasis>& F, const Kernel& kernel) const {
      Profile::Tic("Eval", &comm_);
      Matrix<Real> U_singular;
      Vector<Real> U_direct, U_near_sing;

      Profile::Tic("EvalDirect", &comm_);
      Direct(U_direct, Xt_, elements, F, kernel, order_direct_, comm_);
      Profile::Toc();

      Profile::Tic("EvalSingular", &comm_);
      EvalSingular(U_singular, F, M_singular, kernel.SrcDim(), kernel.TrgDim());
      Profile::Toc();

      Profile::Tic("EvalNearSingular", &comm_);
      EvalNearSingular(U_near_sing, F, M_near_singular, pair_lst, elements.NElem(), Xt_.Dim() / ElemList::CoordDim(), kernel.SrcDim(), kernel.TrgDim(), comm_);
      SCTL_ASSERT(U_near_sing.Dim() == U_direct.Dim());
      Profile::Toc();


      Long Nt = Xt_.Dim() / ElemList::CoordDim();
      const Long dof = U_direct.Dim() / (Nt * kernel.TrgDim());
      SCTL_ASSERT(U_direct.Dim() == Nt * dof * kernel.TrgDim());

      if (U.Dim() != U_direct.Dim()) {
        U.ReInit(U_direct.Dim());
      }
      for (int i = 0; i < U.Dim(); i++) {
        U[i] = (U_direct[i] + U_near_sing[i]) * kernel.template ScaleFactor<Real>();
      }
      if (U_singular.Dim(1)) {
        SCTL_ASSERT(U_singular.Dim(0) == elements.NElem() * dof * kernel.TrgDim());
        const Long Nnodes = U_singular.Dim(1);
        for (int i = 0; i < elements.NElem(); i++) {
          for (int j = 0; j < Nnodes; j++) {
            for (int k = 0; k < dof*kernel.TrgDim(); k++) {
              Real& U_ = U[(i*Nnodes+j)*dof*kernel.TrgDim()+k];
              U_ += U_singular[i*dof*kernel.TrgDim()+k][j] * kernel.template ScaleFactor<Real>();
            }
          }
        }
      }
      Profile::Toc();
    }

    template <Integer ORDER = 5> static void test(Integer order_singular = 10, Integer order_direct = 5, const Comm& comm = Comm::World()) {
      constexpr Integer COORD_DIM = 3;
      constexpr Integer ELEM_DIM = COORD_DIM-1;
      using ElemList = ElemList<COORD_DIM, Basis<Real, ELEM_DIM, ORDER>>;
      using DensityBasis = Basis<Real, ELEM_DIM, ORDER>;
      using PotentialBasis = Basis<Real, ELEM_DIM, ORDER>;

      int np = comm.Size();
      int rank = comm.Rank();
      auto build_torus = [rank,np](ElemList& elements, long Nt, long Np, Real Rmajor, Real Rminor){
        auto nodes = ElemList::CoordBasis::Nodes();
        auto torus = [](Real theta, Real phi, Real Rmajor, Real Rminor) {
          Real R = Rmajor + Rminor * cos<Real>(phi);
          Real X = R * cos<Real>(theta);
          Real Y = R * sin<Real>(theta);
          Real Z = Rminor * sin<Real>(phi);
          return std::make_tuple(X,Y,Z);
        };

        long start = Nt*Np*(rank+0)/np;
        long end   = Nt*Np*(rank+1)/np;
        elements.ReInit(end - start);
        for (long ii = start; ii < end; ii++) {
          long i = ii / Np;
          long j = ii % Np;
          for (int k = 0; k < ElemList::CoordBasis::Size(); k++) {
            Real X, Y, Z;
            Real theta = 2 * const_pi<Real>() * (i + nodes[0][k]) / Nt;
            Real phi   = 2 * const_pi<Real>() * (j + nodes[1][k]) / Np;
            std::tie(X,Y,Z) = torus(theta, phi, Rmajor, Rminor);
            elements(ii-start,0)[k] = X;
            elements(ii-start,1)[k] = Y;
            elements(ii-start,2)[k] = Z;
          }
        }
      };
      ElemList elements_src, elements_trg;
      build_torus(elements_src, 28, 16, 2, 1.0);
      build_torus(elements_trg, 29, 17, 2, 0.99);

      Vector<Real> Xt;
      Vector<PotentialBasis> U_onsurf, U_offsurf;
      Vector<DensityBasis> density_sl, density_dl;
      { // Set Xt, elements_src, elements_trg, density_sl, density_dl, U
        Real X0[COORD_DIM] = {3,2,1};
        std::function<void(Real*,Real*,Real*)> potential = [X0](Real* U, Real* X, Real* Xn) {
          Real dX[COORD_DIM] = {X[0]-X0[0],X[1]-X0[1],X[2]-X0[2]};
          Real Rinv = 1/sqrt(dX[0]*dX[0]+dX[1]*dX[1]+dX[2]*dX[2]);
          U[0] = Rinv;
        };
        std::function<void(Real*,Real*,Real*)> potential_normal_derivative = [X0](Real* U, Real* X, Real* Xn) {
          Real dX[COORD_DIM] = {X[0]-X0[0],X[1]-X0[1],X[2]-X0[2]};
          Real Rinv = 1/sqrt(dX[0]*dX[0]+dX[1]*dX[1]+dX[2]*dX[2]);
          Real RdotN = dX[0]*Xn[0]+dX[1]*Xn[1]+dX[2]*Xn[2];
          U[0] = -RdotN * Rinv*Rinv*Rinv;
        };

        DiscretizeSurfaceFn<COORD_DIM,1>(density_sl, elements_src, potential_normal_derivative);
        DiscretizeSurfaceFn<COORD_DIM,1>(density_dl, elements_src, potential);
        DiscretizeSurfaceFn<COORD_DIM,1>(U_onsurf  , elements_src, potential);
        DiscretizeSurfaceFn<COORD_DIM,1>(U_offsurf , elements_trg, potential);

        for (long i = 0; i < elements_trg.NElem(); i++) { // Set Xt
          for (long j = 0; j < PotentialBasis::Size(); j++) {
            for (int k = 0; k < COORD_DIM; k++) {
              Xt.PushBack(elements_trg(i,k)[j]);
            }
          }
        }
      }

      GenericKernel<Laplace3D_DxU> Laplace_DxU;
      GenericKernel<Laplace3D_FxU> Laplace_FxU;

      Profile::Enable(true);
      if (1) { // Greeen's identity test (Laplace, on-surface)
        Profile::Tic("OnSurface", &comm);
        Quadrature<Real> quadrature_DxU, quadrature_FxU;
        quadrature_FxU.Setup<DensityBasis, PotentialBasis>(elements_src, Laplace_FxU, order_singular, order_direct, -1.0, comm);
        quadrature_DxU.Setup<DensityBasis, PotentialBasis>(elements_src, Laplace_DxU, order_singular, order_direct, -1.0, comm);

        Vector<PotentialBasis> U_sl, U_dl;
        quadrature_FxU.Eval(U_sl, elements_src, density_sl, Laplace_FxU);
        quadrature_DxU.Eval(U_dl, elements_src, density_dl, Laplace_DxU);
        Profile::Toc();

        Real max_err = 0;
        Vector<PotentialBasis> err(U_onsurf.Dim());
        for (long i = 0; i < U_sl.Dim(); i++) {
          for (long j = 0; j < PotentialBasis::Size(); j++) {
            err[i][j] = 0.5*U_onsurf[i][j] - (U_sl[i][j] + U_dl[i][j]);
            max_err = std::max<Real>(max_err, fabs(err[i][j]));
          }
        }
        { // Print error
          Real glb_err;
          comm.Allreduce(Ptr2ConstItr<Real>(&max_err,1), Ptr2Itr<Real>(&glb_err,1), 1, Comm::CommOp::MAX);
          if (!comm.Rank()) std::cout<<"Error = "<<glb_err<<'\n';
        }
        { // Write VTK output
          VTUData vtu;
          vtu.AddElems(elements_src, err, ORDER);
          vtu.WriteVTK("err", comm);
        }
        { // Write VTK output
          VTUData vtu;
          vtu.AddElems(elements_src, U_onsurf, ORDER);
          vtu.WriteVTK("U", comm);
        }
      }
      if (1) { // Greeen's identity test (Laplace, off-surface)
        Profile::Tic("OffSurface", &comm);
        Quadrature<Real> quadrature_DxU, quadrature_FxU;
        quadrature_FxU.Setup<DensityBasis>(elements_src, Xt, Laplace_FxU, order_singular, order_direct, -1.0, comm);
        quadrature_DxU.Setup<DensityBasis>(elements_src, Xt, Laplace_DxU, order_singular, order_direct, -1.0, comm);

        Vector<Real> U_sl, U_dl;
        quadrature_FxU.Eval(U_sl, elements_src, density_sl, Laplace_FxU);
        quadrature_DxU.Eval(U_dl, elements_src, density_dl, Laplace_DxU);
        Profile::Toc();

        Real max_err = 0;
        Vector<PotentialBasis> err(elements_trg.NElem());
        for (long i = 0; i < elements_trg.NElem(); i++) {
          for (long j = 0; j < PotentialBasis::Size(); j++) {
            err[i][j] = U_offsurf[i][j] - (U_sl[i*PotentialBasis::Size()+j] + U_dl[i*PotentialBasis::Size()+j]);
            max_err = std::max<Real>(max_err, fabs(err[i][j]));
          }
        }
        { // Print error
          Real glb_err;
          comm.Allreduce(Ptr2ConstItr<Real>(&max_err,1), Ptr2Itr<Real>(&glb_err,1), 1, Comm::CommOp::MAX);
          if (!comm.Rank()) std::cout<<"Error = "<<glb_err<<'\n';
        }
        { // Write VTK output
          VTUData vtu;
          vtu.AddElems(elements_trg, err, ORDER);
          vtu.WriteVTK("err", comm);
        }
        { // Write VTK output
          VTUData vtu;
          vtu.AddElems(elements_trg, U_offsurf, ORDER);
          vtu.WriteVTK("U", comm);
        }
      }
      Profile::print(&comm);
    }

    static void test1() {
      const Comm& comm = Comm::World();
      constexpr Integer ORDER = 15;
      Integer order_singular = 20;
      Integer order_direct = 20;

      constexpr Integer COORD_DIM = 3;
      constexpr Integer ELEM_DIM = COORD_DIM-1;
      using ElemList = ElemList<COORD_DIM, Basis<Real, ELEM_DIM, ORDER>>;
      using DensityBasis = Basis<Real, ELEM_DIM, ORDER>;
      using PotentialBasis = Basis<Real, ELEM_DIM, ORDER>;

      int np = comm.Size();
      int rank = comm.Rank();
      auto build_sphere = [rank,np](ElemList& elements, Real X, Real Y, Real Z, Real R){
        auto nodes = ElemList::CoordBasis::Nodes();

        long start = 2*COORD_DIM*(rank+0)/np;
        long end   = 2*COORD_DIM*(rank+1)/np;
        elements.ReInit(end - start);
        for (long ii = start; ii < end; ii++) {
          long i = ii / 2;
          long j = ii % 2;
          for (int k = 0; k < ElemList::CoordBasis::Size(); k++) {
            Real coord[COORD_DIM];
            coord[(i+0)%COORD_DIM] = (j ? -1.0 : 1.0);
            coord[(i+1)%COORD_DIM] = 2.0 * nodes[j?1:0][k] - 1.0;
            coord[(i+2)%COORD_DIM] = 2.0 * nodes[j?0:1][k] - 1.0;
            Real R0 = sqrt<Real>(coord[0]*coord[0] + coord[1]*coord[1] + coord[2]*coord[2]);

            elements(ii-start,0)[k] = X + R * coord[0] / R0;
            elements(ii-start,1)[k] = Y + R * coord[1] / R0;
            elements(ii-start,2)[k] = Z + R * coord[2] / R0;
          }
        }
      };

      ElemList elements;
      build_sphere(elements, 0.0, 0.0, 0.0, 1.00);

      Vector<DensityBasis> density_sl;
      { // Set density_sl
        std::function<void(Real*,Real*,Real*)> sigma = [](Real* U, Real* X, Real* Xn) {
          Real R = sqrt(X[0]*X[0]+X[1]*X[1]+X[2]*X[2]);
          Real sinp = sqrt(X[1]*X[1] + X[2]*X[2]) / R;
          Real cosp = -X[0] / R;

          U[0] = -1.5;
          U[1] = 0;
          U[2] = 0;
        };
        DiscretizeSurfaceFn<COORD_DIM,3>(density_sl, elements, sigma);
      }

      GenericKernel<Stokes3D_DxU> Stokes_DxU;
      GenericKernel<Stokes3D_FxU> Stokes_FxU;

      Profile::Enable(true);
      if (1) {
        Vector<PotentialBasis> U;
        Quadrature<Real> quadrature_FxU;
        quadrature_FxU.Setup<DensityBasis, PotentialBasis>(elements, Stokes_FxU, order_singular, order_direct, -1.0, comm);
        quadrature_FxU.Eval(U, elements, density_sl, Stokes_FxU);

        { // Write VTK output
          VTUData vtu;
          vtu.AddElems(elements, U, ORDER);
          vtu.WriteVTK("U", comm);
        }
        { // Write VTK output
          VTUData vtu;
          vtu.AddElems(elements, density_sl, ORDER);
          vtu.WriteVTK("sigma", comm);
        }
      }
      Profile::print(&comm);
    }

  private:

    static void scan(Vector<Long>& dsp, const Vector<Long>& cnt) {
      dsp.ReInit(cnt.Dim());
      if (cnt.Dim()) dsp[0] = 0;
      omp_par::scan(cnt.begin(), dsp.begin(), cnt.Dim());
    }

    template <class Basis> static void eval_basis(Vector<Real>& value, const Vector<Basis> X, Integer dof, Integer Nnds, const typename Basis::EvalOpType& EvalOp) {
      Long Nelem = X.Dim() / dof;
      SCTL_ASSERT(X.Dim() == Nelem * dof);

      value.ReInit(Nelem*Nnds*dof);
      Matrix<Real> X_(Nelem*dof, Nnds, value.begin(),false);
      Basis::Eval(X_, X, EvalOp);
      for (Long j = 0; j < Nelem; j++) { // Rearrange data
        Matrix<Real> X(Nnds, dof, X_[j*dof], false);
        X = Matrix<Real>(dof, Nnds, X_[j*dof], false).Transpose();
      }
    }

    template <int CoordDim, int FnDim, class FnBasis, class ElemList> static void DiscretizeSurfaceFn(Vector<FnBasis>& U, const ElemList& elements, std::function<void(Real*,Real*,Real*)> fn) {
      using CoordBasis = typename ElemList::CoordBasis;
      const long Nelem = elements.NElem();
      U.ReInit(Nelem * FnDim);

      Matrix<Real> X, X_grad;
      { // Set X, X_grad
        Vector<CoordBasis> coord = elements.ElemVector();
        Vector<CoordBasis> coord_grad;
        CoordBasis::Grad(coord_grad, coord);

        const auto Meval = CoordBasis::SetupEval(FnBasis::Nodes());
        CoordBasis::Eval(X, coord, Meval);
        CoordBasis::Eval(X_grad, coord_grad, Meval);
      }

      for (long i = 0; i < Nelem; i++) {
        for (long j = 0; j < FnBasis::Size(); j++) {
          Real X_[CoordDim], Xn[CoordDim], U_[FnDim];
          for (long k = 0; k < CoordDim; k++) {
            X_[k] = X[i*CoordDim+k][j];
          }
          { // Set Xn
            Real Xu[CoordDim], Xv[CoordDim];
            for (long k = 0; k < CoordDim; k++) {
              Xu[k] = X_grad[(i*CoordDim+k)*2+0][j];
              Xv[k] = X_grad[(i*CoordDim+k)*2+1][j];
            }
            Real dA = 0;
            for (long k = 0; k < CoordDim; k++) {
              Xn[k] = Xu[(k+1)%CoordDim] * Xv[(k+2)%CoordDim];
              Xn[k] -= Xv[(k+1)%CoordDim] * Xu[(k+2)%CoordDim];
              dA += Xn[k] * Xn[k];
            }
            dA = sqrt(dA);
            for (long k = 0; k < CoordDim; k++) {
              Xn[k] /= dA;
            }
          }
          fn(U_, X_, Xn);
          for (long k = 0; k < FnDim; k++) {
            U[i*FnDim+k][j] = U_[k];
          }
        }
      }
    }

    Vector<Real> Xt_;
    Matrix<Real> M_singular;
    Matrix<Real> M_near_singular;
    Vector<Pair<Long,Long>> pair_lst;
    Integer order_direct_;

    Real period_length_;
    Comm comm_;
};


template <class Real, Integer ORDER=10> class Stellarator {
  private:
    static constexpr Integer order_singular = 20;
    static constexpr Integer order_direct = 25;
    static constexpr Integer COORD_DIM = 3;
    static constexpr Integer ELEM_DIM = COORD_DIM-1;
    using ElemBasis = Basis<Real, ELEM_DIM, ORDER>;
    using ElemLst = ElemList<COORD_DIM, ElemBasis>;

    struct Laplace3D_dUxF {
      template <class ValueType> static constexpr ValueType ScaleFactor() {
        return 1 / (4 * const_pi<ValueType>());
      }
      template <class ValueType> static void Eval(ValueType (&u)[3][1], const ValueType (&r)[3], const ValueType (&n)[3], void* ctx_ptr) {
        ValueType r2 = r[0]*r[0]+r[1]*r[1]+r[2]*r[2];
        ValueType rinv = (r2>1e-16 ? 1/sqrt<ValueType>(r2) : 0);
        ValueType rinv3 = rinv * rinv * rinv;
        u[0][0] = -r[0] * rinv3;
        u[1][0] = -r[1] * rinv3;
        u[2][0] = -r[2] * rinv3;
      }
    };

    struct BiotSavart3D {
      template <class ValueType> static constexpr ValueType ScaleFactor() {
        return 1 / (4 * const_pi<ValueType>());
      }
      template <class ValueType> static void Eval(ValueType (&u)[3][3], const ValueType (&r)[3], const ValueType (&n)[3], void* ctx_ptr) {
        ValueType r2 = r[0]*r[0]+r[1]*r[1]+r[2]*r[2];
        ValueType rinv = (r2>1e-16 ? 1/sqrt<ValueType>(r2) : 0);
        ValueType rinv3 = rinv * rinv * rinv;
        u[0][0] =   (0) * rinv3; u[1][0] =  r[2] * rinv3; u[2][0] = -r[1] * rinv3;
        u[0][1] = -r[2] * rinv3; u[1][1] =   (0) * rinv3; u[2][1] =  r[0] * rinv3;
        u[0][2] =  r[1] * rinv3; u[1][2] = -r[0] * rinv3; u[2][2] =   (0) * rinv3;
      }
    };

    struct BiotSavartGrad3D {
      template <class ValueType> static constexpr ValueType ScaleFactor() {
        return 1 / (4 * const_pi<ValueType>());
      }
      template <class ValueType> static void Eval(ValueType (&u)[3][9], const ValueType (&r)[3], const ValueType (&n)[3], void* ctx_ptr) {
        ValueType r2 = r[0]*r[0]+r[1]*r[1]+r[2]*r[2];
        ValueType rinv = (r2>1e-16 ? 1/sqrt<ValueType>(r2) : 0);
        ValueType rinv2 = rinv * rinv;
        ValueType rinv3 = rinv2 * rinv;
        ValueType rinv5 = rinv2 * rinv3;

        u[0][0] =                                      0; u[1][0] =              - 3 * r[2] * r[0] * rinv5; u[2][0] =                3 * r[1] * r[0] * rinv5;
        u[0][1] =                                      0; u[1][1] =              - 3 * r[2] * r[1] * rinv5; u[2][1] = -(1) * rinv3 + 3 * r[1] * r[1] * rinv5;
        u[0][2] =                                      0; u[1][2] =  (1) * rinv3 - 3 * r[2] * r[2] * rinv5; u[2][2] =                3 * r[1] * r[2] * rinv5;

        u[0][3] =                3 * r[2] * r[0] * rinv5; u[1][3] =                                      0; u[2][3] =  (1) * rinv3 - 3 * r[0] * r[0] * rinv5;
        u[0][4] =                3 * r[2] * r[1] * rinv5; u[1][4] =                                      0; u[2][4] =              - 3 * r[0] * r[1] * rinv5;
        u[0][5] = -(1) * rinv3 + 3 * r[2] * r[2] * rinv5; u[1][5] =                                      0; u[2][5] =              - 3 * r[0] * r[2] * rinv5;

        u[0][6] =              - 3 * r[1] * r[0] * rinv5; u[1][6] = -(1) * rinv3 + 3 * r[0] * r[0] * rinv5; u[2][6] =                                      0;
        u[0][7] =  (1) * rinv3 - 3 * r[1] * r[1] * rinv5; u[1][7] =                3 * r[0] * r[1] * rinv5; u[2][7] =                                      0;
        u[0][8] =              - 3 * r[1] * r[2] * rinv5; u[1][8] =                3 * r[0] * r[2] * rinv5; u[2][8] =                                      0;
      }
    };

    struct Laplace3D_dUxD {
      template <class ValueType> static constexpr ValueType ScaleFactor() {
        return 1 / (4 * const_pi<ValueType>());
      }
      template <class ValueType> static void Eval(ValueType (&u)[3][3], const ValueType (&r)[3], const ValueType (&n)[3], void* ctx_ptr) {
        ValueType r2 = r[0]*r[0]+r[1]*r[1]+r[2]*r[2];
        ValueType rinv = (r2>1e-16 ? 1/sqrt<ValueType>(r2) : 0);
        ValueType rdotn = r[0]*n[0] + r[1]*n[1] + r[2]*n[2];
        ValueType rinv2 = rinv * rinv;
        ValueType rinv3 = rinv * rinv2;
        ValueType rinv5 = rinv3 * rinv2;

        u[0][0] = -1 * rinv3 + 3 * r[0] * r[0] * rinv5;
        u[0][1] = -0 * rinv3 + 3 * r[0] * r[1] * rinv5;
        u[0][2] = -0 * rinv3 + 3 * r[0] * r[2] * rinv5;

        u[1][0] = -0 * rinv3 + 3 * r[1] * r[0] * rinv5;
        u[1][1] = -1 * rinv3 + 3 * r[1] * r[1] * rinv5;
        u[1][2] = -0 * rinv3 + 3 * r[1] * r[2] * rinv5;

        u[2][0] = -0 * rinv3 + 3 * r[2] * r[0] * rinv5;
        u[2][1] = -0 * rinv3 + 3 * r[2] * r[1] * rinv5;
        u[2][2] = -1 * rinv3 + 3 * r[2] * r[2] * rinv5;
      }
    };

    struct Laplace3D_DxdU {
      template <class ValueType> static constexpr ValueType ScaleFactor() {
        return 1 / (4 * const_pi<ValueType>());
      }
      template <class ValueType> static void Eval(ValueType (&u)[1][3], const ValueType (&r)[3], const ValueType (&n)[3], void* ctx_ptr) {
        ValueType r2 = r[0]*r[0]+r[1]*r[1]+r[2]*r[2];
        ValueType rinv = (r2>1e-16 ? 1/sqrt<ValueType>(r2) : 0);
        ValueType rdotn = r[0]*n[0] + r[1]*n[1] + r[2]*n[2];
        ValueType rinv2 = rinv * rinv;
        ValueType rinv3 = rinv * rinv2;
        ValueType rinv5 = rinv3 * rinv2;
        u[0][0] = -n[0] * rinv3 + 3*rdotn * r[0] * rinv5;
        u[0][1] = -n[1] * rinv3 + 3*rdotn * r[1] * rinv5;
        u[0][2] = -n[2] * rinv3 + 3*rdotn * r[2] * rinv5;
      }
    };

    struct Laplace3D_Fxd2U {
      template <class ValueType> static constexpr ValueType ScaleFactor() {
        return 1 / (4 * const_pi<ValueType>());
      }
      template <class ValueType> static void Eval(ValueType (&u)[1][9], const ValueType (&r)[3], const ValueType (&n)[3], void* ctx_ptr) {
        ValueType r2 = r[0]*r[0]+r[1]*r[1]+r[2]*r[2];
        ValueType rinv = (r2>1e-16 ? 1/sqrt<ValueType>(r2) : 0);
        ValueType rinv2 = rinv * rinv;
        ValueType rinv3 = rinv * rinv2;
        ValueType rinv5 = rinv3 * rinv2;

        u[0][0+3*0] = -1 * rinv3 + 3 * r[0] * r[0] * rinv5;
        u[0][1+3*0] = -0 * rinv3 + 3 * r[0] * r[1] * rinv5;
        u[0][2+3*0] = -0 * rinv3 + 3 * r[0] * r[2] * rinv5;

        u[0][0+3*1] = -0 * rinv3 + 3 * r[1] * r[0] * rinv5;
        u[0][1+3*1] = -1 * rinv3 + 3 * r[1] * r[1] * rinv5;
        u[0][2+3*1] = -0 * rinv3 + 3 * r[1] * r[2] * rinv5;

        u[0][0+3*2] = -0 * rinv3 + 3 * r[2] * r[0] * rinv5;
        u[0][1+3*2] = -0 * rinv3 + 3 * r[2] * r[1] * rinv5;
        u[0][2+3*2] = -1 * rinv3 + 3 * r[2] * r[2] * rinv5;
      }
    };

    static Real max_norm(const Vector<Real>& x) {
      Real err = 0;
      for (const auto& a : x) err = std::max(err, fabs<Real>(a));
      return err;
    }

  public:
    static Vector<ElemBasis> compute_dot_prod(const Vector<ElemBasis>& A, const Vector<ElemBasis>& B) {
      const Long Nelem = A.Dim() / COORD_DIM;
      const Long Nnodes = ElemBasis::Size();
      SCTL_ASSERT(A.Dim() == Nelem * COORD_DIM);
      SCTL_ASSERT(B.Dim() == Nelem * COORD_DIM);
      Vector<ElemBasis> AdotB(Nelem);
      for (Long i = 0; i < Nelem; i++) {
        for (Long j = 0; j < Nnodes; j++) {
          Real a_dot_b = 0;
          a_dot_b += A[i*COORD_DIM+0][j]*B[i*COORD_DIM+0][j];
          a_dot_b += A[i*COORD_DIM+1][j]*B[i*COORD_DIM+1][j];
          a_dot_b += A[i*COORD_DIM+2][j]*B[i*COORD_DIM+2][j];
          AdotB[i][j] = a_dot_b;
        }
      }
      return AdotB;
    }
    static Real compute_inner_prod(const Vector<ElemBasis>& area_elem, const Vector<ElemBasis>& A, const Vector<ElemBasis>& B) {
      const auto& quad_wts = ElemBasis::QuadWts();
      const Long Nnodes = ElemBasis::Size();
      const Long Nelem = area_elem.Dim();
      const Long dof = B.Dim() / Nelem;
      Real sum = 0;
      for (Long i = 0; i < Nelem; i++) {
        for (Long j = 0; j < Nnodes; j++) {
          Real AdotB = 0;
          for (Long k = 0; k < dof; k++) {
            AdotB += A[i*dof+k][j] * B[i*dof+k][j];
          }
          sum += AdotB * area_elem[i][j] * quad_wts[j];
        }
      }
      return sum;
    }
    static void compute_harmonic_vector_potentials(Vector<ElemBasis>& Jt, Vector<ElemBasis>& Jp, const Stellarator<Real,ORDER>& S) {
      Comm comm = Comm::World();
      Real gmres_tol = 1e-10;
      Long max_iter = 100;

      auto cheb2grid = [] (const Vector<ElemBasis>& X, Long Mt, Long Mp, Long Nt, Long Np) {
        const Long dof = X.Dim() / (Mt * Mp);
        SCTL_ASSERT(X.Dim() == Mt * Mp *dof);

        Vector<Real> Xf(dof*Nt*Np); Xf = 0;
        const Long Nnodes = ElemBasis::Size();
        const Matrix<Real>& Mnodes = Basis<Real,1,ORDER>::Nodes();
        for (Long t = 0; t < Nt; t++) {
          for (Long p = 0; p < Np; p++) {
            Real theta = t / (Real)Nt;
            Real phi   = p / (Real)Np;
            Long i = (Long)(theta * Mt);
            Long j = (Long)(phi   * Mp);
            Real x = theta * Mt - i;
            Real y = phi   * Mp - j;
            Long elem_idx = i * Mp + j;

            Vector<Real> Interp0(ORDER);
            Vector<Real> Interp1(ORDER);
            { // Set Interp0, Interp1
              auto node = [&Mnodes] (Long i) {
                return Mnodes[0][i];
              };
              for (Long i = 0; i < ORDER; i++) {
                Real wt_x = 1, wt_y = 1;
                for (Long j = 0; j < ORDER; j++) {
                  if (j != i) {
                    wt_x *= (x - node(j)) / (node(i) - node(j));
                    wt_y *= (y - node(j)) / (node(i) - node(j));
                  }
                  Interp0[i] = wt_x;
                  Interp1[i] = wt_y;
                }
              }
            }

            for (Long ii = 0; ii < ORDER; ii++) {
              for (Long jj = 0; jj < ORDER; jj++) {
                Long node_idx = jj * ORDER + ii;
                for (Long k = 0; k < dof; k++) {
                  Xf[(k*Nt+t)*Np+p] += X[elem_idx*dof+k][node_idx] * Interp0[ii] * Interp1[jj];
                }
              }
            }
          }
        }
        return Xf;
      };
      auto grid2cheb = [] (const Vector<Real>& Xf, Long Nt, Long Np, Long Mt, Long Mp) {
        Long dof = Xf.Dim() / (Nt*Np);
        SCTL_ASSERT(Xf.Dim() == dof*Nt*Np);
        Vector<ElemBasis> X(Mt*Mp*dof);
        constexpr Integer INTERP_ORDER = 12;

        for (Long tt = 0; tt < Mt; tt++) {
          for (Long pp = 0; pp < Mp; pp++) {
            for (Long t = 0; t < ORDER; t++) {
              for (Long p = 0; p < ORDER; p++) {
                Matrix<Real> Mnodes = Basis<Real,1,ORDER>::Nodes();
                Real theta = (tt + Mnodes[0][t]) / Mt;
                Real phi   = (pp + Mnodes[0][p]) / Mp;
                Long i = (Long)(theta * Nt);
                Long j = (Long)(phi   * Np);
                Real x = theta * Nt - i;
                Real y = phi   * Np - j;

                Vector<Real> Interp0(INTERP_ORDER);
                Vector<Real> Interp1(INTERP_ORDER);
                { // Set Interp0, Interp1
                  auto node = [] (Long i) {
                    return (Real)i - (INTERP_ORDER-1)/2;
                  };
                  for (Long i = 0; i < INTERP_ORDER; i++) {
                    Real wt_x = 1, wt_y = 1;
                    for (Long j = 0; j < INTERP_ORDER; j++) {
                      if (j != i) {
                        wt_x *= (x - node(j)) / (node(i) - node(j));
                        wt_y *= (y - node(j)) / (node(i) - node(j));
                      }
                      Interp0[i] = wt_x;
                      Interp1[i] = wt_y;
                    }
                  }
                }

                for (Long k = 0; k < dof; k++) {
                  Real X0 = 0;
                  for (Long ii = 0; ii < INTERP_ORDER; ii++) {
                    for (Long jj = 0; jj < INTERP_ORDER; jj++) {
                      Long idx_i = (i + ii-(INTERP_ORDER-1)/2 + Nt) % Nt;
                      Long idx_j = (j + jj-(INTERP_ORDER-1)/2 + Np) % Np;
                      X0 += Interp0[ii] * Interp1[jj] * Xf[(k*Nt+idx_i)*Np+idx_j];
                    }
                  }
                  Long elem_idx = tt * Mp + pp;
                  Long node_idx = p * ORDER + t;
                  X[elem_idx*dof+k][node_idx] = X0;
                }
              }
            }
          }
        }
        return X;
      };

      Long Nelem = S.NElem();
      if (Jp.Dim() != Nelem * COORD_DIM) Jp.ReInit(Nelem * COORD_DIM);
      if (Jt.Dim() != Nelem * COORD_DIM) Jt.ReInit(Nelem * COORD_DIM);
      for (Long k = 0; k < S.Nsurf(); k++) {
        Long Nt = S.NTor(k)*ORDER*2, Np = S.NPol(k)*ORDER*2;
        const auto& X_ = S.GetElemList().ElemVector();
        Vector<ElemBasis> X(S.NTor(k)*S.NPol(k)*COORD_DIM, (Iterator<ElemBasis>)X_.begin()+S.ElemDsp(k)*COORD_DIM, false);

        biest::Surface<Real> SS(Nt, Np);
        biest::SurfaceOp<Real> surf_op(comm, Nt, Np);
        SS.Coord() = cheb2grid(X, S.NTor(k), S.NPol(k), Nt, Np);

        Vector<Real> dX, d2X;
        surf_op.Grad2D(dX, SS.Coord());
        surf_op.Grad2D(d2X, dX);

        Vector<Real> Jt_(COORD_DIM * Nt * Np);
        Vector<Real> Jp_(COORD_DIM * Nt * Np);
        { // Set Jt_, Jp_
          Vector<Real> DivV, InvLapDivV, GradInvLapDivV;
          for (Long i = 0; i < Nt*Np; i++) { // Set V
            for (Long k =0; k < COORD_DIM; k++) {
              Jt_[k * Nt*Np + i] = dX[(k*2+0) * Nt*Np + i];
              Jp_[k * Nt*Np + i] = dX[(k*2+1) * Nt*Np + i];
            }
          }

          surf_op.SurfDiv(DivV, dX, Jt_);
          surf_op.GradInvSurfLap(GradInvLapDivV, dX, d2X, DivV, gmres_tol * max_norm(Jt_) / max_norm(DivV), max_iter, 1.5);
          Jt_ = Jt_ - GradInvLapDivV;

          surf_op.SurfDiv(DivV, dX, Jp_);
          surf_op.GradInvSurfLap(GradInvLapDivV, dX, d2X, DivV, gmres_tol * max_norm(Jp_) / max_norm(DivV), max_iter, 1.5);
          Jp_ = Jp_ - GradInvLapDivV;
        }
        Vector<ElemBasis> Jt__(S.NTor(k)*S.NPol(k)*COORD_DIM, (Iterator<ElemBasis>)Jt.begin()+S.ElemDsp(k)*COORD_DIM, false);
        Vector<ElemBasis> Jp__(S.NTor(k)*S.NPol(k)*COORD_DIM, (Iterator<ElemBasis>)Jp.begin()+S.ElemDsp(k)*COORD_DIM, false);
        Jt__ = grid2cheb(Jt_, Nt, Np, S.NTor(k), S.NPol(k));
        Jp__ = grid2cheb(Jp_, Nt, Np, S.NTor(k), S.NPol(k));
      }
    }
    static void compute_norm_area_elem(const Stellarator<Real,ORDER>& S, Vector<ElemBasis>& normal, Vector<ElemBasis>& area_elem){ // Set normal, area_elem
      const Vector<ElemBasis>& X = S.GetElemList().ElemVector();
      const Long Nelem = X.Dim() / COORD_DIM;
      const Long Nnodes = ElemBasis::Size();

      Vector<ElemBasis> dX;
      ElemBasis::Grad(dX, X);

      area_elem.ReInit(Nelem);
      normal.ReInit(Nelem * COORD_DIM);
      for (Long i = 0; i < Nelem; i++) {
        for (Long j = 0; j < Nnodes; j++) {
          Tensor<Real,true,COORD_DIM> x, n;
          Tensor<Real,true,COORD_DIM,2> dx;
          x(0) = X[i*COORD_DIM+0][j];
          x(1) = X[i*COORD_DIM+1][j];
          x(2) = X[i*COORD_DIM+2][j];

          dx(0,0) = dX[i*COORD_DIM*2+0][j];
          dx(0,1) = dX[i*COORD_DIM*2+1][j];
          dx(1,0) = dX[i*COORD_DIM*2+2][j];
          dx(1,1) = dX[i*COORD_DIM*2+3][j];
          dx(2,0) = dX[i*COORD_DIM*2+4][j];
          dx(2,1) = dX[i*COORD_DIM*2+5][j];

          n(0) = dx(1,0) * dx(2,1) - dx(2,0) * dx(1,1);
          n(1) = dx(2,0) * dx(0,1) - dx(0,0) * dx(2,1);
          n(2) = dx(0,0) * dx(1,1) - dx(1,0) * dx(0,1);
          Real area_elem_ = sqrt<Real>(n(0)*n(0) + n(1)*n(1) + n(2)*n(2));
          Real ooae = 1 / area_elem_;
          n(0) *= ooae;
          n(1) *= ooae;
          n(2) *= ooae;

          normal[i*COORD_DIM+0][j] = n(0);
          normal[i*COORD_DIM+1][j] = n(1);
          normal[i*COORD_DIM+2][j] = n(2);
          area_elem[i][j] = area_elem_;
        }
      }
    }


    static Vector<ElemBasis> compute_B(const Stellarator<Real,ORDER>& S, const Vector<ElemBasis>& sigma, Real alpha, Real beta) {
      const Long Nelem = S.NElem();
      Vector<ElemBasis> B(S.NElem() * COORD_DIM);
      if (sigma.Dim()) {
        const Long Nnodes = ElemBasis::Size();
        Vector<ElemBasis> normal, area_elem;
        compute_norm_area_elem(S, normal, area_elem);
        if (S.Nsurf() == 2) {
          Long Nelem0 = S.NTor(0)*S.NPol(0);
          for (Long i = 0; i < Nelem0*COORD_DIM; i++) {
            for (Long j = 0; j < Nnodes; j++) {
              normal[i][j] *= -1.0;
            }
          }
        }
        EvalQuadrature(B, S.quadrature_FxdU, S, sigma, S.Laplace_FxdU);
        for (Long i = 0; i < Nelem; i++) {
          for (Long j = 0; j < Nnodes; j++) {
            for (Long k = 0; k < COORD_DIM; k++) {
              B[i*COORD_DIM+k][j] -= 0.5*sigma[i][j]*normal[i*COORD_DIM+k][j];
            }
          }
        }
      } else {
        B = 0;
      }
      if (S.Nsurf() >= 1) B += S.Bt0*alpha;
      if (S.Nsurf() >= 2) B += S.Bp0*beta;
      return B;
    }
    static Vector<ElemBasis> compute_dB(const Stellarator<Real,ORDER>& S, const Vector<ElemBasis>& sigma, Real alpha, Real beta) {
      const Long Nelem = S.NElem();
      Vector<ElemBasis> dB(S.NElem() * COORD_DIM * COORD_DIM);
      if (sigma.Dim()) {
        EvalQuadrature(dB, S.quadrature_Fxd2U, S, sigma, S.Laplace_Fxd2U);
      } else {
        dB = 0;
      }
      if (S.Nsurf() >= 1) dB += S.dBt0*alpha;
      if (S.Nsurf() >= 2) dB += S.dBp0*beta;
      return dB;
    }
    static void compute_flux(Real& flux_tor, Real& flux_pol, const Stellarator<Real,ORDER>& S, const Vector<ElemBasis>& B, const Vector<ElemBasis>& normal) {
      const Long Nelem = S.NElem();
      const Long Nnodes = ElemBasis::Size();
      SCTL_ASSERT(B.Dim() == Nelem*COORD_DIM);
      SCTL_ASSERT(normal.Dim() == Nelem*COORD_DIM);

      Vector<ElemBasis> J(Nelem * COORD_DIM);
      for (Long i = 0; i < Nelem; i++) { // Set J
        for (Long j = 0; j < Nnodes; j++) {
          Tensor<Real,true,COORD_DIM> b, n;
          b(0) = B[i*COORD_DIM+0][j];
          b(1) = B[i*COORD_DIM+1][j];
          b(2) = B[i*COORD_DIM+2][j];

          n(0) = normal[i*COORD_DIM+0][j];
          n(1) = normal[i*COORD_DIM+1][j];
          n(2) = normal[i*COORD_DIM+2][j];

          J[i*COORD_DIM+0][j] = n(1) * b(2) - n(2) * b(1);
          J[i*COORD_DIM+1][j] = n(2) * b(0) - n(0) * b(2);
          J[i*COORD_DIM+2][j] = n(0) * b(1) - n(1) * b(0);
        }
      }

      Vector<ElemBasis> A;
      EvalQuadrature(A, S.quadrature_FxU, S, J, S.Laplace_FxU);

      Vector<Real> circ_pol(S.Nsurf()), circ_tor(S.Nsurf());
      { // compute circ
        Vector<ElemBasis> dX;
        ElemBasis::Grad(dX, S.GetElemList().ElemVector());
        const auto& quad_wts = ElemBasis::QuadWts();
        for (Long k = 0; k < S.Nsurf(); k++) {
          circ_pol[k] = 0;
          circ_tor[k] = 0;
          Long Ndsp = S.ElemDsp(k);
          for (Long i = 0; i < S.NTor(k)*S.NPol(k); i++) {
            for (Long j = 0; j < Nnodes; j++) {
              circ_pol[k] += A[(Ndsp+i)*COORD_DIM+0][j] * dX[(Ndsp+i)*COORD_DIM*2+1][j] * quad_wts[j] / S.NTor(k);
              circ_pol[k] += A[(Ndsp+i)*COORD_DIM+1][j] * dX[(Ndsp+i)*COORD_DIM*2+3][j] * quad_wts[j] / S.NTor(k);
              circ_pol[k] += A[(Ndsp+i)*COORD_DIM+2][j] * dX[(Ndsp+i)*COORD_DIM*2+5][j] * quad_wts[j] / S.NTor(k);

              circ_tor[k] += A[(Ndsp+i)*COORD_DIM+0][j] * dX[(Ndsp+i)*COORD_DIM*2+0][j] * quad_wts[j] / S.NPol(k);
              circ_tor[k] += A[(Ndsp+i)*COORD_DIM+1][j] * dX[(Ndsp+i)*COORD_DIM*2+2][j] * quad_wts[j] / S.NPol(k);
              circ_tor[k] += A[(Ndsp+i)*COORD_DIM+2][j] * dX[(Ndsp+i)*COORD_DIM*2+4][j] * quad_wts[j] / S.NPol(k);
            }
          }
        }
      }
      if (S.Nsurf() == 1) {
        flux_tor = circ_pol[0];
        flux_pol = 0;
      } else if (S.Nsurf() == 2) {
        flux_tor = circ_pol[1] - circ_pol[0];
        flux_pol = circ_tor[0] - circ_tor[1];
      } else {
        SCTL_ASSERT(false);
      }
    }
    static Vector<Real> compute_A(const Stellarator<Real,ORDER>& S, const Vector<Real>& x) {
      const Long Nelem = S.NElem();
      const Long Nnodes = ElemBasis::Size();
      SCTL_ASSERT(x.Dim() == Nelem*Nnodes+S.Nsurf());

      Vector<ElemBasis> normal, area_elem;
      compute_norm_area_elem(S, normal, area_elem);
      if (S.Nsurf() == 2) {
        Long Nelem0 = S.NTor(0)*S.NPol(0);
        for (Long i = 0; i < Nelem0*COORD_DIM; i++) {
          for (Long j = 0; j < Nnodes; j++) {
            normal[i][j] *= -1.0;
          }
        }
      }

      Vector<ElemBasis> sigma(Nelem);
      for (Long i = 0; i < Nelem; i++) {
        for (Long j = 0; j < Nnodes; j++) {
          sigma[i][j] = x[i*Nnodes+j];
        }
      }
      Real alpha = (S.Nsurf() >= 1 ? x[Nelem*Nnodes + 0] : 0);
      Real beta  = (S.Nsurf() >= 2 ? x[Nelem*Nnodes + 1] : 0);

      Vector<ElemBasis> B = compute_B(S, sigma, alpha, beta);
      Vector<ElemBasis> BdotN = compute_dot_prod(B, normal);

      Real flux_tor, flux_pol;
      compute_flux(flux_tor, flux_pol, S, B, normal);

      Vector<Real> Ax(Nelem*Nnodes+S.Nsurf());
      for (Long i = 0; i < Nelem; i++) {
        for (Long j = 0; j < Nnodes; j++) {
          Ax[i*Nnodes+j] = BdotN[i][j];
        }
      }
      if (S.Nsurf() >= 1) Ax[Nelem*Nnodes + 0] = flux_tor;
      if (S.Nsurf() >= 2) Ax[Nelem*Nnodes + 1] = flux_pol;
      return Ax;
    }
    static void compute_invA(Vector<ElemBasis>& sigma, Real& alpha, Real& beta, const Stellarator<Real,ORDER>& S, Vector<ElemBasis>& Bdotn, Real flux_tor, Real flux_pol, const Comm& comm) {
      typename ParallelSolver<Real>::ParallelOp BIOp = [&S](Vector<Real>* Ax, const Vector<Real>& x) {
        (*Ax) = compute_A(S, x);
      };
      const Long Nelem = S.NElem();
      const Long Nnodes = ElemBasis::Size();

      Vector<Real> rhs_(Nelem * Nnodes + S.Nsurf());
      for (Long i = 0; i < Nelem; i++) {
        for (Long j = 0; j < Nnodes; j++) {
          rhs_[i*Nnodes+j] = Bdotn[i][j];
        }
      }
      if (S.Nsurf() >= 1) rhs_[Nelem * Nnodes + 0] = flux_tor;
      if (S.Nsurf() >= 2) rhs_[Nelem * Nnodes + 1] = flux_pol;

      Vector<Real> x_(Nelem * Nnodes + S.Nsurf());
      x_ = 0;
      ParallelSolver<Real> linear_solver(comm, true);
      linear_solver(&x_, BIOp, rhs_, 1e-10, 100);

      sigma.ReInit(Nelem);
      for (Long i = 0; i < Nelem; i++) {
        for (Long j = 0; j < Nnodes; j++) {
          sigma[i][j] = x_[i*Nnodes+j];
        }
      }
      alpha = (S.Nsurf() >= 1 ? x_[Nelem * Nnodes + 0] : 0);
      beta  = (S.Nsurf() >= 2 ? x_[Nelem * Nnodes + 1] : 0);
    };
    static void compute_invA(Vector<ElemBasis>& sigma, Real& alpha, Real& beta, const Stellarator<Real,ORDER>& S, Real flux_tor, Real flux_pol, const Comm& comm) {
      Vector<ElemBasis> Bdotn(S.NElem());
      Bdotn = 0;
      compute_invA(sigma, alpha, beta, S, Bdotn, flux_tor, flux_pol, comm);
    }
    static Vector<Real> compute_Aadj(const Stellarator<Real,ORDER>& S, const Vector<Real>& x) {
      const Long Nelem = S.NElem();
      const Long Nnodes = ElemBasis::Size();
      SCTL_ASSERT(x.Dim() == Nelem*Nnodes+S.Nsurf());

      Vector<ElemBasis> normal, area_elem;
      compute_norm_area_elem(S, normal, area_elem);
      if (S.Nsurf() == 2) {
        Long Nelem0 = S.NTor(0)*S.NPol(0);
        for (Long i = 0; i < Nelem0*COORD_DIM; i++) {
          for (Long j = 0; j < Nnodes; j++) {
            normal[i][j] *= -1.0;
          }
        }
      }

      Vector<ElemBasis> x0(Nelem);
      for (Long i = 0; i < Nelem; i++) {
        for (Long j = 0; j < Nnodes; j++) {
          x0[i][j] = x[i*Nnodes+j];
        }
      }
      Real x1 = (S.Nsurf() >= 1 ? x[Nelem*Nnodes + 0] : 0);
      Real x2 = (S.Nsurf() >= 2 ? x[Nelem*Nnodes + 1] : 0);

      Vector<ElemBasis> Ax0;
      Real Ax1, Ax2;
      { // Set Ax0, Ax1, Ax2
        Vector<ElemBasis> x0_n(Nelem*COORD_DIM);
        for (Long i = 0; i < Nelem; i++) {
          for (Long j = 0; j < Nnodes; j++) {
            x0_n[i*COORD_DIM+0][j] = x0[i][j] * normal[i*COORD_DIM+0][j];
            x0_n[i*COORD_DIM+1][j] = x0[i][j] * normal[i*COORD_DIM+1][j];
            x0_n[i*COORD_DIM+2][j] = x0[i][j] * normal[i*COORD_DIM+2][j];
          }
        }
        EvalQuadrature(Ax0, S.quadrature_dUxF, S, x0_n, S.Laplace_dUxF);
        Ax0 = x0*(-0.5) - Ax0;

        Ax1 = (S.Nsurf() >= 1 ? compute_inner_prod(area_elem, compute_dot_prod(S.Bt0, normal), x0) : 0);
        Ax2 = (S.Nsurf() >= 2 ? compute_inner_prod(area_elem, compute_dot_prod(S.Bp0, normal), x0) : 0);
      }

      // TODO: precompute A21adj, A22adj
      auto compute_A21adj = [&S,&normal,&area_elem] (bool toroidal_flux) {
        const Long Nelem = S.NElem();
        const Long Nnodes = ElemBasis::Size();
        Vector<ElemBasis> density(Nelem * COORD_DIM);
        { // Set density
          Real scal[2];
          if (S.Nsurf() == 1) {
            SCTL_ASSERT(toroidal_flux == true);
            scal[0] = 1.0 / S.NTor(0);
            scal[1] = 0;
          } else if (S.Nsurf() == 2) {
            if (toroidal_flux == true) {
              scal[0] = -1.0 / S.NTor(0);
              scal[1] = 1.0 / S.NTor(1);
            } else {
              scal[0] = 1.0 / S.NPol(0);
              scal[1] = -1.0 / S.NPol(1);
            }
          } else {
            SCTL_ASSERT(false);
          }

          Vector<ElemBasis> dX;
          ElemBasis::Grad(dX, S.GetElemList().ElemVector());
          for (Long k = 0; k < S.Nsurf(); k++) {
            for (Long i_ = 0; i_ < S.NTor(k)*S.NPol(k); i_++) {
              Long i = S.ElemDsp(k) + i_;
              for (Long j = 0; j < Nnodes; j++) {
                Real s = scal[k] / area_elem[i][j];
                density[i*COORD_DIM+0][j] = dX[i*COORD_DIM*2+0+(toroidal_flux?1:0)][j] * s;
                density[i*COORD_DIM+1][j] = dX[i*COORD_DIM*2+2+(toroidal_flux?1:0)][j] * s;
                density[i*COORD_DIM+2][j] = dX[i*COORD_DIM*2+4+(toroidal_flux?1:0)][j] * s;
              }
            }
          }
        }

        Vector<ElemBasis> Gdensity, nxGdensity(Nelem * COORD_DIM), A21adj;
        EvalQuadrature(Gdensity, S.quadrature_FxU, S, density, S.Laplace_FxU);
        for (Long i = 0; i < Nelem; i++) { // Set nxGdensity
          for (Long j = 0; j < Nnodes; j++) {
            Tensor<Real,true,COORD_DIM> Gdensity_, n;
            Gdensity_(0) = Gdensity[i*COORD_DIM+0][j];
            Gdensity_(1) = Gdensity[i*COORD_DIM+1][j];
            Gdensity_(2) = Gdensity[i*COORD_DIM+2][j];

            n(0) = normal[i*COORD_DIM+0][j];
            n(1) = normal[i*COORD_DIM+1][j];
            n(2) = normal[i*COORD_DIM+2][j];

            nxGdensity[i*COORD_DIM+0][j] = n(1) * Gdensity_(2) - n(2) * Gdensity_(1);
            nxGdensity[i*COORD_DIM+1][j] = n(2) * Gdensity_(0) - n(0) * Gdensity_(2);
            nxGdensity[i*COORD_DIM+2][j] = n(0) * Gdensity_(1) - n(1) * Gdensity_(0);
          }
        }
        EvalQuadrature(A21adj, S.quadrature_dUxF, S, nxGdensity, S.Laplace_dUxF);
        return A21adj;
      };
      if (S.Nsurf() >= 1) Ax0 += compute_A21adj( true) * x1;
      if (S.Nsurf() >= 2) Ax0 += compute_A21adj(false) * x2;
      if (S.Nsurf() == 1) { // Add flux part of Ax1, Ax2
        Real flux_tor, flux_pol;
        compute_flux(flux_tor, flux_pol, S, S.Bt0, normal);
        Ax1 += flux_tor * x1;
        Ax2 += 0;
      } else if (S.Nsurf() == 2) {
        Real flux_tor, flux_pol;
        compute_flux(flux_tor, flux_pol, S, S.Bt0, normal);
        Ax1 += flux_tor * x1 + flux_pol * x2;
        compute_flux(flux_tor, flux_pol, S, S.Bp0, normal);
        Ax2 += flux_tor * x1 + flux_pol * x2;
      } else {
        SCTL_ASSERT(false);
      }

      Vector<Real> Ax(Nelem*Nnodes+S.Nsurf());
      for (Long i = 0; i < Nelem; i++) {
        for (Long j = 0; j < Nnodes; j++) {
          Ax[i*Nnodes+j] = Ax0[i][j];
        }
      }
      if (S.Nsurf() >= 1) Ax[Nelem*Nnodes + 0] = Ax1;
      if (S.Nsurf() >= 2) Ax[Nelem*Nnodes + 1] = Ax2;
      return Ax;
    }
    static Vector<Real> compute_invAadj(const Stellarator<Real,ORDER>& S, const Vector<Real>& b, const Comm& comm) {
      typename ParallelSolver<Real>::ParallelOp BIOp = [&S](Vector<Real>* Ax, const Vector<Real>& x) {
        (*Ax) = compute_Aadj(S,x);
      };
      const Long Nelem = S.NElem();
      const Long Nnodes = ElemBasis::Size();

      Vector<Real> x(b.Dim());
      x = 0;
      ParallelSolver<Real> linear_solver(comm, true);
      linear_solver(&x, BIOp, b, 1e-8, 100);
      return x;
    }

    static Vector<ElemBasis> compute_pressure_jump(const Vector<Stellarator<Real,ORDER>>& Svec, const Vector<Real>& pressure) {
      const Long Nnodes = ElemBasis::Size();
      const Long Nsurf = Svec.Dim();
      SCTL_ASSERT(pressure.Dim() == Nsurf);
      Vector<Vector<ElemBasis>> total_pressure(Nsurf);
      for (Long i = 0; i < Nsurf; i++) { // Set total_pressure
        const Long Nelem = Svec[i].NElem();
        const auto& B = Svec[i].B;
        total_pressure[i].ReInit(Nelem);
        for (Long j = 0; j < Nelem; j++) {
          for (Long  k = 0; k < Nnodes; k++) {
            Real B2 = 0;
            B2 += B[j*COORD_DIM+0][k] * B[j*COORD_DIM+0][k];
            B2 += B[j*COORD_DIM+1][k] * B[j*COORD_DIM+1][k];
            B2 += B[j*COORD_DIM+2][k] * B[j*COORD_DIM+2][k];
            total_pressure[i][j][k] = 0.5 * B2 + pressure[i];
          }
        }
      }

      Vector<Long> elem_cnt, elem_dsp;
      for (Long i = 0; i < Nsurf; i++) {
        if (i == 0) {
          elem_cnt.PushBack(Svec[i].NTor(0) * Svec[i].NPol(0));
          elem_dsp.PushBack(0);
        } else {
          elem_cnt.PushBack(Svec[i].NTor(1) * Svec[i].NPol(1));
          elem_dsp.PushBack(elem_dsp[i-1] + elem_cnt[i-1]);
        }
      }
      Vector<ElemBasis> pressure_jump(elem_dsp[Nsurf-1] + elem_cnt[Nsurf-1]);
      pressure_jump = 0;

      for (Long i = 0; i < Nsurf-1; i++) { // Set pressure_jump
        Long Nelem, offset;
        if (i == 0) {
          offset = 0;
          Nelem = Svec[i].NTor(0) * Svec[i].NPol(0);
        } else {
          offset = Svec[i].NTor(0) * Svec[i].NPol(0);
          Nelem = Svec[i].NTor(1) * Svec[i].NPol(1);
        }
        for (Long j = 0; j < Nelem; j++) {
          for (Long k = 0; k < Nnodes; k++) {
            Real T0 = total_pressure[i][offset+j][k];
            Real T1 = (i+1<Nsurf ? total_pressure[i+1][j][k] : 0);
            pressure_jump[elem_dsp[i]+j][k] = T1 - T0;
          }
        }
      }
      return pressure_jump;
    }
    static void compute_gvec(const Vector<Stellarator<Real,ORDER>>& Svec, const Vector<Real>& pressure) {
      Vector<ElemBasis> pressure_jump = compute_pressure_jump(Svec, pressure);
      const Long Nnodes = ElemBasis::Size();
      const Long Nsurf = Svec.Dim();
      Long elem_offset = 0;
      for (Long i = 0; i < Nsurf; i++) { // Allocate
        Svec[i].gvec.ReInit(Svec[i].NElem());
        Svec[i].gvec = 0;
      }
      for (Long i = 0; i < Nsurf-1; i++) { // Set gvec
        Long Nelem, offset;
        if (i == 0) {
          offset = 0;
          Nelem = Svec[i].NTor(0) * Svec[i].NPol(0);
        } else {
          offset = Svec[i].NTor(0) * Svec[i].NPol(0);
          Nelem = Svec[i].NTor(1) * Svec[i].NPol(1);
        }
        for (Long j = 0; j < Nelem; j++) {
          for (Long k = 0; k < Nnodes; k++) {
            Real jump = pressure_jump[elem_offset+j][k];
            Svec[i].gvec[offset+j][k] = 0.5 * jump * jump;
            if (i+1<Nsurf) Svec[i+1].gvec[j][k] = 0.5 * jump * jump;
          }
        }
        elem_offset += Nelem;
      }
    }
    static void compute_dgdB(const Vector<Stellarator<Real,ORDER>>& Svec, const Vector<Real>& pressure) {
      Vector<ElemBasis> pressure_jump = compute_pressure_jump(Svec, pressure);
      const Long Nnodes = ElemBasis::Size();
      const Long Nsurf = Svec.Dim();
      Long elem_offset = 0;
      for (Long i = 0; i < Nsurf; i++) { // Allocate
        Svec[i].dgdB.ReInit(Svec[i].NElem() * COORD_DIM);
        Svec[i].dgdB = 0;
      }
      for (Long i = 0; i < Nsurf-1; i++) { // Set dgdB
        Long Nelem, offset;
        if (i == 0) {
          offset = 0;
          Nelem = Svec[i].NTor(0) * Svec[i].NPol(0);
        } else {
          offset = Svec[i].NTor(0) * Svec[i].NPol(0);
          Nelem = Svec[i].NTor(1) * Svec[i].NPol(1);
        }
        for (Long j = 0; j < Nelem; j++) {
          for (Long k = 0; k < Nnodes; k++) {
            Real jump = pressure_jump[elem_offset+j][k];
            Svec[i].dgdB[(offset+j)*COORD_DIM+0][k] = -jump * Svec[i].B[(offset+j)*COORD_DIM+0][k];
            Svec[i].dgdB[(offset+j)*COORD_DIM+1][k] = -jump * Svec[i].B[(offset+j)*COORD_DIM+1][k];
            Svec[i].dgdB[(offset+j)*COORD_DIM+2][k] = -jump * Svec[i].B[(offset+j)*COORD_DIM+2][k];
            if (i+1<Nsurf) {
              Svec[i+1].dgdB[j*COORD_DIM+0][k] = jump * Svec[i+1].B[j*COORD_DIM+0][k];
              Svec[i+1].dgdB[j*COORD_DIM+1][k] = jump * Svec[i+1].B[j*COORD_DIM+1][k];
              Svec[i+1].dgdB[j*COORD_DIM+2][k] = jump * Svec[i+1].B[j*COORD_DIM+2][k];
            }
          }
        }
        elem_offset += Nelem;
      }
    }
    static Real compute_g(const Vector<Stellarator<Real,ORDER>>& Svec, const Vector<Real>& pressure) {
      Real g = 0;
      compute_gvec(Svec, pressure);
      for (Long i = 0; i < Svec.Dim(); i++) { // Set gvec
        Vector<ElemBasis> normal, area_elem, wt(Svec[i].NElem());
        compute_norm_area_elem(Svec[i], normal, area_elem);
        wt = 0.5;
        if (i == Svec.Dim()-1) {
          Long Nsurf = Svec[i].Nsurf();
          Long Nelem = Svec[i].NTor(Nsurf-1) * Svec[i].NPol(Nsurf-1);
          Long offset = Svec[i].ElemDsp(Nsurf-1);
          for (Long j = 0; j < Nelem; j++) {
            wt[offset + j] = 1.0;
          }
        }
        g += compute_inner_prod(area_elem, Svec[i].gvec, wt);
      }
      return g;
    }

    Stellarator(const Vector<Long>& NtNp = Vector<Long>()) {
      NtNp_ = NtNp;
      Long Nsurf = NtNp_.Dim() / 2;
      SCTL_ASSERT(Nsurf*2 == NtNp_.Dim());

      Long Nelem = 0;
      elem_dsp.ReInit(Nsurf+1);
      elem_dsp[0] = 0;
      for (Long i = 0; i < Nsurf; i++) {
        Nelem += NtNp_[i*2+0]*NtNp_[i*2+1];
        elem_dsp[i+1] = Nelem;
      }
      elements.ReInit(Nelem);
      for (Long i = 0; i < Nsurf; i++) {
        InitSurf(i, this->Nsurf());
      }
    }

    Long ElemIdx(Long s, Long t, Long p) {
      SCTL_ASSERT(0 <= s && s < Nsurf());
      SCTL_ASSERT(0 <= t && t < NtNp_[s*2+0]);
      SCTL_ASSERT(0 <= p && p < NtNp_[s*2+1]);
      return elem_dsp[s] + t*NtNp_[s*2+1] + p;
    }
    ElemBasis& Elem(Long elem, Integer dim) {
      return elements(elem,dim);
    }
    const ElemBasis& Elem(Long elem, Integer dim) const {
      return elements(elem,dim);
    }
    const ElemLst& GetElemList() const {
      return elements;
    }

    Long Nsurf() const {
      return elem_dsp.Dim()-1;
    }
    Long ElemDsp(Long s) const {
      return elem_dsp[s];
    }
    Long NTor(Long s) const {
      return NtNp_[s*2+0];
    }
    Long NPol(Long s) const {
      return NtNp_[s*2+1];
    }
    Long NElem() const {
      return elements.NElem();
    }

    static Vector<ElemBasis> compute_gradient(const Stellarator<Real,ORDER>& S_, const Vector<Real>& pressure, const Vector<Real>& flux_tor_, const Vector<Real>& flux_pol_, Real* g_ptr = nullptr) {
      Comm comm = Comm::World();

      Vector<Stellarator<Real,ORDER>> Svec(S_.Nsurf());
      for (Long i = 0; i < S_.Nsurf(); i++) { // Set Svec[i] (quadratures, B)
        const Long elem_dsp = (i==0 ? 0 : S_.ElemDsp(i-1));
        const Long Nnodes = ElemBasis::Size();
        Stellarator<Real,ORDER>& S = Svec[i];
        if (i == 0) { // Init S
          Vector<Long> NtNp;
          NtNp.PushBack(S_.NTor(i));
          NtNp.PushBack(S_.NPol(i));
          S = Stellarator<Real,ORDER>(NtNp);
        } else {
          Vector<Long> NtNp;
          NtNp.PushBack(S_.NTor(i-1));
          NtNp.PushBack(S_.NPol(i-1));
          NtNp.PushBack(S_.NTor(i));
          NtNp.PushBack(S_.NPol(i));
          S = Stellarator<Real,ORDER>(NtNp);
        }
        for (Long j = 0; j < S.NElem(); j++) { // Set S coordinates
          for (Long k = 0; k < Nnodes; k++) {
            S.Elem(j,0)[k] = S_.Elem(elem_dsp+j,0)[k];
            S.Elem(j,1)[k] = S_.Elem(elem_dsp+j,1)[k];
            S.Elem(j,2)[k] = S_.Elem(elem_dsp+j,2)[k];
          }
        }

        SetupQuadrature(S.quadrature_dBS  , S, S.BiotSavartGrad, order_singular, order_direct, -1.0, comm, -0.01 * pow<-2,Real>(ORDER));
        SetupQuadrature(S.quadrature_BS   , S, S.BiotSavart    , order_singular, order_direct, -1.0, comm, -0.01 * pow<-2,Real>(ORDER));
        SetupQuadrature(S.quadrature_FxU  , S, S.Laplace_FxU   , order_singular, order_direct, -1.0, comm);
        SetupQuadrature(S.quadrature_FxdU , S, S.Laplace_FxdU  , order_singular, order_direct, -1.0, comm);
        SetupQuadrature(S.quadrature_dUxF , S, S.Laplace_dUxF  , order_singular, order_direct, -1.0, comm);
        SetupQuadrature(S.quadrature_dUxD , S, S.Laplace_dUxD  , order_singular, order_direct, -1.0, comm,  0.01 * pow<-2,Real>(ORDER));
        SetupQuadrature(S.quadrature_Fxd2U, S, S.Laplace_Fxd2U , order_singular, order_direct, -1.0, comm, -0.01 * pow<-2,Real>(ORDER));

        { // Set Bt0, Bp0, dBt0, dBp0
          Vector<ElemBasis> Jt, Jp;
          compute_harmonic_vector_potentials(Jt, Jp, S);
          EvalQuadrature(S.Bt0 , S.quadrature_BS , S, Jp, S.BiotSavart);
          EvalQuadrature(S.Bp0 , S.quadrature_BS , S, Jt, S.BiotSavart);
          EvalQuadrature(S.dBt0, S.quadrature_dBS, S, Jp, S.BiotSavartGrad);
          EvalQuadrature(S.dBp0, S.quadrature_dBS, S, Jt, S.BiotSavartGrad);
        }

        compute_invA(S.sigma, S.alpha, S.beta, S, flux_tor_[i], flux_pol_[i], comm);
        S.B = compute_B(S, S.sigma, S.alpha, S.beta);
        if (0) { // Write VTU
          VTUData vtu;
          vtu.AddElems(S.GetElemList(), S.sigma, ORDER);
          vtu.WriteVTK("sigma"+std::to_string(i), comm);
        }
        if (0) { // Write VTU
          VTUData vtu;
          vtu.AddElems(S.GetElemList(), S.B, ORDER);
          vtu.WriteVTK("B"+std::to_string(i), comm);
        }
      }
      compute_gvec(Svec, pressure);
      compute_dgdB(Svec, pressure);
      if (g_ptr != nullptr) g_ptr[0] = compute_g(Svec, pressure);

      auto compute_gradient = [&comm] (const Stellarator<Real,ORDER>& S) {
        const Long Nnodes = ElemBasis::Size();
        const Long Nelem = S.NElem();

        const auto& sigma = S.sigma;
        const auto& alpha = S.alpha;
        const auto& beta = S.beta;
        const auto& B = S.B;

        Vector<ElemBasis> normal, area_elem;
        compute_norm_area_elem(S, normal, area_elem);
        if (S.Nsurf() == 2) {
          Long Nelem0 = S.NTor(0)*S.NPol(0);
          for (Long i = 0; i < Nelem0*COORD_DIM; i++) {
            for (Long j = 0; j < Nnodes; j++) {
              normal[i][j] *= -1.0;
            }
          }
        }
        auto compute_H = [] (const ElemList<COORD_DIM,ElemBasis>& elem_lst, const Vector<ElemBasis>& normal) {
          const Long Nnodes = ElemBasis::Size();
          const Long Nelem = elem_lst.NElem();

          const Vector<ElemBasis> X = elem_lst.ElemVector();
          Vector<ElemBasis> dX, d2X, H(Nelem);
          ElemBasis::Grad(dX, X);
          ElemBasis::Grad(d2X, dX);
          for (Long i = 0; i < Nelem; i++) {
            for (Long j = 0; j < Nnodes; j++) {
              Tensor<Real,true,2,2> I, invI, II;
              for (Long k0 = 0; k0 < 2; k0++) {
                for (Long k1 = 0; k1 < 2; k1++) {
                  I(k0,k1) = 0;
                  I(k0,k1) += dX[(i*COORD_DIM+0)*2+k0][j] * dX[(i*COORD_DIM+0)*2+k1][j];
                  I(k0,k1) += dX[(i*COORD_DIM+1)*2+k0][j] * dX[(i*COORD_DIM+1)*2+k1][j];
                  I(k0,k1) += dX[(i*COORD_DIM+2)*2+k0][j] * dX[(i*COORD_DIM+2)*2+k1][j];

                  II(k0,k1) = 0;
                  II(k0,k1) += d2X[(i*COORD_DIM+0)*4+k0*2+k1][j] * normal[i*COORD_DIM+0][j];
                  II(k0,k1) += d2X[(i*COORD_DIM+1)*4+k0*2+k1][j] * normal[i*COORD_DIM+1][j];
                  II(k0,k1) += d2X[(i*COORD_DIM+2)*4+k0*2+k1][j] * normal[i*COORD_DIM+2][j];
                }
              }
              { // Set invI
                Real detI = I(0,0)*I(1,1)-I(0,1)*I(1,0);
                invI(0,0) = I(1,1) / detI;
                invI(0,1) = -I(0,1) / detI;
                invI(1,0) = -I(1,0) / detI;
                invI(1,1) = I(0,0) / detI;
              }
              { // Set H
                H[i][j] = 0;
                H[i][j] += -0.5 * II(0,0)*invI(0,0);
                H[i][j] += -0.5 * II(0,1)*invI(0,1);
                H[i][j] += -0.5 * II(1,0)*invI(1,0);
                H[i][j] += -0.5 * II(1,1)*invI(1,1);
              }
            }
          }
          return H;
        };
        Vector<ElemBasis> H = compute_H(S.GetElemList(), normal);

        auto compute_dg_dnu = [&S,&normal,&area_elem,&H]() { // dg_dnu = (B*B) 2H - (2 B) \cdot (n \cdnot nabla) \nabla G[sigma] + (2 B) \alpha dB0_dnu \hat{\theta} + sigma (\nabla D)^T [2 B] + (2H) sigma (\nabla G)^T [2 B]
          const Long Nelem = S.NElem();
          const Long Nnodes = ElemBasis::Size();
          const Vector<ElemBasis>& gvec = S.gvec;
          const Vector<ElemBasis>& v    = S.dgdB;
          const auto& sigma = S.sigma;
          const auto& alpha = S.alpha;
          const auto& beta = S.beta;
          const auto& B = S.B;

          Vector<ElemBasis> dg_dnu0(Nelem), dg_dnu1(Nelem), dg_dnu2(Nelem), dg_dnu3(Nelem), dg_dnu4(Nelem);
          dg_dnu0 = 0;
          dg_dnu1 = 0;
          dg_dnu2 = 0;
          dg_dnu3 = 0;
          dg_dnu4 = 0;

          // dg_dnu0 = (B*B) 2H
          for (Long i = 0; i < Nelem; i++) {
            for (Long j = 0; j < Nnodes; j++) {
              dg_dnu0[i][j] = gvec[i][j] * (2.0*H[i][j]) * 0.5;
              // multiplicative factor 0.5 is there so that this term is not
              // counted twice from shape derivative of regions on either side
              // of the domain.
            }
          }

          // dg_dnu1 = (2 B) \cdot (n \cdnot \nabla) B
          Vector<ElemBasis> dB = compute_dB(S, sigma, alpha, beta);
          for (Long i = 0; i < Nelem; i++) {
            for (Long j = 0; j < Nnodes; j++) {
              dg_dnu1[i][j] = 0;
              dg_dnu1[i][j] -= dB[i*9+0][j] * normal[i*COORD_DIM+0][j] * v[i*COORD_DIM+0][j];
              dg_dnu1[i][j] -= dB[i*9+1][j] * normal[i*COORD_DIM+1][j] * v[i*COORD_DIM+0][j];
              dg_dnu1[i][j] -= dB[i*9+2][j] * normal[i*COORD_DIM+2][j] * v[i*COORD_DIM+0][j];

              dg_dnu1[i][j] -= dB[i*9+3][j] * normal[i*COORD_DIM+0][j] * v[i*COORD_DIM+1][j];
              dg_dnu1[i][j] -= dB[i*9+4][j] * normal[i*COORD_DIM+1][j] * v[i*COORD_DIM+1][j];
              dg_dnu1[i][j] -= dB[i*9+5][j] * normal[i*COORD_DIM+2][j] * v[i*COORD_DIM+1][j];

              dg_dnu1[i][j] -= dB[i*9+6][j] * normal[i*COORD_DIM+0][j] * v[i*COORD_DIM+2][j];
              dg_dnu1[i][j] -= dB[i*9+7][j] * normal[i*COORD_DIM+1][j] * v[i*COORD_DIM+2][j];
              dg_dnu1[i][j] -= dB[i*9+8][j] * normal[i*COORD_DIM+2][j] * v[i*COORD_DIM+2][j];
            }
          }

          // dg_dnu3 = (sigma (\nabla D)^T [2 B]
          Vector<ElemBasis> nablaDtv;
          EvalQuadrature(nablaDtv, S.quadrature_dUxD, S, v, S.Laplace_dUxD);
          for (Long i = 0; i < Nelem; i++) {
            for (Long j = 0; j < Nnodes; j++) {
              dg_dnu3[i][j] = 0;
              dg_dnu3[i][j] += sigma[i][j] * nablaDtv[i*COORD_DIM+0][j]*normal[i*COORD_DIM+0][j];
              dg_dnu3[i][j] += sigma[i][j] * nablaDtv[i*COORD_DIM+1][j]*normal[i*COORD_DIM+1][j];
              dg_dnu3[i][j] += sigma[i][j] * nablaDtv[i*COORD_DIM+2][j]*normal[i*COORD_DIM+2][j];
            }
          }

          // dg_dnu4 = (2H) sigma (\nabla G)^T [2 B]
          EvalQuadrature(dg_dnu4, S.quadrature_dUxF, S, v, S.Laplace_dUxF);
          for (Long i = 0; i < Nelem; i++) {
            for (Long j = 0; j < Nnodes; j++) {
              dg_dnu4[i][j] += 0.5 * v[i*COORD_DIM+0][j] * normal[i*COORD_DIM+0][j];
              dg_dnu4[i][j] += 0.5 * v[i*COORD_DIM+1][j] * normal[i*COORD_DIM+1][j];
              dg_dnu4[i][j] += 0.5 * v[i*COORD_DIM+2][j] * normal[i*COORD_DIM+2][j];
              dg_dnu4[i][j] *= 2*H[i][j] * sigma[i][j];
            }
          }

          return dg_dnu0 + dg_dnu1 + dg_dnu3 - dg_dnu4;
        };
        Vector<ElemBasis> dg_dnu = compute_dg_dnu();

        auto compute_dg_dsigma = [&S,&normal,&area_elem] () {
          const Long Nnodes = ElemBasis::Size();
          const Long Nelem = S.NElem();

          const auto& B = S.B;
          const Vector<ElemBasis>& dgdB = S.dgdB;

          auto compute_dg_dsigma = [&S,&B,&dgdB,&normal]() { // dg_dsigma = \int 2 B \cdot (\nabla G + n/2)
            Vector<ElemBasis> B_dot_gradG;
            EvalQuadrature(B_dot_gradG, S.quadrature_dUxF, S, dgdB, S.Laplace_dUxF);
            return B_dot_gradG * (-1.0) + compute_dot_prod(dgdB,normal) * 0.5;
          };
          auto compute_dg_dalpha = [&S,&B,&dgdB,&area_elem] () {
            auto dB_dalpha = compute_B(S, Vector<ElemBasis>(),1,0);
            return compute_inner_prod(area_elem, dgdB,dB_dalpha);
          };
          auto compute_dg_dbeta  = [&S,&B,&dgdB,&area_elem] () {
            auto dB_dalpha = compute_B(S, Vector<ElemBasis>(),0,1);
            return compute_inner_prod(area_elem, dgdB,dB_dalpha);
          };

          Vector<Real> dg_dsigma(Nelem*Nnodes+S.Nsurf());
          Vector<ElemBasis> dg_dsigma_ = compute_dg_dsigma();
          for (Long i = 0; i < Nelem; i++) {
            for (Long j = 0; j < Nnodes; j++) {
              dg_dsigma[i*Nnodes+j] = dg_dsigma_[i][j];
            }
          }
          if (S.Nsurf() >= 1) dg_dsigma[Nelem*Nnodes+0] = compute_dg_dalpha();
          if (S.Nsurf() >= 2) dg_dsigma[Nelem*Nnodes+1] = compute_dg_dbeta ();
          return dg_dsigma;
        };
        Vector<Real> dg_dsigma = compute_dg_dsigma();
        Vector<Real> dg_dsigma_invA = compute_invAadj(S, dg_dsigma, comm);

        ///////////////////////////////////////////////////////////////////////////////////////////////////////////////
        ///////////////////////////////////////////////////////////////////////////////////////////////////////////////

        auto compute_grad_adj = [&S,&area_elem] (const Vector<ElemBasis>& V) {
          const Long Nelem = S.NElem();
          const Long Nnodes = ElemBasis::Size();

          Vector<ElemBasis> du_dX(Nelem*COORD_DIM*2);
          { // Set du_dX
            Vector<ElemBasis> dX;
            ElemBasis::Grad(dX, S.GetElemList().ElemVector());

            auto inv2x2 = [](Tensor<Real, true, 2, 2> M) {
              Tensor<Real, true, 2, 2> Mout;
              Real oodet = 1 / (M(0,0) * M(1,1) - M(0,1) * M(1,0));
              Mout(0,0) =  M(1,1) * oodet;
              Mout(0,1) = -M(0,1) * oodet;
              Mout(1,0) = -M(1,0) * oodet;
              Mout(1,1) =  M(0,0) * oodet;
              return Mout;
            };
            for (Long i = 0; i < Nelem; i++) {
              for (Long j = 0; j < Nnodes; j++) {
                Tensor<Real, true, 3, 2> dX_du;
                dX_du(0,0) = dX[(i*COORD_DIM+0)*2+0][j];
                dX_du(1,0) = dX[(i*COORD_DIM+1)*2+0][j];
                dX_du(2,0) = dX[(i*COORD_DIM+2)*2+0][j];
                dX_du(0,1) = dX[(i*COORD_DIM+0)*2+1][j];
                dX_du(1,1) = dX[(i*COORD_DIM+1)*2+1][j];
                dX_du(2,1) = dX[(i*COORD_DIM+2)*2+1][j];

                Tensor<Real, true, 2, 2> G; // = dX_du.Transpose() * dX_du;
                G(0,0) = dX_du(0,0) * dX_du(0,0) + dX_du(1,0) * dX_du(1,0) + dX_du(2,0) * dX_du(2,0);
                G(0,1) = dX_du(0,0) * dX_du(0,1) + dX_du(1,0) * dX_du(1,1) + dX_du(2,0) * dX_du(2,1);
                G(1,0) = dX_du(0,1) * dX_du(0,0) + dX_du(1,1) * dX_du(1,0) + dX_du(2,1) * dX_du(2,0);
                G(1,1) = dX_du(0,1) * dX_du(0,1) + dX_du(1,1) * dX_du(1,1) + dX_du(2,1) * dX_du(2,1);

                Tensor<Real, true, 2, 2> Ginv = inv2x2(G);
                du_dX[(i*COORD_DIM+0)*2+0][j] = Ginv(0,0) * dX_du(0,0) + Ginv(0,1) * dX_du(0,1);
                du_dX[(i*COORD_DIM+1)*2+0][j] = Ginv(0,0) * dX_du(1,0) + Ginv(0,1) * dX_du(1,1);
                du_dX[(i*COORD_DIM+2)*2+0][j] = Ginv(0,0) * dX_du(2,0) + Ginv(0,1) * dX_du(2,1);
                du_dX[(i*COORD_DIM+0)*2+1][j] = Ginv(1,0) * dX_du(0,0) + Ginv(1,1) * dX_du(0,1);
                du_dX[(i*COORD_DIM+1)*2+1][j] = Ginv(1,0) * dX_du(1,0) + Ginv(1,1) * dX_du(1,1);
                du_dX[(i*COORD_DIM+2)*2+1][j] = Ginv(1,0) * dX_du(2,0) + Ginv(1,1) * dX_du(2,1);
              }
            }
          }

          Vector<ElemBasis> dudX_V(Nelem*2);
          for (Long i = 0; i < Nelem; i++) {
            for (Long j = 0; j < Nnodes; j++) {
              dudX_V[i*2+0][j] = 0;
              dudX_V[i*2+1][j] = 0;

              dudX_V[i*2+0][j] += du_dX[(i*COORD_DIM+0)*2+0][j] * V[i*COORD_DIM+0][j] * area_elem[i][j];
              dudX_V[i*2+0][j] += du_dX[(i*COORD_DIM+1)*2+0][j] * V[i*COORD_DIM+1][j] * area_elem[i][j];
              dudX_V[i*2+0][j] += du_dX[(i*COORD_DIM+2)*2+0][j] * V[i*COORD_DIM+2][j] * area_elem[i][j];

              dudX_V[i*2+1][j] += du_dX[(i*COORD_DIM+0)*2+1][j] * V[i*COORD_DIM+0][j] * area_elem[i][j];
              dudX_V[i*2+1][j] += du_dX[(i*COORD_DIM+1)*2+1][j] * V[i*COORD_DIM+1][j] * area_elem[i][j];
              dudX_V[i*2+1][j] += du_dX[(i*COORD_DIM+2)*2+1][j] * V[i*COORD_DIM+2][j] * area_elem[i][j];
            }
          }

          Vector<ElemBasis> grad_dudX_V;
          ElemBasis::Grad(grad_dudX_V, dudX_V);

          Vector<ElemBasis> grad_adj_V(Nelem);
          for (Long i = 0; i < Nelem; i++) {
            for (Long j = 0; j < Nnodes; j++) {
              grad_adj_V[i][j] = -(grad_dudX_V[(i*2+0)*2+0][j] + grad_dudX_V[(i*2+1)*2+1][j]) / area_elem[i][j];
            }
          }
          return grad_adj_V;
        };
        auto compute_u_dAdnu_v_0 = [&S,&normal,&H,&compute_grad_adj] (const Vector<Real>& u_, const Vector<ElemBasis>& v, Real alpha, Real beta) {
          const Long Nnodes = ElemBasis::Size();
          const Long Nelem = S.NElem();

          Vector<ElemBasis> dAdnu0(Nelem), dAdnu1(Nelem), dAdnu2(Nelem), dAdnu3(Nelem);
          Vector<ElemBasis> u(Nelem), u_n(Nelem*COORD_DIM);
          for (Long i = 0; i < Nelem; i++) {
            for (Long j = 0; j < Nnodes; j++) {
              u[i][j] = u_[i*Nnodes+j];
              u_n[i*COORD_DIM+0][j] = u[i][j] * normal[i*COORD_DIM+0][j];
              u_n[i*COORD_DIM+1][j] = u[i][j] * normal[i*COORD_DIM+1][j];
              u_n[i*COORD_DIM+2][j] = u[i][j] * normal[i*COORD_DIM+2][j];
            }
          }

          // dAdnu0 = u B \cdot grad_nu
          Vector<ElemBasis> B = compute_B(S, v, alpha, beta);
          Vector<ElemBasis> u_B(Nelem*COORD_DIM);
          for (Long i = 0; i < Nelem; i++) {
            for (Long j = 0; j < Nnodes; j++) {
              u_B[i*COORD_DIM+0][j] = u[i][j] * B[i*COORD_DIM+0][j];
              u_B[i*COORD_DIM+1][j] = u[i][j] * B[i*COORD_DIM+1][j];
              u_B[i*COORD_DIM+2][j] = u[i][j] * B[i*COORD_DIM+2][j];
            }
          }
          dAdnu0 = compute_grad_adj(u_B)*(-1.0);

          // dAdnu1 = (u n) \cdot (n \cdnot \nabla) B
          Vector<ElemBasis> dB = compute_dB(S, v, alpha, beta);
          for (Long i = 0; i < Nelem; i++) {
            for (Long j = 0; j < Nnodes; j++) {
              dAdnu1[i][j] = 0;
              dAdnu1[i][j] -= dB[i*9+0][j] * normal[i*COORD_DIM+0][j] * u_n[i*COORD_DIM+0][j];
              dAdnu1[i][j] -= dB[i*9+1][j] * normal[i*COORD_DIM+0][j] * u_n[i*COORD_DIM+1][j];
              dAdnu1[i][j] -= dB[i*9+2][j] * normal[i*COORD_DIM+0][j] * u_n[i*COORD_DIM+2][j];

              dAdnu1[i][j] -= dB[i*9+3][j] * normal[i*COORD_DIM+1][j] * u_n[i*COORD_DIM+0][j];
              dAdnu1[i][j] -= dB[i*9+4][j] * normal[i*COORD_DIM+1][j] * u_n[i*COORD_DIM+1][j];
              dAdnu1[i][j] -= dB[i*9+5][j] * normal[i*COORD_DIM+1][j] * u_n[i*COORD_DIM+2][j];

              dAdnu1[i][j] -= dB[i*9+6][j] * normal[i*COORD_DIM+2][j] * u_n[i*COORD_DIM+0][j];
              dAdnu1[i][j] -= dB[i*9+7][j] * normal[i*COORD_DIM+2][j] * u_n[i*COORD_DIM+1][j];
              dAdnu1[i][j] -= dB[i*9+8][j] * normal[i*COORD_DIM+2][j] * u_n[i*COORD_DIM+2][j];
            }
          }

          // dAdnu2 = (2H) v (I/2 + \nabla G)^T [u n]
          EvalQuadrature(dAdnu2, S.quadrature_dUxF, S, u_n, S.Laplace_dUxF);
          for (Long i = 0; i < Nelem; i++) {
            for (Long j = 0; j < Nnodes; j++) {
              dAdnu2[i][j] += 0.5 * u_n[i*COORD_DIM+0][j] * normal[i*COORD_DIM+0][j];
              dAdnu2[i][j] += 0.5 * u_n[i*COORD_DIM+1][j] * normal[i*COORD_DIM+1][j];
              dAdnu2[i][j] += 0.5 * u_n[i*COORD_DIM+2][j] * normal[i*COORD_DIM+2][j];
              dAdnu2[i][j] *= -2*H[i][j] * v[i][j];
            }
          }

          // dAdnu3 = (v n \cdot \nabla D[u]
          Vector<ElemBasis> nablaDt_u_n;
          EvalQuadrature(nablaDt_u_n, S.quadrature_dUxD, S, u_n, S.Laplace_dUxD);
          for (Long i = 0; i < Nelem; i++) {
            for (Long j = 0; j < Nnodes; j++) {
              dAdnu3[i][j] = 0;
              dAdnu3[i][j] += v[i][j] * nablaDt_u_n[i*COORD_DIM+0][j]*normal[i*COORD_DIM+0][j];
              dAdnu3[i][j] += v[i][j] * nablaDt_u_n[i*COORD_DIM+1][j]*normal[i*COORD_DIM+1][j];
              dAdnu3[i][j] += v[i][j] * nablaDt_u_n[i*COORD_DIM+2][j]*normal[i*COORD_DIM+2][j];
            }
          }

          return dAdnu0 + dAdnu1 + dAdnu2 + dAdnu3;
        };
        auto compute_u_dAdnu_v_1 = [&S,&area_elem,&normal,&H,&compute_grad_adj] (const Vector<ElemBasis>& sigma, Real alpha, Real beta, bool toroidal_flux) {
          const Long Nnodes = ElemBasis::Size();
          const Long Nelem = S.NElem();

          Vector<ElemBasis> B = compute_B(S, sigma, alpha, beta);
          Vector<ElemBasis> gradB = compute_dB(S, sigma, alpha, beta);

          auto compute_v = [&S,&area_elem,&toroidal_flux] (const Vector<ElemBasis>& X) {
            const Long Nelem = S.NElem();
            const Long Nnodes = ElemBasis::Size();

            Real scal[2];
            if (S.Nsurf() == 1) {
              SCTL_ASSERT(toroidal_flux == true);
              scal[0] = 1.0 / S.NTor(0);
              scal[1] = 0;
            } else if (S.Nsurf() == 2) {
              if (toroidal_flux == true) {
                scal[0] = -1.0 / S.NTor(0);
                scal[1] = 1.0 / S.NTor(1);
              } else {
                scal[0] = 1.0 / S.NPol(0);
                scal[1] = -1.0 / S.NPol(1);
              }
            } else {
              SCTL_ASSERT(false);
            }

            Vector<ElemBasis> v(Nelem * COORD_DIM);
            Vector<ElemBasis> dX;
            ElemBasis::Grad(dX, X);
            for (Long k = 0; k < S.Nsurf(); k++) {
              for (Long i_ = 0; i_ < S.NTor(k)*S.NPol(k); i_++) {
                Long i = S.ElemDsp(k) + i_;
                for (Long j = 0; j < Nnodes; j++) {
                  Real s = scal[k] / area_elem[i][j];
                  v[i*COORD_DIM+0][j] = dX[i*COORD_DIM*2+0+(toroidal_flux?1:0)][j] * s;
                  v[i*COORD_DIM+1][j] = dX[i*COORD_DIM*2+2+(toroidal_flux?1:0)][j] * s;
                  v[i*COORD_DIM+2][j] = dX[i*COORD_DIM*2+4+(toroidal_flux?1:0)][j] * s;
                }
              }
            }
            return v;
          };
          auto compute_AxB = [&S] (const Vector<ElemBasis>& A, const Vector<ElemBasis>& B) {
            const Long Nelem = S.NElem();
            const Long Nnodes = ElemBasis::Size();

            Vector<ElemBasis> J(Nelem * COORD_DIM);
            for (Long i = 0; i < Nelem; i++) { // Set J
              for (Long j = 0; j < Nnodes; j++) {
                Tensor<Real,true,COORD_DIM> a, b;
                a(0) = A[i*COORD_DIM+0][j];
                a(1) = A[i*COORD_DIM+1][j];
                a(2) = A[i*COORD_DIM+2][j];

                b(0) = B[i*COORD_DIM+0][j];
                b(1) = B[i*COORD_DIM+1][j];
                b(2) = B[i*COORD_DIM+2][j];

                J[i*COORD_DIM+0][j] = a(1) * b(2) - a(2) * b(1);
                J[i*COORD_DIM+1][j] = a(2) * b(0) - a(0) * b(2);
                J[i*COORD_DIM+2][j] = a(0) * b(1) - a(1) * b(0);
              }
            }
            return J;
          };

          auto compute_dphi_dnu0 = [&S,&normal,&compute_AxB,&compute_v,&B,compute_grad_adj] () {
            const Long Nelem = S.NElem();
            const Long Nnodes = ElemBasis::Size();

            Vector<ElemBasis> Gv;
            Vector<ElemBasis> v = compute_v(S.GetElemList().ElemVector());
            EvalQuadrature(Gv, S.quadrature_FxU, S, v, S.Laplace_FxU);
            Vector<ElemBasis> BxGv = compute_AxB(B,Gv);
            return compute_grad_adj(BxGv)*(-1.0);
          };
          auto compute_dphi_dnu1 = [&S,&normal,&H,&compute_AxB,&compute_v,&B] () {
            const Long Nelem = S.NElem();
            const Long Nnodes = ElemBasis::Size();

            Vector<ElemBasis> Gv;
            Vector<ElemBasis> v = compute_v(S.GetElemList().ElemVector());
            EvalQuadrature(Gv, S.quadrature_FxU, S, v, S.Laplace_FxU);

            Vector<ElemBasis> BxGv = compute_AxB(B,Gv);
            Vector<ElemBasis> n_dot_BxGv = compute_dot_prod(normal,BxGv);

            Vector<ElemBasis> dphi_dnu(Nelem);
            for (Long i = 0; i < Nelem; i++) {
              for (Long j = 0; j < Nnodes; j++) {
                dphi_dnu[i][j] = n_dot_BxGv[i][j] * 2*H[i][j];
              }
            }
            return dphi_dnu;
          };
          auto compute_dphi_dnu2 = [&S,&normal,&H,&compute_AxB,&compute_v,&B] () {
            const Long Nelem = S.NElem();
            const Long Nnodes = ElemBasis::Size();

            Vector<ElemBasis> GnxB;
            Vector<ElemBasis> nxB = compute_AxB(normal,B);
            EvalQuadrature(GnxB, S.quadrature_FxU, S, nxB, S.Laplace_FxU);

            Vector<ElemBasis> v = compute_v(S.GetElemList().ElemVector());
            Vector<ElemBasis> v_dot_GnxB = compute_dot_prod(v,GnxB);

            Vector<ElemBasis> dphi_dnu(Nelem);
            for (Long i = 0; i < Nelem; i++) {
              for (Long j = 0; j < Nnodes; j++) {
                dphi_dnu[i][j] = v_dot_GnxB[i][j] * 2*H[i][j];
              }
            }
            return dphi_dnu;
          };
          auto compute_dphi_dnu3 = [&S,&normal,&area_elem,&H,&compute_AxB,&compute_v,&B] () {
            const Long Nelem = S.NElem();
            const Long Nnodes = ElemBasis::Size();

            Vector<ElemBasis> GnxB;
            Vector<ElemBasis> nxB = compute_AxB(normal,B);
            EvalQuadrature(GnxB, S.quadrature_FxU, S, nxB, S.Laplace_FxU);

            Vector<ElemBasis> dGnxB = compute_v(GnxB);
            Vector<ElemBasis> v = compute_v(S.GetElemList().ElemVector());
            Vector<ElemBasis> dv_dnu1(Nelem), dv_dnu2(Nelem);
            { // Set dv_dnu1, dv_dnu2
              for (Long i = 0; i < Nelem; i++) {
                for (Long j = 0; j < Nnodes; j++) {
                  dv_dnu1[i][j] = 0;
                  dv_dnu1[i][j] += -GnxB[i*COORD_DIM+0][j] * v[i*COORD_DIM+0][j] * 2 * H[i][j];
                  dv_dnu1[i][j] += -GnxB[i*COORD_DIM+1][j] * v[i*COORD_DIM+1][j] * 2 * H[i][j];
                  dv_dnu1[i][j] += -GnxB[i*COORD_DIM+2][j] * v[i*COORD_DIM+2][j] * 2 * H[i][j];

                  dv_dnu2[i][j] = 0;
                  dv_dnu2[i][j] += -dGnxB[i*COORD_DIM+0][j] * normal[i*COORD_DIM+0][j];
                  dv_dnu2[i][j] += -dGnxB[i*COORD_DIM+1][j] * normal[i*COORD_DIM+1][j];
                  dv_dnu2[i][j] += -dGnxB[i*COORD_DIM+2][j] * normal[i*COORD_DIM+2][j];
                }
              }
            }
            return dv_dnu1 + dv_dnu2;
          };
          auto compute_dphi_dnu4 = [&S,&normal,&compute_AxB,&compute_v,&B] () {
            const Long Nelem = S.NElem();
            const Long Nnodes = ElemBasis::Size();

            Vector<ElemBasis> dGnxB;
            Vector<ElemBasis> nxB = compute_AxB(normal,B);
            EvalQuadrature(dGnxB, S.quadrature_FxdU, S, nxB, S.Laplace_FxdU);

            Vector<ElemBasis> v = compute_v(S.GetElemList().ElemVector());
            Vector<ElemBasis> dphi_dnu(Nelem);
            for (Long i = 0; i < Nelem; i++) {
              for (Long j = 0; j < Nnodes; j++) {
                Real dphi_dnu_ = 0;
                dphi_dnu_ += -normal[i*COORD_DIM+0][j] * dGnxB[(i*COORD_DIM+0)*COORD_DIM+0][j] * v[i*COORD_DIM+0][j];
                dphi_dnu_ += -normal[i*COORD_DIM+1][j] * dGnxB[(i*COORD_DIM+0)*COORD_DIM+1][j] * v[i*COORD_DIM+0][j];
                dphi_dnu_ += -normal[i*COORD_DIM+2][j] * dGnxB[(i*COORD_DIM+0)*COORD_DIM+2][j] * v[i*COORD_DIM+0][j];

                dphi_dnu_ += -normal[i*COORD_DIM+0][j] * dGnxB[(i*COORD_DIM+1)*COORD_DIM+0][j] * v[i*COORD_DIM+1][j];
                dphi_dnu_ += -normal[i*COORD_DIM+1][j] * dGnxB[(i*COORD_DIM+1)*COORD_DIM+1][j] * v[i*COORD_DIM+1][j];
                dphi_dnu_ += -normal[i*COORD_DIM+2][j] * dGnxB[(i*COORD_DIM+1)*COORD_DIM+2][j] * v[i*COORD_DIM+1][j];

                dphi_dnu_ += -normal[i*COORD_DIM+0][j] * dGnxB[(i*COORD_DIM+2)*COORD_DIM+0][j] * v[i*COORD_DIM+2][j];
                dphi_dnu_ += -normal[i*COORD_DIM+1][j] * dGnxB[(i*COORD_DIM+2)*COORD_DIM+1][j] * v[i*COORD_DIM+2][j];
                dphi_dnu_ += -normal[i*COORD_DIM+2][j] * dGnxB[(i*COORD_DIM+2)*COORD_DIM+2][j] * v[i*COORD_DIM+2][j];
                dphi_dnu[i][j] = dphi_dnu_;
              }
            }
            return dphi_dnu;
          };
          auto compute_dphi_dnu5 = [&S,&normal,&compute_AxB,&compute_v,&B] () {
            const Long Nelem = S.NElem();
            const Long Nnodes = ElemBasis::Size();

            Vector<ElemBasis> nxB = compute_AxB(normal,B);

            Vector<ElemBasis> dGv;
            Vector<ElemBasis> v = compute_v(S.GetElemList().ElemVector());
            EvalQuadrature(dGv, S.quadrature_FxdU, S, v, S.Laplace_FxdU);

            Vector<ElemBasis> dphi_dnu(Nelem);
            for (Long i = 0; i < Nelem; i++) {
              for (Long j = 0; j < Nnodes; j++) {
                Real dphi_dnu_ = 0;
                dphi_dnu_ += -normal[i*COORD_DIM+0][j] * dGv[(i*COORD_DIM+0)*COORD_DIM+0][j] * nxB[i*COORD_DIM+0][j];
                dphi_dnu_ += -normal[i*COORD_DIM+1][j] * dGv[(i*COORD_DIM+0)*COORD_DIM+1][j] * nxB[i*COORD_DIM+0][j];
                dphi_dnu_ += -normal[i*COORD_DIM+2][j] * dGv[(i*COORD_DIM+0)*COORD_DIM+2][j] * nxB[i*COORD_DIM+0][j];

                dphi_dnu_ += -normal[i*COORD_DIM+0][j] * dGv[(i*COORD_DIM+1)*COORD_DIM+0][j] * nxB[i*COORD_DIM+1][j];
                dphi_dnu_ += -normal[i*COORD_DIM+1][j] * dGv[(i*COORD_DIM+1)*COORD_DIM+1][j] * nxB[i*COORD_DIM+1][j];
                dphi_dnu_ += -normal[i*COORD_DIM+2][j] * dGv[(i*COORD_DIM+1)*COORD_DIM+2][j] * nxB[i*COORD_DIM+1][j];

                dphi_dnu_ += -normal[i*COORD_DIM+0][j] * dGv[(i*COORD_DIM+2)*COORD_DIM+0][j] * nxB[i*COORD_DIM+2][j];
                dphi_dnu_ += -normal[i*COORD_DIM+1][j] * dGv[(i*COORD_DIM+2)*COORD_DIM+1][j] * nxB[i*COORD_DIM+2][j];
                dphi_dnu_ += -normal[i*COORD_DIM+2][j] * dGv[(i*COORD_DIM+2)*COORD_DIM+2][j] * nxB[i*COORD_DIM+2][j];
                dphi_dnu[i][j] = dphi_dnu_;
              }
            }
            return dphi_dnu;
          };
          auto compute_dphi_dnu6 = [&S,&normal,&compute_AxB,&compute_v,&gradB] () {
            const Long Nelem = S.NElem();
            const Long Nnodes = ElemBasis::Size();

            Vector<ElemBasis> Gv;
            Vector<ElemBasis> v = compute_v(S.GetElemList().ElemVector());
            EvalQuadrature(Gv, S.quadrature_FxU, S, v, S.Laplace_FxU);
            Vector<ElemBasis> nxGv = compute_AxB(Gv,normal);

            Vector<ElemBasis> dphi_dnu(Nelem);
            for (Long i = 0; i < Nelem; i++) {
              for (Long j = 0; j < Nnodes; j++) {
                Real dphi_dnu_ = 0;
                dphi_dnu_ += -nxGv[i*COORD_DIM+0][j] * gradB[(i*COORD_DIM+0)*COORD_DIM+0][j] * normal[i*COORD_DIM+0][j];
                dphi_dnu_ += -nxGv[i*COORD_DIM+1][j] * gradB[(i*COORD_DIM+0)*COORD_DIM+1][j] * normal[i*COORD_DIM+0][j];
                dphi_dnu_ += -nxGv[i*COORD_DIM+2][j] * gradB[(i*COORD_DIM+0)*COORD_DIM+2][j] * normal[i*COORD_DIM+0][j];

                dphi_dnu_ += -nxGv[i*COORD_DIM+0][j] * gradB[(i*COORD_DIM+1)*COORD_DIM+0][j] * normal[i*COORD_DIM+1][j];
                dphi_dnu_ += -nxGv[i*COORD_DIM+1][j] * gradB[(i*COORD_DIM+1)*COORD_DIM+1][j] * normal[i*COORD_DIM+1][j];
                dphi_dnu_ += -nxGv[i*COORD_DIM+2][j] * gradB[(i*COORD_DIM+1)*COORD_DIM+2][j] * normal[i*COORD_DIM+1][j];

                dphi_dnu_ += -nxGv[i*COORD_DIM+0][j] * gradB[(i*COORD_DIM+2)*COORD_DIM+0][j] * normal[i*COORD_DIM+2][j];
                dphi_dnu_ += -nxGv[i*COORD_DIM+1][j] * gradB[(i*COORD_DIM+2)*COORD_DIM+1][j] * normal[i*COORD_DIM+2][j];
                dphi_dnu_ += -nxGv[i*COORD_DIM+2][j] * gradB[(i*COORD_DIM+2)*COORD_DIM+2][j] * normal[i*COORD_DIM+2][j];
                dphi_dnu[i][j] = dphi_dnu_;
              }
            }
            return dphi_dnu;
          };
          auto compute_dphi_dnu7 = [&S,&normal,&H,&compute_AxB,&compute_v,&sigma] () {
            const Long Nelem = S.NElem();
            const Long Nnodes = ElemBasis::Size();

            Vector<ElemBasis> Gv;
            Vector<ElemBasis> v = compute_v(S.GetElemList().ElemVector());
            EvalQuadrature(Gv, S.quadrature_FxU, S, v, S.Laplace_FxU);
            Vector<ElemBasis> nxGv = compute_AxB(Gv,normal);

            Vector<ElemBasis> dphi_dnu(Nelem);
            EvalQuadrature(dphi_dnu, S.quadrature_dUxF, S, nxGv, S.Laplace_dUxF);
            for (Long i = 0; i < Nelem; i++) {
              for (Long j = 0; j < Nnodes; j++) {
                dphi_dnu[i][j] *= -2*H[i][j] * sigma[i][j];
              }
            }
            return dphi_dnu;
          };
          auto compute_dphi_dnu8 = [&S,&normal,&H,&compute_AxB,&compute_v,&sigma] () {
            const Long Nelem = S.NElem();
            const Long Nnodes = ElemBasis::Size();

            Vector<ElemBasis> Gv;
            Vector<ElemBasis> v = compute_v(S.GetElemList().ElemVector());
            EvalQuadrature(Gv, S.quadrature_FxU, S, v, S.Laplace_FxU);
            Vector<ElemBasis> nxGv = compute_AxB(Gv,normal);

            Vector<ElemBasis> dphi_dnu(Nelem);
            Vector<ElemBasis> nablaDt_nxGv;
            EvalQuadrature(nablaDt_nxGv, S.quadrature_dUxD, S, nxGv, S.Laplace_dUxD);
            for (Long i = 0; i < Nelem; i++) {
              for (Long j = 0; j < Nnodes; j++) {
                dphi_dnu[i][j] = 0;
                dphi_dnu[i][j] += sigma[i][j] * nablaDt_nxGv[i*COORD_DIM+0][j]*normal[i*COORD_DIM+0][j];
                dphi_dnu[i][j] += sigma[i][j] * nablaDt_nxGv[i*COORD_DIM+1][j]*normal[i*COORD_DIM+1][j];
                dphi_dnu[i][j] += sigma[i][j] * nablaDt_nxGv[i*COORD_DIM+2][j]*normal[i*COORD_DIM+2][j];
              }
            }
            return dphi_dnu;
          };

          auto dphi_dnu0 = compute_dphi_dnu0();
          auto dphi_dnu1 = compute_dphi_dnu1();
          auto dphi_dnu2 = compute_dphi_dnu2();
          auto dphi_dnu3 = compute_dphi_dnu3();
          auto dphi_dnu4 = compute_dphi_dnu4();
          auto dphi_dnu5 = compute_dphi_dnu5();
          auto dphi_dnu6 = compute_dphi_dnu6();
          auto dphi_dnu7 = compute_dphi_dnu7();
          auto dphi_dnu8 = compute_dphi_dnu8();

          return (dphi_dnu0+dphi_dnu1+dphi_dnu2+dphi_dnu3+dphi_dnu4+dphi_dnu5+dphi_dnu6+dphi_dnu7+dphi_dnu8);
        };

        { // Set dg_dnu -= dg_dsigma invA dA_dnu sigma
          dg_dnu -= compute_u_dAdnu_v_0(dg_dsigma_invA, sigma, alpha, beta);
          if (S.Nsurf() >= 1) dg_dnu -= compute_u_dAdnu_v_1(sigma, alpha, beta,  true) * dg_dsigma_invA[Nelem*Nnodes+0];
          if (S.Nsurf() >= 2) dg_dnu -= compute_u_dAdnu_v_1(sigma, alpha, beta, false) * dg_dsigma_invA[Nelem*Nnodes+1];
        }
        return dg_dnu;
      };
      Vector<ElemBasis> dgdnu;
      { // Set dgdnu
        dgdnu.ReInit(S_.NElem());
        dgdnu = 0;
        for (Long i = 0; i < S_.Nsurf(); i++) {
          const Long elem_dsp = (i==0 ? 0 : S_.ElemDsp(i-1));
          const Long Nnodes = ElemBasis::Size();
          auto dgdnu_ = compute_gradient(Svec[i]);
          if (0) { // Write VTU
            VTUData vtu;
            vtu.AddElems(Svec[i].GetElemList(), dgdnu_, ORDER);
            vtu.WriteVTK("dgdnu-"+std::to_string(i), comm);
          }
          for (Long j = 0; j < (i==0?0:Svec[i].NTor(0)*Svec[i].NPol(0)); j++) {
            for (Long k = 0; k < Nnodes; k++) {
              dgdnu[elem_dsp+j][k] -= dgdnu_[j][k];
            }
          }
          for (Long j = (i==0?0:Svec[i].NTor(0)*Svec[i].NPol(0)); j < dgdnu_.Dim(); j++) {
            for (Long k = 0; k < Nnodes; k++) {
              dgdnu[elem_dsp+j][k] += dgdnu_[j][k];
            }
          }
        }
      }
      return dgdnu;
    }

    static Vector<ElemBasis> compute_pressure_jump(const Stellarator<Real,ORDER>& S_, const Vector<Real>& pressure, const Vector<Real>& flux_tor_, const Vector<Real>& flux_pol_, Real* g_ptr = nullptr) {
      Comm comm = Comm::World();

      Vector<Stellarator<Real,ORDER>> Svec(S_.Nsurf());
      for (Long i = 0; i < S_.Nsurf(); i++) { // Set Svec[i] (quadratures, B)
        const Long elem_dsp = (i==0 ? 0 : S_.ElemDsp(i-1));
        const Long Nnodes = ElemBasis::Size();
        Stellarator<Real,ORDER>& S = Svec[i];
        if (i == 0) { // Init S
          Vector<Long> NtNp;
          NtNp.PushBack(S_.NTor(i));
          NtNp.PushBack(S_.NPol(i));
          S = Stellarator<Real,ORDER>(NtNp);
        } else {
          Vector<Long> NtNp;
          NtNp.PushBack(S_.NTor(i-1));
          NtNp.PushBack(S_.NPol(i-1));
          NtNp.PushBack(S_.NTor(i));
          NtNp.PushBack(S_.NPol(i));
          S = Stellarator<Real,ORDER>(NtNp);
        }
        for (Long j = 0; j < S.NElem(); j++) { // Set S coordinates
          for (Long k = 0; k < Nnodes; k++) {
            S.Elem(j,0)[k] = S_.Elem(elem_dsp+j,0)[k];
            S.Elem(j,1)[k] = S_.Elem(elem_dsp+j,1)[k];
            S.Elem(j,2)[k] = S_.Elem(elem_dsp+j,2)[k];
          }
        }

        SetupQuadrature(S.quadrature_BS   , S, S.BiotSavart    , order_singular, order_direct, -1.0, comm);//, -0.01 * pow<-2,Real>(ORDER));
        SetupQuadrature(S.quadrature_FxU  , S, S.Laplace_FxU   , order_singular, order_direct, -1.0, comm);
        SetupQuadrature(S.quadrature_FxdU , S, S.Laplace_FxdU  , order_singular, order_direct, -1.0, comm);

        { // Set Bt0, Bp0, dBt0, dBp0
          Vector<ElemBasis> Jt, Jp;
          compute_harmonic_vector_potentials(Jt, Jp, S);
          EvalQuadrature(S.Bt0 , S.quadrature_BS , S, Jp, S.BiotSavart);
          EvalQuadrature(S.Bp0 , S.quadrature_BS , S, Jt, S.BiotSavart);

          Vector<ElemBasis> normal, area_elem;
          compute_norm_area_elem(S, normal, area_elem);
          if (S.Nsurf() == 2) {
            Long Nelem0 = S.NTor(0)*S.NPol(0);
            for (Long i = 0; i < Nelem0*COORD_DIM; i++) {
              for (Long j = 0; j < Nnodes; j++) {
                normal[i][j] *= -1.0;
              }
            }
          }

          const Long Nelem = S.NElem();
          const Long Nnodes = ElemBasis::Size();
          for (Long j = 0; j < Nelem; j++) {
            for (Long k = 0; k < Nnodes; k++) {
              Real Jxn[COORD_DIM];
              Jxn[0] = Jp[j*COORD_DIM+1][k] * normal[j*COORD_DIM+2][k] - Jp[j*COORD_DIM+2][k] * normal[j*COORD_DIM+1][k];
              Jxn[1] = Jp[j*COORD_DIM+2][k] * normal[j*COORD_DIM+0][k] - Jp[j*COORD_DIM+0][k] * normal[j*COORD_DIM+2][k];
              Jxn[2] = Jp[j*COORD_DIM+0][k] * normal[j*COORD_DIM+1][k] - Jp[j*COORD_DIM+1][k] * normal[j*COORD_DIM+0][k];
              S.Bt0[j*COORD_DIM+0][k] += 0.5 * Jxn[0];
              S.Bt0[j*COORD_DIM+1][k] += 0.5 * Jxn[1];
              S.Bt0[j*COORD_DIM+2][k] += 0.5 * Jxn[2];

              Jxn[0] = Jt[j*COORD_DIM+1][k] * normal[j*COORD_DIM+2][k] - Jt[j*COORD_DIM+2][k] * normal[j*COORD_DIM+1][k];
              Jxn[1] = Jt[j*COORD_DIM+2][k] * normal[j*COORD_DIM+0][k] - Jt[j*COORD_DIM+0][k] * normal[j*COORD_DIM+2][k];
              Jxn[2] = Jt[j*COORD_DIM+0][k] * normal[j*COORD_DIM+1][k] - Jt[j*COORD_DIM+1][k] * normal[j*COORD_DIM+0][k];
              S.Bp0[j*COORD_DIM+0][k] += 0.5 * Jxn[0];
              S.Bp0[j*COORD_DIM+1][k] += 0.5 * Jxn[1];
              S.Bp0[j*COORD_DIM+2][k] += 0.5 * Jxn[2];
            }
          }
        }

        compute_invA(S.sigma, S.alpha, S.beta, S, flux_tor_[i], flux_pol_[i], comm);
        S.B = compute_B(S, S.sigma, S.alpha, S.beta);
      }
      if (g_ptr != nullptr) g_ptr[0] = compute_g(Svec, pressure);
      return compute_pressure_jump(Svec, pressure);
    }




    static void test() {
      Comm comm = Comm::World();
      Profile::Enable(true);

      Long Nsurf = 2;
      Stellarator<Real,ORDER> S;
      Vector<Real> flux_tor(Nsurf), flux_pol(Nsurf), pressure(Nsurf);
      { // Init S, flux_tor, flux_pol, pressure
        Vector<Long> NtNp;
        for (Long i = 0; i < Nsurf; i++) {
          NtNp.PushBack(30);
          NtNp.PushBack(4);
        }
        S = Stellarator<Real,ORDER>(NtNp);
        flux_tor = 1;
        flux_pol = 1;
        pressure = 0;

        //flux_tor[0] = 1; //0.791881512;
        //flux_tor[1] = 1;
        //flux_pol[0] = 0;
        //flux_pol[1] = 0;
        //pressure[0] = 0;
        //pressure[1] = 0;
      }

      { // find equilibrium flux surfaces
        {
        //auto filter = [](const Stellarator<Real,ORDER>& S, Vector<ElemBasis>& f) {
        //  auto cheb2grid = [] (const Vector<ElemBasis>& X, Long Mt, Long Mp, Long Nt, Long Np) {
        //    const Long dof = X.Dim() / (Mt * Mp);
        //    SCTL_ASSERT(X.Dim() == Mt * Mp *dof);

        //    Vector<Real> Xf(dof*Nt*Np); Xf = 0;
        //    const Long Nnodes = ElemBasis::Size();
        //    const Matrix<Real>& Mnodes = Basis<Real,1,ORDER>::Nodes();
        //    for (Long t = 0; t < Nt; t++) {
        //      for (Long p = 0; p < Np; p++) {
        //        Real theta = t / (Real)Nt;
        //        Real phi   = p / (Real)Np;
        //        Long i = (Long)(theta * Mt);
        //        Long j = (Long)(phi   * Mp);
        //        Real x = theta * Mt - i;
        //        Real y = phi   * Mp - j;
        //        Long elem_idx = i * Mp + j;

        //        Vector<Real> Interp0(ORDER);
        //        Vector<Real> Interp1(ORDER);
        //        { // Set Interp0, Interp1
        //          auto node = [&Mnodes] (Long i) {
        //            return Mnodes[0][i];
        //          };
        //          for (Long i = 0; i < ORDER; i++) {
        //            Real wt_x = 1, wt_y = 1;
        //            for (Long j = 0; j < ORDER; j++) {
        //              if (j != i) {
        //                wt_x *= (x - node(j)) / (node(i) - node(j));
        //                wt_y *= (y - node(j)) / (node(i) - node(j));
        //              }
        //              Interp0[i] = wt_x;
        //              Interp1[i] = wt_y;
        //            }
        //          }
        //        }

        //        for (Long ii = 0; ii < ORDER; ii++) {
        //          for (Long jj = 0; jj < ORDER; jj++) {
        //            Long node_idx = jj * ORDER + ii;
        //            for (Long k = 0; k < dof; k++) {
        //              Xf[(k*Nt+t)*Np+p] += X[elem_idx*dof+k][node_idx] * Interp0[ii] * Interp1[jj];
        //            }
        //          }
        //        }
        //      }
        //    }
        //    return Xf;
        //  };
        //  auto grid2cheb = [] (const Vector<Real>& Xf, Long Nt, Long Np, Long Mt, Long Mp) {
        //    Long dof = Xf.Dim() / (Nt*Np);
        //    SCTL_ASSERT(Xf.Dim() == dof*Nt*Np);
        //    Vector<ElemBasis> X(Mt*Mp*dof);
        //    constexpr Integer INTERP_ORDER = 12;

        //    for (Long tt = 0; tt < Mt; tt++) {
        //      for (Long pp = 0; pp < Mp; pp++) {
        //        for (Long t = 0; t < ORDER; t++) {
        //          for (Long p = 0; p < ORDER; p++) {
        //            Matrix<Real> Mnodes = Basis<Real,1,ORDER>::Nodes();
        //            Real theta = (tt + Mnodes[0][t]) / Mt;
        //            Real phi   = (pp + Mnodes[0][p]) / Mp;
        //            Long i = (Long)(theta * Nt);
        //            Long j = (Long)(phi   * Np);
        //            Real x = theta * Nt - i;
        //            Real y = phi   * Np - j;

        //            Vector<Real> Interp0(INTERP_ORDER);
        //            Vector<Real> Interp1(INTERP_ORDER);
        //            { // Set Interp0, Interp1
        //              auto node = [] (Long i) {
        //                return (Real)i - (INTERP_ORDER-1)/2;
        //              };
        //              for (Long i = 0; i < INTERP_ORDER; i++) {
        //                Real wt_x = 1, wt_y = 1;
        //                for (Long j = 0; j < INTERP_ORDER; j++) {
        //                  if (j != i) {
        //                    wt_x *= (x - node(j)) / (node(i) - node(j));
        //                    wt_y *= (y - node(j)) / (node(i) - node(j));
        //                  }
        //                  Interp0[i] = wt_x;
        //                  Interp1[i] = wt_y;
        //                }
        //              }
        //            }

        //            for (Long k = 0; k < dof; k++) {
        //              Real X0 = 0;
        //              for (Long ii = 0; ii < INTERP_ORDER; ii++) {
        //                for (Long jj = 0; jj < INTERP_ORDER; jj++) {
        //                  Long idx_i = (i + ii-(INTERP_ORDER-1)/2 + Nt) % Nt;
        //                  Long idx_j = (j + jj-(INTERP_ORDER-1)/2 + Np) % Np;
        //                  X0 += Interp0[ii] * Interp1[jj] * Xf[(k*Nt+idx_i)*Np+idx_j];
        //                }
        //              }
        //              Long elem_idx = tt * Mp + pp;
        //              Long node_idx = p * ORDER + t;
        //              X[elem_idx*dof+k][node_idx] = X0;
        //            }
        //          }
        //        }
        //      }
        //    }
        //    return X;
        //  };

        //  Long dof = f.Dim() / S.NElem();
        //  SCTL_ASSERT(f.Dim() == S.NElem() * dof);
        //  for (Long i = 0; i < S.Nsurf(); i++) {
        //    const Long Mt = S.NTor(i);
        //    const Long Mp = S.NPol(i);
        //    const Long Nelem = Mt * Mp;
        //    const Long offset = S.ElemDsp(i);
        //    const Long Nt = Mt * ORDER / 5;
        //    const Long Np = Mp * ORDER / 5;

        //    Vector<ElemBasis> f_(Nelem*dof, f.begin() + offset*dof, false);
        //    Vector<Real> f_fourier = cheb2grid(f_, Mt, Mp, Nt, Np);
        //    f_ = grid2cheb(f_fourier, Nt, Np, Mt, Mp);
        //  }
        //};
        }

        auto filter = [](const Stellarator<Real,ORDER>& S, const Comm& comm, Vector<ElemBasis>& f, Real sigma) {
          auto cheb2grid = [] (const Vector<ElemBasis>& X, Long Mt, Long Mp, Long Nt, Long Np) {
            const Long dof = X.Dim() / (Mt * Mp);
            SCTL_ASSERT(X.Dim() == Mt * Mp *dof);

            Vector<Real> Xf(dof*Nt*Np); Xf = 0;
            const Long Nnodes = ElemBasis::Size();
            const Matrix<Real>& Mnodes = Basis<Real,1,ORDER>::Nodes();
            for (Long t = 0; t < Nt; t++) {
              for (Long p = 0; p < Np; p++) {
                Real theta = t / (Real)Nt;
                Real phi   = p / (Real)Np;
                Long i = (Long)(theta * Mt);
                Long j = (Long)(phi   * Mp);
                Real x = theta * Mt - i;
                Real y = phi   * Mp - j;
                Long elem_idx = i * Mp + j;

                Vector<Real> Interp0(ORDER);
                Vector<Real> Interp1(ORDER);
                { // Set Interp0, Interp1
                  auto node = [&Mnodes] (Long i) {
                    return Mnodes[0][i];
                  };
                  for (Long i = 0; i < ORDER; i++) {
                    Real wt_x = 1, wt_y = 1;
                    for (Long j = 0; j < ORDER; j++) {
                      if (j != i) {
                        wt_x *= (x - node(j)) / (node(i) - node(j));
                        wt_y *= (y - node(j)) / (node(i) - node(j));
                      }
                      Interp0[i] = wt_x;
                      Interp1[i] = wt_y;
                    }
                  }
                }

                for (Long ii = 0; ii < ORDER; ii++) {
                  for (Long jj = 0; jj < ORDER; jj++) {
                    Long node_idx = jj * ORDER + ii;
                    for (Long k = 0; k < dof; k++) {
                      Xf[(k*Nt+t)*Np+p] += X[elem_idx*dof+k][node_idx] * Interp0[ii] * Interp1[jj];
                    }
                  }
                }
              }
            }
            return Xf;
          };
          auto grid2cheb = [] (const Vector<Real>& Xf, Long Nt, Long Np, Long Mt, Long Mp) {
            Long dof = Xf.Dim() / (Nt*Np);
            SCTL_ASSERT(Xf.Dim() == dof*Nt*Np);
            Vector<ElemBasis> X(Mt*Mp*dof);
            constexpr Integer INTERP_ORDER = 12;

            for (Long tt = 0; tt < Mt; tt++) {
              for (Long pp = 0; pp < Mp; pp++) {
                for (Long t = 0; t < ORDER; t++) {
                  for (Long p = 0; p < ORDER; p++) {
                    Matrix<Real> Mnodes = Basis<Real,1,ORDER>::Nodes();
                    Real theta = (tt + Mnodes[0][t]) / Mt;
                    Real phi   = (pp + Mnodes[0][p]) / Mp;
                    Long i = (Long)(theta * Nt);
                    Long j = (Long)(phi   * Np);
                    Real x = theta * Nt - i;
                    Real y = phi   * Np - j;

                    Vector<Real> Interp0(INTERP_ORDER);
                    Vector<Real> Interp1(INTERP_ORDER);
                    { // Set Interp0, Interp1
                      auto node = [] (Long i) {
                        return (Real)i - (INTERP_ORDER-1)/2;
                      };
                      for (Long i = 0; i < INTERP_ORDER; i++) {
                        Real wt_x = 1, wt_y = 1;
                        for (Long j = 0; j < INTERP_ORDER; j++) {
                          if (j != i) {
                            wt_x *= (x - node(j)) / (node(i) - node(j));
                            wt_y *= (y - node(j)) / (node(i) - node(j));
                          }
                          Interp0[i] = wt_x;
                          Interp1[i] = wt_y;
                        }
                      }
                    }

                    for (Long k = 0; k < dof; k++) {
                      Real X0 = 0;
                      for (Long ii = 0; ii < INTERP_ORDER; ii++) {
                        for (Long jj = 0; jj < INTERP_ORDER; jj++) {
                          Long idx_i = (i + ii-(INTERP_ORDER-1)/2 + Nt) % Nt;
                          Long idx_j = (j + jj-(INTERP_ORDER-1)/2 + Np) % Np;
                          X0 += Interp0[ii] * Interp1[jj] * Xf[(k*Nt+idx_i)*Np+idx_j];
                        }
                      }
                      Long elem_idx = tt * Mp + pp;
                      Long node_idx = p * ORDER + t;
                      X[elem_idx*dof+k][node_idx] = X0;
                    }
                  }
                }
              }
            }
            return X;
          };
          auto fourier_filter = [](Vector<Real>& X, long Nt_, long Np_, Real sigma, const Comm& comm) {
            long dof = X.Dim() / (Nt_ * Np_);
            SCTL_ASSERT(X.Dim() == dof * Nt_ * Np_);

            FFT<Real> fft_r2c, fft_c2r;
            StaticArray<Long, 2> fft_dim = {Nt_, Np_};
            fft_r2c.Setup(FFT_Type::R2C, 1, Vector<Long>(2, fft_dim, false), omp_get_max_threads());
            fft_c2r.Setup(FFT_Type::C2R, 1, Vector<Long>(2, fft_dim, false), omp_get_max_threads());

            long Nt = Nt_;
            long Np = fft_r2c.Dim(1) / (Nt * 2);
            SCTL_ASSERT(fft_r2c.Dim(1) == Nt * Np * 2);

            //auto filter_fn = [](Real x2, Real sigma) {return exp(-x2/(2*sigma*sigma));};
            auto filter_fn = [](Real x2, Real sigma) {return (x2<sigma*sigma?1.0:0.0);};

            Vector<Real> normal, gradX;
            biest::SurfaceOp<Real> op(comm, Nt_, Np_);
            Vector<Real> coeff(fft_r2c.Dim(1));
            for (long k = 0; k < dof; k++) {
              Vector<Real> X_(Nt_*Np_, X.begin() + k*Nt_*Np_, false);
              fft_r2c.Execute(X_, coeff);
              for (long t = 0; t < Nt; t++) {
                for (long p = 0; p < Np; p++) {
                  Real tt = (t - (t > Nt / 2 ? Nt : 0)) / (Real)(Nt / 2);
                  Real pp = p / (Real)Np;
                  Real f = filter_fn(tt*tt+pp*pp, sigma);
                  coeff[(t * Np + p) * 2 + 0] *= f;
                  coeff[(t * Np + p) * 2 + 1] *= f;
                }
              }
              fft_c2r.Execute(coeff, X_);
            }
          };

          Long dof = f.Dim() / S.NElem();
          SCTL_ASSERT(f.Dim() == S.NElem() * dof);
          for (Long i = 0; i < S.Nsurf(); i++) {
            const Long Mt = S.NTor(i);
            const Long Mp = S.NPol(i);
            const Long Nelem = Mt * Mp;
            const Long offset = S.ElemDsp(i);
            const Long Nt = Mt * ORDER * 4;
            const Long Np = Mp * ORDER * 4;

            Vector<ElemBasis> f_(Nelem*dof, f.begin() + offset*dof, false);
            Vector<Real> f_fourier = cheb2grid(f_, Mt, Mp, Nt, Np);
            fourier_filter(f_fourier, Nt, Np, 0.25 * sigma, comm);
            f_ = grid2cheb(f_fourier, Nt, Np, Mt, Mp);
          }
        };

        Long iter = 0;
        Real dt = 0.1;
        while (1) { // time-step
          Vector<ElemBasis> dgdnu = compute_gradient(S, pressure, flux_tor, flux_pol)*(-1);
          //Vector<ElemBasis> dgdnu = compute_pressure_jump(S, pressure, flux_tor, flux_pol)*(-1);

          Vector<ElemBasis> dXdt(dgdnu.Dim()*COORD_DIM);
          { // Set dXdt
            dXdt = 0;
            const Long Nnodes = ElemBasis::Size();
            Vector<ElemBasis> normal, area_elem;
            compute_norm_area_elem(S, normal, area_elem);
            for (Long i = 0; i < S.ElemDsp(S.Nsurf()-1); i++) {
              for (Long j = 0; j < Nnodes; j++) {
                dXdt[i*COORD_DIM+0][j] = normal[i*COORD_DIM+0][j] * dgdnu[i][j];
                dXdt[i*COORD_DIM+1][j] = normal[i*COORD_DIM+1][j] * dgdnu[i][j];
                dXdt[i*COORD_DIM+2][j] = normal[i*COORD_DIM+2][j] * dgdnu[i][j];
              }
            }
            filter(S, comm, dXdt, 0.1);
          }
          { // Update dt
            const Long Nelem = S.NElem();
            Stellarator<Real,ORDER> S0 = S, S1 = S, S2 = S;
            for (Long i = 0; i < S.NElem(); i++) {
              S0.Elem(i, 0) += dXdt[i*COORD_DIM+0] * 0.0 * dt;
              S0.Elem(i, 1) += dXdt[i*COORD_DIM+1] * 0.0 * dt;
              S0.Elem(i, 2) += dXdt[i*COORD_DIM+2] * 0.0 * dt;

              S1.Elem(i, 0) += dXdt[i*COORD_DIM+0] * 0.5 * dt;
              S1.Elem(i, 1) += dXdt[i*COORD_DIM+1] * 0.5 * dt;
              S1.Elem(i, 2) += dXdt[i*COORD_DIM+2] * 0.5 * dt;

              S2.Elem(i, 0) += dXdt[i*COORD_DIM+0] * 1.0 * dt;
              S2.Elem(i, 1) += dXdt[i*COORD_DIM+1] * 1.0 * dt;
              S2.Elem(i, 2) += dXdt[i*COORD_DIM+2] * 1.0 * dt;
            }
            Real g0, g1, g2;
            compute_pressure_jump(S0, pressure, flux_tor, flux_pol, &g0);
            compute_pressure_jump(S1, pressure, flux_tor, flux_pol, &g1);
            compute_pressure_jump(S2, pressure, flux_tor, flux_pol, &g2);
            { // Calculate optimal step size dt
              Real a = 2*g0 - 4*g1 + 2*g2;
              Real b =-3*g0 + 4*g1 -   g2;
              Real c = g0;
              Real s = -b/(2*a);
              dt *= s;

              Real g_ = a*s*s + b*s + c;
              std::cout<<"g = "<<g_<<' ';
              std::cout<<g0<<' ';
              std::cout<<g1<<' ';
              std::cout<<g2<<' ';
              std::cout<<dt<<'\n';
            }
          }

          { // Write VTU
            VTUData vtu;
            vtu.AddElems(S.GetElemList(), dgdnu*dt, ORDER);
            vtu.WriteVTK("dgdnu"+std::to_string(iter), comm);
          }
          { // Write VTU
            VTUData vtu;
            vtu.AddElems(S.GetElemList(), dXdt*dt, ORDER);
            vtu.WriteVTK("dXdt"+std::to_string(iter), comm);
          }
          { // Write VTU
            Vector<ElemBasis> pressure_jump = compute_pressure_jump(S, pressure, flux_tor, flux_pol);
            VTUData vtu;
            vtu.AddElems(S.GetElemList(), pressure_jump, ORDER);
            vtu.WriteVTK("pressure_jump"+std::to_string(iter), comm);
          }

          { // Update S <-- filter(S + dXdt * dt)
            const Long Nelem = S.NElem();
            Vector<ElemBasis> X(Nelem*COORD_DIM);
            for (Long i = 0; i < S.NElem(); i++) {
              X[i*COORD_DIM+0] = S.Elem(i, 0) + dXdt[i*COORD_DIM+0] * dt;
              X[i*COORD_DIM+1] = S.Elem(i, 1) + dXdt[i*COORD_DIM+1] * dt;
              X[i*COORD_DIM+2] = S.Elem(i, 2) + dXdt[i*COORD_DIM+2] * dt;
            }
            filter(S, comm, X, 0.3);
            for (Long i = 0; i < S.NElem(); i++) {
              S.Elem(i, 0) = X[i*COORD_DIM+0];
              S.Elem(i, 1) = X[i*COORD_DIM+1];
              S.Elem(i, 2) = X[i*COORD_DIM+2];
            }
          }
          iter++;
        }
        return;
      }

      { // Verify using finite difference approximation
        Vector<ElemBasis> dgdnu = compute_gradient(S, pressure, flux_tor, flux_pol);
        { // Write VTU
          VTUData vtu;
          vtu.AddElems(S.GetElemList(), dgdnu, ORDER);
          vtu.WriteVTK("dgdnu", comm);
        }

        Real eps = 1e-4;
        const Long Nnodes = ElemBasis::Size();
        Vector<ElemBasis> normal, area_elem;
        compute_norm_area_elem(S, normal, area_elem);

        Vector<ElemBasis> nu = area_elem;
        for (Long i = S.ElemDsp(S.Nsurf()-1); i < S.NElem(); i++) nu[i] = 0;

        Stellarator<Real,ORDER> S0 = S, S1 = S;
        for (Long i = 0; i < S.NElem(); i++) {
          for (Long j = 0; j < Nnodes; j++) {
            S0.Elem(i, 0)[j] -= 0.5 * eps * normal[i*COORD_DIM+0][j] * nu[i][j];
            S0.Elem(i, 1)[j] -= 0.5 * eps * normal[i*COORD_DIM+1][j] * nu[i][j];
            S0.Elem(i, 2)[j] -= 0.5 * eps * normal[i*COORD_DIM+2][j] * nu[i][j];

            S1.Elem(i, 0)[j] += 0.5 * eps * normal[i*COORD_DIM+0][j] * nu[i][j];
            S1.Elem(i, 1)[j] += 0.5 * eps * normal[i*COORD_DIM+1][j] * nu[i][j];
            S1.Elem(i, 2)[j] += 0.5 * eps * normal[i*COORD_DIM+2][j] * nu[i][j];
          }
        }
        Real g0, g1;
        compute_pressure_jump(S0, pressure, flux_tor, flux_pol, &g0);
        compute_pressure_jump(S1, pressure, flux_tor, flux_pol, &g1);
        std::cout<<"g0 = "<<g0<<";  g1 = "<<g1<<";  dgdnu_ = "<<(g1-g0)/eps<<'\n';
        std::cout<<"dgdnu  = "<<compute_inner_prod(area_elem, dgdnu, nu)<<'\n';
      }
    }

    static void test_() {
      Comm comm = Comm::World();
      Profile::Enable(true);

      Real flux_tor = 1.0, flux_pol = 1.0;
      Stellarator<Real,ORDER> S;
      { // Init S
        Vector<Long> NtNp;
        NtNp.PushBack(20);
        NtNp.PushBack(4);
        //NtNp.PushBack(20);
        //NtNp.PushBack(4);
        S = Stellarator<Real,ORDER>(NtNp);
      }
      if (S.Nsurf() == 1) flux_pol = 0.0;

      Vector<ElemBasis> pressure;
      { // Set pressure
        Vector<ElemBasis> normal, area_elem;
        compute_norm_area_elem(S, normal, area_elem);
        pressure = area_elem*0;
      }

      /////////////////////////////////////////////////////////////////////////////////////////////////////////////////

      SetupQuadrature(S.quadrature_BS  , S, S.BiotSavart  , order_singular, order_direct, -1.0, comm, -0.01 * pow<-2,Real>(ORDER));
      SetupQuadrature(S.quadrature_FxU , S, S.Laplace_FxU , order_singular, order_direct, -1.0, comm);
      SetupQuadrature(S.quadrature_FxdU, S, S.Laplace_FxdU, order_singular, order_direct, -1.0, comm);
      SetupQuadrature(S.quadrature_dUxF, S, S.Laplace_dUxF, order_singular, order_direct, -1.0, comm);

      Vector<ElemBasis> Bt0, Bp0;
      { // Set Bt0, Bp0
        Vector<ElemBasis> Jt, Jp;
        compute_harmonic_vector_potentials(Jt, Jp, S);
        EvalQuadrature(Bt0, S.quadrature_BS, S, Jp, S.BiotSavart);
        EvalQuadrature(Bp0, S.quadrature_BS, S, Jt, S.BiotSavart);
      }

      auto compute_B = [&S,&Bt0,&Bp0] (const Vector<ElemBasis>& sigma, Real alpha, Real beta) {
        const Long Nelem = S.NElem();
        Vector<ElemBasis> B(S.NElem() * COORD_DIM);
        if (sigma.Dim()) {
          const Long Nnodes = ElemBasis::Size();
          Vector<ElemBasis> normal, area_elem;
          compute_norm_area_elem(S, normal, area_elem);
          EvalQuadrature(B, S.quadrature_FxdU, S, sigma, S.Laplace_FxdU);
          for (Long i = 0; i < Nelem; i++) {
            for (Long j = 0; j < Nnodes; j++) {
              for (Long k = 0; k < COORD_DIM; k++) {
                B[i*COORD_DIM+k][j] -= 0.5*sigma[i][j]*normal[i*COORD_DIM+k][j];
              }
            }
          }
        } else {
          B = 0;
        }
        if (S.Nsurf() >= 1) B += Bt0*alpha;
        if (S.Nsurf() >= 2) B += Bp0*beta;
        return B;
      };
      auto compute_flux = [&S] (Real& flux_tor, Real& flux_pol, const Vector<ElemBasis>& B, const Vector<ElemBasis>& normal) {
        const Long Nelem = S.NElem();
        const Long Nnodes = ElemBasis::Size();
        SCTL_ASSERT(B.Dim() == Nelem*COORD_DIM);
        SCTL_ASSERT(normal.Dim() == Nelem*COORD_DIM);

        Vector<ElemBasis> J(Nelem * COORD_DIM);
        for (Long i = 0; i < Nelem; i++) { // Set J
          for (Long j = 0; j < Nnodes; j++) {
            Tensor<Real,true,COORD_DIM> b, n;
            b(0) = B[i*COORD_DIM+0][j];
            b(1) = B[i*COORD_DIM+1][j];
            b(2) = B[i*COORD_DIM+2][j];

            n(0) = normal[i*COORD_DIM+0][j];
            n(1) = normal[i*COORD_DIM+1][j];
            n(2) = normal[i*COORD_DIM+2][j];

            J[i*COORD_DIM+0][j] = n(1) * b(2) - n(2) * b(1);
            J[i*COORD_DIM+1][j] = n(2) * b(0) - n(0) * b(2);
            J[i*COORD_DIM+2][j] = n(0) * b(1) - n(1) * b(0);
          }
        }

        Vector<ElemBasis> A;
        EvalQuadrature(A, S.quadrature_FxU, S, J, S.Laplace_FxU);

        Vector<Real> circ_pol(S.Nsurf()), circ_tor(S.Nsurf());
        { // compute circ
          Vector<ElemBasis> dX;
          ElemBasis::Grad(dX, S.GetElemList().ElemVector());
          const auto& quad_wts = ElemBasis::QuadWts();
          for (Long k = 0; k < S.Nsurf(); k++) {
            circ_pol[k] = 0;
            circ_tor[k] = 0;
            Long Ndsp = S.ElemDsp(k);
            for (Long i = 0; i < S.NTor(k)*S.NPol(k); i++) {
              for (Long j = 0; j < Nnodes; j++) {
                circ_pol[k] += A[(Ndsp+i)*COORD_DIM+0][j] * dX[(Ndsp+i)*COORD_DIM*2+1][j] * quad_wts[j] / S.NTor(k);
                circ_pol[k] += A[(Ndsp+i)*COORD_DIM+1][j] * dX[(Ndsp+i)*COORD_DIM*2+3][j] * quad_wts[j] / S.NTor(k);
                circ_pol[k] += A[(Ndsp+i)*COORD_DIM+2][j] * dX[(Ndsp+i)*COORD_DIM*2+5][j] * quad_wts[j] / S.NTor(k);

                circ_tor[k] += A[(Ndsp+i)*COORD_DIM+0][j] * dX[(Ndsp+i)*COORD_DIM*2+0][j] * quad_wts[j] / S.NPol(k);
                circ_tor[k] += A[(Ndsp+i)*COORD_DIM+1][j] * dX[(Ndsp+i)*COORD_DIM*2+2][j] * quad_wts[j] / S.NPol(k);
                circ_tor[k] += A[(Ndsp+i)*COORD_DIM+2][j] * dX[(Ndsp+i)*COORD_DIM*2+4][j] * quad_wts[j] / S.NPol(k);
              }
            }
          }
        }
        if (S.Nsurf() == 1) {
          flux_tor = circ_pol[0];
          flux_pol = 0;
        } else if (S.Nsurf() == 2) {
          flux_tor = circ_pol[1] - circ_pol[0];
          flux_pol = circ_tor[0] - circ_tor[1];
        } else {
          SCTL_ASSERT(false);
        }
      };
      auto compute_A = [&S,compute_B,&compute_flux] (const Vector<Real>& x) {
        const Long Nelem = S.NElem();
        const Long Nnodes = ElemBasis::Size();
        SCTL_ASSERT(x.Dim() == Nelem*Nnodes+S.Nsurf());

        Vector<ElemBasis> normal, area_elem;
        compute_norm_area_elem(S, normal, area_elem);

        Vector<ElemBasis> sigma(Nelem);
        for (Long i = 0; i < Nelem; i++) {
          for (Long j = 0; j < Nnodes; j++) {
            sigma[i][j] = x[i*Nnodes+j];
          }
        }
        Real alpha = (S.Nsurf() >= 1 ? x[Nelem*Nnodes + 0] : 0);
        Real beta  = (S.Nsurf() >= 2 ? x[Nelem*Nnodes + 1] : 0);

        Vector<ElemBasis> B = compute_B(sigma, alpha, beta);
        Vector<ElemBasis> BdotN = compute_dot_prod(B, normal);

        Real flux_tor, flux_pol;
        compute_flux(flux_tor, flux_pol, B, normal);

        Vector<Real> Ax(Nelem*Nnodes+S.Nsurf());
        for (Long i = 0; i < Nelem; i++) {
          for (Long j = 0; j < Nnodes; j++) {
            Ax[i*Nnodes+j] = BdotN[i][j];
          }
        }
        if (S.Nsurf() >= 1) Ax[Nelem*Nnodes + 0] = flux_tor;
        if (S.Nsurf() >= 2) Ax[Nelem*Nnodes + 1] = flux_pol;
        return Ax;
      };
      auto compute_invA = [&S,&comm,&compute_A] (Vector<ElemBasis>& sigma, Real& alpha, Real& beta, Real flux_tor, Real flux_pol) {
        typename ParallelSolver<Real>::ParallelOp BIOp = [&compute_A](Vector<Real>* Ax, const Vector<Real>& x) {
          (*Ax) = compute_A(x);
        };
        const Long Nelem = S.NElem();
        const Long Nnodes = ElemBasis::Size();

        Vector<Real> rhs_(Nelem * Nnodes + S.Nsurf());
        rhs_ = 0;
        if (S.Nsurf() >= 1) rhs_[Nelem * Nnodes + 0] = flux_tor;
        if (S.Nsurf() >= 2) rhs_[Nelem * Nnodes + 1] = flux_pol;

        Vector<Real> x_(Nelem * Nnodes + S.Nsurf());
        x_ = 0;
        ParallelSolver<Real> linear_solver(comm, true);
        linear_solver(&x_, BIOp, rhs_, 1e-8, 100);

        sigma.ReInit(Nelem);
        for (Long i = 0; i < Nelem; i++) {
          for (Long j = 0; j < Nnodes; j++) {
            sigma[i][j] = x_[i*Nnodes+j];
          }
        }
        alpha = (S.Nsurf() >= 1 ? x_[Nelem * Nnodes + 0] : 0);
        beta  = (S.Nsurf() >= 2 ? x_[Nelem * Nnodes + 1] : 0);
      };

      Vector<ElemBasis> dg_dnu = compute_gradient(S, pressure, flux_tor, flux_pol);
      { // Write VTU
        VTUData vtu;
        vtu.AddElems(S.GetElemList(), dg_dnu, ORDER);
        vtu.WriteVTK("dg_dnu", comm);
      }

      if (1) { // test grad_g
        auto compute_g = [&S,&Bt0,&Bp0,&compute_B,&compute_invA,&comm] (const Vector<ElemBasis>& nu, Real eps, Real flux_tor, Real flux_pol, const Vector<ElemBasis>& pressure) {
          const Long Nelem = S.NElem();
          const Long Nnodes = ElemBasis::Size();

          Vector<ElemBasis> normal, area_elem;
          compute_norm_area_elem(S, normal, area_elem);

          Vector<ElemBasis> X_orig(Nelem*COORD_DIM);
          for (Long i = 0; i < Nelem; i++) {
            for (Long j = 0; j < Nnodes; j++) {
              X_orig[i*COORD_DIM+0][j] = S.Elem(i,0)[j];
              X_orig[i*COORD_DIM+1][j] = S.Elem(i,1)[j];
              X_orig[i*COORD_DIM+2][j] = S.Elem(i,2)[j];
              S.Elem(i,0)[j] += eps*nu[i][j] * normal[i*COORD_DIM+0][j];
              S.Elem(i,1)[j] += eps*nu[i][j] * normal[i*COORD_DIM+1][j];
              S.Elem(i,2)[j] += eps*nu[i][j] * normal[i*COORD_DIM+2][j];
            }
          }
          /////////////////////////////////////////////////////////////////////////////////////////////////////////////

          SetupQuadrature(S.quadrature_BS  , S, S.BiotSavart  , order_singular, order_direct, -1.0, comm, -0.01 * pow<-2,Real>(ORDER));
          SetupQuadrature(S.quadrature_FxU , S, S.Laplace_FxU , order_singular, order_direct, -1.0, comm);
          SetupQuadrature(S.quadrature_FxdU, S, S.Laplace_FxdU, order_singular, order_direct, -1.0, comm);

          Vector<ElemBasis> Jt, Jp;
          compute_harmonic_vector_potentials(Jt, Jp, S);
          EvalQuadrature(Bt0, S.quadrature_BS, S, Jp, S.BiotSavart);
          EvalQuadrature(Bp0, S.quadrature_BS, S, Jt, S.BiotSavart);

          Real alpha, beta;
          Vector<ElemBasis> sigma;
          compute_invA(sigma, alpha, beta, flux_tor, flux_pol);
          Vector<ElemBasis> B = compute_B(sigma, alpha, beta);

          compute_norm_area_elem(S, normal, area_elem);
          Real g = compute_inner_prod(area_elem, compute_gvec(S,B,pressure), area_elem*0+1);

          /////////////////////////////////////////////////////////////////////////////////////////////////////////////
          for (Long i = 0; i < Nelem; i++) {
            for (Long j = 0; j < Nnodes; j++) {
              S.Elem(i,0)[j] = X_orig[i*COORD_DIM+0][j];
              S.Elem(i,1)[j] = X_orig[i*COORD_DIM+1][j];
              S.Elem(i,2)[j] = X_orig[i*COORD_DIM+2][j];
            }
          }

          return g;
        };

        Vector<ElemBasis> normal, area_elem;
        compute_norm_area_elem(S, normal, area_elem);
        const Long Nelem = S.NElem();
        {
          Vector<ElemBasis> nu(Nelem);
          nu = area_elem;

          Real eps = 1e-4;
          Real g0 = compute_g(nu,-eps, flux_tor, flux_pol, pressure);
          Real g1 = compute_g(nu,eps, flux_tor, flux_pol, pressure);
          std::cout<<"g = "<<g0<<"  g = "<<g1<<"  dg_dnu = "<<(g1-g0)/(2*eps)<<'\n';
          std::cout<<"dg_dnu = "<<compute_inner_prod(area_elem,nu, dg_dnu)<<'\n';
        }
        {
          Vector<ElemBasis> nu(Nelem);
          nu = 1;

          Real eps = 1e-4;
          Real g0 = compute_g(nu,-eps, flux_tor, flux_pol, pressure);
          Real g1 = compute_g(nu,eps, flux_tor, flux_pol, pressure);
          std::cout<<"g = "<<g0<<"  g = "<<g1<<"  dg_dnu = "<<(g1-g0)/(2*eps)<<'\n';
          std::cout<<"dg_dnu = "<<compute_inner_prod(area_elem,nu, dg_dnu)<<'\n';
        }
        {
          Vector<ElemBasis> nu(Nelem);
          nu = dg_dnu;

          Real eps = 1e-4;
          Real g0 = compute_g(nu,-eps, flux_tor, flux_pol, pressure);
          Real g1 = compute_g(nu,eps, flux_tor, flux_pol, pressure);
          std::cout<<"g = "<<g0<<"  g = "<<g1<<"  dg_dnu = "<<(g1-g0)/(2*eps)<<'\n';
          std::cout<<"dg_dnu = "<<compute_inner_prod(area_elem,nu, dg_dnu)<<'\n';
        }
      }
    }

    static void test_askham() {
      auto Setup = [] (Stellarator<Real,ORDER>& S, const Comm& comm) { // Set quadratures, Bt0, Bp0, ...
        SetupQuadrature(S.quadrature_dBS  , S, S.BiotSavartGrad, order_singular, order_direct, -1.0, comm, -0.01 * pow<-2,Real>(ORDER));
        SetupQuadrature(S.quadrature_BS   , S, S.BiotSavart    , order_singular, order_direct, -1.0, comm, -0.01 * pow<-2,Real>(ORDER));
        SetupQuadrature(S.quadrature_FxU  , S, S.Laplace_FxU   , order_singular, order_direct, -1.0, comm);
        SetupQuadrature(S.quadrature_FxdU , S, S.Laplace_FxdU  , order_singular, order_direct, -1.0, comm);
        SetupQuadrature(S.quadrature_dUxF , S, S.Laplace_dUxF  , order_singular, order_direct, -1.0, comm);
        SetupQuadrature(S.quadrature_dUxD , S, S.Laplace_dUxD  , order_singular, order_direct, -1.0, comm,  0.01 * pow<-2,Real>(ORDER));
        SetupQuadrature(S.quadrature_Fxd2U, S, S.Laplace_Fxd2U , order_singular, order_direct, -1.0, comm, -0.01 * pow<-2,Real>(ORDER));

        { // Set Bt0, Bp0, dBt0, dBp0
          Vector<ElemBasis> Jt, Jp;
          compute_harmonic_vector_potentials(Jt, Jp, S);
          EvalQuadrature(S.Bt0 , S.quadrature_BS , S, Jp, S.BiotSavart);
          EvalQuadrature(S.Bp0 , S.quadrature_BS , S, Jt, S.BiotSavart);
          EvalQuadrature(S.dBt0, S.quadrature_dBS, S, Jp, S.BiotSavartGrad);
          EvalQuadrature(S.dBp0, S.quadrature_dBS, S, Jt, S.BiotSavartGrad);
        }
      };
      auto compute_grad = [] (const Stellarator<Real,ORDER>& S, const Vector<ElemBasis>& V) {
        const Long Nelem = S.GetElemList().NElem();
        const Long Nnodes = ElemBasis::Size();
        const Long dof = V.Dim() / Nelem;
        SCTL_ASSERT(Nelem * dof == V.Dim());

        Vector<ElemBasis> du_dX(Nelem*COORD_DIM*2);
        { // Set du_dX
          Vector<ElemBasis> dX;
          ElemBasis::Grad(dX, S.GetElemList().ElemVector());

          auto inv2x2 = [](Tensor<Real, true, 2, 2> M) {
            Tensor<Real, true, 2, 2> Mout;
            Real oodet = 1 / (M(0,0) * M(1,1) - M(0,1) * M(1,0));
            Mout(0,0) =  M(1,1) * oodet;
            Mout(0,1) = -M(0,1) * oodet;
            Mout(1,0) = -M(1,0) * oodet;
            Mout(1,1) =  M(0,0) * oodet;
            return Mout;
          };
          for (Long i = 0; i < Nelem; i++) {
            for (Long j = 0; j < Nnodes; j++) {
              Tensor<Real, true, 3, 2> dX_du;
              dX_du(0,0) = dX[(i*COORD_DIM+0)*2+0][j];
              dX_du(1,0) = dX[(i*COORD_DIM+1)*2+0][j];
              dX_du(2,0) = dX[(i*COORD_DIM+2)*2+0][j];
              dX_du(0,1) = dX[(i*COORD_DIM+0)*2+1][j];
              dX_du(1,1) = dX[(i*COORD_DIM+1)*2+1][j];
              dX_du(2,1) = dX[(i*COORD_DIM+2)*2+1][j];

              Tensor<Real, true, 2, 2> G; // = dX_du.Transpose() * dX_du;
              G(0,0) = dX_du(0,0) * dX_du(0,0) + dX_du(1,0) * dX_du(1,0) + dX_du(2,0) * dX_du(2,0);
              G(0,1) = dX_du(0,0) * dX_du(0,1) + dX_du(1,0) * dX_du(1,1) + dX_du(2,0) * dX_du(2,1);
              G(1,0) = dX_du(0,1) * dX_du(0,0) + dX_du(1,1) * dX_du(1,0) + dX_du(2,1) * dX_du(2,0);
              G(1,1) = dX_du(0,1) * dX_du(0,1) + dX_du(1,1) * dX_du(1,1) + dX_du(2,1) * dX_du(2,1);

              Tensor<Real, true, 2, 2> Ginv = inv2x2(G);
              du_dX[(i*COORD_DIM+0)*2+0][j] = Ginv(0,0) * dX_du(0,0) + Ginv(0,1) * dX_du(0,1);
              du_dX[(i*COORD_DIM+1)*2+0][j] = Ginv(0,0) * dX_du(1,0) + Ginv(0,1) * dX_du(1,1);
              du_dX[(i*COORD_DIM+2)*2+0][j] = Ginv(0,0) * dX_du(2,0) + Ginv(0,1) * dX_du(2,1);
              du_dX[(i*COORD_DIM+0)*2+1][j] = Ginv(1,0) * dX_du(0,0) + Ginv(1,1) * dX_du(0,1);
              du_dX[(i*COORD_DIM+1)*2+1][j] = Ginv(1,0) * dX_du(1,0) + Ginv(1,1) * dX_du(1,1);
              du_dX[(i*COORD_DIM+2)*2+1][j] = Ginv(1,0) * dX_du(2,0) + Ginv(1,1) * dX_du(2,1);
            }
          }
        }

        Vector<ElemBasis> dV;
        ElemBasis::Grad(dV, V);

        Vector<ElemBasis> gradV(Nelem*dof*COORD_DIM);
        for (Long i = 0; i < Nelem; i++) {
          for (Long j = 0; j < Nnodes; j++) {
            for (Long k = 0; k < dof; k++) {
              gradV[(i*dof+k)*COORD_DIM+0][j] = dV[(i*dof+k)*2+0][j] * du_dX[(i*COORD_DIM+0)*2+0][j] + dV[(i*dof+k)*2+1][j] * du_dX[(i*COORD_DIM+0)*2+1][j];
              gradV[(i*dof+k)*COORD_DIM+1][j] = dV[(i*dof+k)*2+0][j] * du_dX[(i*COORD_DIM+1)*2+0][j] + dV[(i*dof+k)*2+1][j] * du_dX[(i*COORD_DIM+1)*2+1][j];
              gradV[(i*dof+k)*COORD_DIM+2][j] = dV[(i*dof+k)*2+0][j] * du_dX[(i*COORD_DIM+2)*2+0][j] + dV[(i*dof+k)*2+1][j] * du_dX[(i*COORD_DIM+2)*2+1][j];
            }
          }
        }
        return gradV;
      };
      auto compute_surfdiv = [&compute_grad] (const Stellarator<Real,ORDER>& S, const Vector<ElemBasis>& V) {
        const Long Nelem = S.GetElemList().NElem();
        const Long Nnodes = ElemBasis::Size();
        SCTL_ASSERT(V.Dim() == Nelem* COORD_DIM);

        Vector<ElemBasis> gradV = compute_grad(S, V);
        Vector<ElemBasis> divV(Nelem);
        for (Long i = 0; i < Nelem; i++) {
          for (Long j = 0; j < Nnodes; j++) {
            divV[i][j] = gradV[(i*COORD_DIM+0)*COORD_DIM+0][j] + gradV[(i*COORD_DIM+1)*COORD_DIM+1][j] + gradV[(i*COORD_DIM+2)*COORD_DIM+2][j];
          }
        }
        return divV;
      };
      auto compute_g = [](const Stellarator<Real,ORDER>& S, const Vector<ElemBasis>& B) {
        const Long Nelem = S.NElem();
        const Long Nnodes = ElemBasis::Size();
        Vector<ElemBasis> normal, area_elem;
        compute_norm_area_elem(S, normal, area_elem);
        Vector<ElemBasis> B2(Nelem);
        for (Long i = 0; i < Nelem; i++) {
          for (Long j = 0; j < Nnodes; j++) {
            B2[i][j] = 0;
            B2[i][j] += B[i*COORD_DIM+0][j] * B[i*COORD_DIM+0][j];
            B2[i][j] += B[i*COORD_DIM+1][j] * B[i*COORD_DIM+1][j];
            B2[i][j] += B[i*COORD_DIM+2][j] * B[i*COORD_DIM+2][j];
          }
        }
        return compute_inner_prod(area_elem,B2, B2) * 0.25;
      };
      auto compute_H = [] (const ElemList<COORD_DIM,ElemBasis>& elem_lst, const Vector<ElemBasis>& normal) {
        const Long Nnodes = ElemBasis::Size();
        const Long Nelem = elem_lst.NElem();

        const Vector<ElemBasis> X = elem_lst.ElemVector();
        Vector<ElemBasis> dX, d2X, H(Nelem);
        ElemBasis::Grad(dX, X);
        ElemBasis::Grad(d2X, dX);
        for (Long i = 0; i < Nelem; i++) {
          for (Long j = 0; j < Nnodes; j++) {
            Tensor<Real,true,2,2> I, invI, II;
            for (Long k0 = 0; k0 < 2; k0++) {
              for (Long k1 = 0; k1 < 2; k1++) {
                I(k0,k1) = 0;
                I(k0,k1) += dX[(i*COORD_DIM+0)*2+k0][j] * dX[(i*COORD_DIM+0)*2+k1][j];
                I(k0,k1) += dX[(i*COORD_DIM+1)*2+k0][j] * dX[(i*COORD_DIM+1)*2+k1][j];
                I(k0,k1) += dX[(i*COORD_DIM+2)*2+k0][j] * dX[(i*COORD_DIM+2)*2+k1][j];

                II(k0,k1) = 0;
                II(k0,k1) += d2X[(i*COORD_DIM+0)*4+k0*2+k1][j] * normal[i*COORD_DIM+0][j];
                II(k0,k1) += d2X[(i*COORD_DIM+1)*4+k0*2+k1][j] * normal[i*COORD_DIM+1][j];
                II(k0,k1) += d2X[(i*COORD_DIM+2)*4+k0*2+k1][j] * normal[i*COORD_DIM+2][j];
              }
            }
            { // Set invI
              Real detI = I(0,0)*I(1,1)-I(0,1)*I(1,0);
              invI(0,0) = I(1,1) / detI;
              invI(0,1) = -I(0,1) / detI;
              invI(1,0) = -I(1,0) / detI;
              invI(1,1) = I(0,0) / detI;
            }
            { // Set H
              H[i][j] = 0;
              H[i][j] += -0.5 * II(0,0)*invI(0,0);
              H[i][j] += -0.5 * II(0,1)*invI(0,1);
              H[i][j] += -0.5 * II(1,0)*invI(1,0);
              H[i][j] += -0.5 * II(1,1)*invI(1,1);
            }
          }
        }
        return H;
      };
      auto compute_A_ = [](const Stellarator<Real,ORDER>& S, const Vector<Real>& x) {
        const Long Nelem = S.NElem();
        const Long Nnodes = ElemBasis::Size();
        SCTL_ASSERT(x.Dim() == Nelem*Nnodes+S.Nsurf());

        Vector<ElemBasis> normal, area_elem;
        compute_norm_area_elem(S, normal, area_elem);
        if (S.Nsurf() == 2) {
          Long Nelem0 = S.NTor(0)*S.NPol(0);
          for (Long i = 0; i < Nelem0*COORD_DIM; i++) {
            for (Long j = 0; j < Nnodes; j++) {
              normal[i][j] *= -1.0;
            }
          }
        }

        Vector<ElemBasis> sigma(Nelem);
        for (Long i = 0; i < Nelem; i++) {
          for (Long j = 0; j < Nnodes; j++) {
            sigma[i][j] = x[i*Nnodes+j];
          }
        }
        Real alpha = (S.Nsurf() >= 1 ? x[Nelem*Nnodes + 0] : 0);
        Real beta  = (S.Nsurf() >= 2 ? x[Nelem*Nnodes + 1] : 0);

        Vector<ElemBasis> B = compute_B(S, sigma, alpha, beta);
        Vector<ElemBasis> BdotN = compute_dot_prod(B, normal);

        Real flux_tor = 0, flux_pol = 0;
        //compute_flux(flux_tor, flux_pol, S, B, normal);
        { // compute flux_tor
          SCTL_ASSERT(S.Nsurf() == 1);
          const Long Nelem = S.NElem();
          const Long Nnodes = ElemBasis::Size();
          SCTL_ASSERT(B.Dim() == Nelem*COORD_DIM);
          SCTL_ASSERT(normal.Dim() == Nelem*COORD_DIM);

          Vector<ElemBasis> J(Nelem * COORD_DIM);
          for (Long i = 0; i < Nelem; i++) { // Set J
            for (Long j = 0; j < Nnodes; j++) {
              Tensor<Real,true,COORD_DIM> b, n;
              b(0) = B[i*COORD_DIM+0][j];
              b(1) = B[i*COORD_DIM+1][j];
              b(2) = B[i*COORD_DIM+2][j];

              n(0) = normal[i*COORD_DIM+0][j];
              n(1) = normal[i*COORD_DIM+1][j];
              n(2) = normal[i*COORD_DIM+2][j];

              J[i*COORD_DIM+0][j] = n(1) * b(2) - n(2) * b(1);
              J[i*COORD_DIM+1][j] = n(2) * b(0) - n(0) * b(2);
              J[i*COORD_DIM+2][j] = n(0) * b(1) - n(1) * b(0);
            }
          }

          Vector<ElemBasis> A;
          EvalQuadrature(A, S.quadrature_FxU, S, J, S.Laplace_FxU);

          Vector<Real> circ_pol(S.Nsurf()), circ_tor(S.Nsurf());
          { // compute circ
            const Vector<ElemBasis>& X = S.GetElemList().ElemVector();
            const auto& quad_wts = ElemBasis::QuadWts();
            for (Long k = 0; k < S.Nsurf(); k++) {
              circ_pol[k] = 0;
              circ_tor[k] = 0;
              Long Ndsp = S.ElemDsp(k);
              for (Long i = 0; i < S.NTor(k)*S.NPol(k); i++) {
                for (Long j = 0; j < Nnodes; j++) {
                  Tensor<Real,true,COORD_DIM> x, n, axis, phi_over_R, nxphi_over_R;
                  { // Set nxphi_over_R
                    x(0) = S.Elem(Ndsp+i,0)[j];
                    x(1) = S.Elem(Ndsp+i,1)[j];
                    x(2) = S.Elem(Ndsp+i,2)[j];

                    n(0) = normal[(Ndsp+i)*COORD_DIM+0][j];
                    n(1) = normal[(Ndsp+i)*COORD_DIM+1][j];
                    n(2) = normal[(Ndsp+i)*COORD_DIM+2][j];

                    axis(0) = 0;
                    axis(1) = 0;
                    axis(2) = 1;

                    phi_over_R(0) = axis(1) * x(2) - axis(2) * x(1);
                    phi_over_R(1) = axis(2) * x(0) - axis(0) * x(2);
                    phi_over_R(2) = axis(0) * x(1) - axis(1) * x(0);
                    Real scale = 1 / (phi_over_R(0)*phi_over_R(0) + phi_over_R(1)*phi_over_R(1) + phi_over_R(2)*phi_over_R(2));
                    phi_over_R(0) *= scale;
                    phi_over_R(1) *= scale;
                    phi_over_R(2) *= scale;

                    nxphi_over_R(0) = n(1) * phi_over_R(2) - n(2) * phi_over_R(1);
                    nxphi_over_R(1) = n(2) * phi_over_R(0) - n(0) * phi_over_R(2);
                    nxphi_over_R(2) = n(0) * phi_over_R(1) - n(1) * phi_over_R(0);
                  }

                  circ_pol[k] += A[(Ndsp+i)*COORD_DIM+0][j] * nxphi_over_R(0) * quad_wts[j] * area_elem[i][j] / (2 * const_pi<Real>());
                  circ_pol[k] += A[(Ndsp+i)*COORD_DIM+1][j] * nxphi_over_R(1) * quad_wts[j] * area_elem[i][j] / (2 * const_pi<Real>());
                  circ_pol[k] += A[(Ndsp+i)*COORD_DIM+2][j] * nxphi_over_R(2) * quad_wts[j] * area_elem[i][j] / (2 * const_pi<Real>());

                  //circ_tor[k] += ;
                  //circ_tor[k] += ;
                  //circ_tor[k] += ;
                }
              }
            }
          }
          if (S.Nsurf() == 1) {
            flux_tor = circ_pol[0];
            flux_pol = 0;
          } else if (S.Nsurf() == 2) {
            flux_tor = circ_pol[1] - circ_pol[0];
            flux_pol = circ_tor[0] - circ_tor[1];
          } else {
            SCTL_ASSERT(false);
          }
        }
        { // update flux_tor
          Vector<ElemBasis> G_BdotN(Nelem), phi_dot_N_over_R(Nelem);
          EvalQuadrature(G_BdotN, S.quadrature_FxU, S, BdotN, S.Laplace_FxU);
          for (Long i = 0; i < Nelem; i++) {
            for (Long j = 0; j < Nnodes; j++) {
              Tensor<Real,true,COORD_DIM> x, axis, phi_over_R;
              x(0) = S.Elem(i,0)[j];
              x(1) = S.Elem(i,1)[j];
              x(2) = S.Elem(i,2)[j];

              axis(0) = 0;
              axis(1) = 0;
              axis(2) = 1;
              phi_over_R(0) = axis(1) * x(2) - axis(2) * x(1);
              phi_over_R(1) = axis(2) * x(0) - axis(0) * x(2);
              phi_over_R(2) = axis(0) * x(1) - axis(1) * x(0);
              Real scale = 1 / (phi_over_R(0)*phi_over_R(0) + phi_over_R(1)*phi_over_R(1) + phi_over_R(2)*phi_over_R(2));
              phi_over_R(0) *= scale;
              phi_over_R(1) *= scale;
              phi_over_R(2) *= scale;

              phi_dot_N_over_R[i][j] = 0;
              phi_dot_N_over_R[i][j] += normal[i*COORD_DIM+0][j] * phi_over_R(0);
              phi_dot_N_over_R[i][j] += normal[i*COORD_DIM+1][j] * phi_over_R(1);
              phi_dot_N_over_R[i][j] += normal[i*COORD_DIM+2][j] * phi_over_R(2);
            }
          }
          flux_tor += compute_inner_prod(area_elem, phi_dot_N_over_R, G_BdotN)/(2*const_pi<Real>());
        }

        Vector<Real> Ax(Nelem*Nnodes+S.Nsurf());
        for (Long i = 0; i < Nelem; i++) {
          for (Long j = 0; j < Nnodes; j++) {
            Ax[i*Nnodes+j] = BdotN[i][j];
          }
        }
        if (S.Nsurf() >= 1) Ax[Nelem*Nnodes + 0] = flux_tor;
        if (S.Nsurf() >= 2) Ax[Nelem*Nnodes + 1] = flux_pol;
        return Ax;
      };
      auto compute_invA_ = [&compute_A_](Vector<ElemBasis>& sigma, Real& alpha, Real& beta, const Stellarator<Real,ORDER>& S, Vector<ElemBasis>& Bdotn, Real flux_tor, Real flux_pol, const Comm& comm) {
        typename ParallelSolver<Real>::ParallelOp BIOp = [&S,&compute_A_](Vector<Real>* Ax, const Vector<Real>& x) {
          (*Ax) = compute_A_(S, x);
        };
        const Long Nelem = S.NElem();
        const Long Nnodes = ElemBasis::Size();

        Vector<Real> rhs_(Nelem * Nnodes + S.Nsurf());
        for (Long i = 0; i < Nelem; i++) {
          for (Long j = 0; j < Nnodes; j++) {
            rhs_[i*Nnodes+j] = Bdotn[i][j];
          }
        }
        if (S.Nsurf() >= 1) rhs_[Nelem * Nnodes + 0] = flux_tor;
        if (S.Nsurf() >= 2) rhs_[Nelem * Nnodes + 1] = flux_pol;

        Vector<Real> x_(Nelem * Nnodes + S.Nsurf());
        x_ = 0;
        ParallelSolver<Real> linear_solver(comm, true);
        linear_solver(&x_, BIOp, rhs_, 1e-6, 100);

        sigma.ReInit(Nelem);
        for (Long i = 0; i < Nelem; i++) {
          for (Long j = 0; j < Nnodes; j++) {
            sigma[i][j] = x_[i*Nnodes+j];
          }
        }
        alpha = (S.Nsurf() >= 1 ? x_[Nelem * Nnodes + 0] : 0);
        beta  = (S.Nsurf() >= 2 ? x_[Nelem * Nnodes + 1] : 0);
      };
      Comm comm = Comm::World();
      Profile::Enable(true);

      Long Nsurf = 1;
      Stellarator<Real,ORDER> S;
      Vector<Real> flux_tor(Nsurf), flux_pol(Nsurf);
      { // Init S, flux_tor, flux_pol, pressure
        Vector<Long> NtNp;
        NtNp.PushBack(30);
        NtNp.PushBack(4);
        S = Stellarator<Real,ORDER>(NtNp);
        flux_tor = 1;
        flux_pol = 1;
      }
      Setup(S, comm);

      const Long Nelem = S.NElem();
      const Long Nnodes = ElemBasis::Size();
      Vector<ElemBasis> normal, area_elem;
      compute_norm_area_elem(S, normal, area_elem);

      Vector<ElemBasis> nu(Nelem);
      { // Set nu
        //nu = area_elem;
        //nu = 1;
        //for (Long i = 0; i < Nelem; i++) {
        //  for (Long j = 0; j < Nnodes; j++) {
        //    Tensor<Real,true,COORD_DIM> x;
        //    x(0) = S.Elem(i,0)[j];
        //    x(1) = S.Elem(i,1)[j];
        //    x(2) = S.Elem(i,2)[j];
        //    nu[i][j] = x(2);
        //  }
        //}

        for (Long i = 0; i < Nelem; i++) {
          for (Long j = 0; j < Nnodes; j++) {
            Tensor<Real,true,COORD_DIM> x;
            x(0) = S.Elem(i,0)[j]-8;
            x(1) = S.Elem(i,1)[j]+6;
            x(2) = S.Elem(i,2)[j]-3;
            nu[i][j] = 1/sqrt(x(0)*x(0)+x(1)*x(1)+x(2)*x(2));
          }
        }

        nu = nu * (1.0/sqrt(compute_inner_prod(area_elem, nu, nu)));
      }
      { // Write VTU
        VTUData vtu;
        vtu.AddElems(S.GetElemList(), nu, ORDER);
        vtu.WriteVTK("nu", comm);
      }

      Vector<ElemBasis> B, nu_dBdn, nu_n_dot_dBdn;
      { // Set B, nu_dBdn, nu_n_dot_dBdn
        Real alpha, beta;
        Vector<ElemBasis> sigma;
        compute_invA(sigma, alpha, beta, S, flux_tor[0], flux_pol[0], comm);
        B = compute_B(S, sigma, alpha, beta);
        Vector<ElemBasis> dB = compute_dB(S, sigma, alpha, beta);
        nu_dBdn.ReInit(Nelem * COORD_DIM);
        nu_n_dot_dBdn.ReInit(Nelem);
        for (Long i = 0; i < Nelem; i++) {
          for (Long j = 0; j < Nnodes; j++) {
            Real nu_dBdn_[COORD_DIM] = {0,0,0};
            nu_dBdn_[0] -= dB[(i*COORD_DIM+0)*COORD_DIM+0][j] * normal[i*COORD_DIM+0][j] * nu[i][j];
            nu_dBdn_[0] -= dB[(i*COORD_DIM+0)*COORD_DIM+1][j] * normal[i*COORD_DIM+1][j] * nu[i][j];
            nu_dBdn_[0] -= dB[(i*COORD_DIM+0)*COORD_DIM+2][j] * normal[i*COORD_DIM+2][j] * nu[i][j];
            nu_dBdn_[1] -= dB[(i*COORD_DIM+1)*COORD_DIM+0][j] * normal[i*COORD_DIM+0][j] * nu[i][j];
            nu_dBdn_[1] -= dB[(i*COORD_DIM+1)*COORD_DIM+1][j] * normal[i*COORD_DIM+1][j] * nu[i][j];
            nu_dBdn_[1] -= dB[(i*COORD_DIM+1)*COORD_DIM+2][j] * normal[i*COORD_DIM+2][j] * nu[i][j];
            nu_dBdn_[2] -= dB[(i*COORD_DIM+2)*COORD_DIM+0][j] * normal[i*COORD_DIM+0][j] * nu[i][j];
            nu_dBdn_[2] -= dB[(i*COORD_DIM+2)*COORD_DIM+1][j] * normal[i*COORD_DIM+1][j] * nu[i][j];
            nu_dBdn_[2] -= dB[(i*COORD_DIM+2)*COORD_DIM+2][j] * normal[i*COORD_DIM+2][j] * nu[i][j];
            nu_dBdn[i*COORD_DIM+0][j] = nu_dBdn_[0];
            nu_dBdn[i*COORD_DIM+1][j] = nu_dBdn_[1];
            nu_dBdn[i*COORD_DIM+2][j] = nu_dBdn_[2];

            Real nu_n_dot_dBdn_ = 0;
            nu_n_dot_dBdn_ += nu_dBdn_[0] * normal[i*COORD_DIM+0][j];
            nu_n_dot_dBdn_ += nu_dBdn_[1] * normal[i*COORD_DIM+1][j];
            nu_n_dot_dBdn_ += nu_dBdn_[2] * normal[i*COORD_DIM+2][j];
            nu_n_dot_dBdn[i][j] = nu_n_dot_dBdn_;
          }
        }
      }
      { // Write VTU
        VTUData vtu;
        vtu.AddElems(S.GetElemList(), B, ORDER);
        vtu.WriteVTK("B", comm);
      }

      Real dgdnu;
      Vector<ElemBasis> dBdnu, n_dot_dBdnu;
      { // Set dBdnu, n_dot_dBdnu, dgdnu (finite-difference approximation)
        Real eps = 1e-3;
        Stellarator<Real,ORDER> S0 = S, S1 = S;
        for (Long i = 0; i < Nelem; i++) {
          for (Long j = 0; j < Nnodes; j++) {
            S0.Elem(i, 0)[j] -= 0.5 * eps * normal[i*COORD_DIM+0][j] * nu[i][j];
            S0.Elem(i, 1)[j] -= 0.5 * eps * normal[i*COORD_DIM+1][j] * nu[i][j];
            S0.Elem(i, 2)[j] -= 0.5 * eps * normal[i*COORD_DIM+2][j] * nu[i][j];

            S1.Elem(i, 0)[j] += 0.5 * eps * normal[i*COORD_DIM+0][j] * nu[i][j];
            S1.Elem(i, 1)[j] += 0.5 * eps * normal[i*COORD_DIM+1][j] * nu[i][j];
            S1.Elem(i, 2)[j] += 0.5 * eps * normal[i*COORD_DIM+2][j] * nu[i][j];
          }
        }
        Setup(S0, comm);
        Setup(S1, comm);

        Real alpha0, alpha1, beta0, beta1;
        Vector<ElemBasis> sigma0, sigma1;
        compute_invA(sigma0, alpha0, beta0, S0, flux_tor[0], flux_pol[0], comm);
        compute_invA(sigma1, alpha1, beta1, S1, flux_tor[0], flux_pol[0], comm);
        Vector<ElemBasis> B0 = compute_B(S0, sigma0, alpha0, beta0);
        Vector<ElemBasis> B1 = compute_B(S1, sigma1, alpha1, beta1);
        dBdnu = (B1 - B0) * (1/eps);
        dgdnu = (compute_g(S1,B1) - compute_g(S0,B0)) * (1/eps);

        n_dot_dBdnu.ReInit(Nelem);
        for (Long i = 0; i < Nelem; i++) {
          for (Long j = 0; j < Nnodes; j++) {
            Real n_dot_dBdnu_ = 0;
            n_dot_dBdnu_ += normal[i*COORD_DIM+0][j] * dBdnu[i*COORD_DIM+0][j];
            n_dot_dBdnu_ += normal[i*COORD_DIM+1][j] * dBdnu[i*COORD_DIM+1][j];
            n_dot_dBdnu_ += normal[i*COORD_DIM+2][j] * dBdnu[i*COORD_DIM+2][j];
            n_dot_dBdnu[i][j] = n_dot_dBdnu_;
          }
        }
      }

      Vector<ElemBasis> B_dot_gradnu, nu_surfdivB, surfdivBnu;
      { // Set B_dot_gradnu
        Vector<ElemBasis> gradnu = compute_grad(S, nu);
        B_dot_gradnu.ReInit(Nelem);
        for (Long i = 0; i < Nelem; i++) {
          for (Long j = 0; j < Nnodes; j++) {
            Real B_dot_gradnu_ = 0;
            B_dot_gradnu_ += B[i*COORD_DIM+0][j] * gradnu[i*COORD_DIM+0][j];
            B_dot_gradnu_ += B[i*COORD_DIM+1][j] * gradnu[i*COORD_DIM+1][j];
            B_dot_gradnu_ += B[i*COORD_DIM+2][j] * gradnu[i*COORD_DIM+2][j];
            B_dot_gradnu[i][j] = B_dot_gradnu_;
          }
        }
      }
      { // Set nu_surfdivB
        Vector<ElemBasis> surfdivB = compute_surfdiv(S, B);
        nu_surfdivB.ReInit(Nelem);
        for (Long i = 0; i < Nelem; i++) {
          for (Long j = 0; j < Nnodes; j++) {
            nu_surfdivB[i][j] = nu[i][j] * surfdivB[i][j];
          }
        }
      }
      { // Set surfdivBnu
        Vector<ElemBasis> Bnu(Nelem*COORD_DIM);
        for (Long i = 0; i < Nelem; i++) {
          for (Long j = 0; j < Nnodes; j++) {
            Bnu[i*COORD_DIM+0][j] = B[i*COORD_DIM+0][j] * nu[i][j];
            Bnu[i*COORD_DIM+1][j] = B[i*COORD_DIM+1][j] * nu[i][j];
            Bnu[i*COORD_DIM+2][j] = B[i*COORD_DIM+2][j] * nu[i][j];
          }
        }
        surfdivBnu = compute_surfdiv(S, Bnu);
      }

      // nu_surfdivB == -nu_n_dot_dBdn
      // B_dot_gradnu == n_dot_dBdnu
      // surfdivBnu == B_dot_gradnu - nu_n_dot_dBdn

      Vector<ElemBasis> dBdnu_;
      { // Compute dBdnu_
        Real alpha, beta;
        Real flux_tor = 0, flux_pol = 0;
        { // Set flux_tor, flux_pol
          Vector<ElemBasis> B_dot_phi_over_R(Nelem);
          for (Long i = 0; i < Nelem; i++) {
            for (Long j = 0; j < Nnodes; j++) {
              Tensor<Real,true,COORD_DIM> x, n, axis, phi_over_R;
              { // Set phi_over_R
                x(0) = S.Elem(i,0)[j];
                x(1) = S.Elem(i,1)[j];
                x(2) = S.Elem(i,2)[j];

                n(0) = normal[i*COORD_DIM+0][j];
                n(1) = normal[i*COORD_DIM+1][j];
                n(2) = normal[i*COORD_DIM+2][j];

                axis(0) = 0;
                axis(1) = 0;
                axis(2) = 1;

                phi_over_R(0) = axis(1) * x(2) - axis(2) * x(1);
                phi_over_R(1) = axis(2) * x(0) - axis(0) * x(2);
                phi_over_R(2) = axis(0) * x(1) - axis(1) * x(0);
                Real scale = 1 / (phi_over_R(0)*phi_over_R(0) + phi_over_R(1)*phi_over_R(1) + phi_over_R(2)*phi_over_R(2));
                phi_over_R(0) *= scale;
                phi_over_R(1) *= scale;
                phi_over_R(2) *= scale;
              }
              B_dot_phi_over_R[i][j] = 0;
              B_dot_phi_over_R[i][j] += B[i*COORD_DIM+0][j] * phi_over_R(0);
              B_dot_phi_over_R[i][j] += B[i*COORD_DIM+1][j] * phi_over_R(1);
              B_dot_phi_over_R[i][j] += B[i*COORD_DIM+2][j] * phi_over_R(2);
            }
          }
          flux_tor = -compute_inner_prod(area_elem, B_dot_phi_over_R, nu) / (2 * const_pi<Real>());
        }
        Vector<ElemBasis> sigma, Bdotn = B_dot_gradnu - nu_n_dot_dBdn;
        compute_invA_(sigma, alpha, beta, S, Bdotn, flux_tor, flux_pol, comm);
        dBdnu_ = compute_B(S, sigma, alpha, beta) + nu_dBdn;
      }
      { // Write VTU
        VTUData vtu;
        vtu.AddElems(S.GetElemList(), dBdnu, ORDER);
        vtu.WriteVTK("dBdnu", comm);
      }
      { // Write VTU
        VTUData vtu;
        vtu.AddElems(S.GetElemList(), dBdnu_, ORDER);
        vtu.WriteVTK("dBdnu_", comm);
      }
      { // Write VTU
        VTUData vtu;
        vtu.AddElems(S.GetElemList(), dBdnu_ - dBdnu, ORDER);
        vtu.WriteVTK("err", comm);
      }


      Real dgdnu0, dgdnu1, dgdnu2;
      { // Set dgdnu0 = \int_{Gamma} (B^2 - p) B . B'
        Vector<ElemBasis> dB = dBdnu - nu_dBdn;
        Vector<ElemBasis> B2_p(Nelem), B_dot_dB(Nelem);
        for (Long i = 0; i < Nelem; i++) {
          for (Long j = 0; j < Nnodes; j++) {
            B2_p[i][j] = 0;
            B2_p[i][j] += B[i*COORD_DIM+0][j]*B[i*COORD_DIM+0][j];
            B2_p[i][j] += B[i*COORD_DIM+1][j]*B[i*COORD_DIM+1][j];
            B2_p[i][j] += B[i*COORD_DIM+2][j]*B[i*COORD_DIM+2][j];

            B_dot_dB[i][j] = 0;
            B_dot_dB[i][j] += B[i*COORD_DIM+0][j] * dB[i*COORD_DIM+0][j];
            B_dot_dB[i][j] += B[i*COORD_DIM+1][j] * dB[i*COORD_DIM+1][j];
            B_dot_dB[i][j] += B[i*COORD_DIM+2][j] * dB[i*COORD_DIM+2][j];
          }
        }
        dgdnu0 = compute_inner_prod(area_elem, B2_p, B_dot_dB);
      }
      { // Set dgdnu1 = \int_{Gamma} (B^2-p) B . nu_dBdn
        Vector<ElemBasis> dB = nu_dBdn;
        Vector<ElemBasis> B2_p(Nelem), B_dot_dB(Nelem);
        for (Long i = 0; i < Nelem; i++) {
          for (Long j = 0; j < Nnodes; j++) {
            B2_p[i][j] = 0;
            B2_p[i][j] += B[i*COORD_DIM+0][j]*B[i*COORD_DIM+0][j];
            B2_p[i][j] += B[i*COORD_DIM+1][j]*B[i*COORD_DIM+1][j];
            B2_p[i][j] += B[i*COORD_DIM+2][j]*B[i*COORD_DIM+2][j];

            B_dot_dB[i][j] = 0;
            B_dot_dB[i][j] += B[i*COORD_DIM+0][j] * dB[i*COORD_DIM+0][j];
            B_dot_dB[i][j] += B[i*COORD_DIM+1][j] * dB[i*COORD_DIM+1][j];
            B_dot_dB[i][j] += B[i*COORD_DIM+2][j] * dB[i*COORD_DIM+2][j];
          }
        }
        dgdnu1 = compute_inner_prod(area_elem, B2_p, B_dot_dB);
      }
      { // Set dgdnu2 = \int_{Gamma} 2H(B^2-p)^2 \nu
        Vector<ElemBasis> H = compute_H(S.GetElemList(), normal);
        Vector<ElemBasis> H_B2_p_2(Nelem);
        for (Long i = 0; i < Nelem; i++) {
          for (Long j = 0; j < Nnodes; j++) {
            Real B2_p = 0;
            B2_p += B[i*COORD_DIM+0][j]*B[i*COORD_DIM+0][j];
            B2_p += B[i*COORD_DIM+1][j]*B[i*COORD_DIM+1][j];
            B2_p += B[i*COORD_DIM+2][j]*B[i*COORD_DIM+2][j];
            H_B2_p_2[i][j] = H[i][j] * B2_p*B2_p;
          }
        }
        dgdnu2 = 0.5 * compute_inner_prod(area_elem,H_B2_p_2, nu);
      }

      std::cout<<dgdnu0<<' '<<dgdnu1<<' '<<dgdnu2<<' '<<dgdnu0+dgdnu1+dgdnu2<<'\n';
      std::cout<<dgdnu<<'\n';


#if 0
      Comm comm = Comm::World();
      Profile::Enable(true);

      Real flux_tor = 1.0, flux_pol = 1.0;
      Stellarator<Real,ORDER> S;
      { // Init S
        Vector<Long> NtNp;
        NtNp.PushBack(20);
        NtNp.PushBack(4);
        S = Stellarator<Real,ORDER>(NtNp);
      }
      Vector<ElemBasis> pressure(S.NElem());
      pressure = 0;

      /////////////////////////////////////////////////////////////////////////////////////////////////////////////////

      if (S.Nsurf() == 1) flux_pol = 0.0;
      SetupQuadrature(S.quadrature_dBS , S, S.BiotSavartGrad, order_singular, order_direct, -1.0, comm, -0.01 * pow<-2,Real>(ORDER));
      SetupQuadrature(S.quadrature_BS  , S, S.BiotSavart    , order_singular, order_direct, -1.0, comm, -0.01 * pow<-2,Real>(ORDER));
      SetupQuadrature(S.quadrature_FxU , S, S.Laplace_FxU   , order_singular, order_direct, -1.0, comm);
      SetupQuadrature(S.quadrature_FxdU, S, S.Laplace_FxdU  , order_singular, order_direct, -1.0, comm);
      SetupQuadrature(S.quadrature_dUxF, S, S.Laplace_dUxF  , order_singular, order_direct, -1.0, comm);

      Vector<ElemBasis> Bt0, Bp0;
      Vector<ElemBasis> dBt0, dBp0;
      { // Set Bt0, Bp0
        Vector<ElemBasis> Jt, Jp;
        compute_harmonic_vector_potentials(Jt, Jp, S);
        EvalQuadrature(Bt0, S.quadrature_BS, S, Jp, S.BiotSavart);
        EvalQuadrature(Bp0, S.quadrature_BS, S, Jt, S.BiotSavart);
        EvalQuadrature(dBt0, S.quadrature_dBS, S, Jp, S.BiotSavartGrad);
        EvalQuadrature(dBp0, S.quadrature_dBS, S, Jt, S.BiotSavartGrad);
      }

      auto compute_B = [&S,&Bt0,&Bp0] (const Vector<ElemBasis>& sigma, Real alpha, Real beta) {
        const Long Nelem = S.NElem();
        Vector<ElemBasis> B(S.NElem() * COORD_DIM);
        if (sigma.Dim()) {
          const Long Nnodes = ElemBasis::Size();
          Vector<ElemBasis> normal, area_elem;
          compute_norm_area_elem(S, normal, area_elem);
          EvalQuadrature(B, S.quadrature_FxdU, S, sigma, S.Laplace_FxdU);
          for (Long i = 0; i < Nelem; i++) {
            for (Long j = 0; j < Nnodes; j++) {
              for (Long k = 0; k < COORD_DIM; k++) {
                B[i*COORD_DIM+k][j] -= 0.5*sigma[i][j]*normal[i*COORD_DIM+k][j];
              }
            }
          }
        } else {
          B = 0;
        }
        if (S.Nsurf() >= 1) B += Bt0*alpha;
        if (S.Nsurf() >= 2) B += Bp0*beta;
        return B;
      };
      auto compute_dB = [&S,&dBt0,&dBp0] (const Vector<ElemBasis>& sigma, Real alpha, Real beta) {
        const Long Nelem = S.NElem();
        Vector<ElemBasis> dB(S.NElem() * COORD_DIM * COORD_DIM);
        if (sigma.Dim()) {
          EvalQuadrature(dB, S.quadrature_Fxd2U, S, sigma, S.Laplace_Fxd2U);
        } else {
          dB = 0;
        }
        if (S.Nsurf() >= 1) dB += dBt0*alpha;
        if (S.Nsurf() >= 2) dB += dBp0*beta;
        return dB;
      };
      auto compute_flux = [&S] (Real& flux_tor, Real& flux_pol, const Vector<ElemBasis>& B, const Vector<ElemBasis>& normal) {
        const Long Nelem = S.NElem();
        const Long Nnodes = ElemBasis::Size();
        SCTL_ASSERT(B.Dim() == Nelem*COORD_DIM);
        SCTL_ASSERT(normal.Dim() == Nelem*COORD_DIM);

        Vector<ElemBasis> J(Nelem * COORD_DIM);
        for (Long i = 0; i < Nelem; i++) { // Set J
          for (Long j = 0; j < Nnodes; j++) {
            Tensor<Real,true,COORD_DIM> b, n;
            b(0) = B[i*COORD_DIM+0][j];
            b(1) = B[i*COORD_DIM+1][j];
            b(2) = B[i*COORD_DIM+2][j];

            n(0) = normal[i*COORD_DIM+0][j];
            n(1) = normal[i*COORD_DIM+1][j];
            n(2) = normal[i*COORD_DIM+2][j];

            J[i*COORD_DIM+0][j] = n(1) * b(2) - n(2) * b(1);
            J[i*COORD_DIM+1][j] = n(2) * b(0) - n(0) * b(2);
            J[i*COORD_DIM+2][j] = n(0) * b(1) - n(1) * b(0);
          }
        }

        Vector<ElemBasis> A;
        EvalQuadrature(A, S.quadrature_FxU, S, J, S.Laplace_FxU);

        Vector<Real> circ_pol(S.Nsurf()), circ_tor(S.Nsurf());
        { // compute circ
          Vector<ElemBasis> dX;
          ElemBasis::Grad(dX, S.GetElemList().ElemVector());
          const auto& quad_wts = ElemBasis::QuadWts();
          for (Long k = 0; k < S.Nsurf(); k++) {
            circ_pol[k] = 0;
            circ_tor[k] = 0;
            Long Ndsp = S.ElemDsp(k);
            for (Long i = 0; i < S.NTor(k)*S.NPol(k); i++) {
              for (Long j = 0; j < Nnodes; j++) {
                circ_pol[k] += A[(Ndsp+i)*COORD_DIM+0][j] * dX[(Ndsp+i)*COORD_DIM*2+1][j] * quad_wts[j] / S.NTor(k);
                circ_pol[k] += A[(Ndsp+i)*COORD_DIM+1][j] * dX[(Ndsp+i)*COORD_DIM*2+3][j] * quad_wts[j] / S.NTor(k);
                circ_pol[k] += A[(Ndsp+i)*COORD_DIM+2][j] * dX[(Ndsp+i)*COORD_DIM*2+5][j] * quad_wts[j] / S.NTor(k);

                circ_tor[k] += A[(Ndsp+i)*COORD_DIM+0][j] * dX[(Ndsp+i)*COORD_DIM*2+0][j] * quad_wts[j] / S.NPol(k);
                circ_tor[k] += A[(Ndsp+i)*COORD_DIM+1][j] * dX[(Ndsp+i)*COORD_DIM*2+2][j] * quad_wts[j] / S.NPol(k);
                circ_tor[k] += A[(Ndsp+i)*COORD_DIM+2][j] * dX[(Ndsp+i)*COORD_DIM*2+4][j] * quad_wts[j] / S.NPol(k);
              }
            }
          }
        }
        if (S.Nsurf() == 1) {
          flux_tor = circ_pol[0];
          flux_pol = 0;
        } else if (S.Nsurf() == 2) {
          flux_tor = circ_pol[1] - circ_pol[0];
          flux_pol = circ_tor[0] - circ_tor[1];
        } else {
          SCTL_ASSERT(false);
        }
      };
      auto compute_A = [&S,compute_B,&compute_flux] (const Vector<Real>& x) {
        const Long Nelem = S.NElem();
        const Long Nnodes = ElemBasis::Size();
        SCTL_ASSERT(x.Dim() == Nelem*Nnodes+S.Nsurf());

        Vector<ElemBasis> normal, area_elem;
        compute_norm_area_elem(S, normal, area_elem);

        Vector<ElemBasis> sigma(Nelem);
        for (Long i = 0; i < Nelem; i++) {
          for (Long j = 0; j < Nnodes; j++) {
            sigma[i][j] = x[i*Nnodes+j];
          }
        }
        Real alpha = (S.Nsurf() >= 1 ? x[Nelem*Nnodes + 0] : 0);
        Real beta  = (S.Nsurf() >= 2 ? x[Nelem*Nnodes + 1] : 0);

        Vector<ElemBasis> B = compute_B(sigma, alpha, beta);
        Vector<ElemBasis> BdotN = compute_dot_prod(B, normal);

        Real flux_tor, flux_pol;
        compute_flux(flux_tor, flux_pol, B, normal);

        Vector<Real> Ax(Nelem*Nnodes+S.Nsurf());
        for (Long i = 0; i < Nelem; i++) {
          for (Long j = 0; j < Nnodes; j++) {
            Ax[i*Nnodes+j] = BdotN[i][j];
          }
        }
        if (S.Nsurf() >= 1) Ax[Nelem*Nnodes + 0] = flux_tor;
        if (S.Nsurf() >= 2) Ax[Nelem*Nnodes + 1] = flux_pol;
        return Ax;
      };
      auto compute_invA = [&S,&comm,&compute_A] (Vector<ElemBasis>& sigma, Real& alpha, Real& beta, Real flux_tor, Real flux_pol) {
        typename ParallelSolver<Real>::ParallelOp BIOp = [&compute_A](Vector<Real>* Ax, const Vector<Real>& x) {
          (*Ax) = compute_A(x);
        };
        const Long Nelem = S.NElem();
        const Long Nnodes = ElemBasis::Size();

        Vector<Real> rhs_(Nelem * Nnodes + S.Nsurf());
        rhs_ = 0;
        if (S.Nsurf() >= 1) rhs_[Nelem * Nnodes + 0] = flux_tor;
        if (S.Nsurf() >= 2) rhs_[Nelem * Nnodes + 1] = flux_pol;

        Vector<Real> x_(Nelem * Nnodes + S.Nsurf());
        x_ = 0;
        ParallelSolver<Real> linear_solver(comm, true);
        linear_solver(&x_, BIOp, rhs_, 1e-8, 100);

        sigma.ReInit(Nelem);
        for (Long i = 0; i < Nelem; i++) {
          for (Long j = 0; j < Nnodes; j++) {
            sigma[i][j] = x_[i*Nnodes+j];
          }
        }
        alpha = (S.Nsurf() >= 1 ? x_[Nelem * Nnodes + 0] : 0);
        beta  = (S.Nsurf() >= 2 ? x_[Nelem * Nnodes + 1] : 0);
      };
      auto compute_H = [] (const ElemList<COORD_DIM,ElemBasis>& elem_lst, const Vector<ElemBasis>& normal) {
        const Long Nnodes = ElemBasis::Size();
        const Long Nelem = elem_lst.NElem();

        const Vector<ElemBasis> X = elem_lst.ElemVector();
        Vector<ElemBasis> dX, d2X, H(Nelem);
        ElemBasis::Grad(dX, X);
        ElemBasis::Grad(d2X, dX);
        for (Long i = 0; i < Nelem; i++) {
          for (Long j = 0; j < Nnodes; j++) {
            Tensor<Real,true,2,2> I, invI, II;
            for (Long k0 = 0; k0 < 2; k0++) {
              for (Long k1 = 0; k1 < 2; k1++) {
                I(k0,k1) = 0;
                I(k0,k1) += dX[(i*COORD_DIM+0)*2+k0][j] * dX[(i*COORD_DIM+0)*2+k1][j];
                I(k0,k1) += dX[(i*COORD_DIM+1)*2+k0][j] * dX[(i*COORD_DIM+1)*2+k1][j];
                I(k0,k1) += dX[(i*COORD_DIM+2)*2+k0][j] * dX[(i*COORD_DIM+2)*2+k1][j];

                II(k0,k1) = 0;
                II(k0,k1) += d2X[(i*COORD_DIM+0)*4+k0*2+k1][j] * normal[i*COORD_DIM+0][j];
                II(k0,k1) += d2X[(i*COORD_DIM+1)*4+k0*2+k1][j] * normal[i*COORD_DIM+1][j];
                II(k0,k1) += d2X[(i*COORD_DIM+2)*4+k0*2+k1][j] * normal[i*COORD_DIM+2][j];
              }
            }
            { // Set invI
              Real detI = I(0,0)*I(1,1)-I(0,1)*I(1,0);
              invI(0,0) = I(1,1) / detI;
              invI(0,1) = -I(0,1) / detI;
              invI(1,0) = -I(1,0) / detI;
              invI(1,1) = I(0,0) / detI;
            }
            { // Set H
              H[i][j] = 0;
              H[i][j] += -0.5 * II(0,0)*invI(0,0);
              H[i][j] += -0.5 * II(0,1)*invI(0,1);
              H[i][j] += -0.5 * II(1,0)*invI(1,0);
              H[i][j] += -0.5 * II(1,1)*invI(1,1);
            }
          }
        }
        return H;
      };

      auto compute_grad = [&S,&compute_B,&compute_dB,&compute_invA,&compute_H](Vector<ElemBasis>& pressure, Real flux_tor, Real flux_pol) {
        const Long Nelem = S.NElem();
        const Long Nnodes = ElemBasis::Size();

        Real alpha, beta;
        Vector<ElemBasis> sigma;
        compute_invA(sigma, alpha, beta, flux_tor, flux_pol);
        Vector<ElemBasis> B = compute_B(sigma, alpha, beta);
        Vector<ElemBasis> dB = compute_dB(sigma, alpha, beta);

        Vector<ElemBasis> normal, area_elem;
        compute_norm_area_elem(S, normal, area_elem);
        Vector<ElemBasis> gvec = compute_gvec(S, B, pressure);
        Vector<ElemBasis> dgdB = compute_dgdB(S, B, pressure);
        Vector<ElemBasis> H = compute_H(S.GetElemList(), normal);

        Vector<ElemBasis> dgdnu(Nelem);
        dgdnu = 0;
        for (Long i = 0; i < Nelem; i++) {
          for (Long j = 0; j < Nnodes; j++) {
            Real dgdB_dot_dBdn = 0;
            Real dBdn[COORD_DIM] = {0,0,0};
            for (Long k = 0; k < COORD_DIM; k++) {
              dBdn[0] += dB[(i*COORD_DIM+0)*COORD_DIM+k][j] * normal[i*COORD_DIM+k][j];
              dBdn[1] += dB[(i*COORD_DIM+1)*COORD_DIM+k][j] * normal[i*COORD_DIM+k][j];
              dBdn[2] += dB[(i*COORD_DIM+2)*COORD_DIM+k][j] * normal[i*COORD_DIM+k][j];
            }
            for (Long k = 0; k < COORD_DIM; k++) {
              dgdB_dot_dBdn += dgdB[i*COORD_DIM+k][j] * dBdn[k];
            }
            dgdnu[i][j] = dgdB_dot_dBdn + 2*H[i][j]*gvec[i][j];
          }
        }
        return dgdnu;
      };
      auto dg_dnu0 = compute_gradient(S, pressure, flux_tor, flux_pol);
      auto dg_dnu1 = compute_grad    (   pressure, flux_tor, flux_pol);
      { // Write VTU
        VTUData vtu;
        vtu.AddElems(S.GetElemList(), dg_dnu0, ORDER);
        vtu.WriteVTK("dg_dnu0", comm);
      }
      { // Write VTU
        VTUData vtu;
        vtu.AddElems(S.GetElemList(), dg_dnu1, ORDER);
        vtu.WriteVTK("dg_dnu1", comm);
      }
#endif
    }

  private:

    static void tmp() {
      //if (0) { // Save data
      //  Matrix<Real> M(S.NtNp_[0]*ORDER, S.NtNp_[1]*ORDER);
      //  for (Long tt = 0; tt < S.NtNp_[0]; tt++) {
      //    for (Long pp = 0; pp < S.NtNp_[1]; pp++) {
      //      for (Long t = 0; t < ORDER; t++) {
      //        for (Long p = 0; p < ORDER; p++) {
      //          Long elem_idx = tt * S.NtNp_[1] + pp;
      //          Long node_idx = p * ORDER + t;
      //          M[tt*ORDER+t][pp*ORDER+p] = dg_dnu[elem_idx][node_idx];
      //        }
      //      }
      //    }
      //  }
      //  M.Write("dg_dnu.mat");
      //}
      //if (0) { // filter dg_dnu and write VTU
      //  const Long Nelem = S.NElem();
      //  const Long Nnodes = ElemBasis::Size();
      //  const Integer INTERP_ORDER = 12;
      //  Long Nt = S.NtNp_[0]*ORDER/5, Np = S.NtNp_[1]*ORDER/5;

      //  Matrix<Real> M(Nt, Np); M = 0;
      //  const auto& quad_wts = ElemBasis::QuadWts();
      //  const Matrix<Real>& Mnodes = Basis<Real,1,ORDER>::Nodes();
      //  for (Long tt = 0; tt < S.NtNp_[0]; tt++) {
      //    for (Long pp = 0; pp < S.NtNp_[1]; pp++) {
      //      for (Long t = 0; t < ORDER; t++) {
      //        for (Long p = 0; p < ORDER; p++) {
      //          Real theta = (tt + Mnodes[0][t]) / S.NtNp_[0];
      //          Real phi   = (pp + Mnodes[0][p]) / S.NtNp_[1];
      //          Long i = (Long)(theta * Nt);
      //          Long j = (Long)(phi   * Np);
      //          Real x = theta * Nt - i;
      //          Real y = phi   * Np - j;

      //          Long elem_idx = tt * S.NtNp_[1] + pp;
      //          Long node_idx = p * ORDER + t;

      //          Vector<Real> Interp0(INTERP_ORDER);
      //          Vector<Real> Interp1(INTERP_ORDER);
      //          { // Set Interp0, Interp1
      //            auto node = [] (Long i) {
      //              return (Real)i - (INTERP_ORDER-1)/2;
      //            };
      //            for (Long i = 0; i < INTERP_ORDER; i++) {
      //              Real wt_x = 1, wt_y = 1;
      //              for (Long j = 0; j < INTERP_ORDER; j++) {
      //                if (j != i) {
      //                  wt_x *= (x - node(j)) / (node(i) - node(j));
      //                  wt_y *= (y - node(j)) / (node(i) - node(j));
      //                }
      //                Interp0[i] = wt_x;
      //                Interp1[i] = wt_y;
      //              }
      //            }
      //          }

      //          for (Long ii = 0; ii < INTERP_ORDER; ii++) {
      //            for (Long jj = 0; jj < INTERP_ORDER; jj++) {
      //              Long idx_i = (i + ii-(INTERP_ORDER-1)/2 + Nt) % Nt;
      //              Long idx_j = (j + jj-(INTERP_ORDER-1)/2 + Np) % Np;
      //              M[idx_i][idx_j] += dg_dnu[elem_idx][node_idx] * quad_wts[node_idx] * Interp0[ii] * Interp1[jj] / (S.NtNp_[0] * S.NtNp_[1]) * (Nt * Np);
      //            }
      //          }
      //        }
      //      }
      //    }
      //  }

      //  Vector<ElemBasis> f(Nelem);
      //  for (Long tt = 0; tt < S.NtNp_[0]; tt++) {
      //    for (Long pp = 0; pp < S.NtNp_[1]; pp++) {
      //      for (Long t = 0; t < ORDER; t++) {
      //        for (Long p = 0; p < ORDER; p++) {
      //          Matrix<Real> Mnodes = Basis<Real,1,ORDER>::Nodes();
      //          Real theta = (tt + Mnodes[0][t]) / S.NtNp_[0];
      //          Real phi   = (pp + Mnodes[0][p]) / S.NtNp_[1];
      //          Long i = (Long)(theta * Nt);
      //          Long j = (Long)(phi   * Np);
      //          Real x = theta * Nt - i;
      //          Real y = phi   * Np - j;

      //          Vector<Real> Interp0(INTERP_ORDER);
      //          Vector<Real> Interp1(INTERP_ORDER);
      //          { // Set Interp0, Interp1
      //            auto node = [] (Long i) {
      //              return (Real)i - (INTERP_ORDER-1)/2;
      //            };
      //            for (Long i = 0; i < INTERP_ORDER; i++) {
      //              Real wt_x = 1, wt_y = 1;
      //              for (Long j = 0; j < INTERP_ORDER; j++) {
      //                if (j != i) {
      //                  wt_x *= (x - node(j)) / (node(i) - node(j));
      //                  wt_y *= (y - node(j)) / (node(i) - node(j));
      //                }
      //                Interp0[i] = wt_x;
      //                Interp1[i] = wt_y;
      //              }
      //            }
      //          }

      //          Real f0 = 0;
      //          for (Long ii = 0; ii < INTERP_ORDER; ii++) {
      //            for (Long jj = 0; jj < INTERP_ORDER; jj++) {
      //              Long idx_i = (i + ii-(INTERP_ORDER-1)/2 + Nt) % Nt;
      //              Long idx_j = (j + jj-(INTERP_ORDER-1)/2 + Np) % Np;
      //              f0 += Interp0[ii] * Interp1[jj] * M[idx_i][idx_j];
      //            }
      //          }

      //          Long elem_idx = tt * S.NtNp_[1] + pp;
      //          Long node_idx = p * ORDER + t;
      //          f[elem_idx][node_idx] = f0;
      //        }
      //      }
      //    }
      //  }
      //  { // Write VTU
      //    VTUData vtu;
      //    vtu.AddElems(S.GetElemList(), f, ORDER);
      //    vtu.WriteVTK("dg_dnu_filtered", comm);
      //  }
      //  dg_dnu = f;
      //}
    }

    static void FlipNormal(Vector<ElemBasis>& v) {
      for (Long i = 0; i < v.Dim(); i++) {
        const auto elem = v[i];
        for (Long j0 = 0; j0 < ORDER; j0++) {
          for (Long j1 = 0; j1 < ORDER; j1++) {
            v[i][j0*ORDER+j1] = elem[j0*ORDER+(ORDER-j1-1)];
          }
        }
      }
    }
    template <class Kernel> static void SetupQuadrature(Quadrature<Real>& quadrature, const Stellarator<Real,ORDER>& S, const Kernel& kernel, Integer order_singular, Integer order_direct, Real period_length, const Comm& comm, Real Rqbx = 0) {
      if (S.Nsurf() == 2) {
        Long Nelem0 = S.NTor(0)*S.NPol(0);
        ElemList<COORD_DIM, ElemBasis> elem_lst = S.GetElemList();
        { // Update elem_lst
          Vector<ElemBasis> X = elem_lst.ElemVector();
          Vector<ElemBasis> X0(Nelem0*COORD_DIM, X.begin(), false);
          FlipNormal(X0);
          elem_lst.ReInit(X);
        }
        quadrature.template Setup<ElemBasis, ElemBasis>(elem_lst, kernel, order_singular, order_direct, period_length, comm, Rqbx);
      } else {
        quadrature.template Setup<ElemBasis, ElemBasis>(S.GetElemList(), kernel, order_singular, order_direct, period_length, comm, Rqbx);
      }
    }
    template <class Kernel> static void EvalQuadrature(Vector<ElemBasis>& potential, const Quadrature<Real>& quadrature, const Stellarator<Real,ORDER>& S, const Vector<ElemBasis>& density, const Kernel& kernel) {
      if (S.Nsurf() == 2) {
        Long Nelem0 = S.NTor(0)*S.NPol(0);
        Vector<ElemBasis> potential_, density_ = density;
        ElemList<COORD_DIM, ElemBasis> elem_lst = S.GetElemList();
        { // Update elem_lst
          Vector<ElemBasis> X = elem_lst.ElemVector();
          Vector<ElemBasis> X0(Nelem0*COORD_DIM, X.begin(), false);
          FlipNormal(X0);
          elem_lst.ReInit(X);
        }
        { // Update density_
          Long dof = density_.Dim() / S.NElem();
          Vector<ElemBasis> density0(Nelem0*dof, density_.begin(), false);
          FlipNormal(density0);
        }
        quadrature.Eval(potential_, elem_lst, density_, kernel);
        { // Update potential_
          Long dof = potential_.Dim() / S.NElem();
          Vector<ElemBasis> potential0(Nelem0*dof, potential_.begin(), false);
          FlipNormal(potential0);
        }
        potential = potential_;
      } else {
        quadrature.Eval(potential, S.GetElemList(), density, kernel);
      }
    }

    void InitSurf(Long l, Long Nsurf) {
      const auto& nodes = ElemBasis::Nodes();
      const Long Nt = NTor(l);
      const Long Np = NPol(l);

      for (Long i = 0; i < Nt; i++) {
        for (Long j = 0; j < Np; j++) {
          for (Long k = 0; k < ElemBasis::Size(); k++) {
            Real theta = (i + nodes[0][k]) * 2*const_pi<Real>()/Nt;
            Real phi   = (j + nodes[1][k]) * 2*const_pi<Real>()/Np;
            Real X,Y,Z;
            SurfGeom(X,Y,Z,theta,phi, (2.0+l)/(1.0+Nsurf));
            Elem(ElemIdx(l,i,j),0)[k] = X;
            Elem(ElemIdx(l,i,j),1)[k] = Y;
            Elem(ElemIdx(l,i,j),2)[k] = Z;
          }
        }
      }
    }

    static void SurfGeom(Real& X, Real& Y, Real& Z, Real theta, Real phi, Real s) {
      Integer Nperiod = 5;
#if 0
      Real Aspect_ratio = 10.27932548522949;
      Real coeffmat[21][21] = { 0.00000478813217,  0.00000000000000,  0.00000351611652,  0.00000135354389,  0.00000061357832,  0.00000220091101,  0.00000423862912, -0.00003000058678,  0.00000064187111, -0.00024228452821,  0.00003116775770,  0.00000176210710,  0.00000289141326, -0.00000150300525,  0.00000772853855,  0.00000098855242,  0.00000316606793,  0.00000002168364,  0.00000212047939,  0.00000299016097,  0.00000443224508,
                                0.00000028202930,  0.00000000000000, -0.00000249222421, -0.00000203136278,  0.00000131104809,  0.00000011987446, -0.00000370760154,  0.00004553918916, -0.00007711342914, -0.00004685295062,  0.00011049838213, -0.00000197486270,  0.00000395827146,  0.00000615046474,  0.00000755337123,  0.00000700606006,  0.00000922725030, -0.00000043310337,  0.00000107416383,  0.00000449787694,  0.00000305137178,
                                0.00001226376662,  0.00000000000000,  0.00000270820692,  0.00000208059305,  0.00000521478523,  0.00001779037302,  0.00000846544117,  0.00001120913385, -0.00065816845745, -0.00085107452469, -0.00013171190221, -0.00005540943675, -0.00001835885450,  0.00000101879823,  0.00000209222071,  0.00000091532502, -0.00000521515358, -0.00000209227142, -0.00000678545939, -0.00000034963549, -0.00000015111488,
                                0.00001560274177,  0.00000000000000,  0.00000350691471, -0.00001160475040, -0.00001763036562,  0.00003487367940, -0.00002787247831, -0.00000910982726,  0.00008818832430, -0.00524408789352,  0.00009378376126,  0.00004184526188,  0.00002849263365, -0.00002757280527,  0.00003388467667,  0.00000706207265,  0.00000625263419, -0.00003315929280, -0.00001181772132,  0.00000311426015,  0.00001875682574,
                               -0.00000398287420,  0.00000000000000, -0.00001524541040,  0.00001724056165,  0.00002245173346,  0.00002806861812, -0.00000388776925,  0.00008143573359, -0.00005900909309,  0.00110496615525,  0.00134626252111,  0.00005128383054, -0.00001372421866,  0.00003612563887,  0.00002236580076, -0.00002728391883,  0.00001981237256,  0.00000655450458,  0.00000985319002,  0.00001347597299,  0.00000645987802,
                                0.00003304968050,  0.00000000000000, -0.00000530822217,  0.00001324870937, -0.00003610889689, -0.00005478735329, -0.00005818806312, -0.00037112057908, -0.00017812002625, -0.00093204283621,  0.00115969858598, -0.00033559172880, -0.00010441876657, -0.00001617923044, -0.00000555065844,  0.00007343527250, -0.00004408047607,  0.00000403802142,  0.00001843931204,  0.00001694047933,  0.00001213414362,
                               -0.00000751115658,  0.00000000000000,  0.00005457974839, -0.00000334614515,  0.00005845565465,  0.00015000770509,  0.00021849104087,  0.00002724147635,  0.00167233624961,  0.00011666602222,  0.00276563479565, -0.00085952825611, -0.00030217235326, -0.00008841593808,  0.00000997664119, -0.00015285826521,  0.00002517224675,  0.00003009161810,  0.00001883217556,  0.00002146127554,  0.00001822445302,
                               -0.00004128706860,  0.00000000000000, -0.00003496417776,  0.00001088761655, -0.00000298955979, -0.00005359326315, -0.00019021633489, -0.00017992728681, -0.00347794801928,  0.00064632791327,  0.00449698418379, -0.00017710507382,  0.00006126180233,  0.00018059254216,  0.00002354096432,  0.00008189838991, -0.00010060678323, -0.00017183290038,  0.00019413756672,  0.00021334811754,  0.00011263617489,
                                0.00000853522670, -0.00000000000000, -0.00006544789358,  0.00005424076880, -0.00000679056529, -0.00001249735487, -0.00053082982777,  0.00035396864405, -0.00115020677913,  0.05894451215863,  0.06573092192411,  0.01498018857092,  0.00278125284240,  0.00145188067108,  0.00033717858605,  0.00000800427370, -0.00009335305367,  0.00024286781263, -0.00023916347709,  0.00031213948387,  0.00018134393031,
                               -0.00002521496390, -0.00000000000000, -0.00054337945767,  0.00012690725271,  0.00053313979879,  0.00064233405283, -0.00047686311882,  0.00176536326762,  0.00074157933705, -0.02684566564858,  1.00000000000000,  0.07176169008017,  0.00837037432939, -0.00000381640211,  0.00088998704450, -0.00049218931235, -0.00024546548957, -0.00036608282244,  0.00049480766756,  0.00031158892671,  0.00006898906577,
                                0.00021280418150,  0.00028127161204, -0.00070030166535,  0.00022237010126, -0.00028713891516, -0.00013800295710,  0.00005912094275,  0.00172126013786, -0.00618684850633,  0.03608432412148,    Aspect_ratio  ,  0.49896776676178,  0.00091372377938, -0.00085712829605, -0.00124801427592, -0.00007427225501, -0.00005245858847,  0.00002841771493,  0.00020249813679, -0.00014303345233,  0.00001406490901,
                                0.00023699452868,  0.00008661757602,  0.00025744654704, -0.00022715188970, -0.00076146807987,  0.00055185536621, -0.00012325309217, -0.00072356045712, -0.00160693109501,  0.00246682553552, -0.14175094664097, -0.36207047104836, -0.04089594259858,  0.00060774467420,  0.00088646943914,  0.00004865296432, -0.00041878610500, -0.00023025234987, -0.00009676301852, -0.00000000000000,  0.00008409228758,
                                0.00011432896281, -0.00000707848403,  0.00004698805787, -0.00043642931269,  0.00081384339137, -0.00065635429928, -0.00011831733718,  0.00017413357273,  0.00224463525228,  0.00478497287259,  0.03294761106372,  0.01078986655921,  0.10731782764196,  0.00075034319889, -0.00009241879889,  0.00055023463210,  0.00006596000458,  0.00005045382932,  0.00014874986664,  0.00000000000000, -0.00015369028552,
                                0.00001037383754,  0.00009250180301,  0.00026204055757,  0.00007424291834, -0.00047751804232,  0.00029184055165,  0.00050921301590, -0.00004825839278, -0.00029933769838,  0.00279659987427,  0.00210463814437, -0.00618590926751, -0.02400829829276, -0.02316811867058, -0.00086368201301, -0.00032258985448, -0.00018304496189,  0.00008438774967, -0.00008305341908,  0.00000000000000,  0.00013047417451,
                               -0.00001376930322, -0.00001723831701, -0.00011543079017, -0.00022646733851,  0.00013467084500, -0.00004661652201, -0.00008419520600,  0.00035772417323, -0.00011815709877,  0.00028718306567,  0.00092207465786, -0.00317224999890,  0.00061770365573,  0.01017294172198,  0.00294739892706,  0.00014669894881,  0.00015702951350,  0.00003432080121, -0.00008555022214, -0.00000000000000,  0.00000454909878,
                               -0.00000196001542, -0.00003198397462, -0.00004425687075, -0.00004129848094, -0.00003789070615, -0.00027583551127,  0.00025874207495, -0.00002334945384, -0.00007259396807, -0.00008295358566,  0.00011360697681, -0.00101968157105,  0.00046784928418, -0.00208410434425, -0.00313158822246, -0.00046005158219, -0.00010552268213, -0.00005850767775,  0.00003971093611,  0.00000000000000, -0.00005275657168,
                               -0.00001065901233, -0.00001934838656, -0.00001220186732, -0.00002060524639, -0.00000225423423, -0.00001894621164, -0.00001533334580, -0.00001791087379,  0.00008156246622, -0.00008441298269,  0.00021060956351, -0.00030303673702,  0.00075949780876, -0.00010539998038,  0.00109045265708,  0.00068949378328,  0.00009268362192,  0.00003471063246,  0.00001204656473, -0.00000000000000,  0.00001500743110,
                                0.00000105878155, -0.00000910870767, -0.00000172467264, -0.00000722095228,  0.00000699280463, -0.00002061720625, -0.00000889817693, -0.00001993474507,  0.00000370749740, -0.00000090311920,  0.00002677819793,  0.00043428712524,  0.00210293265991,  0.00018200518389, -0.00009621794743, -0.00035250501242, -0.00012996385340, -0.00002185157609, -0.00001116586463, -0.00000000000000, -0.00000451994811,
                                0.00000424055270, -0.00000463139304,  0.00000301006116, -0.00000123974939,  0.00000632465435, -0.00002090823000,  0.00001773388794,  0.00000121050368,  0.00001886057362, -0.00001043497195, -0.00002269273500, -0.00021979617304, -0.00001043962493, -0.00116343051195, -0.00004193381756,  0.00007944958634,  0.00007301353617,  0.00002082651736, -0.00000119863023, -0.00000000000000, -0.00001440504820,
                               -0.00000391270805, -0.00000490489265, -0.00000504441778, -0.00000904507579, -0.00000111389932,  0.00000597532107,  0.00000047090245, -0.00001553130096, -0.00001524566323, -0.00000522222899, -0.00007707672921, -0.00004165665086,  0.00015764687851,  0.00035649110214,  0.00038701237645,  0.00002386798405, -0.00001946414341, -0.00000913835174, -0.00000489907188,  0.00000000000000,  0.00000172327657,
                               -0.00000015388650, -0.00000603232729, -0.00000397650865,  0.00000280493782,  0.00000463132073, -0.00000788678426, -0.00000471605335, -0.00000283715985, -0.00000422824724,  0.00000366817630, -0.00001159603562, -0.00001625759251,  0.00049116823357,  0.00005048640014, -0.00020234247495, -0.00006341376866, -0.00000807822744,  0.00000070463199,  0.00000014041755,  0.00000000000000, -0.00000718306910};
#else
      Real Aspect_ratio = 5;
      Real coeffmat[21][21] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0,            0, 0    , 0, 0, 0, 0, 0, 0, 0, 0, 0,
                               0, 0, 0, 0, 0, 0, 0, 0, 0, 0,            0, 0    , 0, 0, 0, 0, 0, 0, 0, 0, 0,
                               0, 0, 0, 0, 0, 0, 0, 0, 0, 0,            0, 0    , 0, 0, 0, 0, 0, 0, 0, 0, 0,
                               0, 0, 0, 0, 0, 0, 0, 0, 0, 0,            0, 0    , 0, 0, 0, 0, 0, 0, 0, 0, 0,
                               0, 0, 0, 0, 0, 0, 0, 0, 0, 0,            0, 0    , 0, 0, 0, 0, 0, 0, 0, 0, 0,
                               0, 0, 0, 0, 0, 0, 0, 0, 0, 0,            0, 0    , 0, 0, 0, 0, 0, 0, 0, 0, 0,
                               0, 0, 0, 0, 0, 0, 0, 0, 0, 0,            0, 0    , 0, 0, 0, 0, 0, 0, 0, 0, 0,
                               0, 0, 0, 0, 0, 0, 0, 0, 0, 0,            0, 0    , 0, 0, 0, 0, 0, 0, 0, 0, 0,
                               0, 0, 0, 0, 0, 0, 0, 0, 0, 0,            0, 0    , 0, 0, 0, 0, 0, 0, 0, 0, 0,
                               0, 0, 0, 0, 0, 0, 0, 0, 0, 0,            s, 0    , 0, 0, 0, 0, 0, 0, 0, 0, 0,
                               0, 0, 0, 0, 0, 0, 0, 0, 0, 0, Aspect_ratio, 0    , 0, 0, 0, 0, 0, 0, 0, 0, 0,
                               0, 0, 0, 0, 0, 0, 0, 0, 0, 0,            0, 0.2*s, 0, 0, 0, 0, 0, 0, 0, 0, 0,
                               0, 0, 0, 0, 0, 0, 0, 0, 0, 0,            0, 0    , 0, 0, 0, 0, 0, 0, 0, 0, 0,
                               0, 0, 0, 0, 0, 0, 0, 0, 0, 0,            0, 0    , 0, 0, 0, 0, 0, 0, 0, 0, 0,
                               0, 0, 0, 0, 0, 0, 0, 0, 0, 0,            0, 0    , 0, 0, 0, 0, 0, 0, 0, 0, 0,
                               0, 0, 0, 0, 0, 0, 0, 0, 0, 0,            0, 0    , 0, 0, 0, 0, 0, 0, 0, 0, 0,
                               0, 0, 0, 0, 0, 0, 0, 0, 0, 0,            0, 0    , 0, 0, 0, 0, 0, 0, 0, 0, 0,
                               0, 0, 0, 0, 0, 0, 0, 0, 0, 0,            0, 0    , 0, 0, 0, 0, 0, 0, 0, 0, 0,
                               0, 0, 0, 0, 0, 0, 0, 0, 0, 0,            0, 0    , 0, 0, 0, 0, 0, 0, 0, 0, 0,
                               0, 0, 0, 0, 0, 0, 0, 0, 0, 0,            0, 0    , 0, 0, 0, 0, 0, 0, 0, 0, 0,
                               0, 0, 0, 0, 0, 0, 0, 0, 0, 0,            0, 0    , 0, 0, 0, 0, 0, 0, 0, 0, 0};
#endif
      Z = 0;
      Real R = 0;
      for (long i = -10; i <= 10; i++) {
        for (long j = -10; j <= 10; j++) {
          R += coeffmat[i+10][j+10] * cos(-i*phi + Nperiod*j*theta);
          Z += coeffmat[i+10][j+10] * sin(-i*phi + Nperiod*j*theta);
        }
      }
      X = R * cos(theta);
      Y = R * sin(theta);
    }

    GenericKernel<BiotSavart3D>     BiotSavart ;
    GenericKernel<BiotSavartGrad3D> BiotSavartGrad;
    GenericKernel<Laplace3D_FxU >   Laplace_FxU ;
    GenericKernel<Laplace3D_FxdU>   Laplace_FxdU;
    GenericKernel<Laplace3D_dUxF>   Laplace_dUxF;
    GenericKernel<Laplace3D_dUxD>   Laplace_dUxD;
    GenericKernel<Laplace3D_Fxd2U>  Laplace_Fxd2U;

    mutable Quadrature<Real> quadrature_BS  ;
    mutable Quadrature<Real> quadrature_dBS ;
    mutable Quadrature<Real> quadrature_FxU ;
    mutable Quadrature<Real> quadrature_FxdU;
    mutable Quadrature<Real> quadrature_dUxF;
    mutable Quadrature<Real> quadrature_dUxD;
    mutable Quadrature<Real> quadrature_Fxd2U;
    mutable Vector<ElemBasis> Bt0, Bp0, dBt0, dBp0;
    mutable Vector<ElemBasis> sigma, B, gvec, dgdB;
    mutable Real alpha, beta;

    ElemLst elements;
    Vector<Long> NtNp_;
    Vector<Long> elem_dsp;
};

template <class Real, Integer ORDER=10> class MHDEquilib {
    static constexpr Integer fourier_dim0 = 50*10*4;
    static constexpr Integer fourier_dim1 = 10*10*4;
    //static constexpr Integer fourier_upsample = 4;
    static constexpr Integer COORD_DIM = 3;
    static constexpr Integer ELEM_DIM = COORD_DIM-1;
    using ElemBasis = Basis<Real, ELEM_DIM, ORDER>;

    static Vector<Real> cheb2grid(const Vector<ElemBasis>& X, Long Mt, Long Mp, Long Nt, Long Np) {
      const Long dof = X.Dim() / (Mt * Mp);
      SCTL_ASSERT(X.Dim() == Mt * Mp *dof);

      Vector<Real> Xf(dof*Nt*Np); Xf = 0;
      const Long Nnodes = ElemBasis::Size();
      const Matrix<Real>& Mnodes = Basis<Real,1,ORDER>::Nodes();
      for (Long t = 0; t < Nt; t++) {
        for (Long p = 0; p < Np; p++) {
          Real theta = t / (Real)Nt;
          Real phi   = p / (Real)Np;
          Long i = (Long)(theta * Mt);
          Long j = (Long)(phi   * Mp);
          Real x = theta * Mt - i;
          Real y = phi   * Mp - j;
          Long elem_idx = i * Mp + j;

          Vector<Real> Interp0(ORDER);
          Vector<Real> Interp1(ORDER);
          { // Set Interp0, Interp1
            auto node = [&Mnodes] (Long i) {
              return Mnodes[0][i];
            };
            for (Long i = 0; i < ORDER; i++) {
              Real wt_x = 1, wt_y = 1;
              for (Long j = 0; j < ORDER; j++) {
                if (j != i) {
                  wt_x *= (x - node(j)) / (node(i) - node(j));
                  wt_y *= (y - node(j)) / (node(i) - node(j));
                }
                Interp0[i] = wt_x;
                Interp1[i] = wt_y;
              }
            }
          }

          for (Long ii = 0; ii < ORDER; ii++) {
            for (Long jj = 0; jj < ORDER; jj++) {
              Long node_idx = jj * ORDER + ii;
              for (Long k = 0; k < dof; k++) {
                Xf[(k*Nt+t)*Np+p] += X[elem_idx*dof+k][node_idx] * Interp0[ii] * Interp1[jj];
              }
            }
          }
        }
      }
      return Xf;
    }
    static Vector<ElemBasis> grid2cheb(const Vector<Real>& Xf, Long Nt, Long Np, Long Mt, Long Mp) {
      Long dof = Xf.Dim() / (Nt*Np);
      SCTL_ASSERT(Xf.Dim() == dof*Nt*Np);
      Vector<ElemBasis> X(Mt*Mp*dof);
      constexpr Integer INTERP_ORDER = 12;

      for (Long tt = 0; tt < Mt; tt++) {
        for (Long pp = 0; pp < Mp; pp++) {
          for (Long t = 0; t < ORDER; t++) {
            for (Long p = 0; p < ORDER; p++) {
              Matrix<Real> Mnodes = Basis<Real,1,ORDER>::Nodes();
              Real theta = (tt + Mnodes[0][t]) / Mt;
              Real phi   = (pp + Mnodes[0][p]) / Mp;
              Long i = (Long)(theta * Nt);
              Long j = (Long)(phi   * Np);
              Real x = theta * Nt - i;
              Real y = phi   * Np - j;

              Vector<Real> Interp0(INTERP_ORDER);
              Vector<Real> Interp1(INTERP_ORDER);
              { // Set Interp0, Interp1
                auto node = [] (Long i) {
                  return (Real)i - (INTERP_ORDER-1)/2;
                };
                for (Long i = 0; i < INTERP_ORDER; i++) {
                  Real wt_x = 1, wt_y = 1;
                  for (Long j = 0; j < INTERP_ORDER; j++) {
                    if (j != i) {
                      wt_x *= (x - node(j)) / (node(i) - node(j));
                      wt_y *= (y - node(j)) / (node(i) - node(j));
                    }
                    Interp0[i] = wt_x;
                    Interp1[i] = wt_y;
                  }
                }
              }

              for (Long k = 0; k < dof; k++) {
                Real X0 = 0;
                for (Long ii = 0; ii < INTERP_ORDER; ii++) {
                  for (Long jj = 0; jj < INTERP_ORDER; jj++) {
                    Long idx_i = (i + ii-(INTERP_ORDER-1)/2 + Nt) % Nt;
                    Long idx_j = (j + jj-(INTERP_ORDER-1)/2 + Np) % Np;
                    X0 += Interp0[ii] * Interp1[jj] * Xf[(k*Nt+idx_i)*Np+idx_j];
                  }
                }
                Long elem_idx = tt * Mp + pp;
                Long node_idx = p * ORDER + t;
                X[elem_idx*dof+k][node_idx] = X0;
              }
            }
          }
        }
      }
      return X;
    }
    static void fourier_filter(Vector<Real>& X, long Nt_, long Np_, Real sigma) {
      long dof = X.Dim() / (Nt_ * Np_);
      SCTL_ASSERT(X.Dim() == dof * Nt_ * Np_);

      FFT<Real> fft_r2c, fft_c2r;
      StaticArray<Long, 2> fft_dim = {Nt_, Np_};
      fft_r2c.Setup(FFT_Type::R2C, 1, Vector<Long>(2, fft_dim, false), omp_get_max_threads());
      fft_c2r.Setup(FFT_Type::C2R, 1, Vector<Long>(2, fft_dim, false), omp_get_max_threads());

      long Nt = Nt_;
      long Np = fft_r2c.Dim(1) / (Nt * 2);
      SCTL_ASSERT(fft_r2c.Dim(1) == Nt * Np * 2);

      auto filter_fn = [](Real x2, Real sigma) {return exp(-x2/(2*sigma*sigma));};
      //auto filter_fn = [](Real x2, Real sigma) {return (x2<sigma*sigma?1.0:0.0);};

      Vector<Real> coeff(fft_r2c.Dim(1));
      for (long k = 0; k < dof; k++) {
        Vector<Real> X_(Nt_*Np_, X.begin() + k*Nt_*Np_, false);
        fft_r2c.Execute(X_, coeff);
        for (long t = 0; t < Nt; t++) {
          for (long p = 0; p < Np; p++) {
            Real tt = (t - (t > Nt / 2 ? Nt : 0)) * (2*Np/(Real)Nt);
            Real pp = p;
            Real f = filter_fn(tt*tt+pp*pp, sigma);
            coeff[(t * Np + p) * 2 + 0] *= f;
            coeff[(t * Np + p) * 2 + 1] *= f;
          }
        }
        fft_c2r.Execute(coeff, X_);
      }
    };

    static void fourier_print_spectrum(const Vector<Real>& X, long Nt_, long Np_) {
    //  long dof = X.Dim() / (Nt_ * Np_);
    //  SCTL_ASSERT(X.Dim() == dof * Nt_ * Np_);

    //  FFT<Real> fft_r2c, fft_c2r;
    //  StaticArray<Long, 2> fft_dim = {Nt_, Np_};
    //  fft_r2c.Setup(FFT_Type::R2C, 1, Vector<Long>(2, fft_dim, false), omp_get_max_threads());

    //  long Nt = Nt_;
    //  long Np = fft_r2c.Dim(1) / (Nt * 2);
    //  SCTL_ASSERT(fft_r2c.Dim(1) == Nt * Np * 2);

    //  Vector<Real> coeff(fft_r2c.Dim(1));
    //  Vector<Real> spectrum(200); spectrum = 0;
    //  for (long k = 0; k < dof; k++) {
    //    const Vector<Real> X_(Nt_*Np_, (Iterator<Real>)X.begin() + k*Nt_*Np_, false);
    //    fft_r2c.Execute(X_, coeff);
    //    for (long t = 0; t < Nt; t++) {
    //      for (long p = 0; p < Np; p++) {
    //        Real tt = (t - (t > Nt / 2 ? Nt : 0)) / (Real)(Nt / 2);
    //        Real pp = p / (Real)Np;
    //        Long freq = (Long)(sqrt<Real>(tt*tt+pp*pp)*spectrum.Dim());
    //        if (freq<spectrum.Dim()) spectrum[freq] += coeff[(t*Np+p)*2+0]*coeff[(t*Np+p)*2+0];
    //        if (freq<spectrum.Dim()) spectrum[freq] += coeff[(t*Np+p)*2+1]*coeff[(t*Np+p)*2+1];
    //      }
    //    }
    //  }
    //  for (Long i = 0; i < spectrum.Dim(); i++) {
    //    spectrum[i] = log<Real>(sqrt<Real>(spectrum[i]))/log<Real>(10.0);
    //  }
    //  std::cout<<"spectrum = "<<spectrum<<'\n';
    };
    static void fourier_print_spectrum(std::string var_name, const sctl::Vector<Real>& B0, sctl::Long Nt0, sctl::Long Np0) {
      auto Resample = [](const sctl::Vector<Real>& B, long Nt, long Np, long Nt0, long Np0) {
        sctl::Vector<Real> B0;
        biest::SurfaceOp<Real>::Upsample(B,Nt,Np, B0,Nt0,Np0);
        return B0;
      };
      auto max_rel_err = [](const sctl::Vector<Real> A, const sctl::Vector<Real> B) {
        SCTL_ASSERT(A.Dim() == B.Dim());
        Real max_err = 0;
        Real max_val = 1e-20;
        for (sctl::Long i = 0; i < A.Dim(); i++) {
          max_err = std::max<Real>(max_err, fabs(A[i]-B[i]));
          max_val = std::max<Real>(max_val, fabs(A[i]));
        }
        return max_err/max_val;
      };
      auto max_err = [](const sctl::Vector<Real> A, const sctl::Vector<Real> B) {
        SCTL_ASSERT(A.Dim() == B.Dim());
        Real max_err = 0;
        for (sctl::Long i = 0; i < A.Dim(); i++) {
          max_err = std::max<Real>(max_err, fabs(A[i]-B[i]));
        }
        return max_err;
      };

      std::cout<<var_name<<"=[";
      for (sctl::Long i = 1; i < 140; i++) {
        sctl::Long Nt1 = 6*i, Np1 = i;
        auto B1 = Resample(Resample(B0, Nt0,Np0, Nt1,Np1), Nt1,Np1, Nt0,Np0);
        std::cout<<log(max_err(B0, B1))/log(10)<<' ';
      }
      std::cout<<"];\n";

      //if (0) {
      //  auto B1 = Resample(B0, Nt0,Np0, 70*30,14*30);
      //  B1.Write(var_name.c_str());
      //} else {
      //  auto B1 = Resample(B0, Nt0,Np0, 70*30,14*30);
      //  sctl::Vector<Real> B2; B2.Read(var_name.c_str());

      //  Real max_err = 0, max_val = 0;
      //  for (sctl::Long i = 0; i < B1.Dim(); i++) {
      //    max_err = std::max<Real>(max_err, fabs(B1[i]-B2[i]));
      //    max_val = std::max<Real>(max_val, fabs(B2[i]));
      //  }
      //  std::cout<<"Error "<<var_name<<" = "<<max_err/max_val<<'\n';
      //}
    };
    static void fourier_print_spectrum(std::string var_name, const Vector<Real>& X_, Stellarator<Real,ORDER> S_) {
      Long dof = X_.Dim()/(S_.Nsurf()*fourier_dim0*fourier_dim1);
      SCTL_ASSERT(dof * (S_.Nsurf()*fourier_dim0*fourier_dim1) == X_.Dim());
      for (Long i = 0; i < S_.Nsurf(); i++) {
        const Long Mt = S_.NTor(i);
        const Long Mp = S_.NPol(i);
        const Long offset = S_.ElemDsp(i);
        const Long Nt = fourier_dim0;
        const Long Np = fourier_dim1;
        const Vector<Real> X(dof*Nt*Np, (Iterator<Real>)X_.begin() + dof * i * (fourier_dim0*fourier_dim1), false);
        fourier_print_spectrum(var_name+std::to_string(i),X,Nt,Np);
      }
      std::cout<<"\n";
    }

    static void filter_deprecated(const Stellarator<Real,ORDER>& S, const Comm& comm, Vector<ElemBasis>& f, Real sigma) {
      Long dof = f.Dim() / S.NElem();
      SCTL_ASSERT(f.Dim() == S.NElem() * dof);
      for (Long i = 0; i < S.Nsurf()-1; i++) {
        const Long Mt = S.NTor(i);
        const Long Mp = S.NPol(i);
        const Long Nelem = Mt * Mp;
        const Long offset = S.ElemDsp(i);
        const Long Nt = Mt * ORDER * 4;
        const Long Np = Mp * ORDER * 4;

        Vector<ElemBasis> f_(Nelem*dof, f.begin() + offset*dof, false);
        Vector<Real> f_fourier = cheb2grid(f_, Mt, Mp, Nt, Np);
        fourier_filter(f_fourier, Nt, Np, 0.25 * sigma);
        f_ = grid2cheb(f_fourier, Nt, Np, Mt, Mp);
      }
    };

    static Vector<Real> cheb2grid(const Stellarator<Real,ORDER>& S, const Vector<ElemBasis>& f) {
      const Long Nnodes = ElemBasis::Size();
      const Long Nelem = S.NElem();
      const Long dof = f.Dim() / Nelem;
      SCTL_ASSERT(Nelem * dof == f.Dim());

      Vector<Real> f_fourier(dof * S.Nsurf() * (fourier_dim0*fourier_dim1));
      for (Long i = 0; i < S.Nsurf(); i++) {
        const Long Mt = S.NTor(i);
        const Long Mp = S.NPol(i);
        const Long offset = S.ElemDsp(i);
        const Long Nt = fourier_dim0;
        const Long Np = fourier_dim1;

        const Vector<ElemBasis> f_(Mt*Mp*dof, (Iterator<ElemBasis>)f.begin() + offset*dof, false);
        Vector<Real> f_fourier_(dof*Nt*Np, f_fourier.begin() + dof*i * (fourier_dim0*fourier_dim1), false);
        f_fourier_ = cheb2grid(f_, Mt, Mp, Nt, Np);
        SCTL_ASSERT(f_fourier_.Dim() == dof*Nt*Np);
      }
      return f_fourier;
    }
    static Vector<ElemBasis> grid2cheb(const Stellarator<Real,ORDER>& S, const Vector<Real>& f_fourier) {
      const Long Nnodes = ElemBasis::Size();
      const Long Nelem = S.NElem();
      const Long dof = f_fourier.Dim() / (S.Nsurf() * (fourier_dim0*fourier_dim1));
      SCTL_ASSERT(dof * (S.Nsurf() * (fourier_dim0*fourier_dim1)) == f_fourier.Dim());

      Vector<ElemBasis> f(Nelem * dof);
      for (Long i = 0; i < S.Nsurf(); i++) {
        const Long Mt = S.NTor(i);
        const Long Mp = S.NPol(i);
        const Long offset = S.ElemDsp(i);
        const Long Nt = fourier_dim0;
        const Long Np = fourier_dim1;

        Vector<ElemBasis> f_(Mt*Mp*dof, f.begin() + offset*dof, false);
        const Vector<Real> f_fourier_(dof*Nt*Np, (Iterator<Real>)f_fourier.begin() + dof*i * (fourier_dim0*fourier_dim1), false);
        f_ = grid2cheb(f_fourier_, Nt, Np, Mt, Mp);
        SCTL_ASSERT(f_.Dim() == Mt*Mp*dof);
      }
      return f;
    }

    template <class Real, class GradOp> static Long GradientDescent3(GradOp& grad_fn, Eigen::VectorXd& x, Real& fx, Long max_iter, Real tol) {
      Real dt = 0.1;
      Real step_tol = 1e-1;
      for (Long iter = 0; iter < max_iter; iter++) {
        auto time_step = [](Eigen::VectorXd* x, const Real dt, const Eigen::VectorXd& x0_, GradOp& F, Real* error = nullptr) {
          Long N = x0_.size();
          Eigen::VectorXd F0(N), F1(N), F2(N), F3(N);

          F(x0_, F0);
          F(x0_ - F0 * (0.50*dt), F1);
          F(x0_ - F1 * (0.75*dt), F2);
          F(x0_ - F0 * (2.0/9.0*dt) - F1 * (1.0/3.0*dt) - F2 * (4.0/9.0*dt), F3);
          x[0] = x0_ - F0 * (7.0/24.0*dt) - F1 * (1.0/4.0*dt) - F2 * (1.0/3.0*dt) - F3 * (1.0/8.0*dt);

          Real err = 0;
          const Eigen::VectorXd err_vec = x[0] - (x0_ - F0 * (2.0/9.0*dt) - F1 * (1.0/3.0*dt) - F2 * (4.0/9.0*dt));
          for (Long i = 0; i < err_vec.size(); i++) err = std::max(err, fabs(err_vec[i]));
          if (error != nullptr) error[0] = err / dt;
        };

        Real error;
        Eigen::VectorXd x_(x.size());
        time_step(&x_, dt, x, grad_fn, &error);
        if (error > 2.0 * step_tol) {
          dt *= 0.5;
        } else {
          x = x_;
        }
        if (error < 0.5 * step_tol) dt *= 1.4;
        std::cout<<"Error = "<<error<<"  dt = "<<dt<<'\n';
      }
      return max_iter;
    }

    template <class Real, class GradOp> static Long GradientDescent2(GradOp& grad_fn, Eigen::VectorXd& x, Real& fx, Long max_iter, Real tol) {
      Real dt = 0.1;
      Real step_tol = 1e-1;
      for (Long iter = 0; iter < max_iter; iter++) {
        auto time_step = [](Eigen::VectorXd* x, const Real dt, const Eigen::VectorXd& x0_, GradOp& F, Real* error = nullptr) {
          Long N = x0_.size();
          Eigen::VectorXd F0(N), F1(N), F2(N), F3(N);

          F(x0_, F0);
          F(x0_ - F0 * dt, F1);
          F(x0_ - F0 * (0.5 * dt) - F1 * (0.5 * dt), F2);
          F(x0_ - F0 * (0.5 * dt) - F2 * (0.5 * dt), F3);
          x[0] = x0_ - F0 * (0.5 * dt) - F3 * (0.5 * dt);

          Real err0 = 0, err1 = 0;
          const Eigen::VectorXd err_vec0 = (F1 - F3) * dt;
          const Eigen::VectorXd err_vec1 = (F2 - F3) * dt;
          for (Long i = 0; i < err_vec0.size(); i++) err0 = std::max(err0, fabs(err_vec0[i]));
          for (Long i = 0; i < err_vec1.size(); i++) err1 = std::max(err1, fabs(err_vec1[i]));
          if (error != nullptr) {
            error[0] = err1 / dt;
            if (err1 > err0) error[0] = 1;
          }
        };

        Real error;
        Eigen::VectorXd x_(x.size());
        time_step(&x_, dt, x, grad_fn, &error);
        if (error > 2.0 * step_tol) {
          dt *= 0.5;
        } else {
          x = x_;
        }
        if (error < 0.5 * step_tol) dt *= 1.4;
        std::cout<<"Error = "<<error<<"  dt = "<<dt<<'\n';
      }
      return max_iter;
    }

    template <class Real, class GradOp> static Long GradientDescent2_(GradOp& grad_fn, Eigen::VectorXd& x, Real& fx, Long max_iter, Real tol) {
      Real dt_ = 0, fx_ = 0;
      Eigen::VectorXd x_ = x;

      Real dt = 0.296951;
      for (Long iter = 0; iter < max_iter; iter++) {
        Eigen::VectorXd grad0(x.size()), grad1(x.size()), grad2(x.size());
        fx_ = fx;
        fx = grad_fn(x, grad0);
        if (fx < tol) return iter;
        if (iter && fx > fx_) {
          x = x_;
          fx = fx_;
          dt = 0.5 * dt_;
          continue;
        } else {
          x_ = x;
          dt_ = dt;
        }
        { // Update dt
          Real fx1, fx2;
          fx1 = grad_fn(x - grad0 * dt * 0.50, grad1);
          fx1 = grad_fn(x - grad0 * dt * 0.25 - grad1 * dt * 0.25, grad1);
          fx2 = grad_fn(x - grad1 * dt, grad2);
          fx2 = grad_fn(x - grad0 * dt * (1.0/6.0) - grad1 * dt * (2.0/3.0) - grad2 * dt * (1.0/6.0), grad2);

          Real s;
          { // Calculate optimal step size dt
            Real a = 2*fx - 4*fx1 + 2*fx2;
            Real b =-3*fx + 4*fx1 -   fx2;
            Real c = fx;
            s = -b/(2*a);

            Real fx_ = a*s*s + b*s + c;
            std::cout<<"g = "<<fx_<<' ';
            std::cout<<fx<<' ';
            std::cout<<fx1<<' ';
            std::cout<<fx2<<' ';
            std::cout<<dt<<'\n';
          }
          x = x - grad0 * dt * (2*s*s*s/3 - 3*s*s/2 + s) - grad1 * dt * (-4*s*s*s/3 + 2*s*s) - grad2 * dt * (2*s*s*s/3 - s*s/2);
          dt *= s;
        }
      }
      return max_iter;
    }

    template <class Real, class GradOp> static Long GradientDescent(GradOp& grad_fn, Eigen::VectorXd& x, Real& fx, Long max_iter, Real tol) {
      Real dt = 0.192081;
      for (Long iter = 0; iter < max_iter; iter++) {
        Eigen::VectorXd grad(x.size());
        fx = grad_fn(x, grad);
        { // Update dt
          Eigen::VectorXd grad_(x.size());
          Eigen::VectorXd x1 = x - grad * dt * 0.5;
          Eigen::VectorXd x2 = x - grad * dt * 1.0;
          Real fx1 = grad_fn(x1, grad_);
          Real fx2 = grad_fn(x2, grad_);

          { // Calculate optimal step size dt
            Real a = 2*fx - 4*fx1 + 2*fx2;
            Real b =-3*fx + 4*fx1 -   fx2;
            Real c = fx;
            Real s = -b/(2*a);
            dt *= s;

            Real fx_ = a*s*s + b*s + c;
            std::cout<<"g = "<<fx_<<' ';
            std::cout<<fx<<' ';
            std::cout<<fx1<<' ';
            std::cout<<fx2<<' ';
            std::cout<<dt<<'\n';
          }
        }
        x = x - grad * dt;
        if (fx < tol) return iter;
      }
      return max_iter;
    }

    template <class ValueType> static ValueType QuadraticInterp(Real t, const ValueType& v0, const ValueType& v1, const ValueType& v2, Real t0, Real t1, Real t2) {
      ValueType v = v0 * (((t-t1)*(t-t2))/((t0-t1)*(t0-t2)));
      v += v1 * (((t-t0)*(t-t2))/((t1-t0)*(t1-t2)));
      v += v2 * (((t-t0)*(t-t1))/((t2-t0)*(t2-t1)));
      return v;
    }

    template <class Real, class GradOp> static Long GradientDescentNew(GradOp& grad_fn, Eigen::VectorXd& x0, Real& fx, Long max_iter, Real tol) {
      auto compute_inner_prod = [](const Eigen::VectorXd& A, const Eigen::VectorXd& B) {
        Real sum = 0;
        Long N = A.size();
        SCTL_ASSERT(B.size() == N);
        for (Long i = 0; i < N; i++) sum += A[i] * B[i];
        return sum;
      };
      auto compute_dt_scale = [&compute_inner_prod](const Eigen::VectorXd& x, const Eigen::VectorXd& y){
        Real dot_prod = compute_inner_prod(x, y) / sqrt<Real>(compute_inner_prod(x,x)*compute_inner_prod(y,y));
        return 0.05 / sqrt<Real>(1-dot_prod*dot_prod);
      };

      Eigen::VectorXd grad, x = x0;
      Real g = grad_fn(x, grad);

      Real t = 0, dt = -0.1;
      for (Long j = 0; j < max_iter; j++) {
        Eigen::VectorXd grad0, grad1, grad2;
        Real g0 = grad_fn(x + grad*(dt*0.5) + grad *(dt*0.5), grad0);
        Real g1 = grad_fn(x + grad*(dt*0.5) + grad0*(dt*0.5), grad1);
        Real g2 = grad_fn(x + grad*(dt*0.5) + grad1*(dt*0.5), grad2);


        Eigen::VectorXd grad4, grad3;
        Real c = compute_inner_prod(grad2-grad1,grad2-grad1) / compute_inner_prod(grad1-grad0,grad2-grad1);
        grad3 = (grad2 - grad1*c) * (1/(1-c));
        Real g3 = grad_fn(x + grad*(dt*0.5) + grad3*(dt*0.5), grad4);



        std::cout<<sqrt<Real>(compute_inner_prod(grad -grad4,grad -grad4)/compute_inner_prod(grad4,grad4))<<' ';
        std::cout<<sqrt<Real>(compute_inner_prod(grad -grad3,grad -grad3)/compute_inner_prod(grad4,grad4))<<' ';
        std::cout<<sqrt<Real>(compute_inner_prod(grad -grad2,grad -grad2)/compute_inner_prod(grad4,grad4))<<' ';
        std::cout<<sqrt<Real>(compute_inner_prod(grad -grad1,grad -grad1)/compute_inner_prod(grad4,grad4))<<' ';
        std::cout<<sqrt<Real>(compute_inner_prod(grad -grad0,grad -grad0)/compute_inner_prod(grad4,grad4))<<'\n';
        std::cout<<sqrt<Real>(compute_inner_prod(grad0-grad4,grad0-grad4)/compute_inner_prod(grad4,grad4))<<' ';
        std::cout<<sqrt<Real>(compute_inner_prod(grad0-grad3,grad0-grad3)/compute_inner_prod(grad4,grad4))<<' ';
        std::cout<<sqrt<Real>(compute_inner_prod(grad0-grad2,grad0-grad2)/compute_inner_prod(grad4,grad4))<<' ';
        std::cout<<sqrt<Real>(compute_inner_prod(grad0-grad1,grad0-grad1)/compute_inner_prod(grad4,grad4))<<'\n';
        std::cout<<sqrt<Real>(compute_inner_prod(grad1-grad4,grad1-grad4)/compute_inner_prod(grad4,grad4))<<' ';
        std::cout<<sqrt<Real>(compute_inner_prod(grad1-grad3,grad1-grad3)/compute_inner_prod(grad4,grad4))<<' ';
        std::cout<<sqrt<Real>(compute_inner_prod(grad1-grad2,grad1-grad2)/compute_inner_prod(grad4,grad4))<<'\n';
        std::cout<<sqrt<Real>(compute_inner_prod(grad2-grad4,grad2-grad4)/compute_inner_prod(grad4,grad4))<<' ';
        std::cout<<sqrt<Real>(compute_inner_prod(grad2-grad3,grad2-grad3)/compute_inner_prod(grad4,grad4))<<'\n';
        std::cout<<sqrt<Real>(compute_inner_prod(grad3-grad4,grad3-grad4)/compute_inner_prod(grad4,grad4))<<'\n';




        Real s0 = sqrt<Real>(compute_inner_prod(grad -grad4,grad -grad4)/compute_inner_prod(grad4,grad4));
        Real s1 = sqrt<Real>(compute_inner_prod(grad0-grad1,grad0-grad1)/compute_inner_prod(grad4,grad4));
        Real s2 = sqrt<Real>(compute_inner_prod(grad1-grad2,grad1-grad2)/compute_inner_prod(grad4,grad4));
        Real s3 = sqrt<Real>(compute_inner_prod(grad3-grad4,grad3-grad4)/compute_inner_prod(grad4,grad4));
        Real dt_scale = (0.1/s0); // 0.5*(s1/s2);
        if (s3 > s2) dt_scale = std::min(0.5, dt_scale);

        std::cout<<s0<<' '<<s1<<' '<<s2<<' '<<s3<<"        dt_scale = "<<dt_scale<<'\n'; ////////////////////

        if (dt_scale>0.5) {
          t += dt;
          g = g3;
          x = x + grad *(dt*0.5) + grad3*(dt*0.5);
          grad = grad4;
          std::cout<<"iter = "<<j<<"   t = "<<t<<"   g = "<<g<<"   dt = "<<dt<<'\n';
        } else {
          j--;
        }
        dt *= dt_scale;
      }
      return 0;
    }

    template <class Real, class GradOp> static Long GradientDescentNew__(GradOp& grad_fn, Eigen::VectorXd& x, Real& fx, Long max_iter, Real tol) {
      auto compute_inner_prod = [](const Eigen::VectorXd& A, const Eigen::VectorXd& B) {
        Real sum = 0;
        Long N = A.size();
        SCTL_ASSERT(B.size() == N);
        for (Long i = 0; i < N; i++) sum += A[i] * B[i];
        return sum;
      };
      auto compute_dt_scale = [&compute_inner_prod](const Eigen::VectorXd& x, const Eigen::VectorXd& y){
        Real dot_prod = compute_inner_prod(x, y) / sqrt<Real>(compute_inner_prod(x,x)*compute_inner_prod(y,y));
        return 0.05 / sqrt<Real>(1-dot_prod*dot_prod);
      };

      Eigen::VectorXd x0, grad0, grad_op0;
      Eigen::VectorXd x1, grad1, grad_op1;
      Eigen::VectorXd x2, grad2, grad_op2;

      x0 = x;
      Real t = 0, dt = -0.1;
      for (Long j = 0; j < max_iter; j++) {
        Real g0 = grad_fn(x0, grad0);
        Real g1 = grad_fn(x0 + grad0*(dt*0.1), grad1);
        Real g2 = grad_fn(x0 + grad0*(dt*0.1) + grad1*(dt*0.1), grad2);

        // Fit v = v0 + v1 exp(-alpha * dt)
        auto v1 = (grad1-grad0) * (1/(dt*0.1));
        Real alpha = (sqrt<Real>(compute_inner_prod(grad1-grad0,grad1-grad0))/sqrt<Real>(compute_inner_prod(grad2-grad1,grad2-grad1))-1) / (0.1*dt);
        auto v0 = grad0 - v1;

        x0 = x0 + v0*dt - v1 * ((exp<Real>(-alpha*dt)-1.0)/alpha);
        t += dt;

        std::cout<<"iter = "<<j<<"   t = "<<t<<"   g = "<<g0<<"   dt = "<<dt<<'\n';
      }
      return 0;
    }

    template <class Real, class GradOp> static Long GradientDescentNew_(GradOp& grad_fn, Eigen::VectorXd& x0, Real& fx, Long max_iter, Real tol) {
      constexpr Long order = 3;
      Eigen::VectorXd x[order], grad[order], grad_op[order];
      Real g[order], t[order];

      auto compute_inner_prod = [](const Eigen::VectorXd& A, const Eigen::VectorXd& B) {
        Real sum = 0;
        Long N = A.size();
        SCTL_ASSERT(B.size() == N);
        for (Long i = 0; i < N; i++) sum += A[i] * B[i];
        return sum;
      };
      auto compute_dt_scale = [&compute_inner_prod](const Eigen::VectorXd& x, const Eigen::VectorXd& y){
        Real dot_prod = compute_inner_prod(x, y) / sqrt<Real>(compute_inner_prod(x,x)*compute_inner_prod(y,y));
        return 0.05 / sqrt<Real>(1-dot_prod*dot_prod);
      };

      t[1] = 0;
      x[1] = x0;
      g[1] = grad_fn(x[1], grad[1], &grad_op[1]);

      Real dt = -0.1;
      for (Long j = 1; j < order-1; j++) {
        //t[0] = t[1] + dt;
        //x[0] = x[1] + grad[1] * dt;
        //g[0] = grad_fn(x[0], grad[0], &grad_op[0]);

        //Real dot_prod = compute_inner_prod(grad[0], grad[1]) / sqrt<Real>(compute_inner_prod(grad[0], grad[0])*compute_inner_prod(grad[1], grad[1]));
        //Real dt_scale = 0.1 / sqrt<Real>(1-dot_prod*dot_prod);
        //dt *= dt_scale;

        //if (dt_scale < 0.5 || g[0] >= g[1]) {
        //  j--;
        //  continue;
        //}

        //std::cout<<"iter = "<<j<<"   t = "<<t[0]<<"   g = "<<g[0]<<"   dt = "<<dt<<'\n';
        //for (Long i = order-1; i >= 1; i--) {
        //  x[i] = x[i-1];
        //  grad[i] = grad[i-1];
        //  grad_op[i] = grad_op[i-1];
        //  g[i] = g[i-1];
        //  t[i] = t[i-1];
        //}
      }

      for (Long j = 0; j < max_iter; j++) {
        Eigen::VectorXd x_, grad_, grad_op_;
        Real g_ = grad_fn(x[1] + grad[1]*(dt*0.5), grad_, &grad_op_);
        Real dt_scale = compute_dt_scale(grad_, grad[1]);
        dt *= dt_scale;
        if (dt_scale < 0.5 || g[0] >= g[1]) {
          std::cout<<"dt = "<<dt<<'\n';
          j--;
          continue;
        }

        t[0] = t[1] + dt;
        x[0] = x[1] + grad_*dt;
        g[0] = grad_fn(x[0], grad[0], &grad_op[0]);

        dt_scale = compute_dt_scale(grad[0], grad_);
        dt *= dt_scale;
        if (dt_scale < 0.5 || g[0] >= g[1]) {
          std::cout<<"dt =  "<<dt<<'\n';
          j--;
          continue;
        }

        std::cout<<"iter = "<<j<<"   t = "<<t[0]<<"   g = "<<g[0]<<"   dt = "<<dt<<'\n';
        for (Long i = order-1; i >= 1; i--) {
          x[i] = x[i-1];
          grad[i] = grad[i-1];
          grad_op[i] = grad_op[i-1];
          g[i] = g[i-1];
          t[i] = t[i-1];
        }
      }


      for (Long j = 0; j < max_iter; j++) {
        Eigen::VectorXd x_, grad_, grad_op_;
        Real g_ = grad_fn(x[1] + grad[1]*(dt*0.5), grad_, &grad_op_);
        if (compute_dt_scale(grad_, grad[1])<0.5 ||  g_ > g[1]) {
          dt = dt * compute_dt_scale(grad_, grad[1]);
          std::cout<<"dt = "<<dt<<'\n';
          j--;
          continue;
        }

        t[0] = t[1] + dt;
        x[0] = x[1] + grad_*dt;
        g[0] = grad_fn(x[0], grad[0], &grad_op[0]);

        Real dt_scale = compute_dt_scale(grad[0], grad[1]); /// todo use grad_ instead of grad[1]
        dt *= dt_scale;
        if (dt_scale < 0.5 || g[0] >= g[1]) {
          std::cout<<"dt =  "<<dt<<'\n';
          j--;
          continue;
        }

        std::cout<<"iter = "<<j<<"   t = "<<t[0]<<"   g = "<<g[0]<<"   dt = "<<dt<<'\n';
        for (Long i = order-1; i >= 1; i--) {
          x[i] = x[i-1];
          grad[i] = grad[i-1];
          grad_op[i] = grad_op[i-1];
          g[i] = g[i-1];
          t[i] = t[i-1];
        }
      }

      //for (Long iter = 0; iter < max_iter; iter++) {
      //  t[0] = t[1] + dt;
      //  x[0] = x[1] + grad[1] * dt;
      //  g[0] = grad_fn(x[0], grad[0], &grad_op[0]);

      //  Real dot_prod = compute_inner_prod(grad[0], grad[1]) / sqrt<Real>(compute_inner_prod(grad[0], grad[0])*compute_inner_prod(grad[1], grad[1]));
      //  if (dot_prod < 0.95 || g[0] < g[1]) {
      //    dt = dt * 1.5;
      //  }
      //  if (dot_prod < 0.85 || g[0] >= g[1]) {
      //    dt = dt * 0.5;
      //    continue;
      //  }

      //  for (Long i = order-2; i >= 0; i--) {
      //    x[i+1] = x[i];
      //    grad[i+1] = grad[i];
      //    grad_op[i+1] = grad_op[i];
      //    g[i+1] = g[i];
      //    t[i+1] = t[i];
      //  }
      //}


      //for (Long iter = 0; iter < max_iter; iter++) {
      //  fx = grad_fn(x, grad, &grad_op);
      //  { // Update dt
      //    Eigen::VectorXd grad_(x.size());
      //    Eigen::VectorXd x1 = x - grad * dt * 0.5;
      //    Eigen::VectorXd x2 = x - grad * dt * 1.0;
      //    Real fx1 = grad_fn(x1, grad_);
      //    Real fx2 = grad_fn(x2, grad_);

      //    { // Calculate optimal step size dt
      //      Real a = 2*fx - 4*fx1 + 2*fx2;
      //      Real b =-3*fx + 4*fx1 -   fx2;
      //      Real c = fx;
      //      Real s = -b/(2*a);
      //      dt *= s;

      //      Real fx_ = a*s*s + b*s + c;
      //      std::cout<<"g = "<<fx_<<' ';
      //      std::cout<<fx<<' ';
      //      std::cout<<fx1<<' ';
      //      std::cout<<fx2<<' ';
      //      std::cout<<dt<<'\n';
      //    }
      //  }
      //  x = x - grad * dt;
      //  if (fx < tol) return iter;
      //}
      //return max_iter;
      return 0;
    }


  public:
    MHDEquilib(const Stellarator<Real,ORDER>& S, const Vector<Real>& pressure, const Vector<Real>& flux_tor, const Vector<Real>& flux_pol) {
      S_ = S;
      pressure_ = pressure;
      flux_tor_ = flux_tor;
      flux_pol_ = flux_pol;
      iter = 57;
    }

    Real operator()(const Eigen::VectorXd& x, Eigen::VectorXd& grad_direction, Eigen::VectorXd* grad_op_ptr = nullptr) {
      auto compute_H = [] (const ElemList<COORD_DIM,ElemBasis>& elem_lst, const Vector<ElemBasis>& normal) {
        const Long Nnodes = ElemBasis::Size();
        const Long Nelem = elem_lst.NElem();

        const Vector<ElemBasis> X = elem_lst.ElemVector();
        Vector<ElemBasis> dX, d2X, H(Nelem);
        ElemBasis::Grad(dX, X);
        ElemBasis::Grad(d2X, dX);
        for (Long i = 0; i < Nelem; i++) {
          for (Long j = 0; j < Nnodes; j++) {
            Tensor<Real,true,2,2> I, invI, II;
            for (Long k0 = 0; k0 < 2; k0++) {
              for (Long k1 = 0; k1 < 2; k1++) {
                I(k0,k1) = 0;
                I(k0,k1) += dX[(i*COORD_DIM+0)*2+k0][j] * dX[(i*COORD_DIM+0)*2+k1][j];
                I(k0,k1) += dX[(i*COORD_DIM+1)*2+k0][j] * dX[(i*COORD_DIM+1)*2+k1][j];
                I(k0,k1) += dX[(i*COORD_DIM+2)*2+k0][j] * dX[(i*COORD_DIM+2)*2+k1][j];

                II(k0,k1) = 0;
                II(k0,k1) += d2X[(i*COORD_DIM+0)*4+k0*2+k1][j] * normal[i*COORD_DIM+0][j];
                II(k0,k1) += d2X[(i*COORD_DIM+1)*4+k0*2+k1][j] * normal[i*COORD_DIM+1][j];
                II(k0,k1) += d2X[(i*COORD_DIM+2)*4+k0*2+k1][j] * normal[i*COORD_DIM+2][j];
              }
            }
            { // Set invI
              Real detI = I(0,0)*I(1,1)-I(0,1)*I(1,0);
              invI(0,0) = I(1,1) / detI;
              invI(0,1) = -I(0,1) / detI;
              invI(1,0) = -I(1,0) / detI;
              invI(1,1) = I(0,0) / detI;
            }
            { // Set H
              H[i][j] = 0;
              H[i][j] += -0.5 * II(0,0)*invI(0,0);
              H[i][j] += -0.5 * II(0,1)*invI(0,1);
              H[i][j] += -0.5 * II(1,0)*invI(1,0);
              H[i][j] += -0.5 * II(1,1)*invI(1,1);
            }
          }
        }
        return H;
      };

      const Comm comm = Comm::World();
      const Long Nelem = S_.NElem();
      const Long Nnodes = ElemBasis::Size();

      Vector<Real> X_fourier(x.size());
      for (Long i = 0; i < x.size(); i++) { // Set X_fourier
        X_fourier[i] = x(i);
      }
      { // Write to file
        X_fourier.Write(("X"+std::to_string(iter)).c_str());
      }
      Vector<ElemBasis> X = grid2cheb(S_, X_fourier);
      for (Long i = 0; i < Nelem; i++) { // Set S_
        for (Long j = 0; j < Nnodes; j++) {
          S_.Elem(i,0)[j] = X[i*COORD_DIM+0][j];
          S_.Elem(i,1)[j] = X[i*COORD_DIM+1][j];
          S_.Elem(i,2)[j] = X[i*COORD_DIM+2][j];
        }
      }

      Vector<ElemBasis> normal, area_elem;
      Stellarator<Real,ORDER>::compute_norm_area_elem(S_, normal, area_elem);

      Real g;
      //Vector<ElemBasis> dgdnu = Stellarator<Real,ORDER>::compute_gradient(S_, pressure_, flux_tor_, flux_pol_, &g);
      Vector<ElemBasis> dgdnu = Stellarator<Real,ORDER>::compute_pressure_jump(S_, pressure_, flux_tor_, flux_pol_, &g);

      Vector<Real> grad_direction_, grad_op_, grad_direction_orig_;
      { // Set grad_direction_ and filter
        Vector<ElemBasis> H = compute_H(S_.GetElemList(), normal);
        Vector<Real> H_fourier = cheb2grid(S_, H);

        grad_op_.ReInit(x.size());
        grad_direction_.ReInit(x.size());
        grad_direction_orig_.ReInit(x.size());
        Vector<Real> dgdnu_fourier = cheb2grid(S_, dgdnu);
        for (Long i = 0; i < S_.Nsurf(); i++) { // Init grad_direction_, make it divergence-free, filter
          const Long Mt = S_.NTor(i);
          const Long Mp = S_.NPol(i);
          const Long offset = S_.ElemDsp(i);
          const Long Nt = fourier_dim0;
          const Long Np = fourier_dim1;

          const Vector<Real>    dgdnu(          Nt*Np, (Iterator<Real>)dgdnu_fourier.begin()   +           i * (fourier_dim0*fourier_dim1), false);
          const Vector<Real>        H(          Nt*Np, (Iterator<Real>)H_fourier.begin()       +           i * (fourier_dim0*fourier_dim1), false);
          const Vector<Real>        X(COORD_DIM*Nt*Np, (Iterator<Real>)X_fourier.begin()       + COORD_DIM*i * (fourier_dim0*fourier_dim1), false);
          Vector<Real> grad_direction(COORD_DIM*Nt*Np, (Iterator<Real>)grad_direction_.begin() + COORD_DIM*i * (fourier_dim0*fourier_dim1), false);
          Vector<Real>        grad_op(COORD_DIM*Nt*Np, (Iterator<Real>)grad_op_.begin()        + COORD_DIM*i * (fourier_dim0*fourier_dim1), false);

          Vector<Real> dX, d2X, normal, area_elem;
          biest::SurfaceOp<Real> Sop(comm, Nt, Np);
          Sop.Grad2D(dX, X);
          Sop.Grad2D(d2X, dX);
          Sop.SurfNormalAreaElem(&normal, &area_elem, dX, nullptr);

          for (Long j = 0; j < Nt*Np; j++) { // Set grad_op
            grad_op[0*Nt*Np+j] = normal[0*Nt*Np+j] * dgdnu[j] * area_elem[j];
            grad_op[1*Nt*Np+j] = normal[1*Nt*Np+j] * dgdnu[j] * area_elem[j];
            grad_op[2*Nt*Np+j] = normal[2*Nt*Np+j] * dgdnu[j] * area_elem[j];
          }

          if (i < S_.Nsurf() - 1) { // Set grad_direction
            Vector<Real> F(Nt*Np), GradInvLapF;
            for (Long j = 0; j < Nt*Np; j++) { // Set F <-- 2H * dgdnu
              F[j] = 2*H[j] * dgdnu[j];
            }
            Sop.GradInvSurfLap(GradInvLapF, dX, d2X, F, 1e-8, 100, 1.0);

            for (Long j = 0; j < Nt*Np; j++) { // grad_direction <-- normal * dgdnu - GradInvLapF
              grad_direction[0*Nt*Np+j] = normal[0*Nt*Np+j] * dgdnu[j] - GradInvLapF[0*Nt*Np+j]*0;
              grad_direction[1*Nt*Np+j] = normal[1*Nt*Np+j] * dgdnu[j] - GradInvLapF[1*Nt*Np+j]*0;
              grad_direction[2*Nt*Np+j] = normal[2*Nt*Np+j] * dgdnu[j] - GradInvLapF[2*Nt*Np+j]*0;
            }
            { ////////////////////
              Vector<Real> grad_direction_orig(COORD_DIM*Nt*Np, (Iterator<Real>)grad_direction_orig_.begin() + COORD_DIM*i * (fourier_dim0*fourier_dim1), false);
              grad_direction_orig = grad_direction;
            }
            fourier_filter(grad_direction, Nt, Np, 2);

            for (Long k = 0; k < 50; k++) { // reparameterize
              Real max_resid = 0;
              Vector<Real> correction(COORD_DIM*Nt*Np);
              for (Long i = 0; i < Nt*Np; i++) { // Set correction
                Real resid = dgdnu[i];
                resid -= normal[0*Nt*Np+i] * grad_direction[0*Nt*Np+i];
                resid -= normal[1*Nt*Np+i] * grad_direction[1*Nt*Np+i];
                resid -= normal[2*Nt*Np+i] * grad_direction[2*Nt*Np+i];
                max_resid = std::max(max_resid, fabs(resid));

                correction[0*Nt*Np+i] = normal[0*Nt*Np+i] * resid;
                correction[1*Nt*Np+i] = normal[1*Nt*Np+i] * resid;
                correction[2*Nt*Np+i] = normal[2*Nt*Np+i] * resid;
              }
              std::cout<<max_resid<<' '; //////////////////////////////////
              fourier_filter(correction, Nt, Np, 2);

              Real alpha = 0;
              { // Set alpha <-- (dgdnu - x.n, c.n) / (c.n, c.n)
                Real dgdnu_xncn = 0, cncn = 0, max_c = 0;
                for (Long i = 0; i < Nt*Np; i++) {
                  max_c = std::max<Real>(max_c, correction[0*Nt*Np+i]);
                  max_c = std::max<Real>(max_c, correction[1*Nt*Np+i]);
                  max_c = std::max<Real>(max_c, correction[2*Nt*Np+i]);

                  Real resid = dgdnu[i];
                  resid -= normal[0*Nt*Np+i] * grad_direction[0*Nt*Np+i];
                  resid -= normal[1*Nt*Np+i] * grad_direction[1*Nt*Np+i];
                  resid -= normal[2*Nt*Np+i] * grad_direction[2*Nt*Np+i];

                  Real cn = 0;
                  cn += correction[0*Nt*Np+i] * normal[0*Nt*Np+i];
                  cn += correction[1*Nt*Np+i] * normal[1*Nt*Np+i];
                  cn += correction[2*Nt*Np+i] * normal[2*Nt*Np+i];

                  dgdnu_xncn += resid*cn;
                  cncn += cn*cn;
                }
                alpha = dgdnu_xncn / cncn;
                std::cout<<alpha*max_c<<'\n';
              }
              grad_direction += correction * alpha;
            }
            //fourier_print_spectrum("dgdnu", dgdnu, Nt, Np);
            //fourier_print_spectrum("grad_dir", grad_direction, Nt, Np);
          } else {
            grad_direction = 0;
          }
        }
      }

      /////////////////////////////////////////////////////////////////////////
      fourier_print_spectrum("X", X_fourier, S_);
      fourier_print_spectrum("normal", cheb2grid(S_,normal), S_);

      fourier_print_spectrum("dgdnu", cheb2grid(S_, dgdnu), S_);
      fourier_print_spectrum("grad_dir", grad_direction_, S_);
      fourier_print_spectrum("grad_dir_orig", grad_direction_orig_, S_);
      /////////////////////////////////////////////////////////////////////////

      { // Set grad_direction <-- grad_direction_
        if (grad_direction.size() != grad_direction_.Dim()) {
          grad_direction = Eigen::VectorXd(grad_direction_.Dim());
        }
        for (Long i = 0; i < grad_direction.size(); i++) {
          grad_direction(i) = grad_direction_[i];
        }
      }
      if (grad_op_ptr) { // Set grad_op_ptr
        grad_op_ptr[0] = Eigen::VectorXd(grad_op_.Dim());
        for (Long i = 0; i < grad_op_ptr->size(); i++) {
          grad_op_ptr[0](i) = grad_op_[i];
        }
      }

      if (1) { // Write VTU
        VTUData vtu;
        vtu.AddElems(S_.GetElemList(), dgdnu, ORDER);
        vtu.WriteVTK("dgdnu"+std::to_string(iter), comm);
      }
      if (1) { // Write VTU
        VTUData vtu;
        Vector<ElemBasis> grad_direction = grid2cheb(S_, grad_direction_);
        Vector<ElemBasis> dgdnu_ = Stellarator<Real,ORDER>::compute_dot_prod(normal, grad_direction);
        vtu.AddElems(S_.GetElemList(), dgdnu_, ORDER);
        vtu.WriteVTK("dgdnu-"+std::to_string(iter), comm);
      }
      if (1) { // Write VTU
        VTUData vtu;
        Vector<ElemBasis> grad_direction = grid2cheb(S_, grad_direction_);
        vtu.AddElems(S_.GetElemList(), grad_direction, ORDER);
        vtu.WriteVTK("grad_direction"+std::to_string(iter), comm);
      }
      std::cout<<"iter = "<<iter<<"    g = "<<g<<'\n';

      iter++;
      return g;
    }

    static void ComputeEquilibrium(MHDEquilib& mhd_equilib) {
      Comm comm = Comm::World();
      const Long Nelem = mhd_equilib.S_.NElem();
      const Long Nnodes = ElemBasis::Size();

      // Initial guess
      Eigen::VectorXd x;
      { // Set x
        Vector<ElemBasis> X(Nelem * COORD_DIM);
        for (Long i = 0; i < Nelem; i++) { // Set x
          X[i*COORD_DIM+0] = mhd_equilib.S_.Elem(i,0);
          X[i*COORD_DIM+1] = mhd_equilib.S_.Elem(i,1);
          X[i*COORD_DIM+2] = mhd_equilib.S_.Elem(i,2);
        }
        Vector<Real> X_fourier = cheb2grid(mhd_equilib.S_, X);
        x.resize(X_fourier.Dim());
        X_fourier.Read(("X"+std::to_string(mhd_equilib.iter)+"_").c_str()); // Read from file
        for (Long i = 0; i < X_fourier.Dim(); i++) {
          x(i) = X_fourier[i];
        }
      }

      Real fx;
      if (0) {
        LBFGSpp::LBFGSParam<Real> param;
        param.max_iterations = 200;
        param.epsilon = 1e-16;
        param.m = 20;
        LBFGSpp::LBFGSSolver<Real> solver(param);
        Integer niter = solver.minimize(mhd_equilib, x, fx);
      } else {
        //Integer niter = GradientDescentNew(mhd_equilib, x, fx, 200, 1e-12);
        Integer niter = GradientDescent(mhd_equilib, x, fx, 200, 1e-12);
      }
      { // Set x
        // TODO
      }
    }

    static void test() {
      Comm comm = Comm::World();
      Profile::Enable(true);

      Long Nsurf = 2;
      Stellarator<Real,ORDER> S;
      Vector<Real> flux_tor(Nsurf), flux_pol(Nsurf), pressure(Nsurf);
      { // Init S, flux_tor, flux_pol, pressure
        Vector<Long> NtNp;
        for (Long i = 0; i < Nsurf; i++) {
          NtNp.PushBack(70);
          NtNp.PushBack(14);
          //NtNp.PushBack(30);
          //NtNp.PushBack(4);
        }
        S = Stellarator<Real,ORDER>(NtNp);
        flux_tor = 1;
        flux_pol = 1;
        pressure = 0;

        //flux_tor[0] = 1; //0.791881512;
        //flux_tor[1] = 1;
        //flux_pol[0] = 0;
        //flux_pol[1] = 0;
        //pressure[0] = 0;
        //pressure[1] = 0;
      }

      MHDEquilib mhd_equilib(S, pressure, flux_tor, flux_pol);
      ComputeEquilibrium(mhd_equilib);
    }

    static void test_() {
      Comm comm = Comm::World();
      Profile::Enable(true);

      Long Nsurf = 2;
      Stellarator<Real,ORDER> S;
      Vector<Real> flux_tor(Nsurf), flux_pol(Nsurf), pressure(Nsurf);
      { // Init S, flux_tor, flux_pol, pressure
        Vector<Long> NtNp;
        for (Long i = 0; i < Nsurf; i++) {
          NtNp.PushBack(50);
          NtNp.PushBack(10);
          //NtNp.PushBack(30);
          //NtNp.PushBack(4);
        }
        S = Stellarator<Real,ORDER>(NtNp);
        flux_tor = 1;
        flux_pol = 1;
        pressure = 0;

        //flux_tor[0] = 1; //0.791881512;
        //flux_tor[1] = 1;
        //flux_pol[0] = 0;
        //flux_pol[1] = 0;
        //pressure[0] = 0;
        //pressure[1] = 0;
      }
      MHDEquilib mhd_equilib(S, pressure, flux_tor, flux_pol);

      const Long Nelem = mhd_equilib.S_.NElem();
      const Long Nnodes = ElemBasis::Size();
      Eigen::VectorXd x, grad_direction;
      { // Read x
        Vector<ElemBasis> X(Nelem * COORD_DIM);
        for (Long i = 0; i < Nelem; i++) {
          X[i*COORD_DIM+0] = S.Elem(i,0);
          X[i*COORD_DIM+1] = S.Elem(i,1);
          X[i*COORD_DIM+2] = S.Elem(i,2);
        }
        Vector<Real> X_fourier = cheb2grid(S, X);
        //X_fourier.Read("X_tmp");
        x.resize(X_fourier.Dim());
        for (Long i = 0; i < X_fourier.Dim(); i++) {
          x(i) = X_fourier[i];
        }
      }
      Real g = mhd_equilib(x, grad_direction);
      { // Write grad_direction
        Vector<Real> dXdt(grad_direction.size());
        for (Long i = 0; i < dXdt.Dim(); i++) {
          dXdt[i] = grad_direction(i);
        }
        dXdt.Write("dXdt_tmp");
      }
    }

  private:
    Stellarator<Real,ORDER> S_;
    Vector<Real> pressure_;
    Vector<Real> flux_tor_;
    Vector<Real> flux_pol_;
    Long iter;
};



template <class Real, Integer ORDER=5> class Spheres {
  static constexpr Integer COORD_DIM = 3;
  static constexpr Integer ELEM_DIM = COORD_DIM-1;
  using PotentialBasis = Basis<Real, ELEM_DIM, ORDER>;
  using DensityBasis = Basis<Real, ELEM_DIM, ORDER>;
  using CoordBasis = Basis<Real, ELEM_DIM, ORDER>;
  using ElemLst = ElemList<COORD_DIM, CoordBasis>;

  public:
    Spheres(Long N = 0) {
      Vector<Real> X(N*COORD_DIM);
      Vector<Real> R(N);
      X=0;
      R=1;
      for (Long i = 0; i < N; i++) X[i*COORD_DIM] = (i==0?-1.015:1.015); ///////////
      InitSpheres(X,R);
    }

    const ElemLst& GetElem() const {
      return elements;
    }

    static void test() {
      constexpr Integer order_singular = 35;
      constexpr Integer order_direct = 35;
      Comm comm = Comm::World();
      Profile::Enable(true);

      Long Ns = 2;
      Spheres S(Ns);
      S.quadrature_FxT.template Setup<DensityBasis, PotentialBasis>(S.GetElem(), S.Stokes_FxT, order_singular, order_direct, -1.0, comm);
      S.quadrature_FxU.template Setup<DensityBasis, PotentialBasis>(S.GetElem(), S.Stokes_FxU, order_singular, order_direct, -1.0, comm);
      S.quadrature_DxU.template Setup<DensityBasis, PotentialBasis>(S.GetElem(), S.Stokes_DxU, order_singular, order_direct, -1.0, comm);

      const auto SetMotion = [&S](Vector<DensityBasis>& density, const Vector<Real>& force_avg, const Vector<Real>& torque_avg) {
        Long Nelem = S.GetElem().NElem();
        Long Nsurf = S.elem_cnt.Dim();
        const auto& X = S.GetElem().ElemVector();

        Vector<Real> area, Xc;
        Vector<DensityBasis> one(Nelem);
        for (Long i = 0; i < Nelem; i++) {
          for (Long j = 0; j < DensityBasis::Size(); j++) {
            one[i][j] = 1;
          }
        }
        S.SurfInteg(area, one);
        S.SurfInteg(Xc, S.GetElem().ElemVector());
        for (Long i = 0; i < Nsurf; i++) {
          for (Long k = 0; k < COORD_DIM; k++) {
            Xc[i*COORD_DIM+k] /= area[i];
          }
        }

        if (density.Dim() != Nelem*COORD_DIM) density.ReInit(Nelem*COORD_DIM);
        Long elem_itr = 0;
        for (Long i = 0; i < Nsurf; i++) {
          for (Long j = 0; j < S.elem_cnt[i]; j++) {
            for (Long k = 0; k < DensityBasis::Size(); k++) {
              StaticArray<Real,COORD_DIM> dX;
              dX[0] = (X[elem_itr*COORD_DIM+0][k] - Xc[i*COORD_DIM+0]);
              dX[1] = (X[elem_itr*COORD_DIM+1][k] - Xc[i*COORD_DIM+1]);
              dX[2] = (X[elem_itr*COORD_DIM+2][k] - Xc[i*COORD_DIM+2]);
              density[elem_itr*COORD_DIM+0][k] = force_avg[i*COORD_DIM+0]*(1/area[i]) + (torque_avg[i*COORD_DIM+1] * dX[2] - torque_avg[i*COORD_DIM+2] * dX[1]) / (2*area[i]/3);
              density[elem_itr*COORD_DIM+1][k] = force_avg[i*COORD_DIM+1]*(1/area[i]) + (torque_avg[i*COORD_DIM+2] * dX[0] - torque_avg[i*COORD_DIM+0] * dX[2]) / (2*area[i]/3);
              density[elem_itr*COORD_DIM+2][k] = force_avg[i*COORD_DIM+2]*(1/area[i]) + (torque_avg[i*COORD_DIM+0] * dX[1] - torque_avg[i*COORD_DIM+1] * dX[0]) / (2*area[i]/3);
            }
            elem_itr++;
          }
        }
      };
      const auto GetMotion = [&S](Vector<Real>& force_avg, Vector<Real>& torque_avg, const Vector<DensityBasis>& density) {
        Long Nelem = S.GetElem().NElem();
        Long Nsurf = S.elem_cnt.Dim();
        const auto& X = S.GetElem().ElemVector();

        S.SurfInteg(force_avg, density);

        Vector<Real> area, Xc;
        Vector<DensityBasis> one(Nelem);
        for (Long i = 0; i < Nelem; i++) {
          for (Long j = 0; j < DensityBasis::Size(); j++) {
            one[i][j] = 1;
          }
        }
        S.SurfInteg(area, one);
        S.SurfInteg(Xc, S.GetElem().ElemVector());
        for (Long i = 0; i < Nsurf; i++) {
          for (Long k = 0; k < COORD_DIM; k++) {
            Xc[i*COORD_DIM+k] /= area[i];
          }
        }

        { // Set torque_avg
          Long elem_itr = 0;
          Vector<DensityBasis> torque(Nelem*COORD_DIM);
          for (Long i = 0; i < Nsurf; i++) {
            for (Long j = 0; j < S.elem_cnt[i]; j++) {
              for (Long k = 0; k < DensityBasis::Size(); k++) {
                StaticArray<Real,COORD_DIM> dX;
                dX[0] = (X[elem_itr*COORD_DIM+0][k] - Xc[i*COORD_DIM+0]);
                dX[1] = (X[elem_itr*COORD_DIM+1][k] - Xc[i*COORD_DIM+1]);
                dX[2] = (X[elem_itr*COORD_DIM+2][k] - Xc[i*COORD_DIM+2]);
                torque[elem_itr*COORD_DIM+0][k] = dX[1] * density[elem_itr*COORD_DIM+2][k] - dX[2] * density[elem_itr*COORD_DIM+1][k];
                torque[elem_itr*COORD_DIM+1][k] = dX[2] * density[elem_itr*COORD_DIM+0][k] - dX[0] * density[elem_itr*COORD_DIM+2][k];
                torque[elem_itr*COORD_DIM+2][k] = dX[0] * density[elem_itr*COORD_DIM+1][k] - dX[1] * density[elem_itr*COORD_DIM+0][k];
              }
              elem_itr++;
            }
          }
          S.SurfInteg(torque_avg, torque);
        }
      };
      const auto BIOpL = [&GetMotion,&SetMotion](Vector<DensityBasis>& potential, const Vector<DensityBasis>& density) {
        Vector<Real> force_avg, torque_avg;
        GetMotion(force_avg, torque_avg, density);
        SetMotion(potential, force_avg, torque_avg);
      };
      const auto BIOpK = [&S](Vector<DensityBasis>& potential, const Vector<DensityBasis>& density) {
        Vector<DensityBasis> traction;
        S.quadrature_FxT.Eval(traction, S.GetElem(), density, S.Stokes_FxT);

        Vector<CoordBasis> dX;
        const auto X = S.GetElem().ElemVector();
        CoordBasis::Grad(dX, X);

        Long Nelem = S.GetElem().NElem();
        Long Nnodes = CoordBasis::Size();
        potential.ReInit(Nelem * COORD_DIM);
        for (Long i = 0; i < Nelem; i++) {
          for (Long j = 0; j < Nnodes; j++) {
            StaticArray<Real,COORD_DIM> Xn;
            Xn[0] = dX[i*COORD_DIM*2+2][j]*dX[i*COORD_DIM*2+5][j] - dX[i*COORD_DIM*2+4][j]*dX[i*COORD_DIM*2+3][j];
            Xn[1] = dX[i*COORD_DIM*2+4][j]*dX[i*COORD_DIM*2+1][j] - dX[i*COORD_DIM*2+0][j]*dX[i*COORD_DIM*2+5][j];
            Xn[2] = dX[i*COORD_DIM*2+0][j]*dX[i*COORD_DIM*2+3][j] - dX[i*COORD_DIM*2+2][j]*dX[i*COORD_DIM*2+1][j];
            Real AreaElem = sqrt<Real>(Xn[0]*Xn[0] + Xn[1]*Xn[1] + Xn[2]*Xn[2]);
            Real OOAreaElem = 1 / AreaElem;
            Xn[0] *= OOAreaElem;
            Xn[1] *= OOAreaElem;
            Xn[2] *= OOAreaElem;

            potential[i*COORD_DIM+0][j] = traction[i*COORD_DIM*COORD_DIM+0][j]*Xn[0] + traction[i*COORD_DIM*COORD_DIM+1][j]*Xn[1] + traction[i*COORD_DIM*COORD_DIM+2][j]*Xn[2];
            potential[i*COORD_DIM+1][j] = traction[i*COORD_DIM*COORD_DIM+3][j]*Xn[0] + traction[i*COORD_DIM*COORD_DIM+4][j]*Xn[1] + traction[i*COORD_DIM*COORD_DIM+5][j]*Xn[2];
            potential[i*COORD_DIM+2][j] = traction[i*COORD_DIM*COORD_DIM+6][j]*Xn[0] + traction[i*COORD_DIM*COORD_DIM+7][j]*Xn[1] + traction[i*COORD_DIM*COORD_DIM+8][j]*Xn[2];
          }
        }
      };
      const auto BIOp_half_K_L = [&S,&BIOpK,&BIOpL](Vector<DensityBasis>& potential, const Vector<DensityBasis>& density) {
        Vector<DensityBasis> potential_K;
        Vector<DensityBasis> potential_L;
        BIOpK(potential_K, density);
        BIOpL(potential_L, density);

        if (potential.Dim() != potential_K.Dim()) {
          potential.ReInit(potential_K.Dim());
        }
        for (Long i = 0; i < potential_K.Dim(); i++) {
          for (Long k = 0; k < DensityBasis::Size(); k++) {
            potential[i][k] = -0.5*density[i][k] + potential_K[i][k] + potential_L[i][k];
          }
        }
      };
      const auto BIOp_half_K = [&S,&BIOpK,&BIOpL](Vector<DensityBasis>& potential, const Vector<DensityBasis>& density) {
        Vector<DensityBasis> potential_K;
        BIOpK(potential_K, density);

        if (potential.Dim() != potential_K.Dim()) {
          potential.ReInit(potential_K.Dim());
        }
        for (Long i = 0; i < potential_K.Dim(); i++) {
          for (Long k = 0; k < DensityBasis::Size(); k++) {
            potential[i][k] = -0.5*density[i][k] + potential_K[i][k];
          }
        }
      };
      const auto BIOp_half_S_D = [&S,&BIOpL](Vector<DensityBasis>& potential, const Vector<DensityBasis>& density) {
        Vector<DensityBasis> U;
        S.quadrature_DxU.Eval(U, S.GetElem(), density, S.Stokes_DxU);

        Vector<PotentialBasis> U1;
        Vector<DensityBasis> sigma1;
        BIOpL(sigma1,density);
        S.quadrature_FxU.Eval(U1, S.GetElem(), sigma1, S.Stokes_FxU);

        Long Nelem = S.GetElem().NElem();
        Long Nnodes = CoordBasis::Size();
        potential.ReInit(Nelem * COORD_DIM);
        for (Long i = 0; i < Nelem; i++) {
          for (Long j = 0; j < Nnodes; j++) {
            potential[i*COORD_DIM+0][j] = 0.5*density[i*COORD_DIM+0][j] + U[i*COORD_DIM+0][j] + U1[i*COORD_DIM+0][j];
            potential[i*COORD_DIM+1][j] = 0.5*density[i*COORD_DIM+1][j] + U[i*COORD_DIM+1][j] + U1[i*COORD_DIM+1][j];
            potential[i*COORD_DIM+2][j] = 0.5*density[i*COORD_DIM+2][j] + U[i*COORD_DIM+2][j] + U1[i*COORD_DIM+2][j];
          }
        }
      };

      Vector<PotentialBasis> U;
      { // Rachh
        Vector<DensityBasis> sigma0;
        { // Set sigma0
          srand48(comm.Rank());
          Vector<Real> force(Ns*COORD_DIM), torque(Ns*COORD_DIM);
          //for (auto& x : force) x = drand48();
          //for (auto& x : torque) x = drand48();
          force = 0;
          torque = 0;
          force[0] = 1;
          //force[4] = 1;
          SetMotion(sigma0, force, torque);
        }

        Vector<DensityBasis> rhs;
        BIOp_half_K(rhs, sigma0);

        Vector<DensityBasis> sigma;
        { // Set sigma
          Long Nnode = DensityBasis::Size();
          Long Nelem = S.GetElem().NElem();

          typename ParallelSolver<Real>::ParallelOp A = [&S,&BIOp_half_K_L](Vector<Real>* Ax, const Vector<Real>& x) {
            Long Nnode = DensityBasis::Size();
            Long Nelem = S.GetElem().NElem();
            Ax->ReInit(Nelem*COORD_DIM*Nnode);

            Vector<DensityBasis> x_(Nelem*COORD_DIM), Ax_(Nelem*COORD_DIM);
            for (Long i = 0; i < Nelem*COORD_DIM; i++) { // Set x_
              for (Long k = 0; k < Nnode; k++) {
                x_[i][k] = x[i*Nnode+k];
              }
            }
            BIOp_half_K_L(Ax_, x_);
            for (Long i = 0; i < Nelem*COORD_DIM; i++) { // Set Ax
              for (Long k = 0; k < Nnode; k++) {
                (*Ax)[i*Nnode+k] = Ax_[i][k];
              }
            }
          };

          Vector<Real> sigma_(Nelem*COORD_DIM*Nnode), rhs_(Nelem*COORD_DIM*Nnode);
          for (Long i = 0; i < Nelem*COORD_DIM; i++) {// Set rhs_
            for (Long k = 0; k < Nnode; k++) {
              rhs_[i*Nnode+k] = rhs[i][k];
            }
          }
          sigma_ = 0;

          ParallelSolver<Real> linear_solver(comm, true);
          linear_solver(&sigma_, A, rhs_, 1e-6, 50);
          sigma.ReInit(Nelem * COORD_DIM);
          for (Long i = 0; i < Nelem*COORD_DIM; i++) {// Set sigma
            for (Long k = 0; k < Nnode; k++) {
              sigma[i][k] = sigma_[i*Nnode+k] - sigma0[i][k];
            }
          }
        }

        S.quadrature_FxU.Eval(U, S.GetElem(), sigma, S.Stokes_FxU);
        { // Write VTU
          VTUData vtu_sigma;
          vtu_sigma.AddElems(S.elements, sigma, ORDER);
          vtu_sigma.WriteVTK("sphere-sigma0", comm);

          VTUData vtu_U;
          vtu_U.AddElems(S.elements, U, ORDER);
          vtu_U.WriteVTK("sphere-U0", comm);
        }
      }

      { // Tornberg
        Vector<DensityBasis> rhs;
        BIOpL(rhs, U);

        Vector<DensityBasis> sigma;
        { // Set sigma
          Long Nnode = DensityBasis::Size();
          Long Nelem = S.GetElem().NElem();

          typename ParallelSolver<Real>::ParallelOp A = [&S,&BIOp_half_S_D](Vector<Real>* Ax, const Vector<Real>& x) {
            Long Nnode = DensityBasis::Size();
            Long Nelem = S.GetElem().NElem();
            Ax->ReInit(Nelem*COORD_DIM*Nnode);

            Vector<DensityBasis> x_(Nelem*COORD_DIM), Ax_(Nelem*COORD_DIM);
            for (Long i = 0; i < Nelem*COORD_DIM; i++) { // Set x_
              for (Long k = 0; k < Nnode; k++) {
                x_[i][k] = x[i*Nnode+k];
              }
            }
            BIOp_half_S_D(Ax_, x_);
            for (Long i = 0; i < Nelem*COORD_DIM; i++) { // Set Ax
              for (Long k = 0; k < Nnode; k++) {
                (*Ax)[i*Nnode+k] = Ax_[i][k];
              }
            }
          };

          Vector<Real> sigma_(Nelem*COORD_DIM*Nnode), rhs_(Nelem*COORD_DIM*Nnode);
          for (Long i = 0; i < Nelem*COORD_DIM; i++) {// Set rhs_
            for (Long k = 0; k < Nnode; k++) {
              rhs_[i*Nnode+k] = rhs[i][k];
            }
          }
          sigma_ = 0;

          ParallelSolver<Real> linear_solver(comm, true);
          linear_solver(&sigma_, A, rhs_, 1e-6, 50);
          sigma.ReInit(Nelem * COORD_DIM);
          for (Long i = 0; i < Nelem*COORD_DIM; i++) {// Set sigma
            for (Long k = 0; k < Nnode; k++) {
              sigma[i][k] = sigma_[i*Nnode+k];
            }
          }
        }

        Vector<PotentialBasis> U1;
        BIOp_half_S_D(U1, sigma);
        { // Write VTU
          VTUData vtu_sigma;
          vtu_sigma.AddElems(S.elements, sigma, ORDER);
          vtu_sigma.WriteVTK("sphere-sigma1", comm);

          VTUData vtu_U;
          vtu_U.AddElems(S.elements, U1, ORDER);
          vtu_U.WriteVTK("sphere-U1", comm);
        }
      }

      Profile::print(&comm);
    }

  private:

    template <class FnBasis> void SurfInteg(Vector<Real>& I, const Vector<FnBasis>& f) {
      static_assert(std::is_same<FnBasis,CoordBasis>::value, "FnBasis is different from CoordBasis");
      const Long Nelem = elements.NElem();
      const Long dof = f.Dim() / Nelem;
      SCTL_ASSERT(f.Dim() == Nelem * dof);

      auto nodes = FnBasis::Nodes();
      auto quad_wts = FnBasis::QuadWts();
      const Long Nnodes = FnBasis::Size();
      auto EvalOp = CoordBasis::SetupEval(nodes);

      Vector<CoordBasis> dX;
      const auto& X = elements.ElemVector();
      SCTL_ASSERT(X.Dim() == Nelem * COORD_DIM);
      CoordBasis::Grad(dX, X);

      Matrix<Real> I_(Nelem, dof);
      for (Long i = 0; i < Nelem; i++) {
        for (Long k = 0; k < dof; k++) {
          I_[i][k] = 0;
        }
        for (Long j = 0; j < Nnodes; j++) {
          Real dA = 0;
          StaticArray<Real,COORD_DIM> Xn;
          Xn[0] = dX[i*COORD_DIM*2+2][j] * dX[i*COORD_DIM*2+5][j] - dX[i*COORD_DIM*2+3][j] * dX[i*COORD_DIM*2+4][j];
          Xn[1] = dX[i*COORD_DIM*2+4][j] * dX[i*COORD_DIM*2+1][j] - dX[i*COORD_DIM*2+5][j] * dX[i*COORD_DIM*2+0][j];
          Xn[2] = dX[i*COORD_DIM*2+0][j] * dX[i*COORD_DIM*2+3][j] - dX[i*COORD_DIM*2+1][j] * dX[i*COORD_DIM*2+2][j];
          dA += sqrt<Real>(Xn[0]*Xn[0] + Xn[1]*Xn[1] + Xn[2]*Xn[2]) * quad_wts[j];
          for (Long k = 0; k < dof; k++) {
            I_[i][k] += dA * f[i*dof+k][j];
          }
        }
      }

      Long Ns = elem_cnt.Dim();
      if (I.Dim() != Ns * dof) I.ReInit(Ns * dof);
      I = 0;
      Long elem_itr = 0;
      for (Long i = 0; i < Ns; i++) {
        for (Long j = 0; j < elem_cnt[i]; j++) {
          for (Long k = 0; k < dof; k++) {
            I[i*dof+k] += I_[elem_itr][k];
          }
          elem_itr++;
        }
      }
    }

    void InitSpheres(const Vector<Real> X, const Vector<Real>& R){
      SCTL_ASSERT(X.Dim() == R.Dim() * COORD_DIM);
      Long N = R.Dim();
      elements.ReInit(2*COORD_DIM*N);
      auto nodes = ElemLst::CoordBasis::Nodes();
      for (Long l = 0; l < N; l++) {
        for (Integer i = 0; i < COORD_DIM; i++) {
          for (Integer j = 0; j < 2; j++) {
            for (int k = 0; k < ElemLst::CoordBasis::Size(); k++) {
              Real coord[COORD_DIM];
              coord[(i+0)%COORD_DIM] = (j ? -1.0 : 1.0);
              coord[(i+1)%COORD_DIM] = 2.0 * nodes[j?1:0][k] - 1.0;
              coord[(i+2)%COORD_DIM] = 2.0 * nodes[j?0:1][k] - 1.0;
              Real R0 = sqrt<Real>(coord[0]*coord[0] + coord[1]*coord[1] + coord[2]*coord[2]);

              elements((l*COORD_DIM+i)*2+j,0)[k] = X[l*COORD_DIM+0] + R[l] * coord[0] / R0;
              elements((l*COORD_DIM+i)*2+j,1)[k] = X[l*COORD_DIM+1] + R[l] * coord[1] / R0;
              elements((l*COORD_DIM+i)*2+j,2)[k] = X[l*COORD_DIM+2] + R[l] * coord[2] / R0;
            }
          }
        }
      }
      elem_cnt.ReInit(N);
      elem_cnt = 6;
    }

    GenericKernel<Stokes3D_DxU> Stokes_DxU;
    GenericKernel<Stokes3D_FxU> Stokes_FxU;
    GenericKernel<Stokes3D_FxT> Stokes_FxT;

    Quadrature<Real> quadrature_DxU;
    Quadrature<Real> quadrature_FxU;
    Quadrature<Real> quadrature_FxT;

    ElemLst elements;
    Vector<Long> elem_cnt;
};

}  // end namespace

#endif  //_SCTL_BOUNDARY_QUADRATURE_HPP_