import cv2
import numpy as np
from scipy import ndimage

def compute_motion_boudary(flow_clip):

    mx = np.array([[-1, 0, 1], [-1, 0, 1], [-1, 0, 1]])
    my = np.array([[-1, -1, -1], [0, 0, 0], [1, 1, 1]])
    dx_all = []
    dy_all = []
    mb_x = 0
    mb_y = 0

    for flow_img in flow_clip:
        d_x = ndimage.convolve(flow_img, mx)
        d_y = ndimage.convolve(flow_img, my)

        dx_all.append(d_x)
        dy_all.append(d_y)

        mb_x += d_x
        mb_y += d_y

    dx_all = np.array(dx_all)
    dy_all = np.array(dy_all)

    return dx_all, dy_all, mb_x, mb_y

def zero_boundary(frame_mag):

    frame_mag[:8, :] = 0
    frame_mag[:, :8] = 0
    frame_mag[-8:, :] = 0
    frame_mag[:, -8:] = 0
    
    return frame_mag

def motion_mag_downsample(mag, size, input_size):
    block_size = input_size // size
    mask = np.zeros((size,size))
    for i in range(size):
        for j in range(size):
            x_start = i * block_size
            x_end = x_start + block_size
            y_start = j * block_size
            y_end = y_start + block_size

            tmp_block = mag[x_start:x_end, y_start:y_end]

            block_mean = np.mean(tmp_block)
            mask[i, j] = block_mean
    return mask

def motion_sts(flow_clip, size, input_size):

    dx_all, dy_all, dx_sum, dy_sum = compute_motion_boudary(flow_clip)
    mag, ang = cv2.cartToPolar(dx_sum, dy_sum, angleInDegrees=True)
    mag_down = motion_mag_downsample(mag, size, input_size)

    return mag_down


[((7, 9, 7, 55), 'numpy.array', 'np.array', ({(7, 18, 7, 54): '[[-1, 0, 1], [-1, 0, 1], [-1, 0, 1]]'}, {}), '([[-1, 0, 1], [-1, 0, 1], [-1, 0, 1]])', True, 'import numpy as np\n'), ((8, 9, 8, 55), 'numpy.array', 'np.array', ({(8, 18, 8, 54): '[[-1, -1, -1], [0, 0, 0], [1, 1, 1]]'}, {}), '([[-1, -1, -1], [0, 0, 0], [1, 1, 1]])', True, 'import numpy as np\n'), ((24, 13, 24, 29), 'numpy.array', 'np.array', ({(24, 22, 24, 28): 'dx_all'}, {}), '(dx_all)', True, 'import numpy as np\n'), ((25, 13, 25, 29), 'numpy.array', 'np.array', ({(25, 22, 25, 28): 'dy_all'}, {}), '(dy_all)', True, 'import numpy as np\n'), ((40, 11, 40, 32), 'numpy.zeros', 'np.zeros', ({(40, 20, 40, 31): '(size, size)'}, {}), '((size, size))', True, 'import numpy as np\n'), ((57, 15, 57, 67), 'cv2.cartToPolar', 'cv2.cartToPolar', (), '', False, 'import cv2\n'), ((15, 14, 15, 44), 'scipy.ndimage.convolve', 'ndimage.convolve', ({(15, 31, 15, 39): 'flow_img', (15, 41, 15, 43): 'mx'}, {}), '(flow_img, mx)', False, 'from scipy import ndimage\n'), ((16, 14, 16, 44), 'scipy.ndimage.convolve', 'ndimage.convolve', ({(16, 31, 16, 39): 'flow_img', (16, 41, 16, 43): 'my'}, {}), '(flow_img, my)', False, 'from scipy import ndimage\n'), ((50, 25, 50, 43), 'numpy.mean', 'np.mean', ({(50, 33, 50, 42): 'tmp_block'}, {}), '(tmp_block)', True, 'import numpy as np\n')]