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 = "IMG_1734/"
render_dir = "IMG_1734/flowwarp/"

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

os.makedirs(render_dir)
os.makedirs(render_dir + 'prev/')
os.makedirs(render_dir + 'next/')

def warp_flow(img, flow):
  h, w = flow.shape[:2]
  flow = -flow
  flow[:,:,0] += np.arange(w)
  flow[:,:,1] += np.arange(h)[:,np.newaxis]
  res = cv2.remap(img, flow, None, cv2.INTER_LINEAR)
  return res

hsv = []
prev = None
prev_bgr = None
gray = None
prev_gray = None

for i,fn in enumerate(sorted(os.listdir(work_dir))):
  if os.path.basename(fn).startswith('.') or not os.path.isfile(work_dir + fn):
    continue

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

  # store first frame
  if prev_gray is None:
    prev = im
    prev_gray = gray
    #hsv = np.zeros((256,512,3))
    #hsv[...,1] = 255
    continue

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

  #mag, ang = cv2.cartToPolar(flow[...,0], flow[...,1])
  #hsv[...,0] = ang * 180 / np.pi / 2
  #hsv[...,2] = cv2.normalize(mag, None, 0, 255, cv2.NORM_L2)
  #bgr = cv2.cvtColor(hsv.astype('uint8'), cv2.COLOR_HSV2BGR)

  #if prev_bgr is None:
  #  prev_bgr = np.copy(bgr)
  #  continue

  #if (i % 10) == 3:
  #  wd = "test/"
  #elif (i % 10) == 6:
  #  wd = "val/"
  #else:
  #  wd = "train/"

  out_prev = warp_flow(prev, flow)
  #out_next = warp_flow(im, flow)
  ren = "frame_{:05d}.png".format(i)

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