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


[((106, 152), 'numpy.array', 'np.array', (['[[-1, 0, 1], [-1, 0, 1], [-1, 0, 1]]'], {}), '([[-1, 0, 1], [-1, 0, 1], [-1, 0, 1]])\n', (114, 152), True, 'import numpy as np\n'), ((162, 208), 'numpy.array', 'np.array', (['[[-1, -1, -1], [0, 0, 0], [1, 1, 1]]'], {}), '([[-1, -1, -1], [0, 0, 0], [1, 1, 1]])\n', (170, 208), True, 'import numpy as np\n'), ((499, 515), 'numpy.array', 'np.array', (['dx_all'], {}), '(dx_all)\n', (507, 515), True, 'import numpy as np\n'), ((529, 545), 'numpy.array', 'np.array', (['dy_all'], {}), '(dy_all)\n', (537, 545), True, 'import numpy as np\n'), ((843, 865), 'numpy.zeros', 'np.zeros', (['(size, size)'], {}), '((size, size))\n', (851, 865), True, 'import numpy as np\n'), ((1366, 1418), 'cv2.cartToPolar', 'cv2.cartToPolar', (['dx_sum', 'dy_sum'], {'angleInDegrees': '(True)'}), '(dx_sum, dy_sum, angleInDegrees=True)\n', (1381, 1418), False, 'import cv2\n'), ((313, 343), 'scipy.ndimage.convolve', 'ndimage.convolve', (['flow_img', 'mx'], {}), '(flow_img, mx)\n', (329, 343), False, 'from scipy import ndimage\n'), ((358, 388), 'scipy.ndimage.convolve', 'ndimage.convolve', (['flow_img', 'my'], {}), '(flow_img, my)\n', (374, 388), False, 'from scipy import ndimage\n'), ((1162, 1180), 'numpy.mean', 'np.mean', (['tmp_block'], {}), '(tmp_block)\n', (1169, 1180), True, 'import numpy as np\n')]