#!/Users/user/anaconda/envs/cv/bin/python

import os

import sys
sys.path.insert(0, '/Users/user/neural/pix2pix')
from options.dataset_options import DatasetOptions

from shutil import move, copyfile
from PIL import Image, ImageOps
from shutil import copyfile, rmtree
import glob
import numpy as np
import subprocess
import cv2

opt = DatasetOptions().parse()

in_dir = opt.in_dir
out_dir = opt.out_dir
out_base = os.path.basename(out_dir)

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

os.makedirs(out_dir)
if opt.split is True:
  os.makedirs(out_dir + "A/")
  os.makedirs(out_dir + "A/train/")
  os.makedirs(out_dir + "A/test/")
  os.makedirs(out_dir + "A/val/")
  if opt.ab is True:
    os.makedirs(out_dir + "B/")
    os.makedirs(out_dir + "B/train/")
    os.makedirs(out_dir + "B/test/")
    os.makedirs(out_dir + "B/val/")

file = open(os.path.join(out_dir, "opt.txt"), "w") 
for arg in vars(opt):
  file.write("{}: {}\n".format(arg, getattr(opt, arg)))
file.close()

images = sorted(glob.glob(os.path.join(in_dir, '*.*g')))
image_count = len(images)
print("{}, {} images => {}".format(in_dir, image_count, out_base))
for i, fn in enumerate(images):
  pil_image = Image.open(fn).convert('RGB')
  img = np.array(pil_image)
  img = img[:, :, ::-1].copy()

  out_file = "frame_{:05d}.png".format(i)

  if i > 0 and (i % 100) == 0:
    print("{}...".format(i))

  lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
  l, a, b = cv2.split(lab)

  if opt.clahe is True:
    clahe = cv2.createCLAHE(clipLimit=opt.clip_limit, tileGridSize=(8,8))
    l = clahe.apply(l)

  if opt.brightness_gradient is True:
    l = np.add(l.astype('float64'), ((i / image_count) - 0.5) * opt.brightness_sigma)
    np.clip(l, 0, 255, out=l)
    l = l.astype('uint8')

  if opt.brightness_gradient is True or opt.clahe is True:
    limg = cv2.merge((l,a,b))
    img = cv2.cvtColor(limg, cv2.COLOR_LAB2BGR)

  if opt.posterize is True:
    img = cv2.pyrMeanShiftFiltering(img, opt.spatial_window, opt.color_window)
  if opt.grayscale is True:
    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
  if opt.blur != 0:
    img = cv2.GaussianBlur(img, (opt.blur, opt.blur), opt.blur_sigma)
  if opt.canny is True:
    img = cv2.Canny(img, opt.canny_lo, opt.canny_hi)

  if opt.split is True:
    if (i % 10) == 3:
      wd = "test/"
    elif (i % 10) == 6:
      wd = "val/"
    else:
      wd = "train/"

    cv2.imwrite(out_dir + "A/" + wd + out_file, img)
    if opt.ab is True:
      copyfile(fn, out_dir + "B/" + wd + out_file)

  else:
    cv2.imwrite(os.path.join(out_dir, out_file), img)

print("{}...".format(image_count))

if opt.mov:
  print("ffmpeg...")
  mov_file = "{}.mp4".format(out_base)
  cmd = ("ffmpeg",
    "-loglevel", "quiet",
    "-i", os.path.join(out_dir, "frame_%05d.png"),
    "-y", "-c:v", "libx264", "-vf", "fps=30", "-pix_fmt", "yuv420p",
    mov_file)
  process = subprocess.Popen(cmd, stdout=subprocess.PIPE)
  output, error = process.communicate()

  if opt.scp:
    print("scp...")
    cmd = ("scp", mov_file, opt.scp)
    process = subprocess.Popen(cmd, stdout=subprocess.PIPE)
    output, error = process.communicate()

    print("https://asdf.us/neural/" + mov_file)

  # cmd = ("mplayer", mov_file)
  # process = subprocess.Popen(cmd, stdout=subprocess.PIPE)
  # output, error = process.communicate()