#!/usr/bin/env python3
import numpy as np
import cv2
import argparse

parser = argparse.ArgumentParser()
parser.add_argument("-tl")
parser.add_argument("-tr")
parser.add_argument("-bl")
parser.add_argument("-br")
parser.add_argument("-out")
args = parser.parse_args()

if(args.tl is None or args.tr is None or args.bl is None or args.br is None or args.out is None):
    print("Usage: quad -tl <top-left-clip> -tr <top-right-clip> -bl <bottom-left-clip> -br <bottom-right-clip> -out <out>")
    exit()

# currently uses MP4V codec
outcc = cv2.VideoWriter_fourcc('M', 'P', '4', 'V')

print("Layout:")
print("{}\t\t{}".format(args.tl, args.tr))
print("{}\t\t{}".format(args.bl, args.br))
print("Out: {}".format(args.out))

ctl = cv2.VideoCapture(args.tl)
ctr = cv2.VideoCapture(args.tr)
cbl = cv2.VideoCapture(args.bl)
cbr = cv2.VideoCapture(args.br)

if(ctl is None or ctr is None or cbl is None or cbr is None):
    exit()

fpstl = int(ctl.get(cv2.CAP_PROP_FPS))
fpstr = int(ctr.get(cv2.CAP_PROP_FPS))
fpsbl = int(cbl.get(cv2.CAP_PROP_FPS))
fpsbr = int(cbr.get(cv2.CAP_PROP_FPS))
ntl = ctl.get(cv2.CAP_PROP_FRAME_COUNT)
ntr = ctr.get(cv2.CAP_PROP_FRAME_COUNT)
nbl = cbl.get(cv2.CAP_PROP_FRAME_COUNT)
nbr = cbr.get(cv2.CAP_PROP_FRAME_COUNT)
n = max(ntl, ntr, nbl, nbr)
fw = int(ctl.get(cv2.CAP_PROP_FRAME_WIDTH))
fh = int(ctl.get(cv2.CAP_PROP_FRAME_HEIGHT))
fw2 = int(fw*2)
fh2 = int(fh*2)

if(fpstl != fpstr or fpstr != fpsbl or fpsbl != fpsbr):
    print("All frame rates must match! found {}, {}, {}, {}".format(fpstl, fpstr, fpsbl, fpsbr))
    exit()

if(fw is 0 or fh is 0):
    print("Invalid frame size: {} x {}".format(fw, fh))
    exit()

print("Frame size: {} x {}".format(fw, fh))
out = cv2.VideoWriter(args.out, outcc, fpstl, (fw, fh), True)

k = 0
while(True):
    if(k > n):
        break
    print("Frame {} of {}\r".format(k, n), end='')
    k += 1
    # read a frame from each clip
    ret, ftl = ctl.read()
    ret, ftr = ctr.read()
    ret, fbl = cbl.read()
    ret, fbr = cbr.read()

    # write each frame into a double size buffer
    frame = np.zeros((fh2, fw2, 3), dtype="uint8")

    if ftl is not None:
        for i in range(0, fh):
            frame[i][0:fw] = ftl[i]
    if ftr is not None:
        for i in range(0, fh):
            frame[i+fh][0:fw] = ftr[i]

    if fbl is not None:
        for i in range(0, fh):
            frame[i][fw:fw2] = fbl[i]

    if fbr is not None:
        for i in range(0, fh):
            frame[i+fh][fw:fw2] = fbr[i]

    pframe = cv2.resize(frame, (fw, fh))
    out.write(pframe)

out.release()
print("\nDone")