import numpy as np

try:
    import dask.array as da
except ImportError as e:
    opt_import_err = e
else:
    opt_import_err = None

from africanus.gridding.perleypolyhedron.gridder import (
    gridder as np_gridder)
from africanus.gridding.perleypolyhedron.degridder import (
    degridder as np_degridder)
from africanus.gridding.perleypolyhedron.degridder import (
    degridder_serial as np_degridder_serial)
from africanus.gridding.perleypolyhedron.policies import (
    stokes_conversion_policies)
from africanus.util.requirements import requires_optional


def __degrid(uvw,
             gridstack,
             lambdas,
             chanmap,
             image_centres,
             phase_centre,
             cell=None,
             convolution_kernel=None,
             convolution_kernel_width=None,
             convolution_kernel_oversampling=None,
             baseline_transform_policy=None,
             phase_transform_policy=None,
             stokes_conversion_policy=None,
             convolution_policy=None,
             vis_dtype=np.complex128,
             rowparallel=False):
    image_centres = image_centres[0][0]
    if image_centres.ndim != 2:
        raise ValueError(
            "Image centres for DASK wrapper expects list of image centres, "
            "one per facet in radec radians"
        )
    if image_centres.shape[1] != 2:
        raise ValueError("Image centre must be a list of tuples")
    uvw = uvw[0]
    if uvw.ndim != 2 or uvw.shape[1] != 3:
        raise ValueError("UVW array must be nrow x 3")
    gridstack = gridstack[0][0][0][0]
    if gridstack.ndim != 4:
        raise ValueError("Gridstack must be nfacet x nband x ny x nx")
    lambdas = lambdas
    chanmap = chanmap
    if chanmap.size != lambdas.size:
        raise ValueError(
            "Chanmap and corresponding lambdas must match in shape")
    nchan = lambdas.size
    nrow = uvw.shape[0]
    ncorr = stokes_conversion_policies.ncorr_outpy(
        policy_type=stokes_conversion_policy)()
    vis = np.zeros((nrow, nchan, ncorr), dtype=vis_dtype)
    degridcall = np_degridder_serial if not rowparallel else np_degridder
    for fi, f in enumerate(image_centres):
        # add contributions from all facets
        vis[:, :, :] += \
            degridcall(uvw,
                       gridstack[fi, :, :, :],
                       lambdas,
                       chanmap,
                       cell,
                       image_centres,
                       phase_centre,
                       convolution_kernel,
                       convolution_kernel_width,
                       convolution_kernel_oversampling,
                       baseline_transform_policy,
                       phase_transform_policy,
                       stokes_conversion_policy,
                       convolution_policy,
                       vis_dtype=vis_dtype)
    return vis


