from numba import jit
import numpy as np
import math

# from numba import gdb


def ndim_list(shape):
    return [ndim_list(shape[1:]) if len(shape) > 1 else None for _ in range(shape[0])]


@jit(nopython=True, cache=True, nogil=True)
def _aperture_grid_jit(
    grid,
    sum_weight,
    uvw,
    freq_chan,
    chan_map,
    pol_map,
    cf_baseline_map,
    cf_chan_map,
    cf_pol_map,
    imaging_weight,
    conv_kernel_convolved,
    n_uv,
    delta_lm,
    weight_support,
    oversampling,
    phase_gradient,
):
    c = 299792458.0
    uv_scale = np.zeros((2, len(freq_chan)), dtype=np.double)
    uv_scale[0, :] = -(freq_chan * delta_lm[0] * n_uv[0]) / c
    uv_scale[1, :] = -(freq_chan * delta_lm[1] * n_uv[1]) / c

    uv_center = n_uv // 2

    n_time = uvw.shape[0]
    n_baseline = uvw.shape[1]

    n_chan = len(chan_map)
    n_pol = len(pol_map)

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

    u_center = uv_center[0]
    v_center = uv_center[1]

    max_support_center = np.max(weight_support)

    conv_v_center = conv_kernel_convolved.shape[-1] // 2
    conv_u_center = conv_kernel_convolved.shape[-2] // 2

    # print(phase_gradient.shape)
    # print(conv_kernel_convolved.shape)

    conv_kernel_convolved_phase_gradient = (
        conv_kernel_convolved * phase_gradient[0, :, :]
    )

    for i_time in range(n_time):
        for i_baseline in range(n_baseline):
            cf_baseline = cf_baseline_map[i_baseline]
            for i_chan in range(n_chan):
                cf_chan = cf_chan_map[i_chan]
                a_chan = chan_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 + max_support_center < n_u)
                        and (v_center_indx + max_support_center < n_v)
                        and (u_center_indx - max_support_center >= 0)
                        and (v_center_indx - max_support_center >= 0)
                    ):
                        # u_offset = u_center_indx - u_pos
                        # u_center_offset_indx = math.floor(u_offset * oversampling[0] + 0.5)
                        # v_offset = v_center_indx - v_pos
                        # v_center_offset_indx = math.floor(v_offset * oversampling[1] + 0.5)

                        for i_pol in range(n_pol):
                            weighted_data = imaging_weight[
                                i_time, i_baseline, i_chan, i_pol
                            ]

                            if ~np.isnan(weighted_data) and (weighted_data != 0.0):
                                cf_pol = cf_pol_map[i_pol]
                                a_pol = pol_map[i_pol]
                                norm = 0.0

                                """
                                support = weight_support[cf_baseline,cf_chan,cf_pol,:]
                                #support = np.array([13,13])
                                support_center = support // 2
                                start_support = - support_center
                                end_support = support - support_center # end_support is larger by 1 so that python range() gives correct indices
                                """

                                support_u = weight_support[
                                    cf_baseline, cf_chan, cf_pol, 0
                                ]
                                support_v = weight_support[
                                    cf_baseline, cf_chan, cf_pol, 1
                                ]

                                support_center_u = support_u // 2
                                support_center_v = support_v // 2

                                start_support_u = -support_center_u
                                start_support_v = -support_center_v

                                end_support_u = support_u - support_center_u
                                end_support_v = support_v - support_center_v

                                for i_v in range(start_support_v, end_support_v):
                                    v_indx = v_center + i_v
                                    cf_v_indx = oversampling[1] * i_v + conv_v_center

                                    for i_u in range(start_support_u, end_support_u):
                                        u_indx = u_center + i_u
                                        cf_u_indx = (
                                            oversampling[0] * i_u + conv_u_center
                                        )

                                        conv = conv_kernel_convolved_phase_gradient[
                                            cf_baseline,
                                            cf_chan,
                                            cf_pol,
                                            cf_u_indx,
                                            cf_v_indx,
                                        ]

                                        grid[a_chan, a_pol, u_indx, v_indx] = (
                                            grid[a_chan, a_pol, u_indx, v_indx]
                                            + conv * weighted_data
                                        )
                                        norm = norm + conv

                                # print("The norm",norm)
                                sum_weight[a_chan, a_pol] = sum_weight[
                                    a_chan, a_pol
                                ] + weighted_data * np.real(norm)

    return


