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
