##  Copyright (c) 2020 The WebM project authors. All Rights Reserved.
##
##  Use of this source code is governed by a BSD-style license
##  that can be found in the LICENSE file in the root of the source
##  tree. An additional intellectual property rights grant can be found
##  in the file PATENTS.  All contributing project authors may
##  be found in the AUTHORS file in the root of the source tree.
##

import argparse
from os import listdir, path
from PIL import Image
import sys

parser = argparse.ArgumentParser()
parser.add_argument("--frame_path", default="../data/frame/", type=str)
parser.add_argument("--frame_rate", default="25:1", type=str)
parser.add_argument("--interlacing", default="Ip", type=str)
parser.add_argument("--pix_ratio", default="0:0", type=str)
parser.add_argument("--color_space", default="4:2:0", type=str)
parser.add_argument("--output", default="output.y4m", type=str)


def generate(args, frames):
  if len(frames) == 0:
    return
  #sort the frames based on the frame index
  frames = sorted(frames, key=lambda x: x[0])
  #convert the frames to YUV form
  frames = [f.convert("YCbCr") for _, f in frames]
  #write the header
  header = "YUV4MPEG2 W%d H%d F%s %s A%s" % (frames[0].width, frames[0].height,
                                             args.frame_rate, args.interlacing,
                                             args.pix_ratio)
  cs = args.color_space.split(":")
  header += " C%s%s%s\n" % (cs[0], cs[1], cs[2])
  #estimate the sample step based on subsample value
  subsamples = [int(c) for c in cs]
  r_step = [1, int(subsamples[2] == 0) + 1, int(subsamples[2] == 0) + 1]
  c_step = [1, 4 // subsamples[1], 4 // subsamples[1]]
  #write in frames
  with open(args.output, "wb") as y4m:
    y4m.write(header)
    for f in frames:
      y4m.write("FRAME\n")
      px = f.load()
      for k in xrange(3):
        for i in xrange(0, f.height, r_step[k]):
          for j in xrange(0, f.width, c_step[k]):
            yuv = px[j, i]
            y4m.write(chr(yuv[k]))


if __name__ == "__main__":
  args = parser.parse_args()
  frames = []
  frames_mv = []
  for filename in listdir(args.frame_path):
    name, ext = filename.split(".")
    if ext == "png":
      name_parse = name.split("_")
      idx = int(name_parse[-1])
      img = Image.open(path.join(args.frame_path, filename))
      if name_parse[-2] == "mv":
        frames_mv.append((idx, img))
      else:
        frames.append((idx, img))
  if len(frames) == 0:
    print("No frames in directory: " + args.frame_path)
    sys.exit()
  print("----------------------Y4M Info----------------------")
  print("width:  %d" % frames[0][1].width)
  print("height: %d" % frames[0][1].height)
  print("#frame: %d" % len(frames))
  print("frame rate: %s" % args.frame_rate)
  print("interlacing: %s" % args.interlacing)
  print("pixel ratio: %s" % args.pix_ratio)
  print("color space: %s" % args.color_space)
  print("----------------------------------------------------")

  print("Generating ...")
  generate(args, frames)
  if len(frames_mv) != 0:
    args.output = args.output.replace(".y4m", "_mv.y4m")
    generate(args, frames_mv)