/*
 *  This file is part of nifty_gridder.
 *
 *  nifty_gridder is free software; you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation; either version 2 of the License, or
 *  (at your option) any later version.
 *
 *  nifty_gridder is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with nifty_fridder; if not, write to the Free Software
 *  Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
 */

/* Copyright (C) 2019 Max-Planck-Society
   Author: Martin Reinecke */

#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <iostream>
#include <algorithm>

#define POCKETFFT_OPENMP
#include "pocketfft_hdronly.h"

#ifdef __GNUC__
#define RESTRICT __restrict__
#define NOINLINE __attribute__ ((noinline))
#else
#define RESTRICT
#endif

using namespace std;

namespace py = pybind11;

namespace {

//
// basic utilities
//

void myassert(bool cond, const char *msg)
  {
  if (cond) return;
  throw runtime_error(msg);
  }

/*! Returns the remainder of the division \a v1/v2.
    The result is non-negative.
    \a v1 can be positive or negative; \a v2 must be positive. */
template<typename T> inline T fmodulo (T v1, T v2)
  {
  if (v1>=0)
    return (v1<v2) ? v1 : fmod(v1,v2);
  T tmp=fmod(v1,v2)+v2;
  return (tmp==v2) ? T(0) : tmp;
  }

//
// Utilities for Gauss-Legendre quadrature
//

static inline double one_minus_x2 (double x)
  { return (fabs(x)>0.1) ? (1.+x)*(1.-x) : 1.-x*x; }

void legendre_prep(int n, vector<double> &x, vector<double> &w)
  {
  constexpr double pi = 3.141592653589793238462643383279502884197;
  constexpr double eps = 3e-14;
  int m = (n+1)>>1;
  x.resize(m);
  w.resize(m);

  double t0 = 1 - (1-1./n) / (8.*n*n);
  double t1 = 1./(4.*n+2.);

#pragma omp parallel
{
  int i;
#pragma omp for schedule(dynamic,100)
  for (i=1; i<=m; ++i)
    {
    double x0 = cos(pi * ((i<<2)-1) * t1) * t0;

    int dobreak=0;
    int j=0;
    double dpdx;
    while(1)
      {
      double P_1 = 1.0;
      double P0 = x0;
      double dx, x1;

      for (int k=2; k<=n; k++)
        {
        double P_2 = P_1;
        P_1 = P0;
//        P0 = ((2*k-1)*x0*P_1-(k-1)*P_2)/k;
        P0 = x0*P_1 + (k-1.)/k * (x0*P_1-P_2);
        }

      dpdx = (P_1 - x0*P0) * n / one_minus_x2(x0);

      /* Newton step */
      x1 = x0 - P0/dpdx;
      dx = x0-x1;
      x0 = x1;
      if (dobreak) break;

      if (abs(dx)<=eps) dobreak=1;
      if (++j>=100) throw runtime_error("convergence problem");
      }

    x[m-i] = x0;
    w[m-i] = 2. / (one_minus_x2(x0) * dpdx * dpdx);
    }
} // end of parallel region
  }

//
// Start of real gridder functionality
//

template<typename T>
  using pyarr = py::array_t<T>;
template<typename T>
  using pyarr_c = py::array_t<T, py::array::c_style | py::array::forcecast>;

template<typename T> pyarr_c<T> makeArray(const vector<size_t> &shape)
  { return pyarr_c<T>(shape); }

size_t get_w(double epsilon)
  {
  static const vector<double> maxmaperr { 1e8, 0.32, 0.021, 6.2e-4,
    1.08e-5, 1.25e-7, 8.25e-10, 5.70e-12, 1.22e-13, 2.48e-15, 4.82e-17,
    6.74e-19, 5.41e-21, 4.41e-23, 7.88e-25, 3.9e-26 };

  double epssq = epsilon*epsilon;

  for (size_t i=1; i<maxmaperr.size(); ++i)
    if (epssq>maxmaperr[i]) return i;
  throw runtime_error("requested epsilon too small - minimum is 2e-13");
  }

void checkArray(const py::array &arr, const char *aname,
  const vector<size_t> &shape)
  {
  if (size_t(arr.ndim())!=shape.size())
    {
    cerr << "Array '" << aname << "' has " << arr.ndim() << " dimensions; "
            "expected " << shape.size() << endl;
    throw runtime_error("bad dimensionality");
    }
  for (size_t i=0; i<shape.size(); ++i)
    if ((shape[i]!=0) && (size_t(arr.shape(i))!=shape[i]))
      {
      cerr << "Dimension " << i << " of array '" << aname << "' has size "
           << arr.shape(i) << "; expected " << shape[i] << endl;
      throw runtime_error("bad array size");
      }
  }

template<typename T> pyarr_c<T> provideArray(py::object &in,
  const vector<size_t> &shape)
  {
  if (in.is(py::none()))
    {
    auto tmp_ = makeArray<T>(shape);
    size_t sz = size_t(tmp_.size());
    auto tmp = tmp_.mutable_data();
    for (size_t i=0; i<sz; ++i)
      tmp[i] = T(0);
    return tmp_;
    }
  auto tmp_ = in.cast<pyarr_c<T>>();
  checkArray(tmp_, "temporary", shape);
  return tmp_;
  }

template<typename T> pyarr_c<T> complex2hartley
  (const pyarr_c<complex<T>> &grid_)
  {
  checkArray(grid_, "grid", {0,0});
  size_t nu = size_t(grid_.shape(0)), nv = size_t(grid_.shape(1));
  auto grid = grid_.data();

  auto res = makeArray<T>({nu,nv});
  auto grid2 = res.mutable_data();
  {
  py::gil_scoped_release release;
#pragma omp parallel for
  for (size_t u=0; u<nu; ++u)
    {
    size_t xu = (u==0) ? 0 : nu-u;
    for (size_t v=0; v<nv; ++v)
      {
      size_t xv = (v==0) ? 0 : nv-v;
      size_t i1 = u*nv+v;
      size_t i2 = xu*nv+xv;
      grid2[i1] = T(0.5)*(grid[i1].real()+grid[i1].imag()+
                          grid[i2].real()-grid[i2].imag());
      }
    }
  }
  return res;
  }

template<typename T> pyarr_c<complex<T>> hartley2complex
  (const pyarr_c<T> &grid_)
  {
  checkArray(grid_, "grid", {0, 0});
  size_t nu = size_t(grid_.shape(0)), nv = size_t(grid_.shape(1));
  auto grid = grid_.data();

  auto res=makeArray<complex<T>>({nu, nv});
  auto grid2 = res.mutable_data();
  {
  py::gil_scoped_release release;
#pragma omp parallel for
  for (size_t u=0; u<nu; ++u)
    {
    size_t xu = (u==0) ? 0 : nu-u;
    for (size_t v=0; v<nv; ++v)
      {
      size_t xv = (v==0) ? 0 : nv-v;
      size_t i1 = u*nv+v;
      size_t i2 = xu*nv+xv;
      T v1 = T(0.5)*grid[i1];
      T v2 = T(0.5)*grid[i2];
      grid2[i1] = complex<T>(v1+v2, v1-v2);
      }
    }
  }
  return res;
  }

template<typename T> void hartley2_2D(const pyarr_c<T> &in, pyarr_c<T> &out)
  {
  size_t nu=in.shape(0), nv=in.shape(1);
  pocketfft::stride_t s_i{in.strides(0), in.strides(1)},
                      s_o{out.strides(0), out.strides(1)};
  auto d_i = in.data();
  auto ptmp = out.mutable_data();
  {
  py::gil_scoped_release release;
  pocketfft::r2r_hartley({nu, nv}, s_i, s_o, {0,1}, d_i, ptmp, T(1), 0);
#pragma omp parallel for
  for(size_t i=1; i<(nu+1)/2; ++i)
    for(size_t j=1; j<(nv+1)/2; ++j)
       {
       T a = ptmp[i*nv+j];
       T b = ptmp[(nu-i)*nv+j];
       T c = ptmp[i*nv+nv-j];
       T d = ptmp[(nu-i)*nv+nv-j];
       ptmp[i*nv+j] = T(0.5)*(a+b+c-d);
       ptmp[(nu-i)*nv+j] = T(0.5)*(a+b+d-c);
       ptmp[i*nv+nv-j] = T(0.5)*(a+c+d-b);
       ptmp[(nu-i)*nv+nv-j] = T(0.5)*(b+c+d-a);
       }
  }
  }

/* Compute correction factors for the ES gridding kernel
   This implementation follows eqs. (3.8) to (3.10) of Barnett et al. 2018 */
vector<double> correction_factors (size_t n, size_t nval, size_t w)
  {
  constexpr double pi = 3.141592653589793238462643383279502884197;
  auto beta = 2.3*w;
  auto p = int(1.5*w+2);
  double alpha = pi*w/n;
  vector<double> x, wgt;
  legendre_prep(2*p,x,wgt);
  auto psi = x;
  for (auto &v:psi)
    v = exp(beta*(sqrt(1-v*v)-1.));
  vector<double> res(nval);
#pragma omp parallel for schedule(static)
  for (size_t k=0; k<nval; ++k)
    {
    double tmp=0;
    for (int i=0; i<p; ++i)
      tmp += wgt[i]*psi[i]*cos(alpha*k*x[i]);
    res[k] = 1./(w*tmp);
    }
  return res;
  }

template<typename T> struct UVW
  {
  T u, v, w;
  UVW () {}
  UVW (T u_, T v_, T w_) : u(u_), v(v_), w(w_) {}
  UVW operator* (T fct) const
    { return UVW(u*fct, v*fct, w*fct); }
  };

template<typename T> class Baselines
  {
  private:
    vector<UVW<T>> coord;
    vector<T> f_over_c;
    size_t nrows, nchan;

  public:
    Baselines(const pyarr_c<T> &coord_, const pyarr_c<T> &freq_)
      {
      constexpr double speedOfLight = 299792458.;
      checkArray(coord_, "coord", {0, 3});
      checkArray(freq_, "freq", {0});
      nrows = coord_.shape(0);
      nchan = freq_.shape(0);
      myassert(nrows*nchan<(size_t(1)<<32), "too many entries in MS");
      auto freq = freq_.data();
      auto cood = coord_.data();
      {
      py::gil_scoped_release release;
      f_over_c.resize(nchan);
      for (size_t i=0; i<nchan; ++i)
        f_over_c[i] = freq[i]/speedOfLight;
      coord.resize(nrows);
      for (size_t i=0; i<coord.size(); ++i)
        coord[i] = UVW<T>(cood[3*i], cood[3*i+1], cood[3*i+2]);
      }
      }

    UVW<T> effectiveCoord(uint32_t index) const
      {
      size_t irow = index/nchan;
      size_t ichan = index-nchan*irow;
      return coord[irow]*f_over_c[ichan];
      }
    UVW<T> effectiveCoord(size_t irow, size_t ichan) const
      { return coord[irow]*f_over_c[ichan]; }
    size_t Nrows() const { return nrows; }
    size_t Nchannels() const { return nchan; }

    template<typename T2> pyarr_c<T2> ms2vis(const pyarr_c<T2> &ms_,
      const pyarr_c<uint32_t> &idx_) const
      {
      checkArray(idx_, "idx", {0});
      checkArray(ms_, "ms", {nrows, nchan});
      size_t nvis = size_t(idx_.shape(0));
      auto idx = idx_.data();
      auto ms = ms_.data();

      auto res=makeArray<T2>({nvis});
      auto vis = res.mutable_data();
      {
      py::gil_scoped_release release;
#pragma omp parallel for
      for (size_t i=0; i<nvis; ++i)
        vis[i] = ms[idx[i]];
      }
      return res;
      }

    template<typename T2> pyarr_c<T2> vis2ms(const pyarr_c<T2> &vis_,
      const pyarr_c<uint32_t> &idx_, py::object &ms_in) const
      {
      checkArray(vis_, "vis", {0});
      size_t nvis = size_t(vis_.shape(0));
      checkArray(idx_, "idx", {nvis});
      auto idx = idx_.data();
      auto vis = vis_.data();

      auto res = provideArray<T2>(ms_in, {nrows, nchan});
      auto ms = res.mutable_data();
      {
      py::gil_scoped_release release;
#pragma omp parallel for
      for (size_t i=0; i<nvis; ++i)
        ms[idx[i]] = vis[i];
      }
      return res;
      }
  };

constexpr int logsquare=4;

template<typename T> class GridderConfig
  {
  private:
    size_t nx_dirty, ny_dirty;
    double eps, psx, psy;
    size_t w, nsafe, nu, nv;
    T beta;
    vector<T> cfu, cfv;

  public:
    GridderConfig(size_t nxdirty, size_t nydirty, double epsilon,
      double pixsize_x, double pixsize_y)
      : nx_dirty(nxdirty), ny_dirty(nydirty), eps(epsilon),
        psx(pixsize_x), psy(pixsize_y),
        w(get_w(epsilon)), nsafe((w+1)/2),
        nu(max(2*nsafe,2*nx_dirty)), nv(max(2*nsafe,2*ny_dirty)),
        beta(2.3*w),
        cfu(nx_dirty), cfv(ny_dirty)
      {
      {
      py::gil_scoped_release release;
      myassert((nx_dirty&1)==0, "nx_dirty must be even");
      myassert((ny_dirty&1)==0, "ny_dirty must be even");
      myassert(epsilon>0, "epsilon must be positive");
      myassert(pixsize_x>0, "pixsize_x must be positive");
      myassert(pixsize_y>0, "pixsize_y must be positive");

      auto tmp = correction_factors(nu, nx_dirty/2+1, w);
      cfu[nx_dirty/2]=tmp[0];
      cfu[0]=tmp[nx_dirty/2];
      for (size_t i=1; i<nx_dirty/2; ++i)
        cfu[nx_dirty/2-i] = cfu[nx_dirty/2+i] = tmp[i];
      tmp = correction_factors(nv, ny_dirty/2+1, w);
      cfv[ny_dirty/2]=tmp[0];
      cfv[0]=tmp[ny_dirty/2];
      for (size_t i=1; i<ny_dirty/2; ++i)
        cfv[ny_dirty/2-i] = cfv[ny_dirty/2+i] = tmp[i];
      }
      }
    size_t Nxdirty() const { return nx_dirty; }
    size_t Nydirty() const { return ny_dirty; }
    double Epsilon() const { return eps; }
    double Pixsize_x() const { return psx; }
    double Pixsize_y() const { return psy; }
    size_t Nu() const { return nu; }
    size_t Nv() const { return nv; }
    size_t W() const { return w; }
    size_t Nsafe() const { return nsafe; }
    T Beta() const { return beta; }
    pyarr_c<T> grid2dirty(const pyarr_c<T> &grid) const
      {
      checkArray(grid, "grid", {nu, nv});
      auto tmp = makeArray<T>({nu, nv});
      auto ptmp = tmp.mutable_data();
      hartley2_2D<T>(grid, tmp);
      auto res = makeArray<T>({nx_dirty, ny_dirty});
      auto pout = res.mutable_data();
      {
      py::gil_scoped_release release;
      for (size_t i=0; i<nx_dirty; ++i)
        for (size_t j=0; j<ny_dirty; ++j)
          {
          size_t i2 = nu-nx_dirty/2+i;
          if (i2>=nu) i2-=nu;
          size_t j2 = nv-ny_dirty/2+j;
          if (j2>=nv) j2-=nv;
          pout[ny_dirty*i + j] = ptmp[nv*i2+j2]*cfu[i]*cfv[j];
          }
      }
      return res;
      }
    pyarr_c<complex<T>> grid2dirty_c(const pyarr_c<complex<T>> &grid) const
      {
      checkArray(grid, "grid", {nu, nv});
      auto tmp = makeArray<complex<T>>({nu, nv});
      auto ptmp = tmp.mutable_data();
      pocketfft::c2c({nu,nv},{grid.strides(0),grid.strides(1)},
        {tmp.strides(0), tmp.strides(1)}, {0,1}, pocketfft::BACKWARD,
        grid.data(), tmp.mutable_data(), T(1), 0);
      auto res = makeArray<complex<T>>({nx_dirty, ny_dirty});
      auto pout = res.mutable_data();
      {
      py::gil_scoped_release release;
      for (size_t i=0; i<nx_dirty; ++i)
        for (size_t j=0; j<ny_dirty; ++j)
          {
          size_t i2 = nu-nx_dirty/2+i;
          if (i2>=nu) i2-=nu;
          size_t j2 = nv-ny_dirty/2+j;
          if (j2>=nv) j2-=nv;
          pout[ny_dirty*i + j] = ptmp[nv*i2+j2]*cfu[i]*cfv[j];
          }
      }
      return res;
      }
    pyarr_c<T> dirty2grid(const pyarr_c<T> &dirty) const
      {
      checkArray(dirty, "dirty", {nx_dirty, ny_dirty});
      auto pdirty = dirty.data();
      auto tmp = makeArray<T>({nu, nv});
      auto ptmp = tmp.mutable_data();
      {
      py::gil_scoped_release release;
      for (size_t i=0; i<nu*nv; ++i)
        ptmp[i] = 0.;
      for (size_t i=0; i<nx_dirty; ++i)
        for (size_t j=0; j<ny_dirty; ++j)
          {
          size_t i2 = nu-nx_dirty/2+i;
          if (i2>=nu) i2-=nu;
          size_t j2 = nv-ny_dirty/2+j;
          if (j2>=nv) j2-=nv;
          ptmp[nv*i2+j2] = pdirty[ny_dirty*i + j]*cfu[i]*cfv[j];
          }
      }
      hartley2_2D<T>(tmp, tmp);
      return tmp;
      }
    pyarr_c<complex<T>> dirty2grid_c(const pyarr_c<complex<T>> &dirty) const
      {
      checkArray(dirty, "dirty", {nx_dirty, ny_dirty});
      auto pdirty = dirty.data();
      auto tmp = makeArray<complex<T>>({nu, nv});
      auto ptmp = tmp.mutable_data();
      pocketfft::stride_t strides{tmp.strides(0),tmp.strides(1)};
      {
      py::gil_scoped_release release;
      for (size_t i=0; i<nu*nv; ++i)
        ptmp[i] = 0.;
      for (size_t i=0; i<nx_dirty; ++i)
        for (size_t j=0; j<ny_dirty; ++j)
          {
          size_t i2 = nu-nx_dirty/2+i;
          if (i2>=nu) i2-=nu;
          size_t j2 = nv-ny_dirty/2+j;
          if (j2>=nv) j2-=nv;
          ptmp[nv*i2+j2] = pdirty[ny_dirty*i + j]*cfu[i]*cfv[j];
          }
      pocketfft::c2c({nu,nv}, strides, strides, {0,1}, pocketfft::FORWARD,
        ptmp, ptmp, T(1), 0);
      }
      return tmp;
      }
    inline void getpix(T u_in, T v_in, T &u, T &v, int &iu0, int &iv0) const
      {
      u=fmodulo(u_in*psx, T(1))*nu,
      iu0 = int(u-w*0.5 + 1 + nu) - nu;
      if (iu0+w>nu+nsafe) iu0 = nu+nsafe-w;
      v=fmodulo(v_in*psy, T(1))*nv;
      iv0 = int(v-w*0.5 + 1 + nv) - nv;
      if (iv0+w>nv+nsafe) iv0 = nv+nsafe-w;
      }
  };

template<typename T> class Helper
  {
  private:
    const GridderConfig<T> &gconf;
    int nu, nv, nsafe, w;
    T beta;
    complex<T> *grid;
    bool write;
    int su;
  public:
    int sv;
  private:
    int iu0, iv0; // start index of the current visibility
    int bu0, bv0; // start index of the current buffer

    vector<complex<T>> data;

    void dump() const
      {
      if (bu0<-nsafe) return; // nothing written into buffer yet

#pragma omp critical (gridder_writing_to_grid)
{
      int idxu = (bu0+nu)%nu;
      int idxv0 = (bv0+nv)%nv;
      for (int iu=0; iu<su; ++iu)
        {
        int idxv = idxv0;
        for (int iv=0; iv<sv; ++iv)
          {
          grid[idxu*nv + idxv] += data[iu*sv + iv];
          if (++idxv>=nv) idxv=0;
          }
        if (++idxu>=nu) idxu=0;
        }
}
      }

    void load()
      {
      int idxu = (bu0+nu)%nu;
      int idxv0 = (bv0+nv)%nv;
      for (int iu=0; iu<su; ++iu)
        {
        int idxv = idxv0;
        for (int iv=0; iv<sv; ++iv)
          {
          data[iu*sv + iv] = grid[idxu*nv + idxv];
          if (++idxv>=nv) idxv=0;
          }
        if (++idxu>=nu) idxu=0;
        }
      }

  public:
    complex<T> *p0;
    vector<T> kernel;

    Helper(const GridderConfig<T> &gconf_, const complex<T> *grid_, bool write_)
      : gconf(gconf_), nu(gconf.Nu()), nv(gconf.Nv()), nsafe(gconf.Nsafe()), w(gconf.W()),
        beta(gconf.Beta()),
        grid(const_cast<complex<T> *>(grid_)), write(write_),
        su(2*nsafe+(1<<logsquare)), sv(2*nsafe+(1<<logsquare)),
        bu0(-1000000), bv0(-1000000), data(su*sv,T(0)), kernel(2*w)
      {}
    ~Helper() { if (write) dump(); }

    void prep(T u_in, T v_in)
      {
      T u, v;
      gconf.getpix(u_in, v_in, u, v, iu0, iv0);
      T xw=T(2)/w;
      auto x0 = xw*(iu0-u);
      auto y0 = xw*(iv0-v);
      for (int i=0; i<w; ++i)
        {
        auto x = x0+i*xw;
        kernel[i  ] = beta*sqrt(T(1)-x*x);
        auto y = y0+i*xw;
        kernel[i+w] = beta*sqrt(T(1)-y*y);
        }
      for (auto &k : kernel)
        k = exp(k);

      if ((iu0<bu0) || (iv0<bv0) || (iu0+w>bu0+su) || (iv0+w>bv0+sv))
        {
        if (write) { dump(); fill(data.begin(), data.end(), T(0)); }
        bu0=((((iu0+nsafe)>>logsquare)<<logsquare))-nsafe;
        bv0=((((iv0+nsafe)>>logsquare)<<logsquare))-nsafe;
        if (!write) load();
        }
      p0 = data.data() + sv*(iu0-bu0) + iv0-bv0;
      }
  };

template<typename T> pyarr_c<complex<T>> vis2grid_c(
  const Baselines<T> &baselines, const GridderConfig<T> &gconf,
  const pyarr_c<uint32_t> &idx_, const pyarr_c<complex<T>> &vis_)
  {
  checkArray(vis_, "vis", {0});
  size_t nvis = size_t(vis_.shape(0));
  checkArray(idx_, "idx", {nvis});
  auto vis=vis_.data();
  auto idx = idx_.data();

  size_t nu=gconf.Nu(), nv=gconf.Nv();
  auto res = makeArray<complex<T>>({nu, nv});
  auto grid = res.mutable_data();
  {
  py::gil_scoped_release release;
  for (size_t i=0; i<nu*nv; ++i) grid[i] = 0.;
  T beta = gconf.Beta();
  size_t w = gconf.W();

#pragma omp parallel
{
  Helper<T> hlp(gconf, grid, true);
  T emb = exp(-2*beta);
  const T * RESTRICT ku = hlp.kernel.data();
  const T * RESTRICT kv = hlp.kernel.data()+w;

  // Loop over sampling points
#pragma omp for schedule(guided,100)
  for (size_t ipart=0; ipart<nvis; ++ipart)
    {
    UVW<T> coord = baselines.effectiveCoord(idx[ipart]);
    hlp.prep(coord.u, coord.v);
    auto * RESTRICT ptr = hlp.p0;
    auto v(vis[ipart]*emb);
    for (size_t cu=0; cu<w; ++cu)
      {
      complex<T> tmp(v*ku[cu]);
      for (size_t cv=0; cv<w; ++cv)
        ptr[cv] += tmp*kv[cv];
      ptr+=hlp.sv;
      }
    }
} // end of parallel region
  }
  return res;
  }

template<typename T> pyarr_c<T> vis2grid(const Baselines<T> &baselines,
  const GridderConfig<T> &gconf, const pyarr_c<uint32_t> &idx_,
  const pyarr_c<complex<T>> &vis_)
  { return complex2hartley(vis2grid_c(baselines, gconf, idx_, vis_)); }

template<typename T> pyarr_c<complex<T>> ms2grid_c(
  const Baselines<T> &baselines, const GridderConfig<T> &gconf,
  const pyarr_c<uint32_t> &idx_, const pyarr_c<complex<T>> &ms_)
  {
  auto nrows = baselines.Nrows();
  auto nchan = baselines.Nchannels();
  checkArray(ms_, "ms", {nrows, nchan});
  checkArray(idx_, "idx", {0});
  size_t nvis = size_t(idx_.shape(0));
  auto ms = ms_.data();
  auto idx = idx_.data();

  size_t nu=gconf.Nu(), nv=gconf.Nv();
  auto res = makeArray<complex<T>>({nu, nv});
  auto grid = res.mutable_data();
  {
  py::gil_scoped_release release;
  for (size_t i=0; i<nu*nv; ++i) grid[i] = 0.;
  T beta = gconf.Beta();
  size_t w = gconf.W();

#pragma omp parallel
{
  Helper<T> hlp(gconf, grid, true);
  T emb = exp(-2*beta);
  const T * RESTRICT ku = hlp.kernel.data();
  const T * RESTRICT kv = hlp.kernel.data()+w;

  // Loop over sampling points
#pragma omp for schedule(guided,100)
  for (size_t ipart=0; ipart<nvis; ++ipart)
    {
    UVW<T> coord = baselines.effectiveCoord(idx[ipart]);
    hlp.prep(coord.u, coord.v);
    auto * RESTRICT ptr = hlp.p0;
    auto v(ms[idx[ipart]]*emb);
    for (size_t cu=0; cu<w; ++cu)
      {
      complex<T> tmp(v*ku[cu]);
      for (size_t cv=0; cv<w; ++cv)
        ptr[cv] += tmp*kv[cv];
      ptr+=hlp.sv;
      }
    }
} // end of parallel region
  }
  return res;
  }

template<typename T> pyarr_c<T> ms2grid(const Baselines<T> &baselines,
  const GridderConfig<T> &gconf, const pyarr_c<uint32_t> &idx_,
  const pyarr_c<complex<T>> &ms_)
  { return complex2hartley(ms2grid_c(baselines, gconf, idx_, ms_)); }

template<typename T> pyarr_c<complex<T>> grid2vis_c(const Baselines<T> &baselines,
  const GridderConfig<T> &gconf, const pyarr_c<uint32_t> &idx_,
  const pyarr_c<complex<T>> &grid_)
  {
  size_t nu=gconf.Nu(), nv=gconf.Nv();
  checkArray(idx_, "idx", {0});
  auto grid = grid_.data();
  checkArray(grid_, "grid", {nu, nv});
  size_t nvis = size_t(idx_.shape(0));
  auto idx = idx_.data();

  auto res = makeArray<complex<T>>({nvis});
  auto vis = res.mutable_data();
  {
  py::gil_scoped_release release;
  T beta = gconf.Beta();
  size_t w = gconf.W();

  // Loop over sampling points
#pragma omp parallel
{
  Helper<T> hlp(gconf, grid, false);
  T emb = exp(-2*beta);
  const T * RESTRICT ku = hlp.kernel.data();
  const T * RESTRICT kv = hlp.kernel.data()+w;

#pragma omp for schedule(guided,100)
  for (size_t ipart=0; ipart<nvis; ++ipart)
    {
    UVW<T> coord = baselines.effectiveCoord(idx[ipart]);
    hlp.prep(coord.u, coord.v);
    complex<T> r = 0;
    auto * RESTRICT ptr = hlp.p0;
    for (size_t cu=0; cu<w; ++cu)
      {
      complex<T> tmp(0);
      for (size_t cv=0; cv<w; ++cv)
        tmp += ptr[cv] * kv[cv];
      r += tmp*ku[cu];
      ptr += hlp.sv;
      }
    vis[ipart] = r*emb;
    }
}
  }
  return res;
  }

template<typename T> pyarr_c<complex<T>> grid2vis(const Baselines<T> &baselines,
  const GridderConfig<T> &gconf, const pyarr_c<uint32_t> &idx_,
  const pyarr_c<T> &grid_)
  { return grid2vis_c(baselines, gconf, idx_, hartley2complex(grid_)); }

template<typename T> pyarr_c<complex<T>> grid2ms_c(const Baselines<T> &baselines,
  const GridderConfig<T> &gconf, const pyarr_c<uint32_t> &idx_,
  const pyarr_c<complex<T>> &grid_, py::object &ms_in)
  {
  size_t nu=gconf.Nu(), nv=gconf.Nv();
  checkArray(idx_, "idx", {0});
  auto grid = grid_.data();
  checkArray(grid_, "grid", {nu, nv});
  size_t nvis = size_t(idx_.shape(0));
  auto idx = idx_.data();

  auto res = provideArray<complex<T>>(ms_in,
    {baselines.Nrows(), baselines.Nchannels()});
  auto ms = res.mutable_data();
  {
  py::gil_scoped_release release;
  T beta = gconf.Beta();
  size_t w = gconf.W();

  // Loop over sampling points
#pragma omp parallel
{
  Helper<T> hlp(gconf, grid, false);
  T emb = exp(-2*beta);
  const T * RESTRICT ku = hlp.kernel.data();
  const T * RESTRICT kv = hlp.kernel.data()+w;

#pragma omp for schedule(guided,100)
  for (size_t ipart=0; ipart<nvis; ++ipart)
    {
    UVW<T> coord = baselines.effectiveCoord(idx[ipart]);
    hlp.prep(coord.u, coord.v);
    complex<T> r = 0;
    auto * RESTRICT ptr = hlp.p0;
    for (size_t cu=0; cu<w; ++cu)
      {
      complex<T> tmp(0);
      for (size_t cv=0; cv<w; ++cv)
        tmp += ptr[cv] * kv[cv];
      r += tmp*ku[cu];
      ptr += hlp.sv;
      }
    ms[idx[ipart]] += r*emb;
    }
}
  }
  return res;
  }

template<typename T> pyarr_c<complex<T>> grid2ms(const Baselines<T> &baselines,
  const GridderConfig<T> &gconf, const pyarr_c<uint32_t> &idx_,
  const pyarr_c<T> &grid_, py::object &ms_in)
  { return grid2ms_c(baselines, gconf, idx_, hartley2complex(grid_), ms_in); }

template<typename T> pyarr_c<uint32_t> getIndices(const Baselines<T> &baselines,
  const GridderConfig<T> &gconf, const pyarr_c<bool> &flags_, int chbegin,
  int chend, T wmin, T wmax)
  {
  size_t nrow=baselines.Nrows(),
         nchan=baselines.Nchannels(),
         nsafe=gconf.Nsafe();
  if (chbegin<0) chbegin=0;
  if (chend<0) chend=nchan;
  myassert(chend>chbegin, "empty channel range selected");
  myassert(chend<=int(nchan), "chend too large");
  myassert(wmax>wmin, "empty w range selected");
  checkArray(flags_, "flags", {nrow, nchan});
  auto flags = flags_.data();
  constexpr int side=1<<logsquare;
  size_t nbu = (gconf.Nu()+1+side-1) >> logsquare,
         nbv = (gconf.Nv()+1+side-1) >> logsquare;
  vector<uint32_t> acc(nbu*nbv+1, 0);
  vector<uint32_t> tmp(nrow*(chend-chbegin));
  {
  py::gil_scoped_release release;
  for (size_t irow=0, idx=0; irow<nrow; ++irow)
    for (int ichan=chbegin; ichan<chend; ++ichan)
      if (!flags[irow*nchan+ichan])
        {
        auto uvw = baselines.effectiveCoord(irow, ichan);
        if ((uvw.w>=wmin) && (uvw.w<wmax))
          {
          T u, v;
          int iu0, iv0;
          gconf.getpix(uvw.u, uvw.v, u, v, iu0, iv0);
          iu0 = (iu0+nsafe)>>logsquare;
          iv0 = (iv0+nsafe)>>logsquare;
          ++acc[nbv*iu0 + iv0 + 1];
          tmp[idx++] = nbv*iu0 + iv0;
          }
        }

  for (size_t i=1; i<acc.size(); ++i)
    acc[i] += acc[i-1];
  }
  auto res = makeArray<uint32_t>({acc.back()});
  auto iout = res.mutable_data();
  {
  py::gil_scoped_release release;
  for (size_t irow=0, idx=0; irow<nrow; ++irow)
    for (int ichan=chbegin; ichan<chend; ++ichan)
      if (!flags[irow*nchan+ichan])
        {
        auto uvw = baselines.effectiveCoord(irow, ichan);
        if ((uvw.w>=wmin) && (uvw.w<wmax))
          iout[acc[tmp[idx++]]++] = irow*nchan+ichan;
        }
  }
  return res;
  }

const char *Baselines_DS = R"""(
Class storing UVW coordinates and channel information.

Parameters
==========
coord: np.array((nrows, 3), dtype=np.float)
    u, v and w coordinates for each row
freq: np.array((nchannels,), dtype=np.float)
    frequency for each individual channel (in Hz)
)""";

const char *BL_ms2vis_DS = R"""(
Extracts visibility data from a measurement for the provided indices.

Parameters
==========
ms: np.array((nrows, nchannels), dtype=np.complex)
    the measurement set's visibility data
idx: np.array((nvis,), dtype=np.uint32)
    the indices to be extracted

Returns
=======
np.array((nvis,), dtype=np.complex)
    The visibility data for the index array
)""";

const char *BL_vis2ms_DS = R"""(
Produces a new MS with the provided visibilities set.

Parameters
==========
vis: np.array((nvis,), dtype=np.complex)
    The visibility data for the index array
idx: np.array((nvis,), dtype=np.uint32)
    the indices to be inserted

Returns
=======
np.array((nrows, nchannels), dtype=np.complex)
    the measurement set's visibility data (0 where not covered by idx)
)""";

const char *GridderConfig_DS = R"""(
Class storing information related to the gridding/degridding process.

Parameters
==========
nxdirty: int
    x resolution of the dirty image; must be even
nydirty: int
    y resolution of the dirty image; must be even
epsilon: float
    required accuracy for the gridding/degridding step
    Must be >= 2e-13.
pixsize_x: float
    Pixel size in x direction (radians)
pixsize_y: float
    Pixel size in y direction (radians)
)""";

const char *grid2dirty_DS = R"""(
Converts from UV grid to dirty image (FFT, cropping, correction)

Parameters
==========
grid: np.array((nu, nv), dtype=np.float64)
    gridded UV data

Returns
=======
nd.array((nxdirty, nydirty), dtype=np.float64)
    the dirty image
)""";

const char *dirty2grid_DS = R"""(
Converts from a dirty image to a UV grid (correction, padding, FFT)

Parameters
==========
dirty: nd.array((nxdirty, nydirty), dtype=np.float64)
    the dirty image

Returns
=======
np.array((nu, nv), dtype=np.float64)
    gridded UV data
)""";

const char *getIndices_DS = R"""(
Selects a subset of entries from a `Baselines` object.

Parameters
==========
baselines: Baselines
    the Baselines object
gconf: GridderConf
    the GridderConf object to be used with the returned indices.
    (used to optimize the ordering of the indices)
flags: np.array((nrows, nchannels), dtype=np.bool)
    "True" indicates that the value should not be used
chbegin: int
    first channel to use (-1: start with the first available channel)
chend: int
    one-past last channel to use (-1: one past the last available channel)
wmin: float
    only select entries with w>=wmin
wmax: float
    only select entries with w<wmax

Returns
=======
np.array((nvis,), dtype=np.uint32)
    the compressed indices for all entries which match the selected criteria
    and are not flagged.
)""";

const char *vis2grid_DS = R"""(
Grids visibilities onto a UV grid

Parameters
==========
baselines: Baselines
    the Baselines object
gconf: GridderConf
    the GridderConf object to be used
    (used to optimize the ordering of the indices)
idx: np.array((nvis,), dtype=np.uint32)
    the indices for the entries to be gridded
vis: np.array((nvis,), dtype=np.complex)
    The visibility data for the index array

Returns
=======
np.array((nu,nv), dtype=np.float64):
    the gridded visibilities (made real by making use of Hermitian symmetry)
)""";

const char *grid2vis_DS = R"""(
Degrids visibilities from a UV grid

Parameters
==========
baselines: Baselines
    the Baselines object
gconf: GridderConf
    the GridderConf object to be used
    (used to optimize the ordering of the indices)
idx: np.array((nvis,), dtype=np.uint32)
    the indices for the entries to be degridded
grid: np.array((nu,nv), dtype=np.float64):
    the gridded visibilities (made real by making use of Hermitian symmetry)
vis: np.array((nvis,), dtype=np.complex)
    The visibility data for the index array

Returns
=======
np.array((nvis,), dtype=np.complex)
    The degridded visibility data
)""";
} // unnamed namespace