def _mosaic_grid_numpy_wrap(
    vis_data,
    uvw,
    imaging_weight,
    field,
    cf_baseline_map,
    cf_chan_map,
    cf_pol_map,
    conv_kernel,
    weight_support,
    phase_gradient,
    freq_chan,
    grid_parms,
):
    # print('imaging_weight ', imaging_weight.shape)
    import time

    n_chan = imaging_weight.shape[2]
    if grid_parms["chan_mode"] == "cube":
        n_imag_chan = n_chan
        chan_map = (np.arange(0, n_chan)).astype(np.int)
    else:  # continuum
        n_imag_chan = 1  # Making only one continuum image.
        chan_map = (np.zeros(n_chan)).astype(np.int)

    n_imag_pol = imaging_weight.shape[3]
    pol_map = (np.arange(0, n_imag_pol)).astype(np.int)

    n_uv = grid_parms["image_size_padded"]
    delta_lm = grid_parms["cell_size"]
    oversampling = grid_parms["oversampling"]

    if grid_parms["complex_grid"]:
        grid = np.zeros(
            (n_imag_chan, n_imag_pol, n_uv[0], n_uv[1]), dtype=np.complex128
        )
    else:
        grid = np.zeros((n_imag_chan, n_imag_pol, n_uv[0], n_uv[1]), dtype=np.double)
    sum_weight = np.zeros((n_imag_chan, n_imag_pol), dtype=np.double)

    do_psf = grid_parms["do_psf"]
    field_id = grid_parms["field_id"]

    # print('vis_data', vis_data.shape , 'grid ', grid.shape, 'sum_weight', sum_weight.shape, 'cf_chan_map ', cf_chan_map.shape, ' cf_baseline_map', cf_baseline_map.shape, 'cf_pol_map', cf_pol_map.shape, ' conv_kernel',  conv_kernel.shape, 'phase_gradient', phase_gradient.shape, 'field', field.shape,  )

    # start = time.time()
    _aperture_grid_jit(
        grid,
        sum_weight,
        do_psf,
        vis_data,
        uvw,
        freq_chan,
        chan_map,
        pol_map,
        cf_baseline_map,
        cf_chan_map,
        cf_pol_map,
        imaging_weight,
        conv_kernel,
        n_uv,
        delta_lm,
        weight_support,
        oversampling,
        field,
        field_id,
        phase_gradient,
    )
    # time_to_grid = time.time() - start
    # print("time to grid ", time_to_grid)

    return grid, sum_weight


