import math

import numpy as np
from numba import jit, njit, prange
import xarray
from astropy import constants


@jit(nopython=True, cache=True, nogil=True)  # fastmath=True
def prolate_spheroidal_degrid_jit(
    grid,
    vis_data,
    uvw,
    frequency_coord,
    frequency_map,
    time_map,
    pol_map,
    cgk_1D,
    n_uv,
    delta_lm,
    support,
    oversampling,
):
    """Sample (degrid) a model UV grid at visibility coordinates to produce predicted visibilities.

    This is the inverse operation of :func:`prolate_spheroidal_grid_jit`.
    Rather than accumulating visibilities onto a grid, it interpolates the
    model grid at each visibility's UV coordinate using the prolate spheroidal
    convolution kernel, and writes the result back into ``vis_data`` in place.

    For each non-flagged sample the interpolated value is::

        vis_data[i_time, i_baseline, i_chan, i_pol] =
            sum_over_support(conv * grid[a_time, a_chan, a_pol, v_indx, u_indx]) / norm

    where ``norm`` is the sum of convolution kernel values over the support
    window.

    Time and frequency mapping allow the input visibility axes to be collapsed
    or reordered from the image axes.

    Convention note — axis prefixes:
      * ``m_`` denotes an image-domain axis size.
      * ``n_`` denotes a visibility-domain axis size.

    Parameters
    ----------
    grid : complex ndarray, shape (m_time, m_chan, m_pol, m_v, m_u)
        Model UV grid to be sampled.  Read only.  Note that the spatial axes
        are ordered ``(m_v, m_u)`` rather than ``(m_u, m_v)`` as in
        :func:`prolate_spheroidal_grid_jit`.
    vis_data : complex ndarray, shape (n_time, n_baseline, n_vis_chan, n_pol)
        On input, the sole skip condition is ``NaN``: samples where
        ``vis_data`` is ``NaN`` are not degridded and are left unchanged.
        Flagged visibilities must therefore be set to ``NaN`` by the caller
        before invoking this function.
        On output, each processed sample is overwritten with the
        convolution-interpolated model value divided by the kernel norm.
        Modified in place.
    uvw : float ndarray, shape (n_time, n_baseline, 3)
        Baseline coordinates in metres.  The third axis is ordered (U, V, W).
    frequency_coord : float ndarray, shape (n_vis_chan,)
        Channel centre frequencies in Hz, used to convert UVW from metres to
        wavelengths.
    frequency_map : int ndarray, shape (n_vis_chan,)
        Maps each visibility channel index to its source image channel index
        ``a_chan`` in ``grid``.
    time_map : int ndarray, shape (n_time,)
        Maps each visibility time index to its source image time index
        ``a_time`` in ``grid``.
    pol_map : int ndarray, shape (n_pol,)
        Maps each visibility polarization index to its source image
        polarization index ``a_pol`` in ``grid``.
    cgk_1D : float ndarray, shape (oversampling * (support // 2 + 1),)
        Oversampled 1-D prolate spheroidal wave function (PSWF) kernel.
        The 2-D kernel is formed as the outer product of two 1-D lookups.
    n_uv : int ndarray, shape (2,)
        Padded UV-grid dimensions ``[n_u, n_v]`` in pixels.
    delta_lm : float ndarray, shape (2,)
        Image cell size in radians ``[delta_l, delta_m]``.
    support : int
        Full convolution support width in grid pixels.  The kernel spans
        ``[-support//2, support - support//2)`` pixels around each sample.
    oversampling : int
        Oversampling factor of the 1-D kernel ``cgk_1D``.  A higher value
        reduces interpolation error at the cost of a larger kernel table.

    Returns
    -------
    None
        ``vis_data`` is modified in place; nothing is returned.

    Notes
    -----
    - **Flags must be applied before calling**: this function has no flag
      array parameter.  The only skip mechanism is a ``NaN`` check on
      ``vis_data``.  The caller is responsible for setting flagged samples
      to ``NaN`` in ``vis_data`` prior to the call; otherwise flagged
      baselines will receive a (potentially spurious) model prediction.
    - Visibilities whose U or V coordinate is ``NaN`` are skipped entirely.
    - Samples where ``vis_data`` is ``NaN`` are skipped and left unchanged.
    - The sub-pixel offset from the nearest grid point is rounded using
      ``int(x + 0.5)`` rather than ``numpy.round`` to match Fortran/C++
      rounding behaviour (round-half-away-from-zero).
    - Hardcoding ``support`` and ``oversampling`` as literals (currently
      commented out) allows Numba to unroll the inner support loops,
      yielding a significant speed-up.
    - The ``grid`` spatial axes are ordered ``(m_v, m_u)``, which is the
      reverse of the gridding function's ``(m_u, m_v)`` order.
    """
    # By hardcoding the support and oversampling values, the innermost for loops can be unrolled by the compiler leading to significantly faster code.
    # support = 7
    # oversampling = 100

    c = 299792458.0
    uv_scale = np.zeros((2, len(frequency_coord)), dtype=np.double)
    uv_scale[0, :] = -(frequency_coord * delta_lm[0] * n_uv[0]) / c
    uv_scale[1, :] = -(frequency_coord * delta_lm[1] * n_uv[1]) / c

    # oversampling_center = int(oversampling // 2)
    support_center = int(support // 2)
    uv_center = n_uv // 2

    start_support = -support_center
    end_support = (
        support - support_center
    )  # end_support is larger by 1 so that python range() gives correct indices

    n_time = uvw.shape[0]
    n_baseline = uvw.shape[1]
    n_chan = len(frequency_map)
    n_pol = len(pol_map)

    n_u = n_uv[0]
    n_v = n_uv[1]

    for i_time in range(n_time):
        a_time = time_map[i_time]
        for i_baseline in range(n_baseline):
            for i_chan in range(n_chan):
                a_chan = frequency_map[i_chan]
                u = uvw[i_time, i_baseline, 0] * uv_scale[0, i_chan]
                v = uvw[i_time, i_baseline, 1] * uv_scale[1, i_chan]

                if ~np.isnan(u) and ~np.isnan(v):
                    u_pos = u + uv_center[0]
                    v_pos = v + uv_center[1]

                    # Do not use numpy round
                    u_center_indx = int(u_pos + 0.5)
                    v_center_indx = int(v_pos + 0.5)

                    if (
                        (u_center_indx + support_center < n_u)
                        and (v_center_indx + support_center < n_v)
                        and (u_center_indx - support_center >= 0)
                        and (v_center_indx - support_center >= 0)
                    ):
                        u_offset = u_center_indx - u_pos
                        u_center_offset_indx = math.floor(u_offset * oversampling + 0.5)
                        v_offset = v_center_indx - v_pos
                        v_center_offset_indx = math.floor(v_offset * oversampling + 0.5)

                        for i_pol in range(n_pol):

                            a_pol = pol_map[i_pol]
                            norm = 0.0
                            degrid_value = 0.0

                            if ~np.isnan(vis_data[i_time, i_baseline, i_chan, i_pol]):
                                for i_v in range(start_support, end_support):
                                    v_indx = v_center_indx + i_v
                                    v_offset_indx = np.abs(
                                        oversampling * i_v + v_center_offset_indx
                                    )
                                    conv_v = cgk_1D[v_offset_indx]

                                    for i_u in range(start_support, end_support):
                                        u_indx = u_center_indx + i_u
                                        u_offset_indx = np.abs(
                                            oversampling * i_u + u_center_offset_indx
                                        )
                                        conv_u = cgk_1D[u_offset_indx]
                                        conv = conv_u * conv_v
                                        degrid_value += (
                                            conv
                                            * grid[
                                                a_time, a_chan, a_pol, v_indx, u_indx
                                            ]
                                        )
                                        norm = norm + conv

                                if norm > 0.0:
                                    vis_data[i_time, i_baseline, i_chan, i_pol] = (
                                        degrid_value / norm
                                    )
    return


# def sgrid(
#     uvw: np.ndarray,
#     dphase: float,
#     freq: float,
#     c: float,
#     scale: np.ndarray,
#     offset: np.ndarray,
#     sampling: int,
#     pos: np.ndarray,
#     loc,
#     off,
# ):
#     """
#     uvw is a 3 element array
#     """
#     pos[:2] = scale[:2] * uvw[:2] * freq / c + (offset[:2])
#     loc[:2] = np.round(pos[:2]).astype(int)
#     off[:2] = np.round((loc[:2] - pos[:2]) * sampling).astype(int)
#     phase = -2.0 * np.pi * dphase * freq / c
#     phasor = complex(np.cos(phase), np.sin(phase))
#     return phasor


# ### This is a translation of fgridft::dgrid
# def dgrid(
#     uvw: np.ndarray,
#     dphase: np.ndarray,
#     values: np.ndarray,
#     flag,
#     scale,
#     offset,
#     grid,
#     freq,
#     c,
#     support,
#     sampling,
#     convFunc,
#     chanmap,
#     polmap,
# ):
#     """
#     uvw should have the shape (nt, nb, 3)
#     values has shape (nt, nb, nvispol, nvischan) # Corrected understanding of shape
#     flag has shape (nt, nb, nvispol, nvischan)   # Corrected understanding of shape
#     dphase should have the shape (nt,nb,)
#     freq should have shape (nvischan,) # Corrected understanding of shape
#     scale (3) matching u,v,w
#     offset(3) where is u,v,w =0 on the pixel grid
#     convfunc (support//2 * sampling)
#     chanmap (nvischan)
#     polmap (nvispol)
#     """
#     nvispol = values.shape[3]
#     nvischan = values.shape[2]
#     nchan = grid.shape[0]
#     nx = grid.shape[2]
#     ny = grid.shape[3]
#     npol = grid.shape[1]
#     nt = values.shape[0]
#     nb = values.shape[1]
#     supp_centre = support // 2
#     supp_beg = -supp_centre
#     supp_end = supp_beg + support

#     for t in range(nt):
#         for b in range(nb):
#             # Corrected loop order to match values and flag shape
#             for ipol in range(nvispol):
#                 apol = polmap[ipol]
#                 if 0 <= apol < npol:
#                     for ichan in range(nvischan):
#                         achan = chanmap[ichan]
#                         if 0 <= achan < nchan:
#                             pos = np.zeros(2)
#                             loc = np.zeros(2, dtype=int)
#                             off = np.zeros(2, dtype=int)

#                             # Pass the correct frequency for the current channel
#                             phasor = sgrid(
#                                 uvw[t, b],
#                                 dphase[t, b],
#                                 freq[ichan],
#                                 c,
#                                 scale,
#                                 offset,
#                                 sampling,
#                                 pos,
#                                 loc,
#                                 off,
#                             )

#                             # Corrected indexing for flag and values
#                             if not flag[
#                                 t, b, ichan, ipol
#                             ]:  # Access flag using ipol then ichan
#                                 ogrid = (
#                                     ((loc[0] + supp_beg) >= 0)
#                                     and ((loc[0] + supp_end - 1) < nx)
#                                     and ((loc[1] + supp_beg) >= 0)
#                                     and ((loc[1] + supp_end - 1) < ny)
#                                 )

#                                 if ogrid:
#                                     nvalue = 0.0
#                                     norm = 0.0
#                                     for iy in range(supp_beg, supp_end):
#                                         iloc2 = abs(sampling * iy + off[1])
#                                         # Ensure iloc2 is within bounds for convFunc
#                                         if iloc2 < len(convFunc):
#                                             wty = convFunc[iloc2]
#                                             for ix in range(supp_beg, supp_end):
#                                                 iloc1 = abs(sampling * ix + off[0])
#                                                 # Ensure iloc1 is within bounds for convFunc
#                                                 if iloc1 < len(convFunc):
#                                                     wtx = convFunc[iloc1]
#                                                     wt = wtx * wty
#                                                     norm += wt
#                                                     # Access grid using achan then apol, and corrected loc indexing
#                                                     grid_val = grid[
#                                                         achan,
#                                                         apol,
#                                                         loc[0] + ix,
#                                                         loc[1] + iy,
#                                                     ]
#                                                     nvalue += wt * grid_val
#                                     if norm != 0.0:
#                                         # Access values using ipol then ichan
#                                         values[t, b, ichan, ipol] += (
#                                             nvalue * np.conj(phasor)
#                                         ) / norm
#     return


###optimizied with help of Gemini


# @njit
# def dgrid_optimized(
#     uvw,
#     dphase,
#     values,
#     flag,
#     scale,
#     offset,
#     grid,
#     freq,
#     c,
#     support,
#     sampling,
#     convFunc,
#     chanmap,
#     polmap,
# ):

#     nvispol = values.shape[3]
#     nvischan = values.shape[2]
#     nchan = grid.shape[0]
#     nx = grid.shape[2]
#     ny = grid.shape[3]
#     npol = grid.shape[1]
#     nt = values.shape[0]
#     nb = values.shape[1]

#     supp_centre = support // 2
#     supp_beg = -supp_centre
#     supp_end = supp_beg + support
#     # Pre-calculate sgrid components for all time and baseline points
#     # This replaces the loop over t and b and the call to sgrid inside.
#     # Expanding sgrid logic:
#     # pos = (
#     #    scale[:2] * uvw[:, :, :2] * (freq[0] / c) + offset[:2]
#     # )  # Using freq[0] as a placeholder, need to handle freq dependency per channel later
#     # loc = np.round(pos).astype(np.int64)
#     # off = np.round((loc - pos) * sampling).astype(np.int64)

#     # Recalculate pos, loc, off, and phasor per channel dependency outside the inner loops.
#     pos_chan = np.empty((nt, nb, nvischan, 2), dtype=np.float64)
#     loc_chan = np.empty((nt, nb, nvischan, 2), dtype=np.int64)
#     off_chan = np.empty((nt, nb, nvischan, 2), dtype=np.int64)
#     phasor_chan = np.empty((nt, nb, nvischan), dtype=np.complex128)

#     for ichan in range(nvischan):
#         achan = chanmap[ichan]
#         if 0 <= achan < nchan:
#             # sgrid logic applied across nt and nb
#             pos_chan[:, :, ichan, :] = (
#                 scale[:2] * uvw[:, :, :2] * (freq[ichan] / c) + offset[:2]
#             )
#             loc_chan[:, :, ichan, :] = (np.round(pos_chan[:, :, ichan, :])).astype(
#                 np.int64
#             )
#             off_chan[:, :, ichan, :] = (
#                 np.round(
#                     (loc_chan[:, :, ichan, :] - pos_chan[:, :, ichan, :]) * sampling
#                 )
#             ).astype(np.int64)
#             phase = -2.0 * np.pi * dphase * freq[ichan] / c
#             phasor_chan[:, :, ichan] = np.cos(phase) + 1j * np.sin(phase)

#     # Pre-calculate convolution kernel weights for all possible offsets
#     # The offsets for x and y dimensions are independent, so we can calculate 1D kernels and combine them.
#     max_offset_abs = (
#         support // 2 + 1
#     ) * sampling  # Maximum possible absolute offset in sampling units
#     conv_kernel_1d = np.empty(max_offset_abs + 1, dtype=convFunc.dtype)
#     for i in range(max_offset_abs):
#         conv_kernel_1d[i] = convFunc[i]

#     # Loop over time, baseline, channel, and polarization
#     for t in range(nt):
#         for b in range(nb):
#             for ichan in range(nvischan):
#                 achan = chanmap[ichan]
#                 if 0 <= achan < nchan:
#                     loc_tbi = loc_chan[t, b, ichan, :]
#                     off_tbi = off_chan[t, b, ichan, :]
#                     phasor_tbi = phasor_chan[t, b, ichan]

#                     ogrid = (
#                         ((loc_tbi[0] + supp_beg) >= 0)
#                         and ((loc_tbi[0] + supp_end - 1) < nx)
#                         and ((loc_tbi[1] + supp_beg) >= 0)
#                         and ((loc_tbi[1] + supp_end - 1) < ny)
#                     )

#                     if ogrid:
#                         for ipol in range(nvispol):
#                             apol = polmap[ipol]
#                             if (not flag[t, b, ichan, ipol]) and (0 <= apol < npol):
#                                 # Vectorized convolution
#                                 # Extract grid patch
#                                 grid_patch = grid[
#                                     achan,
#                                     apol,
#                                     loc_tbi[0] + supp_beg : loc_tbi[0] + supp_end,
#                                     loc_tbi[1] + supp_beg : loc_tbi[1] + supp_end,
#                                 ]

#                                 # Create 2D kernel from pre-calculated 1D kernels based on offset
#                                 offset_x = off_tbi[0]
#                                 offset_y = off_tbi[1]
#                                 kernel_x = conv_kernel_1d[
#                                     np.abs(
#                                         sampling * np.arange(supp_beg, supp_end)
#                                         + offset_x
#                                     )
#                                 ]
#                                 kernel_y = conv_kernel_1d[
#                                     np.abs(
#                                         sampling * np.arange(supp_beg, supp_end)
#                                         + offset_y
#                                     )
#                                 ]
#                                 conv_kernel_2d = kernel_x.reshape(
#                                     -1, 1
#                                 ) * kernel_y.reshape(1, -1)

#                                 # Perform element-wise multiplication and sum for nvalue and norm
#                                 nvalue = np.sum(conv_kernel_2d * grid_patch)
#                                 norm = np.sum(
#                                     conv_kernel_2d
#                                 )  # The sum of weights is the norm

#                                 if norm != 0.0:
#                                     values[t, b, ichan, ipol] += (
#                                         nvalue * np.conj(phasor_tbi)
#                                     ) / norm
#     return


# some manual optimization of dgrid


# def dgrid2(
#     uvw,
#     dphase,
#     values,
#     flag,
#     scale,
#     offset,
#     grid,
#     freq,
#     c,
#     support,
#     sampling,
#     convFunc,
#     chanmap,
#     polmap,
# ):
#     nvispol = values.shape[3]
#     nvischan = values.shape[2]
#     nchan, npol, nx, ny = grid.shape
#     nt, nb = values.shape[:2]
#     supp_centre = support // 2
#     supp_beg = -supp_centre
#     supp_end = supp_beg + support

#     for t in range(nt):
#         for b in range(nb):
#             for ipol in range(nvispol):
#                 apol = polmap[ipol]
#                 if 0 <= apol < npol:
#                     for ichan in range(nvischan):
#                         achan = chanmap[ichan]
#                         if 0 <= achan < nchan and not flag[t, b, ichan, ipol]:
#                             pos = np.empty(2)
#                             loc = np.empty(2, dtype=int)
#                             off = np.empty(2, dtype=int)

#                             phasor = sgrid(
#                                 uvw[t, b],
#                                 dphase[t, b],
#                                 freq[ichan],
#                                 c,
#                                 scale,
#                                 offset,
#                                 sampling,
#                                 pos,
#                                 loc,
#                                 off,
#                             )

#                             x0, y0 = loc
#                             if (
#                                 (x0 + supp_beg >= 0)
#                                 and (x0 + supp_end < nx)
#                                 and (y0 + supp_beg >= 0)
#                                 and (y0 + supp_end < ny)
#                             ):
#                                 nvalue = 0.0
#                                 norm = 0.0
#                                 for iy in range(supp_beg, supp_end):
#                                     iloc2 = abs(sampling * iy + off[1])
#                                     if iloc2 < len(convFunc):
#                                         wty = convFunc[iloc2]
#                                         for ix in range(supp_beg, supp_end):
#                                             iloc1 = abs(sampling * ix + off[0])
#                                             if iloc1 < len(convFunc):
#                                                 wtx = convFunc[iloc1]
#                                                 wt = wtx * wty
#                                                 norm += wt
#                                                 nvalue += (
#                                                     wt
#                                                     * grid[
#                                                         achan,
#                                                         apol,
#                                                         x0 + ix,
#                                                         y0 + iy,
#                                                     ]
#                                                 )
#                                 if norm != 0.0:
#                                     values[t, b, ichan, ipol] += (
#                                         nvalue * np.conj(phasor)
#                                     ) / norm


# jit version of dgrid2


# @njit
# def sgrid_numba(uvw, dphase, freq, c, scale, offset, sampling):
#     pos = np.empty(2, dtype=np.float64)
#     loc = np.empty(2, dtype=np.int32)
#     off = np.empty(2, dtype=np.int32)

#     pos[0] = scale[0] * uvw[0] * freq / c + offset[0]
#     pos[1] = scale[1] * uvw[1] * freq / c + offset[1]

#     loc[0] = int(np.round(pos[0]))
#     loc[1] = int(np.round(pos[1]))

#     off[0] = int(np.round((float(loc[0]) - pos[0]) * float(sampling)))
#     off[1] = int(np.round((float(loc[1]) - pos[1]) * float(sampling)))

#     phase = -2.0 * np.pi * dphase * freq / c
#     phasor = np.cos(phase) + 1j * np.sin(phase)

#     return phasor, loc, off


# @njit
# def dgrid_numba(
#     uvw,
#     dphase,
#     values,
#     flag,
#     scale,
#     offset,
#     grid,
#     freq,
#     c,
#     support,
#     sampling,
#     convFunc,
#     chanmap,
#     polmap,
# ):
#     nvispol = values.shape[3]
#     nvischan = values.shape[2]
#     nchan, npol, ny, nx = grid.shape
#     nt, nb = values.shape[:2]
#     supp_centre = support // 2
#     supp_beg = -supp_centre
#     supp_end = supp_beg + support
#     pos = np.empty(2, dtype=np.float64)
#     loc = np.empty(2, dtype=np.int32)
#     off = np.empty(2, dtype=np.int32)

#     for t in prange(nt):
#         for b in range(nb):
#             for ipol in range(nvispol):
#                 apol = polmap[ipol]
#                 if 0 <= apol < npol:
#                     for ichan in range(nvischan):
#                         achan = chanmap[ichan]
#                         if 0 <= achan < nchan and not flag[t, b, ichan, ipol]:
#                             # phasor, loc, off = sgrid_numba(
#                             #    uvw[t, b],
#                             #    dphase[t, b],
#                             #    freq[ichan],
#                             #    c,
#                             #    scale,
#                             #    offset,
#                             #    sampling,
#                             # )
#                             # Inlined sgrid

#                             pos[0] = (
#                                 scale[0] * uvw[t, b, 0] * freq[ichan] / c + offset[0]
#                             )
#                             pos[1] = (
#                                 scale[1] * uvw[t, b, 1] * freq[ichan] / c + offset[1]
#                             )

#                             loc[0] = int(np.round(pos[0]))
#                             loc[1] = int(np.round(pos[1]))

#                             off[0] = int(
#                                 np.round((float(loc[0]) - pos[0]) * float(sampling))
#                             )
#                             off[1] = int(
#                                 np.round((float(loc[1]) - pos[1]) * float(sampling))
#                             )

#                             phase = -2.0 * np.pi * dphase[t, b] * freq[ichan] / c
#                             phasor = np.cos(phase) + 1j * np.sin(phase)
#                             # End of inlined sgrid logic

#                             x0, y0 = loc
#                             if (
#                                 (x0 + supp_beg >= 0)
#                                 and ((x0 + supp_end - 1) < nx)
#                                 and (y0 + supp_beg >= 0)
#                                 and ((y0 + supp_end - 1) < ny)
#                             ):
#                                 nvalue = 0.0
#                                 norm = 0.0
#                                 for ix in range(supp_beg, supp_end):
#                                     iloc1 = abs(sampling * ix + off[0])
#                                     wtx = convFunc[iloc1]
#                                     for iy in range(supp_beg, supp_end):
#                                         iloc2 = abs(sampling * iy + off[1])
#                                         wty = convFunc[iloc2]
#                                         wt = wtx * wty
#                                         norm += wt
#                                         nvalue += (
#                                             wt
#                                             * grid[
#                                                 achan,
#                                                 apol,
#                                                 x0 + ix,
#                                                 y0 + iy,
#                                             ]
#                                         )
#                                 if norm != 0.0:
#                                     values[t, b, ichan, ipol] += (
#                                         nvalue * np.conj(phasor)
#                                     ) / norm


# def degrid_spheroid_ms4(
#     vis: xarray.core.datatree.DataTree,
#     grid: np.ndarray,
#     pixelincr: np.ndarray,
#     support: int = 7,
#     sampling: int = 100,
#     incremental: bool = False,
#     whichFunc=0,
# ):
#     """

#     Parameters
#     ----------
#     vis : xarray.core.datatree.DataTree
#         ms4 xradio datatree.
#     grid : np.ndarray
#         model gridded visibilities.
#     pixelincr : np.ndarray
#         [dx,dy] 2 values in radians
#     support : int, optional
#         Support of the prolate spheroidal conv func.
#         The default is 7.
#     sampling : int, optional
#         oversampling of prolate spheroidal conv func .
#         The default is 100.
#     incremental : bool, optional
#         if True add model visibilities to existing VISIBILITY_MODEL.
#         The default is False.
#     whichFunc : TYPE, optional
#         0 use degrid loop and jit
#         1 use degrid vectorized and jit
#         2 use degrid as translated from fortran

#         The default is 0.

#     TODO: The chanmap and polmap may be totally wrong

#     Returns
#     -------
#     None.

#     """
#     func = dgrid_numba
#     if whichFunc == 1:
#         func = dgrid_optimized
#     elif whichFunc == 2:
#         func = dgrid
#     uvw = vis.UVW.data
#     dims = vis.dims
#     # as we are not gridding away from phase_center
#     dphase = np.zeros((dims["time"], dims["baseline_id"]), dtype=float)

#     if not incremental or "VISIBILITY_MODEL" not in vis:
#         modvis = np.zeros(
#             (
#                 dims["time"],
#                 dims["baseline_id"],
#                 dims["frequency"],
#                 dims["polarization"],
#             ),
#             dtype=np.complex128(),
#         )
#     else:
#         modvis = vis.VISIBILITY_MODEL.data
#     flag = vis.FLAG.data
#     # lets flag NaN UVW
#     ###testoo
#     # flag[:, :, :, :] = False
#     ######
#     uvwmask = np.isnan(vis.UVW)
#     # print("number of nan UVW", np.sum(uvwmask) / 3)
#     nan_it, nan_ib = np.where(uvwmask.any(dim="uvw_label"))
#     # print("sum of flag bef", np.sum(flag))
#     # flag[nan_it, nan_ib, :, :] = True
#     for k in range(len(nan_it)):
#         flag[nan_it[k], nan_ib[k], :, :] = True
#         uvw[nan_it[k], nan_ib[k], :] = (
#             1e9  # avoid dealing with nan in some loops later so larger than earth diameter
#         )
#     # print("sum of flags aft", np.sum(flag))
#     nx, ny = grid.shape[-2:]
#     # for sake of completeness making size 3 ofr uvw scales and offset
#     # using same convention of in standard_grid ..negating scale to
#     # deal with -u, -v flip from handedness of uvw in MSv2 onwards
#     scale = np.array([-pixelincr[0] * nx, -pixelincr[1] * ny, 0.0]).astype(np.float64)
#     offset = np.array([nx / 2, ny / 2, 0.0]).astype(np.float64)
#     freq = vis.frequency.data
#     c = constants.c.value
#     convFunc = create_prolate_spheroidal_kernel_1D(sampling, support)
#     nchan = grid.shape[0]
#     npol = grid.shape[1]
#     ## we need to pass chan/pol maps or determine them from info
#     chanmap = (np.arange(dims["frequency"]) / dims["frequency"] * nchan).astype(int)
#     # Channel mapping
#     polmap = np.round(np.arange(dims["polarization"]) % npol).astype(
#         int
#     )  # this is wrong most probably

#     #    dgrid_numba(
#     func(
#         uvw,
#         dphase,
#         modvis,
#         flag,
#         scale,
#         offset,
#         grid,
#         freq,
#         c,
#         support,
#         sampling,
#         convFunc,
#         chanmap,
#         polmap,
#     )
#     datdims = vis.VISIBILITY.dims
#     datcoords = vis.VISIBILITY.coords
#     # print("MODVIS", np.max(modvis))
#     modvis_da = xarray.DataArray(modvis, coords=datcoords, dims=datdims)
#     vis["VISIBILITY_MODEL"] = modvis_da
#     # print("post mod", np.max(vis.VISIBILITY_MODEL))
