import os
import sys
from shutil import move, copyfile
from PIL import Image, ImageOps
from shutil import copyfile, rmtree
import numpy as np
import cv2

work_dir = "woodsized/"
render_dir = "woodflow/"

if os.path.exists(render_dir):
    rmtree(render_dir)

os.makedirs(render_dir)
os.makedirs(render_dir + "A/")
os.makedirs(render_dir + "A/train/")
os.makedirs(render_dir + "A/test/")
os.makedirs(render_dir + "A/val/")
os.makedirs(render_dir + "B/")
os.makedirs(render_dir + "B/train/")
os.makedirs(render_dir + "B/test/")
os.makedirs(render_dir + "B/val/")

hsv = []

for i,fn in enumerate(sorted(os.listdir(work_dir))):
  # load image and convert to grayscale
  pil_image = Image.open(work_dir + fn).convert('RGB')
  im = np.array(pil_image)
  im = im[:, :, ::-1].copy()
  im = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY)

  # store first frame
  if i == 0:
    prev = im
    hsv = np.zeros((256,256,3))
    hsv[...,1] = 255
    continue

  # compute optical flow
  flow = cv2.calcOpticalFlowFarneback(prev, im, None, 0.5, 3, 15, 3, 5, 1.2, 0)

  # turn into magnitude/angle
  mag, ang = cv2.cartToPolar(flow[...,0], flow[...,1])

  # store angle as hue
  hsv[...,0] = ang * 180 / np.pi / 2

  # store magnitude as lum
  hsv[...,2] = cv2.normalize(mag, None, 0, 255, cv2.NORM_MINMAX)

  # convert this HSL to BGR
  bgr = cv2.cvtColor(hsv.astype('uint8'), cv2.COLOR_HSV2BGR)

  ren = "frame_{:05d}.png".format(i)
  if (i % 10) == 3:
    wd = "test/"
  elif (i % 10) == 6:
    wd = "val/"
  else:
    wd = "train/"

  if i == 1:
    prev_bgr = np.copy(bgr)
    continue

  cv2.imwrite(render_dir + "A/" + wd + ren, prev_bgr)
  cv2.imwrite(render_dir + "B/" + wd + ren, bgr)
  # copyfile(work_dir + fn, render_dir + "B/" + wd + ren)
  prev = im
  prev_bgr = np.copy(bgr)