import os
from options.test_options import TestOptions
from options.dataset_options import TestOptions
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
import time

import subprocess
from time import sleep

blur = 3
sigma = 0
canny_lo = 10
canny_hi = 220
frac_a = 0.99
frac_b = 1 - frac_a

if __name__ == '__main__':
    opt = TestOptions().parse()
    data_opt = DatasetOptions().parse()
    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 = opt.start_img.split("/")[-1].split(".")[0]

    render_dir = opt.results_dir + opt.name + "/exp:" + opt.experiment + "/"

    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()

    #copyfile(opt.start_img, render_dir + "frame_00000.png")

    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()
        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)

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

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

                frac_a = data_opt.recursive_frac
                frac_b = 1.0 - frac_a

                array_a = np.multiply(im.astype('float64'), frac_a)
                array_b = np.multiply(last_im.astype('float64'), frac_b)
                im = np.add(array_a, array_b).astype('uint8')
                # print(im.shape, im.dtype)
                last_im = np.roll(tmp_im, 1, axis=1)
            else:
                last_im = im.copy().astype('uint8')
                tmp_im = im.copy().astype('uint8')
                #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 = im[:, :, ::-1].copy()

            if data_opt.clahe is 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:
              img = cv2.pyrMeanShiftFiltering(img, data_opt.spatial_window, data_opt.color_window)
            if data_opt.grayscale is True:
              img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            if data_opt.blur is True:
              img = cv2.GaussianBlur(img, (data_opt.blur_radius, data_opt.blur_radius), data_opt.blur_sigma)
            if data_opt.canny is True:
              img = cv2.Canny(img, data_opt.canny_lo, data_opt.canny_hi)

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


    webpage.save()

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

    t = time.time()
    t /= 60
    t %= 525600
    video_fn = "{}_{}_{}_mogrify.mp4".format(
        opt.name, opt.experiment,
        # opt.how_many, frac_a,
        # blur, sigma, canny_lo, canny_hi,
        int(t))

    cmd = ("/usr/bin/ffmpeg", "-i", render_dir + "ren_%05d.png", "-y", "-c:v", "libx264", "-vf", "fps=30", "-pix_fmt", "yuv420p", 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)