import os
import cv2
import math
import time
import numpy as np

AVERAGING_LENGTH = 7
FPS = 60.0
OFFSET_TO_ANALYZE = 0
FRACTION_TO_ANALYZE = 1

NUM_SINGLE = 3022
NUM_DOUBLE = 2723
NUM_QUADRO = 3983
single_images_0_path = sorted([img for img in os.listdir('images_single')])
double_images_0_path = sorted([img for img in os.listdir('images_double/0')])
double_images_1_path = sorted([img for img in os.listdir('images_double/1')])
quadro_images_0_path = sorted([img for img in os.listdir('images_quadro/0')])
quadro_images_1_path = sorted([img for img in os.listdir('images_quadro/1')])
quadro_images_2_path = sorted([img for img in os.listdir('images_quadro/2')])
quadro_images_3_path = sorted([img for img in os.listdir('images_quadro/3')])

def map_by_layer(image, f, mask):
    mask_3d = np.zeros((mask.shape[0], mask.shape[1], 3))
    mask_3d[:, :, 0] = mask
    mask_3d[:, :, 1] = mask
    mask_3d[:, :, 2] = mask
    return np.where(mask_3d, f(image), image)

def image_difference_mask(prev, curr):
    difference = cv2.subtract(prev, curr)
    b, g, r = cv2.split(difference)
    (_, mask) = cv2.threshold(cv2.bitwise_or(b, cv2.bitwise_or(g, r)), 1, 255, cv2.THRESH_BINARY)
    return mask

def images_difference(images):
    mask = np.zeros((images[0].shape[0], images[0].shape[1]), dtype=np.uint8)
    prev = images[0]
    for i in range(1, len(images)):
        mask = cv2.bitwise_or(mask, image_difference_mask(prev, images[i]))
        prev = images[i]

    return map_by_layer(images[-1], lambda x: x // 3, np.logical_not(mask))

def image_difference(prev, curr):
    mask = image_difference_mask(prev, curr)
    return cv2.bitwise_and(curr, curr, mask=mask)

def load_image(i):
    if i < NUM_SINGLE:
        image = cv2.imread(os.path.join('images_single', single_images_0_path[i]))
        image = np.concatenate((image, 255*np.ones(image.shape, dtype=np.uint8)), axis=0)
        image = np.concatenate((image, 255*np.ones(image.shape, dtype=np.uint8)), axis=1)
    elif i < NUM_SINGLE + NUM_DOUBLE:
        image_0 = cv2.imread(os.path.join('images_double/0', double_images_0_path[i - NUM_SINGLE]))
        image_1 = cv2.imread(os.path.join('images_double/1', double_images_1_path[i - NUM_SINGLE]))
        image = np.concatenate((image_0, image_1), axis=1)
        image = np.concatenate((image, 255*np.ones(image.shape, dtype=np.uint8)), axis=0)
    elif i < NUM_SINGLE + NUM_DOUBLE + NUM_QUADRO:
        image_0 = cv2.imread(os.path.join('images_quadro/0', quadro_images_0_path[i - NUM_SINGLE - NUM_DOUBLE]))
        image_1 = cv2.imread(os.path.join('images_quadro/1', quadro_images_1_path[i - NUM_SINGLE - NUM_DOUBLE]))
        image_2 = cv2.imread(os.path.join('images_quadro/2', quadro_images_2_path[i - NUM_SINGLE - NUM_DOUBLE]))
        image_3 = cv2.imread(os.path.join('images_quadro/3', quadro_images_3_path[i - NUM_SINGLE - NUM_DOUBLE]))
        image_top = np.concatenate((image_0, image_1), axis=1)
        image_bottom = np.concatenate((image_2, image_3), axis=1)
        image = np.concatenate((image_top, image_bottom), axis=0)
    else:
        print("Unsupported")

    return image


start = time.time()

video_name = 'video1.mp4'
frame = load_image(0)
height, width, layers = frame.shape
fourcc = cv2.VideoWriter_fourcc(*'MP4V')
video = cv2.VideoWriter(video_name, fourcc, FPS, (width,height))

next_images = [None] + [load_image(i) for i in range(0, AVERAGING_LENGTH-1)]
max_range = OFFSET_TO_ANALYZE + math.floor((NUM_SINGLE + NUM_DOUBLE + NUM_QUADRO)*FRACTION_TO_ANALYZE) - AVERAGING_LENGTH
for i in range(OFFSET_TO_ANALYZE, max_range):
    if i % 100 == 0:
        print(i)
    next_images.append(load_image(i+AVERAGING_LENGTH-1))
    del next_images[0]
    next_frame = images_difference(next_images)
    video.write(next_frame)

cv2.destroyAllWindows()
video.release()
end = time.time()
print(end - start)