@requires_optional("dask", opt_import_err)
def degridder(uvw,
              gridstack,
              lambdas,
              chanmap,
              cell,
              image_centres,
              phase_centre,
              convolution_kernel,
              convolution_kernel_width,
              convolution_kernel_oversampling,
              baseline_transform_policy,
              phase_transform_policy,
              stokes_conversion_policy,
              convolution_policy,
              vis_dtype=np.complex128,
              rowparallel=False):
    """
    2D Convolutional degridder, discrete to contiguous
    @uvw: value coordinates, (nrow, 3)
    @gridstack: complex gridded data, (nband, npix, npix)
    @lambdas: wavelengths of data channels
    @chanmap: MFS band mapping per channel
    @cell: cell_size in degrees
    @image_centre: new phase centre of image (radians, ra, dec)
    @phase_centre: original phase centre of data (radians, ra, dec)
    @convolution_kernel: packed kernel as generated by kernels package
    @convolution_kernel_width: number of taps in kernel
    @convolution_kernel_oversampling: number of oversampled points in kernel
    @baseline_transform_policy: any accepted policy in
                                .policies.baseline_transform_policies,
                                can be used to tilt image planes for
                                polyhedron faceting
    @phase_transform_policy: any accepted policy in
                             .policies.phase_transform_policies,
                             can be used to facet at provided
                             facet @image_centre
    @stokes_conversion_policy: any accepted correlation to stokes
                               conversion policy in
                               .policies.stokes_conversion_policies
    @convolution_policy: any accepted convolution policy in
                         .policies.convolution_policies
    @vis_dtype: accumulation vis dtype (default complex 128)
    @rowparallel: adds additional threading per row per chunk. This may be
                  necessary for cases where there are few facets and few chunks
                  to get optimal performance. Requires TBB to be installed
                  from your distribution package management system. See numba
                  documentation
                  http://numba.pydata.org/numba-doc/0.46.0/user/threading-layer.html
                  Must set:
                  from numba import config # have to be openmp to support
                  nested parallelism config.THREADING_LAYER = 'threadsafe'
                  before calling this function
    """
    if image_centres.ndim != 2:
        raise ValueError(
            "Image centres for DASK wrapper expects list of "
            "image centres, one per facet in radec radians"
        )
    if image_centres.shape[1] != 2:
        raise ValueError("Image centre must be a list of tuples")
    if gridstack.ndim != 4 or gridstack.shape[0] != image_centres.shape[0]:
        raise ValueError(
            "Grid stack must be nfacet x nband x yy x xx and match number "
            "of image centres"
        )
    vis = da.blockwise(
        __degrid, ("row", "chan", "corr"),
        uvw, ("row", "uvw"),
        gridstack, ("nfacet", "nband", "y", "x"),
        lambdas, ("chan", ),
        chanmap, ("chan", ),
        image_centres, ("nfacet", "coord"),
        convolution_kernel=convolution_kernel,
        convolution_kernel_width=convolution_kernel_width,
        convolution_kernel_oversampling=convolution_kernel_oversampling,
        baseline_transform_policy=baseline_transform_policy,
        phase_transform_policy=phase_transform_policy,
        stokes_conversion_policy=stokes_conversion_policy,
        convolution_policy=convolution_policy,
        cell=cell,
        phase_centre=phase_centre,
        vis_dtype=vis_dtype,
        # goes to one set of grids per row chunk
        adjust_chunks={"row": 1},
        new_axes={
            "corr":
            stokes_conversion_policies.ncorr_outpy(
                policy_type=stokes_conversion_policy)()
        },
        dtype=vis_dtype,
        meta=np.empty(
            (0, 0, 0),
            dtype=vis_dtype)  # row, chan, correlation product as per MSv2 spec
    )
    return vis


def __grid(uvw,
           vis,
           image_centres,
           lambdas=None,
           chanmap=None,
           convolution_kernel=None,
           convolution_kernel_width=None,
           convolution_kernel_oversampling=None,
           baseline_transform_policy=None,
           phase_transform_policy=None,
           stokes_conversion_policy=None,
           convolution_policy=None,
           npix=None,
           cell=None,
           phase_centre=None,
           grid_dtype=np.complex128,
           do_normalize=False):
    image_centres = image_centres[0]
    if image_centres.ndim != 2:
        raise ValueError(
            "Image centres for DASK wrapper expects list of image "
            "centres, one per facet in radec radians"
        )
    if image_centres.shape[1] != 2:
        raise ValueError("Image centre must be a list of tuples")

    uvw = uvw[0]
    vis = vis[0][0]
    lambdas = lambdas[0]
    chanmap = chanmap[0]
    grid_stack = np.zeros(
        (1, image_centres.shape[0], 1, np.max(chanmap) + 1, npix, npix),
        dtype=grid_dtype)
    for fi, f in enumerate(image_centres):
        grid_stack[0, fi, 0, :, :, :] = \
            np_gridder(uvw, vis, lambdas, chanmap, npix, cell, f, phase_centre,
                       convolution_kernel, convolution_kernel_width,
                       convolution_kernel_oversampling,
                       baseline_transform_policy, phase_transform_policy,
                       stokes_conversion_policy,
                       convolution_policy, grid_dtype, do_normalize)
    return grid_stack