PYBIND11_MODULE(nifty_gridder, m)
  {
  using namespace pybind11::literals;

  py::class_<Baselines<double>> (m, "Baselines", Baselines_DS)
    .def(py::init<const pyarr_c<double> &, const pyarr_c<double> &>(),
      "coord"_a, "freq"_a)
    .def ("Nrows",&Baselines<double>::Nrows)
    .def ("Nchannels",&Baselines<double>::Nchannels)
    .def ("ms2vis",&Baselines<double>::ms2vis<complex<double>>, BL_ms2vis_DS,
      "ms"_a, "idx"_a)
    .def ("vis2ms",&Baselines<double>::vis2ms<complex<double>>, BL_vis2ms_DS,
      "vis"_a, "idx"_a, "ms_in"_a=py::none());
  py::class_<GridderConfig<double>> (m, "GridderConfig", GridderConfig_DS)
    .def(py::init<size_t, size_t, double, double, double>(),"nxdirty"_a,
      "nydirty"_a, "epsilon"_a, "pixsize_x"_a, "pixsize_y"_a)
    .def("Nxdirty", &GridderConfig<double>::Nxdirty)
    .def("Nydirty", &GridderConfig<double>::Nydirty)
    .def("Epsilon", &GridderConfig<double>::Epsilon)
    .def("Pixsize_x", &GridderConfig<double>::Pixsize_x)
    .def("Pixsize_y", &GridderConfig<double>::Pixsize_y)
    .def("Nu", &GridderConfig<double>::Nu)
    .def("Nv", &GridderConfig<double>::Nv)
    .def("grid2dirty", &GridderConfig<double>::grid2dirty, grid2dirty_DS,
      "grid"_a)
    .def("grid2dirty_c", &GridderConfig<double>::grid2dirty_c, "grid"_a)
    .def("dirty2grid", &GridderConfig<double>::dirty2grid, dirty2grid_DS,
      "dirty"_a)
    .def("dirty2grid_c", &GridderConfig<double>::dirty2grid_c, "dirty"_a);
  m.def("getIndices", getIndices<double>, getIndices_DS, "baselines"_a,
    "gconf"_a, "flags"_a, "chbegin"_a=-1, "chend"_a=-1,
    "wmin"_a=-1e30, "wmax"_a=1e30);
  m.def("vis2grid",&vis2grid<double>, vis2grid_DS, "baselines"_a, "gconf"_a,
    "idx"_a, "vis"_a);
  m.def("ms2grid",&ms2grid<double>, "baselines"_a, "gconf"_a, "idx"_a, "ms"_a);
  m.def("grid2vis",&grid2vis<double>, grid2vis_DS, "baselines"_a, "gconf"_a,
    "idx"_a, "grid"_a);
  m.def("grid2ms",&grid2ms<double>, "baselines"_a, "gconf"_a, "idx"_a,
    "grid"_a, "ms_in"_a=py::none());
  m.def("vis2grid_c",&vis2grid_c<double>, "baselines"_a, "gconf"_a, "idx"_a,
    "vis"_a);
  m.def("ms2grid_c",&ms2grid_c<double>, "baselines"_a, "gconf"_a, "idx"_a,
    "ms"_a);
  m.def("grid2vis_c",&grid2vis_c<double>, "baselines"_a, "gconf"_a, "idx"_a,
    "grid"_a);
  m.def("grid2ms_c",&grid2ms_c<double>, "baselines"_a, "gconf"_a, "idx"_a,
    "grid"_a, "ms_in"_a=py::none());
  }