def _mosaic_psf_grid_numpy_wrap(
    uvw,
    imaging_weight,
    field,
    cf_baseline_map,
    cf_chan_map,
    cf_pol_map,
    conv_kernel,
    weight_support,
    phase_gradient,
    freq_chan,
    grid_parms,
):
    # print('imaging_weight ', imaging_weight.shape)
    import time

    n_chan = imaging_weight.shape[2]
    if grid_parms["chan_mode"] == "cube":
        n_imag_chan = n_chan
        chan_map = (np.arange(0, n_chan)).astype(np.int)
    else:  # continuum
        n_imag_chan = 1  # Making only one continuum image.
        chan_map = (np.zeros(n_chan)).astype(np.int)

    n_imag_pol = imaging_weight.shape[3]
    pol_map = (np.arange(0, n_imag_pol)).astype(np.int)

    n_uv = grid_parms["image_size_padded"]
    delta_lm = grid_parms["cell_size"]
    oversampling = grid_parms["oversampling"]

    if grid_parms["complex_grid"]:
        grid = np.zeros(
            (n_imag_chan, n_imag_pol, n_uv[0], n_uv[1]), dtype=np.complex128
        )
    else:
        grid = np.zeros((n_imag_chan, n_imag_pol, n_uv[0], n_uv[1]), dtype=np.double)
    sum_weight = np.zeros((n_imag_chan, n_imag_pol), dtype=np.double)

    do_psf = grid_parms["do_psf"]
    field_id = grid_parms["field_id"]

    # print('vis_data', vis_data.shape , 'grid ', grid.shape, 'sum_weight', sum_weight.shape, 'cf_chan_map ', cf_chan_map.shape, ' cf_baseline_map', cf_baseline_map.shape, 'cf_pol_map', cf_pol_map.shape, ' conv_kernel',  conv_kernel.shape, 'phase_gradient', phase_gradient.shape, 'field', field.shape,  )

    vis_data = np.zeros((1, 1, 1, 1), dtype=np.bool)
    # start = time.time()
    _mosaic_grid_jit(
        grid,
        sum_weight,
        do_psf,
        vis_data,
        uvw,
        freq_chan,
        chan_map,
        pol_map,
        cf_baseline_map,
        cf_chan_map,
        cf_pol_map,
        imaging_weight,
        conv_kernel,
        n_uv,
        delta_lm,
        weight_support,
        oversampling,
        field,
        field_id,
        phase_gradient,
    )
    # time_to_grid = time.time() - start
    # print("time to grid ", time_to_grid)

    return grid, sum_weight