@requires_optional("dask", opt_import_err)
def gridder(uvw,
            vis,
            lambdas,
            chanmap,
            npix,
            cell,
            image_centres,
            phase_centre,
            convolution_kernel,
            convolution_kernel_width,
            convolution_kernel_oversampling,
            baseline_transform_policy,
            phase_transform_policy,
            stokes_conversion_policy,
            convolution_policy,
            grid_dtype=np.complex128,
            do_normalize=False):
    """
    2D Convolutional gridder, contiguous to discrete
    @uvw: value coordinates, (nrow, 3)
    @vis: complex data, (nrow, nchan, ncorr)
    @lambdas: wavelengths of data channels
    @chanmap: MFS band mapping
    @npix: number of pixels per axis
    @cell: cell_size in degrees
    @image_centres: new phase centres of images (nfacet, (radians ra, dec))
    @phase_centre: original phase centre of data (radians, ra, dec)
    @convolution_kernel: packed kernel as generated by kernels package
    @convolution_kernel_width: number of taps in kernel
    @convolution_kernel_oversampling: number of oversampled points in kernel
    @baseline_transform_policy: any accepted policy in
                                .policies.baseline_transform_policies,
                                can be used to tilt image planes for
                                polyhedron faceting
    @phase_transform_policy: any accepted policy in
                             .policies.phase_transform_policies,
                             can be used to facet at provided
                             facet @image_centre
    @stokes_conversion_policy: any accepted correlation to stokes
                               conversion policy in
                               .policies.stokes_conversion_policies
    @convolution_policy: any accepted convolution policy in
                        .policies.convolution_policies
    @grid_dtype: accumulation grid dtype (default complex 128)
    @do_normalize: normalize grid by convolution weights
    """
    if len(vis.chunks) != 3 or lambdas.chunks[0] != vis.chunks[1]:
        raise ValueError(
            "Visibility frequency chunking does not match "
            "lambda frequency chunking"
        )
    if len(vis.chunks) != 3 or chanmap.chunks[0] != vis.chunks[1]:
        raise ValueError(
            "Visibility frequency chunking does not match chanmap "
            "frequency chunking"
        )
    if len(vis.chunks) != 3 or len(
            uvw.chunks) != 2 or vis.chunks[0] != uvw.chunks[0]:
        raise ValueError(
            "Visibility row chunking does not match uvw row chunking")
    grids = da.blockwise(
        __grid, ("row", "nfacet", "nstokes", "nband", "y", "x"),
        uvw, ("row", "uvw"),
        vis, ("row", "chan", "corr"),
        image_centres, ("nfacet", "coord"),
        lambdas, ("chan", ),
        chanmap, ("chan", ),
        convolution_kernel=convolution_kernel,
        convolution_kernel_width=convolution_kernel_width,
        convolution_kernel_oversampling=convolution_kernel_oversampling,
        baseline_transform_policy=baseline_transform_policy,
        phase_transform_policy=phase_transform_policy,
        stokes_conversion_policy=stokes_conversion_policy,
        convolution_policy=convolution_policy,
        npix=npix,
        cell=cell,
        phase_centre=phase_centre,
        grid_dtype=grid_dtype,
        do_normalize=do_normalize,
        # goes to one set of grids per row chunk
        adjust_chunks={"row": 1},
        new_axes={
            "nband": np.max(chanmap) + 1,
            # for now will need to be modified if
            # multi-stokes cubes are supported
            "nstokes": 1,
            "y": npix,
            "x": npix
        },
        dtype=grid_dtype,
        meta=np.empty((0, 0, 0, 0, 0, 0), dtype=grid_dtype))

    # Parallel reduction over row dimension
    return grids.mean(axis=0, split_every=2)
