import os
import numpy as np
import moviepy.editor as mpy
import argparse
import logging
from PIL import Image
import pdb

import cv2
def capture_opencv():
  VIDEO_PREFIX = "/path/to/video/video{}.mp4"
  for vid in range(0, 1000):
    vid_dir = "TrainValFrame/video{}".format(vid)
      if not os.path.exists(vid_dir):
        os.mkdir(vid_dir)

    cap = cv2.VideoCapture(VIDEO_PREFIX.format(vid))
    ret = True
    frame_num = cap.get(cv2.cv.CV_CAP_PROP_FRAME_COUNT)
    for frame_id in range(int(frame_num)):
      ret, frame = cap.read(frame_id)
      cv2.imwrite("TrainValFrame/video{}/frame{}.jpg".format(vid, frame_id), frame)
      if ret == False:
        break


def sample_capture(stat_file, video_path, frame_path, start_id, end_id, frames_per_video):
  """
  Capture frames from video.

  frames_per_video: The amount of frames captured from one video.
  """
  frame_amount_stat = open(stat_file, 'a')

  for video_id in range(start_id, end_id):
    video_filename = 'video{}.mp4'.format(video_id)
    clip = mpy.VideoFileClip(os.path.join(video_path, video_filename)

    frame_amount = int(clip.fps * clip.duration)
    frame_amount_stat.write(video_filename + ',{},{},{}\n'.format(frame_amount, clip.w, clip.h))

    frame_dir = os.path.join(frame_path, 'video{:04d}/'.format(video_id))
    if not os.path.exists(frame_dir):
      os.mkdir(frame_dir)  # video7010

    gap = clip.duration / frames_per_video
    for i in range(frames_per_video):
      # t is time in second.
      clip.save_frame(frame_dir + '{}.jpg'.format(i), t=(i * gap))
    print('Video{}: Done!'.format(video_id))

  # frame_amount_stat.close()

def show_statistic_info(split_name, stat_file):
  with open(stat_file) as f:
    lines = f.readlines()

  # Get the minimum and maximum frame amount, width and height
  video_amount = 0
  total_frames = 0
  total_width = 0
  total_height = 0
  min_frames, max_frames = 10000, 0
  min_w, max_w = 10000, 0
  min_h, max_h = 10000, 0

  for line in lines[1:]:
    video_amount += 1
    cur_frames = int(line.split(',')[1])
    total_frames += cur_frames
    cur_w = int(line.split(',')[2])
    total_width += cur_w
    cur_h = int(line.split(',')[3])
    total_height += cur_h

    if cur_frames < min_frames:
      min_frames = cur_frames
    if cur_frames > max_frames:
      max_frames = cur_frames

    if cur_w < min_w:
      min_w = cur_w
    if cur_w > max_w:
      max_w = cur_w

    if cur_h < min_h:
      min_h = cur_h
    if cur_h > max_h:
      max_h = cur_h

  print("Minimum frames of {} set is: {}".format(split_name, min_frames))
  print("Maximum frames of {} set is: {}".format(split_name, max_frames))
  print("Average frames amount of {} set is: {}".format(split_name, total_frames/ video_amount))

  print("Minimum width of {} set is: {}".format(split_name, min_w))
  print("Maximum width of {} set is: {}".format(split_name, max_w))
  print("Average width of {} set is: {}".format(split_name, total_width / video_amount))

  print("Minimum height of {} set is: {}".format(split_name, min_h))
  print("Maximum height of {} set is: {}".format(split_name, max_h))
  print("Average height of {} set is: {}".format(split_name, total_height / video_amount))


def capture(video_path, start_id, end_id, frame_path):
  for video_id in range(start_id, end_id):
    video_filename = 'video{}.mp4'.format(video_id)
    clip = mpy.VideoFileClip(os.path.join(video_path, video_filename))

    frame_dir = os.path.join(frame_path, 'video{:04d}/'.format(video_id))
    if not os.path.exists(frame_dir):
      os.mkdir(frame_dir)  # example: video7010

    i = 0
    for frame in clip.iter_frames(dtype="uint8"):
      frame = Image.fromarray(frame)
      frame.save("{}{:04d}.jpg".format(frame_dir, i))
      i += 1

    if video_id % 100 == 0:
      print("Current Video: ", video_id)


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

    parser.add_argument("--action", type=str)
    parser.add_argument("--video_root", type=str)
    parser.add_argument("--frame_root", type=str)
    parser.add_argument("--start_video", type=int)
    parser.add_argument("--end_video", type=int)

    parser.add_argument("--stat_file", type=str, help="Path to the statistic file.")
    parser.add_argument("--frames_per_video", type=int, help="If sample, the amount of frames for one video.")

    args = parser.parse_args()

    if args.action == "sample_capture":
      sample_capture(stat_file=args.stat_file,
                     video_path=args.video_root,
                     frame_path=args.frame_root,
                     start_id=args.start_video,
                     end_id=args.end_video,
                     frames_per_video=args.frames_per_video)
    if args.action == "show_statistic_info":
      show_statistic_info('train', 'statistic_train.txt')