/*
 *  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();

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

template<typename T> bool isPytype(const py::array &arr)
  {
  auto t1=arr.dtype();
  auto t2=pybind11::dtype::of<T>();
  auto k1=t1.kind();
  auto k2=t2.kind();
  auto s1=t1.itemsize();
  auto s2=t2.itemsize();
  return (k1==k2)&&(s1==s2);
  }
template<typename T> pyarr<T> getPyarr(const py::array &arr, const string &name)
  {
  auto t1=arr.dtype();
  auto t2=pybind11::dtype::of<T>();
  auto k1=t1.kind();
  auto k2=t2.kind();
  auto s1=t1.itemsize();
  auto s2=t2.itemsize();
  myassert((k1==k2)&&(s1==s2),
    "type mismatch for array '", name, "': expected '", k2, s2,
    "', but got '", k1, s1, "'.");
  return arr.cast<pyarr<T>>();
  }

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

void checkArray(const py::array &arr, const string &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 string &name, 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_ = getPyarr<T>(in.cast<py::array>(), name);
  checkArray(tmp_, name, shape);
  return tmp_;
  }

template<typename T> pyarr<T> providePotentialArray(const py::object &in,
  const string &name, const vector<size_t> &shape)
  {
  if (in.is_none())
    return makeArray<T>(vector<size_t>(shape.size(), 0));
  return getPyarr<T>(in.cast<py::array>(), name);
  }

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.float64)
    u, v and w coordinates for each row
freq: np.array((nchannels,), dtype=np.float64)
    frequency for each individual channel (in Hz)
)""";
class PyBaselines: public Baselines
  {
  public:
    using Baselines::Baselines;
    template<typename T> PyBaselines(const pyarr<T> &coord, const pyarr<T> &freq)
      : Baselines(make_const_mav<2>(coord), make_const_mav<1>(freq))
      {}

    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
    )""";

    template<typename T> 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::effectiveUVW(idx,res);
      }
      return res_;
      }

    template<typename T> pyarr<T> ms2vis(const pyarr<T> &ms_,
      const pyarr<uint32_t> &idx_, size_t nthreads) 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<T>({nvis});
      auto vis = make_mav<1>(res);
      {
      py::gil_scoped_release release;
      Baselines::ms2vis(ms, idx, vis, nthreads);
      }
      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 T> pyarr<T> vis2ms(const pyarr<T> &vis_,
      const pyarr<uint32_t> &idx_, py::object &ms_in, size_t nthreads) const
      {
      auto vis=make_const_mav<1>(vis_);
      auto idx=make_const_mav<1>(idx_);
      auto res = provideArray<T>(ms_in, "ms_in", {nrows, nchan});
      auto ms = make_mav<2>(res);
      {
      py::gil_scoped_release release;
      Baselines::vis2ms(vis, idx, ms, nthreads);
      }
      return res;
      }
  };

constexpr auto grid2dirty_DS = R"""(
Converts from UV grid to dirty image (FFT, cropping, correction)

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

Returns
=======
nd.array((nxdirty, nydirty), same dtype as `grid`)
    the dirty image
)""";

constexpr auto dirty2grid_DS = R"""(
Converts from a dirty image to a UV grid (correction, padding, FFT)

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

Returns
=======
np.array((nu, nv), same dtype as `dirty`)
    gridded UV data
)""";

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

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

Returns
=======
np.array((nxdirty, nydirty), same dtype as `img`)
    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.complex64 or 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), same dtype as 'dirty')
    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)
nthreads: int
    the number of threads to use for all calculations involving this object.
)""";
class PyGridderConfig: public GridderConfig
  {
  public:
    using GridderConfig::GridderConfig;
    PyGridderConfig(size_t nxdirty, size_t nydirty, double epsilon,
      double pixsize_x, double pixsize_y, size_t nthreads)
      : GridderConfig(nxdirty, nydirty, epsilon, pixsize_x, pixsize_y, nthreads) {}

    template<typename T> pyarr<T> apply_taper2(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::apply_taper(img2, res2, divide);
      }
      return res;
      }
    py::array apply_taper(const py::array &img, bool divide) const
      {
      if (isPytype<complex<float>>(img))
        return apply_taper2<float>(img, divide);
      if (isPytype<complex<double>>(img))
        return apply_taper2<double>(img, divide);
      myfail("type matching failed: 'img' has neither type 'f4' nor 'f8'");
      }
    template<typename T> pyarr<T> grid2dirty2(const pyarr<T> &grid) const
      {
      auto res = makeArray<T>({nx_dirty, ny_dirty});
      auto grid2=make_const_mav<2>(grid);
      auto res2=make_mav<2>(res);
      {
      py::gil_scoped_release release;
      GridderConfig::grid2dirty(grid2,res2);
      }
      return res;
      }
    py::array grid2dirty(const py::array &grid) const
      {
      if (isPytype<float>(grid))
        return grid2dirty2<float>(grid);
      if (isPytype<double>(grid))
        return grid2dirty2<double>(grid);
      myfail("type matching failed: 'grid' has neither type 'f4' nor 'f8'");
      }
    template<typename T> pyarr<complex<T>> grid2dirty_c2
      (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::grid2dirty_c(grid2,res2);
      }
      return res;
      }
    py::array grid2dirty_c(const py::array &grid) const
      {
      if (isPytype<complex<float>>(grid))
        return grid2dirty_c2<float>(grid);
      if (isPytype<complex<double>>(grid))
        return grid2dirty_c2<double>(grid);
      myfail("type matching failed: 'grid' has neither type 'c8' nor 'c16'");
      }

    template<typename T> pyarr<T> dirty2grid2(const pyarr<T> &dirty) const
      {
      auto dirty2 = make_const_mav<2>(dirty);
      auto grid = makeArray<T>({nu, nv});
      auto grid2=make_mav<2>(grid);
      {
      py::gil_scoped_release release;
      GridderConfig::dirty2grid(dirty2, grid2);
      }
      return grid;
      }
    py::array dirty2grid(const py::array &dirty) const
      {
      if (isPytype<float>(dirty))
        return dirty2grid2<float>(dirty);
      if (isPytype<double>(dirty))
        return dirty2grid2<double>(dirty);
      myfail("type matching failed: 'dirty' has neither type 'f4' nor 'f8'");
      }
    template<typename T> pyarr<complex<T>> dirty2grid_c2(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::dirty2grid_c(dirty2, grid2);
      }
      return grid;
      }
    py::array dirty2grid_c(const py::array &dirty) const
      {
      if (isPytype<complex<float>>(dirty))
        return dirty2grid_c2<float>(dirty);
      if (isPytype<complex<double>>(dirty))
        return dirty2grid_c2<double>(dirty);
      myfail("type matching failed: 'dirty' has neither type 'c8' nor 'c16'");
      }
    template<typename T> pyarr<complex<T>> apply_wscreen2
      (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::apply_wscreen(dirty2, res2, w, adjoint);
      }
      return res;
      }
    py::array apply_wscreen(const py::array &dirty, double w, bool adjoint) const
      {
      if (isPytype<complex<float>>(dirty))
        return apply_wscreen2<float>(dirty, w, adjoint);
      if (isPytype<complex<double>>(dirty))
        return apply_wscreen2<double>(dirty, w, adjoint);
      myfail("type matching failed: 'dirty' has neither type 'c8' nor 'c16'");
      }
  };


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 &baselines, const PyGridderConfig &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_, "wgt", {nvis});
  auto wgt2 = make_const_mav<1>(wgt);
  auto res = provideArray<complex<T>>(grid_in, "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)
)""";
template<typename T> pyarr<T> Pyvis2grid(const PyBaselines &baselines,
  const PyGridderConfig &gconf, const pyarr<uint32_t> &idx_,
  const pyarr<complex<T>> &vis_, py::object &grid_in, const py::object &wgt_)
  {
  auto tmp=Pyvis2grid_c(baselines, gconf, idx_, vis_, None, wgt_);
  auto grd=provideArray<T>(grid_in, "grid_in", {gconf.Nu(), gconf.Nv()});
  {
  py::gil_scoped_release release;
  gridder::detail::complex2hartley(make_const_mav<2>(tmp), make_mav<2>(grd), gconf.Nthreads());
  }
  return grd;
  }

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 &baselines, const PyGridderConfig &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_, "wgt", {nrows, nchan});
  auto wgt2 = make_const_mav<2>(wgt);
  auto res = provideArray<complex<T>>(grid_in, "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<T> Pyms2grid(
  const PyBaselines &baselines, const PyGridderConfig &gconf,
  const pyarr<uint32_t> &idx_, const pyarr<complex<T>> &ms_,
  py::object &grid_in, const py::object &wgt_)
  {
  auto tmp = Pyms2grid_c(baselines, gconf, idx_, ms_, None, wgt_);
  auto res_ = provideArray<T>(grid_in, "grid_in", {gconf.Nu(), gconf.Nv()});
  auto res = make_mav<2>(res_);
  {
  py::gil_scoped_release release;
  gridder::detail::complex2hartley(make_const_mav<2>(tmp), res, gconf.Nthreads());
  }
  return res_;
  }

template<typename T> pyarr<complex<T>> Pygrid2vis_c(
  const PyBaselines &baselines, const PyGridderConfig &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_, "wgt", {nvis});
  auto wgt2 = make_const_mav<1>(wgt);
  auto res = makeArray<complex<T>>({nvis});
  auto vis = make_mav<1>(res);
  vis.fill(0);
  {
  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>> Pygrid2vis(const PyBaselines &baselines,
  const PyGridderConfig &gconf, const pyarr<uint32_t> &idx_,
  const pyarr<T> &grid_, const py::object &wgt_)
  {
  auto tmp=makeArray<complex<T>>({gconf.Nu(), gconf.Nv()});
  gridder::detail::hartley2complex(make_const_mav<2>(grid_),make_mav<2>(tmp), gconf.Nthreads());
  return Pygrid2vis_c(baselines, gconf, idx_, tmp, wgt_);
  }

template<typename T> pyarr<complex<T>> Pygrid2ms_c(
  const PyBaselines &baselines, const PyGridderConfig &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_, "wgt", {nrows, nchan});
  auto wgt2 = make_const_mav<2>(wgt);
  auto res = provideArray<complex<T>>(ms_in, "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>> Pygrid2ms(const PyBaselines &baselines,
  const PyGridderConfig &gconf, const pyarr<uint32_t> &idx_,
  const pyarr<T> &grid_, py::object &ms_in, const py::object &wgt_)
  {
  auto grid2_ = makeArray<complex<T>>({gconf.Nu(), gconf.Nv()});
  auto grid2 = make_mav<2>(grid2_);
  gridder::detail::hartley2complex(make_const_mav<2>(grid_), grid2, gconf.Nthreads());
  return Pygrid2ms_c(baselines, gconf, idx_, grid2_, ms_in, wgt_);
  }

template<typename T> pyarr<complex<T>> apply_holo2(
  const PyBaselines &baselines, const PyGridderConfig &gconf,
  const py::array &idx_, const py::array &grid_, const py::object &wgt_)
  {
  auto idx = getPyarr<uint32_t>(idx_, "idx");
  auto idx2 = make_const_mav<1>(idx);
  auto grid = getPyarr<complex<T>>(grid_, "grid");
  auto grid2 = make_const_mav<2>(grid);
  auto wgt = providePotentialArray<T>(wgt_, "wgt", {idx2.shape(0)});
  auto wgt2 = make_const_mav<1>(wgt);
  auto res = makeArray<complex<T>>({grid2.shape(0),grid2.shape(1)});
  auto res2 = make_mav<2>(res);
  {
  py::gil_scoped_release release;
  apply_holo(baselines, gconf, idx2, grid2, res2, wgt2);
  }
  return res;
  }
py::array Pyapply_holo(
  const PyBaselines &baselines, const PyGridderConfig &gconf,
  const py::array &idx, const py::array &grid, const py::object &wgt)
  {
  if (isPytype<complex<float>>(grid))
    return apply_holo2<float>(baselines, gconf, idx, grid, wgt);
  if (isPytype<complex<double>>(grid))
    return apply_holo2<double>(baselines, gconf, idx, grid, wgt);
  myfail("type matching failed: 'grid' has neither type 'c8' nor 'c16'");
  }

template<typename T> pyarr<T> Pyget_correlations(
  const PyBaselines &baselines, const PyGridderConfig &gconf,
  const pyarr<uint32_t> &idx_, int du, int dv, const py::object &wgt_)
  {
  auto idx = make_const_mav<1>(idx_);
  pyarr<T> wgt2 = providePotentialArray<T>(wgt_, "wgt", {idx.shape(0)});
  auto wgt=make_const_mav<1>(wgt2);

  auto res = makeArray<T>({gconf.Nu(),gconf.Nv()});
  auto ogrid = make_mav<2>(res);
  {
  py::gil_scoped_release release;
  get_correlations(baselines, gconf, idx, du, dv, ogrid, wgt);
  }
  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.
)""";
pyarr<uint32_t> PygetIndices(const PyBaselines &baselines,
  const PyGridderConfig &gconf, const pyarr<bool> &flags_, int chbegin,
  int chend, double wmin, double wmax)
  {
  size_t nidx;
  auto flags = make_const_mav<2>(flags_);
  {
  py::gil_scoped_release release;
  nidx = getIdxSize(baselines, flags, chbegin, chend, wmin, wmax, gconf.Nthreads());
  }
  auto res = makeArray<uint32_t>({nidx});
  auto res2 = make_mav<1>(res);
  {
  py::gil_scoped_release release;
  fillIdx(baselines, gconf, flags, chbegin, chend, wmin, wmax, res2);
  }
  return res;
  }

template<typename T> pyarr<T> vis2dirty2(const PyBaselines &baselines,
  const PyGridderConfig &gconf, const py::array &idx_,
  const py::array &vis_, const py::object &wgt_, bool do_wstacking)
  {
  auto idx = getPyarr<uint32_t>(idx_, "idx");
  auto idx2 = make_const_mav<1>(idx);
  auto dirty = makeArray<T>({gconf.Nxdirty(), gconf.Nydirty()});
  auto dirty2 = make_mav<2>(dirty);
  auto vis = getPyarr<complex<T>>(vis_, "vis");
  auto vis2 = make_const_mav<1>(vis);
  auto wgt = providePotentialArray<T>(wgt_, "wgt", {idx2.shape(0)});
  auto wgt2 = make_const_mav<1>(wgt);
  {
  py::gil_scoped_release release;
  vis2dirty<T>(baselines, gconf, idx2, vis2, wgt2, dirty2, do_wstacking);
  }
  return dirty;
  }

constexpr auto vis2dirty_DS = R"""(
Converts an array of visibilities to a dirty image.

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 provided visibilities
vis: np.array(nvis,), dtype=np.complex64 or np.complex128)
    The input visibilities
    Its data type determines the precision in which the calculation is carried
    out.
wgt: np.array((nvis,), dtype=float with same precision as `vis`, optional
    If present, visibilities are multiplied by the corresponding entries.
do_wstacking: bool
    if True, the full improved w-stacking algorithm is carried out, otherwise
    the w values are assumed to be zero.

Returns
=======
np.array((nxdirty, nydirty), dtype=float of same precision as `vis`.)
    The dirty image
)""";

py::array Pyvis2dirty(const PyBaselines &baselines,
  const PyGridderConfig &gconf, const py::array &idx,
  const py::array &vis, const py::object &wgt, bool do_wstacking)
  {
  if (isPytype<complex<float>>(vis))
    return vis2dirty2<float>(baselines, gconf, idx, vis, wgt, do_wstacking);
  if (isPytype<complex<double>>(vis))
    return vis2dirty2<double>(baselines, gconf, idx, vis, wgt, do_wstacking);
  myfail("type matching failed: 'vis' has neither type 'c8' nor 'c16'");
  }

template<typename T> pyarr<complex<T>> dirty2vis2(const PyBaselines &baselines,
  const PyGridderConfig &gconf, const pyarr<uint32_t> &idx_,
  const pyarr<T> &dirty_, const py::object &wgt_, bool do_wstacking)
  {
  auto idx = getPyarr<uint32_t>(idx_, "idx");
  auto idx2 = make_const_mav<1>(idx);
  auto dirty = getPyarr<T>(dirty_, "dirty");
  auto dirty2 = make_const_mav<2>(dirty_);
  auto wgt = providePotentialArray<T>(wgt_, "wgt", {idx2.shape(0)});
  auto wgt2 = make_const_mav<1>(wgt);
  auto vis = makeArray<complex<T>>({idx2.shape(0)});
  auto vis2 = make_mav<1>(vis);
  vis2.fill(0);
  {
  py::gil_scoped_release release;
  vis2.fill(0);
  dirty2vis<T>(baselines, gconf, idx2, dirty2, wgt2, vis2, do_wstacking);
  }
  return vis;
  }

constexpr auto dirty2vis_DS = R"""(
Converts a dirty image into a 1D array of visibilities.

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 visibilities to be computed
dirty: np.array((nxdirty, nydirty), dtype=np.float32 or np.float64)
    dirty image
    Its data type determines the precision in which the calculation is carried
    out.
wgt: np.array((nvis,), same dtype as `dirty`, optional
    If present, visibilities are multiplied by the corresponding entries.
do_wstacking: bool
    if True, the full improved w-stacking algorithm is carried out, otherwise
    the w values are assumed to be zero.

Returns
=======
np.array((nvis,), dtype=complex of same precision as `dirty`.)
    The visibility data
)""";
py::array Pydirty2vis(const PyBaselines &baselines,
  const PyGridderConfig &gconf, const py::array &idx, const py::array &dirty,
  const py::object &wgt, bool do_wstacking)
  {
  if (isPytype<float>(dirty))
    return dirty2vis2<float>(baselines, gconf, idx, dirty, wgt, do_wstacking);
  if (isPytype<double>(dirty))
    return dirty2vis2<double>(baselines, gconf, idx, dirty, wgt, do_wstacking);
  myfail("type matching failed: 'dirty' has neither type 'f4' nor 'f8'");
  }

template<typename T> py::array ms2dirty2(const py::array &uvw_,
  const py::array &freq_, const py::array &ms_, const py::object &wgt_,
  size_t npix_x, size_t npix_y, double pixsize_x, double pixsize_y, double epsilon,
  bool do_wstacking, size_t nthreads, size_t verbosity)
  {
  auto uvw = getPyarr<double>(uvw_, "uvw");
  auto uvw2 = make_const_mav<2>(uvw);
  auto freq = getPyarr<double>(freq_, "freq");
  auto freq2 = make_const_mav<1>(freq);
  auto ms = getPyarr<complex<T>>(ms_, "ms");
  auto ms2 = make_const_mav<2>(ms);
  auto wgt = providePotentialArray<T>(wgt_, "wgt", {ms2.shape(0),ms2.shape(1)});
  auto wgt2 = make_const_mav<2>(wgt);
  auto dirty = makeArray<T>({npix_x,npix_y});
  auto dirty2 = make_mav<2>(dirty);
  {
  py::gil_scoped_release release;
  ms2dirty(uvw2,freq2,ms2,wgt2,pixsize_x,pixsize_y,epsilon,do_wstacking,
    nthreads,dirty2,verbosity);
  }
  return dirty;
  }

constexpr auto ms2dirty_DS = R"""(
Converts an MS object to dirty image.

Parameters
==========
uvw: np.array((nrows, 3), dtype=np.float64)
    UVW coordinates from the measurement set
freq: np.array((nchan,), dtype=np.float64)
    channel frequencies
ms: np.array((nrows, nchan,), dtype=np.complex64 or np.complex128)
    the input measurement set data.
    Its data type determines the precision in which the calculation is carried
    out.
wgt: np.array((nrows, nchan), float with same precision as `ms`), optional
    If present, its values are multiplied to the output
npix_x, npix_y: int
    dimensions of the dirty image
pixsize_x, pixsize_y: float
    angular pixel size (in radians) of the dirty image
epsilon: float
    accuracy at which the computation should be done. Must be larger than 2e-13.
    If `ms` has type np.complex64, it must be larger than 1e-5.
do_wstacking: bool
    if True, the full improved w-stacking algorithm is carried out, otherwise
    the w values are assumed to be zero.
nthreads: int
    number of threads to use for the calculation
verbosity: int
    0: no output
    1: some output
    2: detailed output

Returns
=======
np.array((nxdirty, nydirty), dtype=float of same precision as `ms`)
    the dirty image
)""";
py::array Pyms2dirty(const py::array &uvw,
  const py::array &freq, const py::array &ms, const py::object &wgt,
  size_t npix_x, size_t npix_y, double pixsize_x, double pixsize_y, double epsilon,
  bool do_wstacking, size_t nthreads, size_t verbosity)
  {
  if (isPytype<complex<float>>(ms))
    return ms2dirty2<float>(uvw, freq, ms, wgt, npix_x, npix_y,
      pixsize_x, pixsize_y, epsilon, do_wstacking, nthreads, verbosity);
  if (isPytype<complex<double>>(ms))
    return ms2dirty2<double>(uvw, freq, ms, wgt, npix_x, npix_y,
      pixsize_x, pixsize_y, epsilon, do_wstacking, nthreads, verbosity);
  myfail("type matching failed: 'ms' has neither type 'c8' nor 'c16'");
  }

template<typename T> py::array dirty2ms2(const py::array &uvw_,
  const py::array &freq_, const py::array &dirty_, const py::object &wgt_,
  double pixsize_x, double pixsize_y, double epsilon,
  bool do_wstacking, size_t nthreads, size_t verbosity)
  {
  auto uvw = getPyarr<double>(uvw_, "uvw");
  auto uvw2 = make_const_mav<2>(uvw);
  auto freq = getPyarr<double>(freq_, "freq");
  auto freq2 = make_const_mav<1>(freq);
  auto dirty = getPyarr<T>(dirty_, "dirty");
  auto dirty2 = make_const_mav<2>(dirty);
  auto wgt = providePotentialArray<T>(wgt_, "wgt", {uvw2.shape(0),freq2.shape(0)});
  auto wgt2 = make_const_mav<2>(wgt);
  auto ms = makeArray<complex<T>>({uvw2.shape(0),freq2.shape(0)});
  auto ms2 = make_mav<2>(ms);
  {
  py::gil_scoped_release release;
  dirty2ms(uvw2,freq2,dirty2,wgt2,pixsize_x,pixsize_y,epsilon,do_wstacking,
    nthreads,ms2,verbosity);
  }
  return ms;
  }

constexpr auto dirty2ms_DS = R"""(
Converts a dirty image to an MS object.

Parameters
==========
uvw: np.array((nrows, 3), dtype=np.float64)
    UVW coordinates from the measurement set
freq: np.array((nchan,), dtype=np.float64)
    channel frequencies
dirty: np.array((nxdirty, nydirty), dtype=np.float32 or np.float64)
    dirty image
    Its data type determines the precision in which the calculation is carried
    out.
wgt: np.array((nrows, nchan), same dtype as `dirty`), optional
    If present, its values are multiplied to the output
pixsize_x, pixsize_y: float
    angular pixel size (in radians) of the dirty image
epsilon: float
    accuracy at which the computation should be done. Must be larger than 2e-13.
    If `dirty` has type np.float32, it must be larger than 1e-5.
do_wstacking: bool
    if True, the full improved w-stacking algorithm is carried out, otherwise
    the w values are assumed to be zero.
nthreads: int
    number of threads to use for the calculation
verbosity: int
    0: no output
    1: some output
    2: detailed output

Returns
=======
np.array((nrows, nchan,), dtype=complex of same precision as `dirty`)
    the measurement set data.
)""";
py::array Pydirty2ms(const py::array &uvw,
  const py::array &freq, const py::array &dirty, const py::object &wgt,
  double pixsize_x, double pixsize_y, double epsilon,
  bool do_wstacking, size_t nthreads, size_t verbosity)
  {
  if (isPytype<float>(dirty))
    return dirty2ms2<float>(uvw, freq, dirty, wgt,
      pixsize_x, pixsize_y, epsilon, do_wstacking, nthreads, verbosity);
  if (isPytype<double>(dirty))
    return dirty2ms2<double>(uvw, freq, dirty, wgt,
      pixsize_x, pixsize_y, epsilon, do_wstacking, nthreads, verbosity);
  myfail("type matching failed: 'dirty' has neither type 'f4' nor 'f8'");
  }

} // unnamed namespace

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

  py::class_<PyBaselines> (m, "Baselines", PyBaselines_DS)
    .def(py::init<const pyarr<double> &, const pyarr<double> &>(),
      "coord"_a, "freq"_a)
    .def ("Nrows",&PyBaselines::Nrows)
    .def ("Nchannels",&PyBaselines::Nchannels)
    .def ("ms2vis",&PyBaselines::ms2vis<complex<double>>,
      PyBaselines::ms2vis_DS, "ms"_a, "idx"_a, "nthreads"_a=1)
    .def ("effectiveuvw",&PyBaselines::effectiveuvw<double>, "idx"_a)
    .def ("vis2ms",&PyBaselines::vis2ms<complex<double>>,
      PyBaselines::vis2ms_DS, "vis"_a, "idx"_a, "ms_in"_a=None, "nthreads"_a=1);
  py::class_<PyGridderConfig> (m, "GridderConfig", GridderConfig_DS)
    .def(py::init<size_t, size_t, double, double, double, size_t>(),"nxdirty"_a,
      "nydirty"_a, "epsilon"_a, "pixsize_x"_a, "pixsize_y"_a, "nthreads"_a=1)
    .def("Nxdirty", &PyGridderConfig::Nxdirty)
    .def("Nydirty", &PyGridderConfig::Nydirty)
    .def("Epsilon", &PyGridderConfig::Epsilon)
    .def("Pixsize_x", &PyGridderConfig::Pixsize_x)
    .def("Pixsize_y", &PyGridderConfig::Pixsize_y)
    .def("Nu", &PyGridderConfig::Nu)
    .def("Nv", &PyGridderConfig::Nv)
    .def("Supp", &PyGridderConfig::Supp)
    .def("apply_taper", &PyGridderConfig::apply_taper, apply_taper_DS,
      "img"_a, "divide"_a=false)
     .def("grid2dirty", &PyGridderConfig::grid2dirty,
        grid2dirty_DS, "grid"_a)
    .def("grid2dirty_c", &PyGridderConfig::grid2dirty_c, "grid"_a)
    .def("dirty2grid", &PyGridderConfig::dirty2grid,
       dirty2grid_DS, "dirty"_a)
    .def("dirty2grid_c", &PyGridderConfig::dirty2grid_c, "dirty"_a)
    .def("apply_wscreen", &PyGridderConfig::apply_wscreen,
      apply_wscreen_DS, "dirty"_a, "w"_a, "adjoint"_a)

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

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

        }));
  m.def("getIndices", PygetIndices, getIndices_DS, "baselines"_a,
    "gconf"_a, "flags"_a, "chbegin"_a=-1, "chend"_a=-1,
    "wmin"_a=-1e30, "wmax"_a=1e30);
  m.def("vis2grid",&Pyvis2grid<double>, vis2grid_DS, "baselines"_a, "gconf"_a,
    "idx"_a, "vis"_a, "grid_in"_a=None, "wgt"_a=None);
  m.def("ms2grid",&Pyms2grid<double>, "baselines"_a, "gconf"_a, "idx"_a, "ms"_a,
    "grid_in"_a=None, "wgt"_a=None);
  m.def("grid2vis",&Pygrid2vis<double>, grid2vis_DS, "baselines"_a, "gconf"_a,
    "idx"_a, "grid"_a, "wgt"_a=None);
  m.def("grid2ms",&Pygrid2ms<double>, "baselines"_a, "gconf"_a, "idx"_a,
    "grid"_a, "ms_in"_a=None, "wgt"_a=None);
  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",&Pyapply_holo, "baselines"_a, "gconf"_a, "idx"_a,
    "grid"_a, "wgt"_a=None);
  m.def("get_correlations", &Pyget_correlations<double>, "baselines"_a, "gconf"_a,
    "idx"_a, "du"_a, "dv"_a, "wgt"_a=None);
  m.def("vis2dirty",&Pyvis2dirty, vis2dirty_DS, "baselines"_a, "gconf"_a,
    "idx"_a, "vis"_a, "wgt"_a=None, "do_wstacking"_a=false);
  m.def("dirty2vis",&Pydirty2vis, "baselines"_a, dirty2vis_DS, "gconf"_a,
    "idx"_a, "dirty"_a, "wgt"_a=None, "do_wstacking"_a=false);

  m.def("ms2dirty", &Pyms2dirty, ms2dirty_DS, "uvw"_a, "freq"_a, "ms"_a,
    "wgt"_a=None, "npix_x"_a, "npix_y"_a, "pixsize_x"_a, "pixsize_y"_a,
    "epsilon"_a, "do_wstacking"_a=false, "nthreads"_a=1, "verbosity"_a=0);
  m.def("dirty2ms", &Pydirty2ms, dirty2ms_DS, "uvw"_a, "freq"_a, "dirty"_a,
    "wgt"_a=None, "pixsize_x"_a, "pixsize_y"_a, "epsilon"_a,
    "do_wstacking"_a=false, "nthreads"_a=1, "verbosity"_a=0);
  }