/*
 *  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_gridder; 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 "gridder_cxx.h"

using namespace std;
using namespace gridder;

namespace py = pybind11;

namespace {

auto None = py::none();

//
// basic utilities
//

//
// Start of real gridder functionality
//
constexpr auto set_nthreads_DS = R"""(
Specifies the number of threads to be used by the module

Parameters
==========
nthreads: int
    the number of threads to be used. Must be >=1.
)""";
constexpr auto get_nthreads_DS = R"""(
Returns the number of threads used by the module

Returns
=======
int : the number of threads used by the module
)""";

template<typename T>
  using pyarr = py::array_t<T, 0>;

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

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<T> provideArray(const py::object &in,
  const vector<size_t> &shape)
  {
  if (in.is_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<T>>();
  checkArray(tmp_, "temporary", shape);
  return tmp_;
  }

template<typename T> pyarr<T> providePotentialArray(const py::object &in,
  const vector<size_t> &shape)
  {
  if (in.is_none())
    return makeArray<T>(vector<size_t>(shape.size(), 0));
  auto tmp_ = in.cast<pyarr<T>>();
  checkArray(tmp_, "temporary", shape);
  return tmp_;
  }

template<typename T> pyarr<T> provideCArray(py::object &in,
  const vector<size_t> &shape)
  {
  if (in.is_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<T>>();
  checkArray(tmp_, "temporary", shape);
  return tmp_;
  }

template<size_t ndim, typename T> mav<T,ndim> make_mav(pyarr<T> &in)
  {
  myassert(ndim==in.ndim(), "dimension mismatch");
  array<size_t,ndim> dims;
  array<ptrdiff_t,ndim> str;
  for (size_t i=0; i<ndim; ++i)
    {
    dims[i]=in.shape(i);
    str[i]=in.strides(i)/sizeof(T);
    myassert(str[i]*ptrdiff_t(sizeof(T))==in.strides(i), "weird strides");
    }
  return mav<T, ndim>(in.mutable_data(),dims,str);
  }
template<size_t ndim, typename T> const_mav<T,ndim> make_const_mav(const pyarr<T> &in)
  {
  myassert(ndim==in.ndim(), "dimension mismatch");
  array<size_t,ndim> dims;
  array<ptrdiff_t,ndim> str;
  for (size_t i=0; i<ndim; ++i)
    {
    dims[i]=in.shape(i);
    str[i]=in.strides(i)/sizeof(T);
    myassert(str[i]*ptrdiff_t(sizeof(T))==in.strides(i), "weird strides");
    }
  return const_mav<T, ndim>(in.data(),dims,str);
  }

constexpr auto PyBaselines_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)
)""";
template<typename T> class PyBaselines: public Baselines<T>
  {
  protected:
    using Baselines<T>::coord;
    using Baselines<T>::f_over_c;
    using Baselines<T>::nrows;
    using Baselines<T>::nchan;

  public:
    using Baselines<T>::Baselines;
    PyBaselines(const pyarr<T> &coord, const pyarr<T> &freq)
      : Baselines<T>(make_const_mav<2>(coord), make_const_mav<1>(freq))
      {}

    using Baselines<T>::effectiveCoord;
    using Baselines<T>::Nrows;
    using Baselines<T>::Nchannels;

    static constexpr auto 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
    )""";

   // using Baselines<T>::effectiveUVW;
    pyarr<T> effectiveuvw(const pyarr<uint32_t> &idx_) const
      {
      size_t nvis = size_t(idx_.shape(0));
      auto idx=make_const_mav<1>(idx_);
      auto res_=makeArray<T>({nvis, 3});
      auto res=make_mav<2>(res_);
      {
      py::gil_scoped_release release;
      Baselines<T>::effectiveUVW(idx,res);
      }
      return res_;
      }

    template<typename T2> pyarr<T2> ms2vis(const pyarr<T2> &ms_,
      const pyarr<uint32_t> &idx_) const
      {
      auto idx=make_const_mav<1>(idx_);
      size_t nvis = size_t(idx.shape(0));
      auto ms = make_const_mav<2>(ms_);
      auto res=makeArray<T2>({nvis});
      auto vis = make_mav<1>(res);
      {
      py::gil_scoped_release release;
      Baselines<T>::ms2vis(ms, idx, vis);
      }
      return res;
      }

    static constexpr auto 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
    ms_in: np.array((nrows, nchannels), dtype=np.complex), optional
        input measurement set to which the visibilities are added.

    Returns
    =======
    np.array((nrows, nchannels), dtype=np.complex)
        the measurement set's visibility data (0 where not covered by idx)
    )""";
    template<typename T2> pyarr<T2> vis2ms(const pyarr<T2> &vis_,
      const pyarr<uint32_t> &idx_, py::object &ms_in) const
      {
      auto vis=make_const_mav<1>(vis_);
      auto idx=make_const_mav<1>(idx_);
      auto res = provideArray<T2>(ms_in, {nrows, nchan});
      auto ms = make_mav<2>(res);
      {
      py::gil_scoped_release release;
      Baselines<T>::vis2ms(vis, idx, ms);
      }
      return res;
      }
  };

constexpr auto apply_taper_DS = R"""(
Applies the taper (or its inverse) to an image

Parameters
==========
img: nd.array((nxdirty, nydirty), dtype=np.float64)
    the image
divide: bool
    if True, the routine dividex by the taper, otherwise it multiplies by it

Returns
=======
np.array((nxdirty, nydirty), dtype=np.float64)
    the image with the taper applied
)""";

constexpr auto apply_wscreen_DS = R"""(
Applies the w screen to an image

Parameters
==========
dirty: nd.array((nxdirty, nydirty), dtype=np.complex128)
    the image
w : float
    the w value to use
adjoint: bool
    if True, apply the complex conjugate of the w screen

Returns
=======
np.array((nxdirty, nydirty), dtype=np.complex128)
    the image with the w screen applied
)""";

constexpr auto 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)
)""";
template<typename T> class PyGridderConfig: public GridderConfig<T>
  {
  protected:
    using GridderConfig<T>::nx_dirty;
    using GridderConfig<T>::ny_dirty;
    using GridderConfig<T>::eps;
    using GridderConfig<T>::psx;
    using GridderConfig<T>::psy;
    using GridderConfig<T>::supp;
    using GridderConfig<T>::nsafe;
    using GridderConfig<T>::nu;
    using GridderConfig<T>::nv;
    using GridderConfig<T>::beta;
    using GridderConfig<T>::cfu;
    using GridderConfig<T>::cfv;
    using GridderConfig<T>::wscreen;

  public:
    using GridderConfig<T>::GridderConfig;
    PyGridderConfig(size_t nxdirty, size_t nydirty, double epsilon,
      double pixsize_x, double pixsize_y)
      : GridderConfig<T>(nxdirty, nydirty, epsilon, pixsize_x, pixsize_y) {}
    using GridderConfig<T>::Nxdirty;
    using GridderConfig<T>::Nydirty;
    using GridderConfig<T>::Epsilon;
    using GridderConfig<T>::Pixsize_x;
    using GridderConfig<T>::Pixsize_y;
    using GridderConfig<T>::Nu;
    using GridderConfig<T>::Nv;
    using GridderConfig<T>::Supp;
    using GridderConfig<T>::Nsafe;
    using GridderConfig<T>::Beta;

    pyarr<T> apply_taper(const pyarr<T> &img, bool divide) const
      {
      auto res = makeArray<T>({nx_dirty, ny_dirty});
      auto img2 = make_const_mav<2>(img);
      auto res2 = make_mav<2>(res);
      {
      py::gil_scoped_release release;
      GridderConfig<T>::apply_taper(img2, res2, divide);
      }
      return res;
      }
    pyarr<complex<T>> grid2dirty_c(const pyarr<complex<T>> &grid) const
      {
      auto res = makeArray<complex<T>>({nx_dirty, ny_dirty});
      auto grid2=make_const_mav<2>(grid);
      auto res2=make_mav<2>(res);
      {
      py::gil_scoped_release release;
      GridderConfig<T>::grid2dirty_c(grid2,res2);
      }
      return res;
      }

    pyarr<complex<T>> dirty2grid_c(const pyarr<complex<T>> &dirty) const
      {
      auto dirty2 = make_const_mav<2>(dirty);
      auto grid = makeArray<complex<T>>({nu, nv});
      auto grid2=make_mav<2>(grid);
      {
      py::gil_scoped_release release;
      GridderConfig<T>::dirty2grid_c(dirty2, grid2);
      }
      return grid;
      }
    pyarr<complex<T>> apply_wscreen(const pyarr<complex<T>> &dirty, double w, bool adjoint) const
      {
      auto dirty2 = make_const_mav<2>(dirty);
      auto res = makeArray<complex<T>>({nx_dirty, ny_dirty});
      auto res2 = make_mav<2>(res);
      {
      py::gil_scoped_release release;
      GridderConfig<T>::apply_wscreen(dirty2, res2, w, adjoint);
      }
      return res;
      }
  };


constexpr auto vis2grid_c_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
grid_in: np.array((nu,nv), dtype=np.complex128), optional
    If present, the result is added to this array.
wgt: np.array((nvis,), dtype=np.float64), optional
    If present, visibilities are multiplied by the corresponding entries.

Returns
=======
np.array((nu,nv), dtype=np.complex128):
    the gridded visibilities
)""";
template<typename T> pyarr<complex<T>> Pyvis2grid_c(
  const PyBaselines<T> &baselines, const PyGridderConfig<T> &gconf,
  const pyarr<uint32_t> &idx_, const pyarr<complex<T>> &vis_,
  py::object &grid_in, const py::object &wgt_)
  {
  auto vis2 = make_const_mav<1>(vis_);
  size_t nvis = vis2.shape(0);
  auto idx2 = make_const_mav<1>(idx_);
  pyarr<T> wgt = providePotentialArray<T>(wgt_, {nvis});
  auto wgt2 = make_const_mav<1>(wgt);
  auto res = provideCArray<complex<T>>(grid_in, {gconf.Nu(), gconf.Nv()});
  auto grid = make_mav<2>(res);
  {
  py::gil_scoped_release release;
  vis2grid_c<T>(baselines, gconf, idx2, vis2, grid, wgt2);
  }
  return res;
  }

constexpr auto 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
grid_in: np.array((nu,nv), dtype=np.float64), optional
    If present, the result is added to this array.
wgt: np.array((nvis,), dtype=np.float64), optional
    If present, visibilities are multiplied by the corresponding entries.

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

constexpr auto ms2grid_c_DS = R"""(
Grids measurement set data 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
ms: np.array((nrows, nchannels), dtype=np.complex128)
    the measurement set.
grid_in: np.array((nu,nv), dtype=np.complex128), optional
    If present, the result is added to this array.
wgt: np.array((nrows, nchannels), dtype=np.float64), optional
    If present, visibilities are multiplied by the corresponding entries.

Returns
=======
np.array((nu,nv), dtype=np.complex128):
    the gridded visibilities
)""";
template<typename T> pyarr<complex<T>> Pyms2grid_c(
  const PyBaselines<T> &baselines, const PyGridderConfig<T> &gconf,
  const pyarr<uint32_t> &idx_, const pyarr<complex<T>> &ms_,
  py::object &grid_in, const py::object &wgt_)
  {
  auto nrows = baselines.Nrows();
  auto nchan = baselines.Nchannels();
  auto ms2 = make_const_mav<2>(ms_);
  auto idx2 = make_const_mav<1>(idx_);
  pyarr<T> wgt = providePotentialArray<T>(wgt_, {nrows, nchan});
  auto wgt2 = make_const_mav<2>(wgt);
  auto res = provideCArray<complex<T>>(grid_in, {gconf.Nu(), gconf.Nv()});
  auto grid = make_mav<2>(res);
  {
  py::gil_scoped_release release;
  ms2grid_c<T>(baselines, gconf, idx2, ms2, grid, wgt2);
  }
  return res;
  }

template<typename T> pyarr<complex<T>> Pygrid2vis_c(
  const PyBaselines<T> &baselines, const PyGridderConfig<T> &gconf,
  const pyarr<uint32_t> &idx_, const pyarr<complex<T>> &grid_,
  const py::object &wgt_)
  {
  auto grid2 = make_const_mav<2>(grid_);
  auto idx2 = make_const_mav<1>(idx_);
  size_t nvis = idx2.shape(0);
  pyarr<T> wgt = providePotentialArray<T>(wgt_, {nvis});
  auto wgt2 = make_const_mav<1>(wgt);
  auto res = makeArray<complex<T>>({nvis});
  auto vis = make_mav<1>(res);
  {
  py::gil_scoped_release release;
  grid2vis_c<T>(baselines, gconf, idx2, grid2, vis, wgt2);
  }
  return res;
  }

constexpr auto 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)
wgt: np.array((nvis,), dtype=np.float64), optional
    If present, visibilities are multiplied by the corresponding entries.

Returns
=======
np.array((nvis,), dtype=np.complex)
    The degridded visibility data
)""";
template<typename T> pyarr<complex<T>> Pygrid2ms_c(
  const PyBaselines<T> &baselines, const PyGridderConfig<T> &gconf,
  const pyarr<uint32_t> &idx_, const pyarr<complex<T>> &grid_,
  py::object &ms_in, const py::object &wgt_)
  {
  auto nrows = baselines.Nrows();
  auto nchan = baselines.Nchannels();
  auto grid2 = make_const_mav<2>(grid_);
  auto idx2 = make_const_mav<1>(idx_);
  pyarr<T> wgt = providePotentialArray<T>(wgt_, {nrows, nchan});
  auto wgt2 = make_const_mav<2>(wgt);
  auto res = provideCArray<complex<T>>(ms_in, {nrows, nchan});
  auto ms = make_mav<2>(res);
  {
  py::gil_scoped_release release;
  grid2ms_c<T>(baselines, gconf, idx2, grid2, ms, wgt2);
  }
  return res;
  }

template<typename T> pyarr<complex<T>> apply_holo(
  const PyBaselines<T> &baselines, const PyGridderConfig<T> &gconf,
  const pyarr<uint32_t> &idx_, const pyarr<complex<T>> &grid_,
  const py::object &wgt_)
  {
  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_.template unchecked<1>();
  pyarr<T> wgt2 = providePotentialArray<T>(wgt_, {nvis});
  bool have_wgt = wgt2.size()>0;
  auto wgt = wgt2.template unchecked<1>();

  auto res = makeArray<complex<T>>({nu, nv});
  auto ogrid = res.mutable_data();
  {
  py::gil_scoped_release release;

  for (size_t i=0; i<nu*nv; ++i) ogrid[i] = T(0);
  T beta = gconf.Beta();
  size_t supp = gconf.Supp();

  // Loop over sampling points
#pragma omp parallel num_threads(nthreads)
{
  Helper<T> hlp(gconf, grid, ogrid);
  T emb = exp(-2*beta);
  int jump = hlp.lineJump();
  const T * ku = hlp.kernel.data();
  const T * kv = hlp.kernel.data()+supp;

#pragma omp for schedule(guided,100)
  for (size_t ipart=0; ipart<nvis; ++ipart)
    {
    auto tidx = idx(ipart);
    UVW<T> coord = baselines.effectiveCoord(tidx);
    hlp.prep(coord.u, coord.v);
    complex<T> r = 0;
    const auto * ptr = hlp.p0r;
    for (size_t cu=0; cu<supp; ++cu)
      {
      complex<T> tmp(0);
      for (size_t cv=0; cv<supp; ++cv)
        tmp += ptr[cv] * kv[cv];
      r += tmp*ku[cu];
      ptr += jump;
      }
    r*=emb*emb;
    if (have_wgt)
      {
      auto twgt = wgt(ipart);
      r*=twgt*twgt;
      }
    auto * wptr = hlp.p0w;
    for (size_t cu=0; cu<supp; ++cu)
      {
      complex<T> tmp(r*ku[cu]);
      for (size_t cv=0; cv<supp; ++cv)
        wptr[cv] += tmp*kv[cv];
      wptr += jump;
      }
    }
}
  }
  return res;
  }

template<typename T> pyarr<T> get_correlations(
  const PyBaselines<T> &baselines, const PyGridderConfig<T> &gconf,
  const pyarr<uint32_t> &idx_, int du, int dv, const py::object &wgt_)
  {
  size_t nu=gconf.Nu(), nv=gconf.Nv();
  size_t supp = gconf.Supp();
  myassert(size_t(abs(du))<supp, "|du| must be smaller than Supp");
  myassert(size_t(abs(dv))<supp, "|dv| must be smaller than Supp");
  checkArray(idx_, "idx", {0});
  size_t nvis = size_t(idx_.shape(0));
  auto idx = idx_.template unchecked<1>();
  pyarr<T> wgt2 = providePotentialArray<T>(wgt_, {nvis});
  bool have_wgt = wgt2.size()>0;
  auto wgt = wgt2.template unchecked<1>();

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

  size_t u0, u1, v0, v1;
  if (du>=0)
    { u0=0; u1=supp-du; }
  else
    { u0=-du; u1=supp; }
  if (dv>=0)
    { v0=0; v1=supp-dv; }
  else
    { v0=-dv; v1=supp; }

  // Loop over sampling points
#pragma omp parallel num_threads(nthreads)
{
  Helper<T,T> hlp(gconf, nullptr, ogrid);
  T emb = exp(-2*beta);
  int jump = hlp.lineJump();
  const T * ku = hlp.kernel.data();
  const T * kv = hlp.kernel.data()+supp;

#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 * wptr = hlp.p0w + u0*jump;
    auto f0 = emb*emb;
    if (have_wgt)
      {
      auto twgt = wgt(ipart);
      f0*=twgt*twgt;
      }
    for (size_t cu=u0; cu<u1; ++cu)
      {
      auto f1=ku[cu]*ku[cu+du]*f0;
      for (size_t cv=v0; cv<v1; ++cv)
        wptr[cv] += f1*kv[cv]*kv[cv+dv];
      wptr += jump;
      }
    }
}
  }
  return res;
  }

constexpr auto 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.
)""";
template<typename T> pyarr<uint32_t> getIndices(const PyBaselines<T> &baselines,
  const PyGridderConfig<T> &gconf, const pyarr<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;
  }

template<typename T> pyarr<complex<T>> Pyvis2dirty_wstack(const PyBaselines<T> &baselines,
  const PyGridderConfig<T> &gconf, const pyarr<uint32_t> &idx_,
  const pyarr<complex<T>> &vis_)
  {
  auto nx_dirty=gconf.Nxdirty();
  auto ny_dirty=gconf.Nydirty();
  auto idx2=make_const_mav<1>(idx_);
  auto vis2=make_const_mav<1>(vis_);
  auto dirty = makeArray<complex<T>>({nx_dirty, ny_dirty});
  auto dirty2=make_mav<2>(dirty);
  {
  py::gil_scoped_release release;
  vis2dirty_wstack<T>(baselines, gconf, idx2, vis2, dirty2);
  }
  return dirty;
  }

} // unnamed namespace

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

  py::class_<PyBaselines<double>> (m, "Baselines", PyBaselines_DS)
    .def(py::init<const pyarr<double> &, const pyarr<double> &>(),
      "coord"_a, "freq"_a)
    .def ("Nrows",&PyBaselines<double>::Nrows)
    .def ("Nchannels",&PyBaselines<double>::Nchannels)
    .def ("ms2vis",&PyBaselines<double>::ms2vis<complex<double>>,
      PyBaselines<double>::ms2vis_DS, "ms"_a, "idx"_a)
    .def ("effectiveuvw",&PyBaselines<double>::effectiveuvw, "idx"_a)
    .def ("vis2ms",&PyBaselines<double>::vis2ms<complex<double>>,
      PyBaselines<double>::vis2ms_DS, "vis"_a, "idx"_a, "ms_in"_a=None);
  py::class_<PyGridderConfig<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", &PyGridderConfig<double>::Nxdirty)
    .def("Nydirty", &PyGridderConfig<double>::Nydirty)
    .def("Epsilon", &PyGridderConfig<double>::Epsilon)
    .def("Pixsize_x", &PyGridderConfig<double>::Pixsize_x)
    .def("Pixsize_y", &PyGridderConfig<double>::Pixsize_y)
    .def("Nu", &PyGridderConfig<double>::Nu)
    .def("Nv", &PyGridderConfig<double>::Nv)
    .def("Supp", &PyGridderConfig<double>::Supp)
    .def("apply_taper", &PyGridderConfig<double>::apply_taper, apply_taper_DS,
      "img"_a, "divide"_a=false)
    .def("grid2dirty_c", &PyGridderConfig<double>::grid2dirty_c, "grid"_a)
    .def("dirty2grid_c", &PyGridderConfig<double>::dirty2grid_c, "dirty"_a)
    .def("apply_wscreen", &PyGridderConfig<double>::apply_wscreen,
      apply_wscreen_DS, "dirty"_a, "w"_a, "adjoint"_a)

    // pickle support
    .def(py::pickle(
        // __getstate__
        [](const PyGridderConfig<double> & gc) {
          // Encode object state in tuple
          return py::make_tuple(gc.Nxdirty(), gc.Nydirty(), gc.Epsilon(),
                                gc.Pixsize_x(), gc.Pixsize_y());
        },
        // __setstate__
        [](py::tuple t) {
          myassert(t.size()==5,"Invalid state");

          // Reconstruct from tuple
          return PyGridderConfig<double>(t[0].cast<size_t>(), t[1].cast<size_t>(),
                                       t[2].cast<double>(), t[3].cast<double>(),
                                       t[4].cast<double>());

        }));
  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_c",&Pyvis2grid_c<double>, vis2grid_c_DS, "baselines"_a,
    "gconf"_a, "idx"_a, "vis"_a, "grid_in"_a=None, "wgt"_a=None);
  m.def("ms2grid_c",&Pyms2grid_c<double>, ms2grid_c_DS, "baselines"_a, "gconf"_a,
    "idx"_a, "ms"_a, "grid_in"_a=None, "wgt"_a=None);
  m.def("grid2vis_c",&Pygrid2vis_c<double>, "baselines"_a, "gconf"_a, "idx"_a,
    "grid"_a, "wgt"_a=None);
  m.def("grid2ms_c",&Pygrid2ms_c<double>, "baselines"_a, "gconf"_a, "idx"_a,
    "grid"_a, "ms_in"_a=None, "wgt"_a=None);
  m.def("apply_holo",&apply_holo<double>, "baselines"_a, "gconf"_a, "idx"_a,
    "grid"_a, "wgt"_a=None);
  m.def("get_correlations", &get_correlations<double>, "baselines"_a, "gconf"_a,
    "idx"_a, "du"_a, "dv"_a, "wgt"_a=None);
  m.def("set_nthreads",&set_nthreads, set_nthreads_DS, "nthreads"_a);
  m.def("get_nthreads",&get_nthreads, get_nthreads_DS);
  m.def("vis2dirty_wstack",&Pyvis2dirty_wstack<double>, "baselines"_a, "gconf"_a,
    "idx"_a, "vis"_a);
  }