# Important changes to be made https://github.com/numba/numba/issues/4261
# debug=True and gdb()
@jit(nopython=True, cache=True, nogil=True)
def _mosaic_grid_jit(
    grid,
    sum_weight,
    do_psf,
    vis_data,
    uvw,
    freq_chan,
    chan_map,
    pol_map,
    cf_baseline_map,
    cf_chan_map,
    cf_pol_map,
    imaging_weight,
    conv_kernel,
    n_uv,
    delta_lm,
    weight_support,
    oversampling,
    phase_gradient,
):
    c = 299792458.0
    uv_scale = np.zeros((2, len(freq_chan)), dtype=np.double)
    uv_scale[0, :] = -(freq_chan * delta_lm[0] * n_uv[0]) / c
    uv_scale[1, :] = -(freq_chan * delta_lm[1] * n_uv[1]) / c

    uv_center = n_uv // 2

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

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

    u_center = uv_center[0]
    v_center = uv_center[1]

    max_support_center = np.max(weight_support)

    conv_v_center = conv_kernel.shape[-1] // 2
    conv_u_center = conv_kernel.shape[-2] // 2

    conv_size = np.array(conv_kernel.shape[-2:])

    # print('conv_size',conv_size)

    # print('sizes ',conv_kernel.shape, conv_u_center, conv_v_center)

    # print(phase_gradient.shape)
    # print(conv_kernel_convolved.shape)

    conv_kernel_phase_gradient = conv_kernel * phase_gradient[0, :, :]

    for i_time in range(n_time):
        for i_baseline in range(n_baseline):
            cf_baseline = cf_baseline_map[i_baseline]
            for i_chan in range(n_chan):
                cf_chan = cf_chan_map[i_chan]
                a_chan = chan_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 + max_support_center < n_u)
                        and (v_center_indx + max_support_center < n_v)
                        and (u_center_indx - max_support_center >= 0)
                        and (v_center_indx - max_support_center >= 0)
                    ):
                        u_offset = u_center_indx - u_pos
                        u_center_offset_indx = (
                            math.floor(u_offset * oversampling[0] + 0.5) + conv_u_center
                        )
                        v_offset = v_center_indx - v_pos
                        v_center_offset_indx = (
                            math.floor(v_offset * oversampling[1] + 0.5) + conv_v_center
                        )

                        for i_pol in range(n_pol):
                            if do_psf:
                                weighted_data = imaging_weight[
                                    i_time, i_baseline, i_chan, i_pol
                                ]
                            else:
                                weighted_data = (
                                    vis_data[i_time, i_baseline, i_chan, i_pol]
                                    * imaging_weight[i_time, i_baseline, i_chan, i_pol]
                                )

                            if ~np.isnan(weighted_data) and (weighted_data != 0.0):
                                cf_pol = cf_pol_map[i_pol]
                                a_pol = pol_map[i_pol]
                                norm = 0.0

                                """
                                support = weight_support[cf_baseline,cf_chan,cf_pol,:]
                                #support = np.array([13,13])
                                support_center = support // 2
                                start_support = - support_center
                                end_support = support - support_center # end_support is larger by 1 so that python range() gives correct indices
                                """

                                support_u = weight_support[
                                    cf_baseline, cf_chan, cf_pol, 0
                                ]
                                support_v = weight_support[
                                    cf_baseline, cf_chan, cf_pol, 1
                                ]

                                support_center_u = support_u // 2
                                support_center_v = support_v // 2

                                start_support_u = -support_center_u
                                start_support_v = -support_center_v

                                end_support_u = support_u - support_center_u
                                end_support_v = support_v - support_center_v

                                # print(support)
                                ###############
                                #                                resized_conv_size = (support  + 1)*oversampling
                                #                                start_indx = conv_size//2 - resized_conv_size//2
                                #                                end_indx = start_indx + resized_conv_size
                                #                                normalize_factor = np.real(np.sum(conv_kernel[cf_baseline,cf_chan,cf_pol,start_indx[0]:end_indx[0],start_indx[1]:end_indx[1]])/(oversampling[0]*oversampling[1]))
                                #
                                #                                conv_kernel_phase_gradient = conv_kernel*phase_gradient[field[i_time],:,:]/normalize_factor
                                #                                print(normalize_factor)
                                ##############

                                for i_v in range(start_support_v, end_support_v):
                                    v_indx = v_center_indx + i_v
                                    cf_v_indx = (
                                        oversampling[1] * i_v + v_center_offset_indx
                                    )

                                    for i_u in range(start_support_u, end_support_u):
                                        u_indx = u_center_indx + i_u
                                        cf_u_indx = (
                                            oversampling[0] * i_u + u_center_offset_indx
                                        )

                                        conv = conv_kernel_phase_gradient[
                                            cf_baseline,
                                            cf_chan,
                                            cf_pol,
                                            cf_u_indx,
                                            cf_v_indx,
                                        ]

                                        grid[a_chan, a_pol, u_indx, v_indx] = (
                                            grid[a_chan, a_pol, u_indx, v_indx]
                                            + conv * weighted_data
                                        )
                                        norm = norm + conv
                                if do_psf:
                                    sum_weight[a_chan, a_pol] = sum_weight[
                                        a_chan, a_pol
                                    ] + imaging_weight[
                                        i_time, i_baseline, i_chan, i_pol
                                    ] * np.real(
                                        norm
                                    )
                                else:
                                    sum_weight[a_chan, a_pol] = sum_weight[
                                        a_chan, a_pol
                                    ] + imaging_weight[
                                        i_time, i_baseline, i_chan, i_pol
                                    ] * np.real(
                                        norm**2
                                    )  # *np.real(norm**2)#* np.real(norm) #np.abs(norm**2) #**2 term is needed since the pb is in the image twice (one naturally and another from the gcf)

    return


def _tree_sum_list(list_to_sum):
    import dask.array as da

    while len(list_to_sum) > 1:
        new_list_to_sum = []
        for i in range(0, len(list_to_sum), 2):
            if i < len(list_to_sum) - 1:
                lazy = da.add(list_to_sum[i], list_to_sum[i + 1])
            else:
                lazy = list_to_sum[i]
            new_list_to_sum.append(lazy)
        list_to_sum = new_list_to_sum
    return list_to_sum
