import os
from options.test_options import TestOptions
from options.dataset_options import DatasetOptions
from data import CreateRecursiveDataLoader
from models import create_model
from util.visualizer import Visualizer
from util.util import mkdirs, crop_image
from util import html
from shutil import move, copyfile
from PIL import Image, ImageOps
from skimage.transform import resize
from scipy.misc import imresize
from shutil import copyfile, rmtree
import numpy as np
import cv2
from datetime import datetime
import re
import sys
import math

import subprocess
from time import sleep

if __name__ == '__main__':
  opt = TestOptions().parse()
  data_opt = DatasetOptions().parse(opt.unknown)
  opt.nThreads = 1   # test code only supports nThreads = 1
  opt.batchSize = 1  # test code only supports batchSize = 1
  opt.serial_batches = True  # no shuffle
  opt.no_flip = True  # no flip
  # opt.experiment = data_opt.experiment # opt.start_img.split("/")[-1].split(".")[0]

  if data_opt.tag == '':
    d = datetime.now()
    tag = data_opt.tag = "{}_{}_{}".format(
      opt.name,
      opt.experiment,
      d.strftime('%Y%m%d%H%M')
    )
  else:
    tag = data_opt.tag

  opt.render_dir = render_dir = opt.results_dir + opt.name + "/" + tag + "/"
  A_offset = 0
  A_im = None
  A_dir = None

  def load_first_frame():
    start_img_path = os.path.join(render_dir, "frame_00000.png")
    if data_opt.just_copy:
      copyfile(opt.start_img, start_img_path)
    else:
      print("preload {}".format(opt.start_img))
      A_img = Image.open(opt.start_img).convert('RGB')
      A_im = np.asarray(A_img)
      A = process_image(A_im)
      cv2.imwrite(start_img_path, A)

    numz = re.findall(r'\d+', os.path.basename(opt.start_img))
    print(numz)
    if len(numz) > 0:
      A_offset = int(numz[0])
      print(A_offset)
      if A_offset:
        print(">> starting offset: {}".format(A_offset))
        A_dir = opt.start_img.replace(numz[0], "{:05d}")
        print(A_dir)
      else:
        print("Sequence not found")
    return A_offset, A_im, A_dir

  def process_image(im):
    img = im[:, :, ::-1].copy()
    processed = False

    if data_opt.clahe is True:
      processed = True
      lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
      l, a, b = cv2.split(lab)
      clahe = cv2.createCLAHE(clipLimit=data_opt.clip_limit, tileGridSize=(8,8))
      l = clahe.apply(l)
      limg = cv2.merge((l,a,b))
      img = cv2.cvtColor(limg, cv2.COLOR_LAB2BGR)
    if data_opt.posterize is True:
      processed = True
      img = cv2.pyrMeanShiftFiltering(img, data_opt.spatial_window, data_opt.color_window)
    if data_opt.grayscale is True:
      processed = True
      img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    if data_opt.blur is True:
      processed = True
      img = cv2.GaussianBlur(img, (data_opt.blur_radius, data_opt.blur_radius), data_opt.blur_sigma)
    if data_opt.canny is True:
      processed = True
      img = cv2.Canny(img, data_opt.canny_lo, data_opt.canny_hi)
      img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)

    if processed is False or data_opt.process_frac == 0:
      return img

    src_img = im[:, :, ::-1].copy()
    frac_a = data_opt.process_frac
    frac_b = 1.0 - frac_a
    array_a = np.multiply(src_img.astype('float64'), frac_a)
    array_b = np.multiply(img.astype('float64'), frac_b)
    img = np.add(array_a, array_b).astype('uint8')

    return img

  def render_video():
    print(opt.render_dir)
    if data_opt.render_frames:
      frame_fn = "frame_%05d.png"
      video_fn = tag + "_mog_frame.mp4"
    else:
      frame_fn = "ren_%05d.png"
      video_fn = tag + "_mogrify.mp4"

    # if opt.mov_file is not None and opt.mov_file != 'video.mp4':
    #   video_fn = opt.mov_file + '.mp4'

    cmd = ("ffmpeg", "-i", render_dir + frame_fn, "-y", "-c:v", "libx264", "-vf", "fps=30", "-pix_fmt", "yuv420p", "-s", "456x256", render_dir + video_fn)
    process = subprocess.Popen(cmd, stdout=subprocess.PIPE)
    output, error = process.communicate()

    print("________")

    cmd = ("scp", render_dir + video_fn, "jules@asdf.us:asdf/neural/")
    process = subprocess.Popen(cmd, stdout=subprocess.PIPE)
    output, error = process.communicate()

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

  if data_opt.mov:
    render_video()
    sys.exit(1)
  else:
    print("create render_dir: {}".format(render_dir))
    if os.path.exists(render_dir):
        rmtree(render_dir)
    mkdirs(render_dir)

  # cmd = ("convert", opt.start_img, '-canny', '0x1+10%+30%', render_dir + "frame_00000.png")
  # process = subprocess.Popen(cmd, stdout=subprocess.PIPE)
  # output, error = process.communicate()

  if data_opt.sequence:
    frac_a = data_opt.recursive_frac
    frac_b = data_opt.sequence_frac
    frac_c = 1.0 - frac_a - frac_b
    print("rec: {}, seq: {}, pix: {}".format(frac_a, frac_b, frac_c))
  else:
    frac_a = data_opt.recursive_frac
    frac_b = 1.0 - frac_a
    print("rec: {}, pix: {}".format(frac_a, frac_b))

  A_offset, A_im, A_dir = load_first_frame()

  data_loader = CreateRecursiveDataLoader(opt)
  dataset = data_loader.load_data()
  ds = dataset.dataset
  model = create_model(opt)
  visualizer = Visualizer(opt)
  # create website
  # web_dir = os.path.join(opt.results_dir, opt.name, '%s_%s' % (opt.phase, opt.which_epoch))
  # webpage = html.HTML(web_dir, 'Experiment = %s, Phase = %s, Epoch = %s' % (opt.name, opt.phase, opt.which_epoch))
  # test
  last_im = None
  for i, data in enumerate(data_loader):
    if i >= opt.how_many:
        break
    model.set_input(data)
    model.test()
    visuals = model.get_current_visuals()
    img_path = model.get_image_paths()

    if (i % 100) == 0:
      print('%04d: process image... %s' % (i, img_path))
    # ims = visualizer.save_images(webpage, visuals, img_path, aspect_ratio=opt.aspect_ratio)

    im = visuals['fake_B']
    tmp_path = render_dir + "frame_{:05d}_tmp.png".format(i+1)
    edges_path = render_dir + "frame_{:05d}.png".format(i+1)
    render_path = render_dir + "ren_{:05d}.png".format(i+1)
    if A_dir is not None:
      sequence_path = A_dir.format(A_offset+i+1)
      if not os.path.exists(sequence_path) and opt.how_many > 99999:
        break
    # A_offset

    # save rendered image
    image_pil = Image.fromarray(im, mode='RGB')
    image_pil.save(tmp_path)
    os.rename(tmp_path, render_path)

    def clamp(n,a,b):
      return max(a, min(n, b))
    def lerp(n,a,b):
      return (b-a)*n+a

    if dataset.name() == 'RecursiveDatasetDataLoader':
      if data_opt.recursive and last_im is not None:
        tmp_im = im.copy()

        if data_opt.sequence and A_dir is not None:
          A_img = Image.open(sequence_path).convert('RGB')
          A_im = np.asarray(A_img)
          t = lerp(math.sin(i / data_opt.transition_period * math.pi * 2 ) / 2 + 1, data_opt.transition_min, 1.0)
          frac_a = data_opt.recursive_frac * (1 - t)
          frac_b = data_opt.sequence_frac * (1 - t)
          frac_c = 1.0 - frac_a - frac_b
          array_a = np.multiply(last_im.astype('float64'), frac_a)
          array_b = np.multiply(A_im.astype('float64'), frac_b)
          array_c = np.multiply(im.astype('float64'), frac_c)
          array_ab = np.add(array_a, array_b)
          array_abc = np.add(array_ab, array_c)
          next_im = array_abc.astype('uint8')

        else:
          frac_a = data_opt.recursive_frac
          frac_b = 1.0 - frac_a
          array_a = np.multiply(last_im.astype('float64'), frac_a)
          array_b = np.multiply(im.astype('float64'), frac_b)
          next_im = np.add(array_a, array_b).astype('uint8')

        if data_opt.recurse_roll != 0:
          last_im = np.roll(tmp_im, data_opt.recurse_roll, axis=data_opt.recurse_roll_axis)
        else:
          last_im = next_im.copy().astype('uint8')

      else:
        last_im = im.copy().astype('uint8')
        tmp_im = im.copy().astype('uint8')
        next_im = im
        #print(im.shape, im.dtype)

      # image_pil = Image.fromarray(im, mode='RGB')
      # im = np.asarray(image_pil).astype('uint8')
      #print(im.shape, im.dtype)

      img = process_image(next_im)

      cv2.imwrite(tmp_path, img)
      os.rename(tmp_path, edges_path)

  # webpage.save()

  # os.remove(render_dir + "frame_00000.png")

  render_video()