import argparse
import cv2 as cv
from itertools import count
import numpy as np

def duration(cap):
    """
    return the duration of the capture in seconds
    """
    nframes = int(cap.get(cv.CAP_PROP_FRAME_COUNT))
    fps = cap.get(cv.CAP_PROP_FPS)

    return nframes / fps


def capture_shape(cap):
    """
    return the shape of the video capture
    """
    nframes = int(cap.get(cv.CAP_PROP_FRAME_COUNT))
    height = int(cap.get(cv.CAP_PROP_FRAME_HEIGHT))
    width = int(cap.get(cv.CAP_PROP_FRAME_WIDTH))

    return (nframes, height, width)


def transposed_shape(shape, axis):
    """
    return the shape of the video capture after transposing the given axis to
    become the new time axis
    """
    nframes, height, width = shape
    if axis == 1:
        new_shape = (height, nframes, width)
    elif axis == 2:
        new_shape = (width, height, nframes)

    return new_shape


def transpose_frame(infile, n, axis):
    """
    compute the nth frame of the input file after transposing the given axis to
    become the new time axis
    """
    cap = cv.VideoCapture(infile)

    shape = capture_shape(cap)
    _, height, width = transposed_shape(shape, axis)
    tframe = np.empty((height, width, 3), np.dtype('uint8'))

    for i in count(0):
        # NOTE: reading each frame of the input to produce one output frame is
        # CPU-intensive, but to read the frames fewer times, we would have to
        # store the frames in memory, which is memory-intensive
        ret, frame = cap.read()
        if not ret:
            break

        if axis == 1:
            tframe[i] = frame[n]
        elif axis == 2:
            tframe[:, i] = frame[:, n]

    cap.release()
    return tframe


def transpose_capture(infile, outfile, axis):
    """
    transpose the input file along the given axis, storing the output in
    outfile
    """
    cap = cv.VideoCapture(infile)

    seconds = duration(cap)
    shape = capture_shape(cap)
    nframes, height, width = transposed_shape(shape, axis)
    fps = nframes / seconds

    cap.release()

    fourcc = cv.VideoWriter.fourcc(*'h264')
    out = cv.VideoWriter(outfile, fourcc, fps, (width, height))

    for i in range(0, nframes):
        tframe = transpose_frame(infile, i, axis)
        out.write(tframe)

    out.release()
    cv.destroyAllWindows()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("source")
    parser.add_argument("dest")

    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--height", action='store_true')
    group.add_argument("--width", action='store_true')

    args = parser.parse_args()

    if args.height:
        axis = 1
    elif args.width:
        axis = 2

    transpose_capture(args.source, args.dest, axis)