import numpy as np
from numba import literally, prange
from africanus.util.numba import jit

from africanus.gridding.perleypolyhedron.policies import (
    baseline_transform_policies as btp)
from africanus.gridding.perleypolyhedron.policies import (
    phase_transform_policies as ptp)
from africanus.gridding.perleypolyhedron.policies import (
    convolution_policies as cp)
from africanus.gridding.perleypolyhedron.policies import (
    stokes_conversion_policies as scp)


@jit(nopython=True, nogil=True, fastmath=True, inline="always")
def degridder_row_kernel(uvw,
                         gridstack,
                         wavelengths,
                         chanmap,
                         cell,
                         image_centre,
                         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,
                         nband=0,
                         nrow=0,
                         npix=0,
                         nvischan=0,
                         ncorr=0,
                         vis=None,
                         scale_factor=0,
                         r=0):
    ra0, dec0 = phase_centre
    ra, dec = image_centre
    btp.policy(uvw[r, :], ra, dec, ra0, dec0,
               baseline_transform_policy)
    for c in range(nvischan):
        scaled_u = uvw[r, 0] * scale_factor / wavelengths[c]
        scaled_v = uvw[r, 1] * scale_factor / wavelengths[c]
        scaled_w = uvw[r, 2] * scale_factor / wavelengths[c]
        grid = gridstack[chanmap[c], :, :]
        cp.policy(scaled_u,
                  scaled_v,
                  scaled_w,
                  npix,
                  grid,
                  vis,
                  r,
                  c,
                  convolution_kernel,
                  convolution_kernel_width,
                  convolution_kernel_oversampling,
                  stokes_conversion_policy,
                  policy_type=convolution_policy)
    ptp.policy(vis[r, :, :],
               uvw[r, :],
               wavelengths,
               ra0,
               dec0,
               ra,
               dec,
               policy_type=phase_transform_policy,
               phasesign=-1.0)


@jit(nopython=True, nogil=True, fastmath=True, parallel=True)
def degridder(uvw,
              gridstack,
              wavelengths,
              chanmap,
              cell,
              image_centre,
              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):
    """
    2D Convolutional degridder, discrete to contiguous
    @uvw: value coordinates, (nrow, 3)
    @gridstack: complex gridded data, (nband, npix, npix)
    @wavelengths: wavelengths of data channels
    @chanmap: MFS band mapping per channel
    @cell: cell_size in degrees
    @image_centres: 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)
    """

    if chanmap.size != wavelengths.size:
        raise ValueError(
            "Chanmap and corresponding wavelengths must match in shape")
    chanmap = chanmap.ravel()
    wavelengths = wavelengths.ravel()
    nband = np.max(chanmap) + 1
    nrow = uvw.shape[0]
    npix = gridstack.shape[1]
    if gridstack.shape[1] != gridstack.shape[2]:
        raise ValueError("Grid must be square")
    nvischan = wavelengths.size
    ncorr = scp.ncorr_out(policy_type=literally(stokes_conversion_policy))
    if gridstack.shape[0] < nband:
        raise ValueError(
            "Not enough channel bands in grid stack to match mfs band mapping")
    if uvw.shape[1] != 3:
        raise ValueError("UVW array must be array of tripples")
    if uvw.shape[0] != nrow:
        raise ValueError(
            "UVW array must have same number of rows as vis array")
    if nvischan != wavelengths.size:
        raise ValueError("Chanmap must correspond to visibility channels")

    vis = np.zeros((nrow, nvischan, ncorr), dtype=vis_dtype)

    # scale the FOV using the simularity theorem
    scale_factor = npix * cell / 3600.0 * np.pi / 180.0
    for r in prange(nrow):
        degridder_row_kernel(uvw,
                             gridstack,
                             wavelengths,
                             chanmap,
                             cell,
                             image_centre,
                             phase_centre,
                             convolution_kernel,
                             convolution_kernel_width,
                             convolution_kernel_oversampling,
                             literally(baseline_transform_policy),
                             literally(phase_transform_policy),
                             literally(stokes_conversion_policy),
                             literally(convolution_policy),
                             vis_dtype=vis_dtype,
                             nband=nband,
                             nrow=nrow,
                             npix=npix,
                             nvischan=nvischan,
                             ncorr=ncorr,
                             vis=vis,
                             scale_factor=scale_factor,
                             r=r)
    return vis


@jit(nopython=True, nogil=True, fastmath=True, parallel=False)
def degridder_serial(uvw,
                     gridstack,
                     wavelengths,
                     chanmap,
                     cell,
                     image_centre,
                     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):
    """
    2D Convolutional degridder, discrete to contiguous
    @uvw: value coordinates, (nrow, 3)
    @gridstack: complex gridded data, (nband, npix, npix)
    @wavelengths: wavelengths of data channels
    @chanmap: MFS band mapping per channel
    @cell: cell_size in degrees
    @image_centres: 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)
    """

    if chanmap.size != wavelengths.size:
        raise ValueError(
            "Chanmap and corresponding wavelengths must match in shape")
    chanmap = chanmap.ravel()
    wavelengths = wavelengths.ravel()
    nband = np.max(chanmap) + 1
    nrow = uvw.shape[0]
    npix = gridstack.shape[1]
    if gridstack.shape[1] != gridstack.shape[2]:
        raise ValueError("Grid must be square")
    nvischan = wavelengths.size
    ncorr = scp.ncorr_out(policy_type=literally(stokes_conversion_policy))
    if gridstack.shape[0] < nband:
        raise ValueError(
            "Not enough channel bands in grid stack to match mfs band mapping")
    if uvw.shape[1] != 3:
        raise ValueError("UVW array must be array of tripples")
    if uvw.shape[0] != nrow:
        raise ValueError(
            "UVW array must have same number of rows as vis array")
    if nvischan != wavelengths.size:
        raise ValueError("Chanmap must correspond to visibility channels")

    vis = np.zeros((nrow, nvischan, ncorr), dtype=vis_dtype)

    # scale the FOV using the simularity theorem
    scale_factor = npix * cell / 3600.0 * np.pi / 180.0
    for r in range(nrow):
        degridder_row_kernel(uvw,
                             gridstack,
                             wavelengths,
                             chanmap,
                             cell,
                             image_centre,
                             phase_centre,
                             convolution_kernel,
                             convolution_kernel_width,
                             convolution_kernel_oversampling,
                             literally(baseline_transform_policy),
                             literally(phase_transform_policy),
                             literally(stokes_conversion_policy),
                             literally(convolution_policy),
                             vis_dtype=vis_dtype,
                             nband=nband,
                             nrow=nrow,
                             npix=npix,
                             nvischan=nvischan,
                             ncorr=ncorr,
                             vis=vis,
                             scale_factor=scale_factor,
                             r=r)
    return vis
