#Simple 1D Cases
#Airy Disk dish, blockage, freq
#Gaussian halfWidth
#Poly
#Cos Polyp
#Inverse Poly coeff

#Formula for obscured airy pattern found in https://en.wikipedia.org/wiki/Airy_disk (see Obscured Airy pattern section)
# If ipower is 1 the voltage pattern is returned if ipower is 2 the primary beam is returned.
def _airy_disk(freq_chan,pol,pb_parms,grid_parms):
    '''
    Does not yet handle beam squint
    dish_diameters : list of int
    blockage_diameters : list of int
    frequencies : list of number
    '''

    import numpy as np
    import scipy.constants
    from scipy.special import jn
    
    cell = grid_parms['cell_size']
    image_size = grid_parms['image_size']
    image_center = grid_parms['image_center']
    
    list_dish_diameters = pb_parms['list_dish_diameters']
    list_blockage_diameters = pb_parms['list_blockage_diameters']
    ipower = pb_parms['ipower']
    
    c = scipy.constants.c #299792458
    k = (2*np.pi*freq_chan)/c
    
        
    x = np.arange(-image_center[0], image_size[0]-image_center[0])*cell[0]
    y = np.arange(-image_center[1], image_size[1]-image_center[1])*cell[1]
    
    airy_disk_size = (image_size[0],image_size[1],len(freq_chan),1,len(list_blockage_diameters)) #len(pol) is set initially to 1. For now, the PB is assumed the same. This will change.
    airy_disk =  np.zeros(airy_disk_size)
    
    for i, (dish_diameter, blockage_diameter) in enumerate(zip(list_dish_diameters, list_blockage_diameters)):
        
        aperture = dish_diameter/2
        x_grid, y_grid = np.meshgrid(x,y,indexing='ij')
        
        #r_grid = (np.sin(np.sqrt(x_grid**2 + y_grid**2))[:,:,None]*k*aperture) #d0 x d1 x chan
        r_grid = (np.sqrt(x_grid**2 + y_grid**2)[:,:,None]*k*aperture) #d0 x d1 x chan
        r_grid[image_center[0],image_center[1],:] = 1.0 #Avoid the 0/0 for the centre value.
        
        if blockage_diameter==0.0:
            airy_disk[:,:,:,0,i] = (2.0*jn(1,r_grid)/r_grid)**ipower
        else:
            e = blockage_diameter/dish_diameter
            airy_disk[:,:,:,0,i] = (( 2.0 * jn(1,r_grid)/r_grid   - 2.0 * e * jn(1, r_grid * e)/r_grid )/(1.0 - e**2))**ipower
    
    airy_disk[image_center[0],image_center[1],:,0,:] = 1.0 #Fix centre value
    airy_disk = np.tile(airy_disk,(1,1,1,len(pol),1))
    
    return airy_disk
    

#Formula for obscured airy pattern found in casa6/casa5/code/synthesis/TransformMachines/PBMath1DAiry.cc/h
# If ipower is 1 the voltage pattern is returned if ipower is 2 the primary beam is returned.
def _casa_airy_disk(freq_chan,pol,pb_parms,grid_parms):
    '''
    Does not yet handle beam squint
    dish_diameters : list of int
    blockage_diameters : list of int
    frequencies : list of number
    '''

    import numpy as np
    import scipy.constants
    from scipy.special import jn
    
    cell = grid_parms['cell_size']
    image_size = grid_parms['image_size']
    image_center = grid_parms['image_center']
    
    list_dish_diameters = pb_parms['list_dish_diameters']
    list_blockage_diameters = pb_parms['list_blockage_diameters']
    ipower = pb_parms['ipower']
    
    c = scipy.constants.c #299792458
    k = (2*np.pi*freq_chan)/c
    
        
    x = np.arange(-image_center[0], image_size[0]-image_center[0])*cell[0]
    y = np.arange(-image_center[1], image_size[1]-image_center[1])*cell[1]
    
    airy_disk_size = (image_size[0],image_size[1],len(freq_chan),1,len(list_blockage_diameters)) #len(pol) is set initially to 1. For now, the PB is assumed the same. This will change.
    airy_disk =  np.zeros(airy_disk_size)
    
    for i, (dish_diameter, blockage_diameter) in enumerate(zip(list_dish_diameters, list_blockage_diameters)):
        
        aperture = dish_diameter/2
        x_grid, y_grid = np.meshgrid(x,y,indexing='ij')
        
        #r_grid = (np.sin(np.sqrt(x_grid**2 + y_grid**2))[:,:,None]*k*aperture) #d0 x d1 x chan
        r_grid = (np.sqrt(x_grid**2 + y_grid**2)[:,:,None]*k*aperture) #d0 x d1 x chan
        r_grid[image_center[0],image_center[1],:] = 1.0 #Avoid the 0/0 for the centre value.
        
        if blockage_diameter==0.0:
            airy_disk[:,:,:,0,i] = (2.0*jn(1,r_grid)/r_grid)**ipower
        else:
            area_ratio = (dish_diameter/blockage_diameter)**2
            length_ratio = (dish_diameter/blockage_diameter)
            airy_disk[:,:,:,0,i] = ((area_ratio * 2.0 * jn(1,r_grid)/r_grid   - 2.0 * jn(1, r_grid * length_ratio)/(r_grid * length_ratio) )/(area_ratio - 1.0))**ipower
    
    airy_disk[image_center[0],image_center[1],:,0,:] = 1.0 #Fix centre value
    airy_disk = np.tile(airy_disk,(1,1,1,len(pol),1))
    
    return airy_disk
    


