import numpy as np

'''
A vectorized algorithm to calculate the 2D SF in the physical domain. 
This algorithm is significantly quicker than brute force calculation
without the sampling limitations of the frequency domain approach.

Input: data -2D array

Output: SF - 2D SF
'''

def vectorized_sf(data):

    # Calculate the length of half of each axis
    half_y = np.ma.size(data,0) / 2
    half_x = np.ma.size(data,1) / 2
     
    # Initialize bin array and bin count array
    bin_values = np.zeros((np.ma.size(data,0),np.ma.size(data,1)))
    bin_counts = np.copy(bin_values)
    
    # Loop through each pixel in data                                                                        
    for i in range(len(data)):
        for j in range(len(data[i])):

            # If a pixel is NaN, then skip it
            if not (np.isnan(data[i,j])): 
                
                # Get velocity difference squared between current pixel 
                # and every other pixel
                vdiff = (data - data[i,j])**2
                
                # Find the distances needed to center the current pixel 
                y_roll = half_y - i            
                x_roll = half_x - j
                
                # Ensure that wrapped values are set to NaN. This is done
                # because it is assumed that the input array is not periodic
                if (x_roll > 0 and y_roll > 0):
                    vdiff[(i + half_y):,(j + half_x):] = np.nan
                elif (x_roll > 0 and y_roll < 0):
                    vdiff[0:(i - half_y),(j + half_x):] = np.nan
                elif (x_roll < 0 and y_roll < 0):
                    vdiff[0:(i - half_y),0:(j - half_x)] = np.nan
                elif (x_roll < 0 and y_roll > 0):
                    vdiff[(i + half_y):,0:(j - half_x)] = np.nan   
    
                # Center the current pixel
                vdiff = np.roll(vdiff,(y_roll,x_roll),(0,1))
                
                # Find which bin counts need to be incremented
                counts = np.where(np.isnan(vdiff),0,1)
    
                # Replace NaN values with 0
                np.place(vdiff,np.isnan(vdiff),0)
                
                # Add vdiff values for current pixel to bin values array
                bin_values += vdiff

                # Add bin counts for current pixel to bin count array
                bin_counts += counts

    # Get average SF value for each bin           
    SF = bin_values/bin_counts  

    # Return SF
    return SF