import numpy as np
import scipy
from scipy import constants
from numba import jit
import numba
import xarray as xr
from astroviper._domain._imaging._check_imaging_parameters import _check_grid_parms
from astroviper._domain._imaging._imaging_utils._mosaic_grid import _aperture_grid_jit
from astroviper._utils._parm_utils._check_parms import _check_sel_parms
import copy

def _make_aperture_grid(ms_xds,gcf_xds,img_xds,vis_sel_parms,img_sel_parms,grid_parms):

    #Deep copy so that inputs are not modified
    _vis_sel_parms = copy.deepcopy(vis_sel_parms)
    _img_sel_parms = copy.deepcopy(img_sel_parms)
    _grid_parms = copy.deepcopy(grid_parms)

    _img_sel_parms["overwrite"] = True
    assert(_check_grid_parms(_grid_parms)), "######### ERROR: grid_parms checking failed"
    
    ms_data_group_in, ms_data_group_out = _check_sel_parms(ms_xds,_vis_sel_parms,skip_data_group_out=True)
    img_data_group_in, img_data_group_out = _check_sel_parms(img_xds,_img_sel_parms,default_data_group_out={"mosaic":{"aperture_grid":"APERTURE_GRID","aperture_grid_sum_weight":"APERTURE_GRID_SUM_WEIGHT"}})
    
    #print(ms_data_group_in, ms_data_group_out)
    
    weight_imaging = ms_xds[ms_data_group_in['weight_imaging']].values
    n_chan = weight_imaging.shape[2]
    
    if _grid_parms['chan_mode'] == 'cube':
        n_imag_chan = n_chan
        chan_map = (np.arange(0, n_chan)).astype(int)
    else:  # continuum
        n_imag_chan = 1  # Making only one continuum image.
        chan_map = (np.zeros(n_chan)).astype(int)
        
    n_imag_pol = weight_imaging.shape[3]
    pol_map = (np.arange(0, n_imag_pol)).astype(int)
    
    n_uv = _grid_parms['image_size_padded']
    delta_lm = _grid_parms['cell_size']
    oversampling = gcf_xds.attrs["oversampling"]
    #_grid_parms['oversampling']
    
    _grid_parms['complex_grid'] = True
    if img_data_group_out["aperture_grid"] not in img_xds:
        img_xds[img_data_group_out["aperture_grid"]] = xr.DataArray(np.zeros((n_imag_chan, n_imag_pol, n_uv[0], n_uv[1]), dtype=np.complex128),dims=["frequency","polarization","u","v"])
        img_xds[img_data_group_out["aperture_grid_sum_weight"]] = xr.DataArray(np.zeros((n_imag_chan, n_imag_pol), dtype=np.double),dims=["frequency","polarization"])
        
    grid = img_xds[img_data_group_out["aperture_grid"]].values
    sum_weight = img_xds[img_data_group_out["aperture_grid_sum_weight"]].values
    
    uvw = ms_xds[ms_data_group_in["uvw"]].values
    freq_chan = ms_xds.frequency.values
    imaging_weight = ms_xds[ms_data_group_in["weight_imaging"]].values
    
    #print(img_xds)
    cf_baseline_map = gcf_xds['CF_BASELINE_MAP'].values
    cf_chan_map = gcf_xds['CF_CHAN_MAP'].values
    cf_pol_map = gcf_xds['CF_POL_MAP'].values
    conv_kernel_convolved = gcf_xds['CONV_KERNEL_CONVOLVED'].values
    weight_support = gcf_xds['SUPPORT'].values
    phase_gradient = gcf_xds['PHASE_GRADIENT'].values
    
#    import matplotlib.pyplot as plt
#    plt.figure()
#    plt.plot(img_xds[img_data_group_out["aperture_grid_sum_weight"]].values)
#    plt.title('1')
 
    _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)
    
#    plt.figure()
#    plt.plot(img_xds[img_data_group_out["sum_weight"]].values)
#    plt.title('2')
#    plt.figure()
#    plt.imshow(np.real(phase_gradient[0,:,:]))
#    plt.colorbar()
    
#    plt.figure()
#    plt.imshow(np.imag(phase_gradient[0,:,:]))
#    plt.colorbar()
#    plt.show()
#
#    print(phase_gradient[0,:,:])
#
#    plt.figure()
#    plt.imshow(np.abs(conv_kernel_convolved[0,0,0,:,:]))
#    plt.colorbar()
#
#
#    a = np.real(np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(img_xds[img_data_group_out["aperture_grid"]].values, axes=(2, 3)), axes=(2, 3)), axes=(2, 3)))
#
#    plt.figure()
#    plt.imshow(np.abs(a[80,0,:,:]))
#    plt.colorbar()
#
#    plt.figure()
#    plt.imshow(np.abs(grid[80,0,:,:]))
#    plt.colorbar()
#
#    s = np.sum(np.abs(imaging_weight),axis=(0,1))
#    plt.figure()
#    plt.plot(s)
#
#    plt.show()