#Functions used during the creatiuon of the gridding convolution functions.
#Formula for obscured airy pattern found in https://en.wikipedia.org/wiki/Airy_disk (see Obscured Airy pattern section)
# If ipower is 1 the voltage pattern is returned if ipower is 2 the primary beam is returned.
def _airy_disk_rorder(freq_chan,pol,pb_parms,grid_parms):
    '''
    Does not yet handle beam squint
    dish_diameters : list of int
    blockage_diameters : list of int
    frequencies : list of number
    '''

    import numpy as np
    import scipy.constants
    from scipy.special import jn
    
    cell = grid_parms['cell_size']
    image_size = grid_parms['image_size']
    image_center = grid_parms['image_center']
    
    list_dish_diameters = pb_parms['list_dish_diameters']
    list_blockage_diameters = pb_parms['list_blockage_diameters']
    ipower = pb_parms['ipower']
    
    c = scipy.constants.c #299792458
    k = (2*np.pi*freq_chan)/c
    
    x = np.arange(-image_center[0], image_size[0]-image_center[0])*cell[0]
    y = np.arange(-image_center[1], image_size[1]-image_center[1])*cell[1]
    
    airy_disk_size = (len(list_blockage_diameters),len(freq_chan),1,image_size[0],image_size[1]) #len(pol) is set initially to 1. For now, the PB is assumed the same. This will change.
    airy_disk =  np.zeros(airy_disk_size)
    
    for i, (dish_diameter, blockage_diameter) in enumerate(zip(list_dish_diameters, list_blockage_diameters)):
        
        aperture = dish_diameter/2
        x_grid, y_grid = np.meshgrid(x,y,indexing='ij')
        
        #r_grid = (np.sin(np.sqrt(x_grid**2 + y_grid**2))[:,:,None]*k*aperture) #d0 x d1 x chan
        r_grid = np.moveaxis((np.sqrt(x_grid**2 + y_grid**2)[:,:,None]*k*aperture),2,0) #chan x d0 x d1
        r_grid[:,image_center[0],image_center[1]] = 1.0 #Avoid the 0/0 for the centre value.
        
        if blockage_diameter==0.0:
            airy_disk[i,:,0,:,:] = (2.0*jn(1,r_grid)/r_grid)**ipower
        else:
            e = blockage_diameter/dish_diameter
            airy_disk[i,:,0,:,:] = (( 2.0 * jn(1,r_grid)/r_grid   - 2.0 * e * jn(1, r_grid * e)/r_grid )/(1.0 - e**2))**ipower
    
    airy_disk[:,:,0,image_center[0],image_center[1]] = 1.0 #Fix centre value
    #airy_disk[airy_disk<pb_limit] = 0.0
    airy_disk = np.tile(airy_disk,(1,1,len(pol),1,1))
    
    return airy_disk

#Formula for obscured airy pattern found in casa6/casa5/code/synthesis/TransformMachines/PBMath1DAiry.cc/h
# If ipower is 1 the voltage pattern is returned if ipower is 2 the primary beam is returned.
def _casa_airy_disk_rorder(freq_chan,pol,pb_parms,grid_parms):
    '''
    Does not yet handle beam squint
    dish_diameters : list of int
    blockage_diameters : list of int
    frequencies : list of number
    '''

    import numpy as np
    import scipy.constants
    from scipy.special import jn
    
    cell = grid_parms['cell_size']
    image_size = grid_parms['image_size']
    image_center = grid_parms['image_center']
    
    list_dish_diameters = pb_parms['list_dish_diameters']
    list_blockage_diameters = pb_parms['list_blockage_diameters']
    ipower = pb_parms['ipower']
    
    c = scipy.constants.c #299792458
    k = (2*np.pi*freq_chan)/c
    
    x = np.arange(-image_center[0], image_size[0]-image_center[0])*cell[0]
    y = np.arange(-image_center[1], image_size[1]-image_center[1])*cell[1]
    
    airy_disk_size = (len(list_blockage_diameters),len(freq_chan),1,image_size[0],image_size[1]) #len(pol) is set initially to 1. For now, the PB is assumed the same. This will change.
    airy_disk =  np.zeros(airy_disk_size)
    
    for i, (dish_diameter, blockage_diameter) in enumerate(zip(list_dish_diameters, list_blockage_diameters)):
        
        aperture = dish_diameter/2
        x_grid, y_grid = np.meshgrid(x,y,indexing='ij')
        
        #r_grid = (np.sin(np.sqrt(x_grid**2 + y_grid**2))[:,:,None]*k*aperture) #d0 x d1 x chan
        r_grid = np.moveaxis((np.sqrt(x_grid**2 + y_grid**2)[:,:,None]*k*aperture),2,0) #chan x d0 x d1
        r_grid[:,image_center[0],image_center[1]] = 1.0 #Avoid the 0/0 for the centre value.
        
        if blockage_diameter==0.0:
            airy_disk[i,:,0,:,:] = (2.0*jn(1,r_grid)/r_grid)**ipower
        else:
            area_ratio = (dish_diameter/blockage_diameter)**2
            length_ratio = (dish_diameter/blockage_diameter)
            airy_disk[i,:,0,:,:] = ((area_ratio * 2.0 * jn(1,r_grid)/r_grid   - 2.0 * jn(1, r_grid * length_ratio)/(r_grid * length_ratio) )/(area_ratio - 1.0))**ipower
    
    airy_disk[:,:,0,image_center[0],image_center[1]] = 1.0 #Fix centre value
    airy_disk = np.tile(airy_disk,(1,1,len(pol),1,1))
    
    return airy_disk
