#学習結果から画像を生成するプログラム
#!/usr/bin/env python

# python train_facade.py -g 0 -i ./facade/base --out result_facade --snapshot_interval 10000

from __future__ import print_function
import argparse
import os
import os.path
import glob
import math
from PIL import Image
import cv2

import chainer
from chainer.serializers.npz import NpzDeserializer
import numpy as np

from net import Encoder
from net import Decoder
import scipy as sp

import chainer
import chainer.cuda
from chainer import Variable

import time

"""各種設定"""
#入力画像サイズ
w_in = 256
h_in = 256
#生成したい画像サイズ
w_out = 256
h_out = 256
#保存したパラメータの読み込み
iter_first = 5000
iter_last = 380000
iter_inter = 5000
#パラメータの保存先
param_dir = "./pix2pix_param2/"
#入力画像の保存先
path_man = "./dataset/"
#GPU番号
gpu_num = 0
man_fld = [os.path.basename(x) for x in glob.glob(path_man+'*')]

def out_image(img, man, files, enc, dec, rows, cols, seed, iter_num, gpu):
    def save_image(x, name, mode=None):
        _, C, H, W = x.shape
        x = x.reshape((rows, cols, C, H, W))
        x = x.transpose(0, 3, 1, 4, 2)
        if C==1:
            x = x.reshape((rows*H, cols*W))
        else:
            x = x.reshape((rows*H, cols*W, C))

        img_fld, _ = os.path.splitext(files)
        preview_dir = "./generate_image/"+man+"/"+img_fld+"/"
        if not os.path.exists(preview_dir):
            os.makedirs(preview_dir)
        preview_path =preview_dir+"iter_"+iter_num+".jpg"
        Image.fromarray(x, mode=mode).convert('RGB').save(preview_path)

    np.random.seed(seed)
    n_images = rows * cols
    xp = enc.xp

    in_ch = 3
    out_ch = 3
    batchsize = 1

    in_all = np.zeros((n_images, in_ch, h_in, w_in)).astype("f")
    gen_all = np.zeros((n_images, out_ch, h_out, w_out)).astype("f")

    for it in range(n_images):

        x_in = xp.zeros((batchsize, in_ch, h_in, w_in)).astype("f")
        t_out = xp.zeros((batchsize, out_ch, h_out, w_out)).astype("f")

        x_in[0,:] = xp.asarray(img)
        x_in = Variable(x_in)

        z = enc(x_in)
        x_out = dec(z)

        if gpu >= 0:
            in_all[it,:] = x_in.data.get()[0,:]
            gen_all[it,:] = x_out.data.get()[0,:]
        else:
            in_all[it,:] = x_in.data[0,:]
            gen_all[it,:] = x_out.data[0,:]

    x = np.asarray(np.clip(gen_all * 128 + 128, 0.0, 255.0), dtype=np.uint8)
    save_image(x, "gen")

    #x = np.asarray(np.clip(in_all * 128+128, 0.0, 255.0), dtype=np.uint8)
    #save_image(x, "in")

all_time = time.time()

for man in man_fld:
    path_img = path_man+man+"/"
    img_file = [os.path.basename(x) for x in glob.glob(path_img+'*')]

    for files in img_file:

        start = time.time()

        for num in range(iter_first,iter_last+1,iter_inter):
            iter_num = str(num)
            # trained model
            ENC_W = param_dir+"enc_iter_"+iter_num+".npz"
            DEC_W = param_dir+"dec_iter_"+iter_num+".npz"

            parser = argparse.ArgumentParser(description='chainer implementation of pix2pix')
            parser.add_argument('--gpu', '-g', type=int, default=gpu_num,
                                help='GPU ID (negative value indicates CPU)')
            parser.add_argument('--img', '-i', help='Input image')
            parser.add_argument('--out', '-o', default='result_dehighlight',
                                help='Directory to output the result')
            args = parser.parse_args()


            # Set up a neural network to train
            enc = Encoder(in_ch=3)
            dec = Decoder(out_ch=3)

            chainer.serializers.load_npz(ENC_W, enc)
            chainer.serializers.load_npz(DEC_W, dec)

            if args.gpu >= 0:
                chainer.cuda.get_device(args.gpu).use()  # Make a specified GPU current
                enc.to_gpu()  # Copy the model to the GPU
                dec.to_gpu()

            #入力画像の読み込み
            img_src = Image.open(path_img+files)
            w,h = img_src.size
            #r = 286/min(w,h)
            # resize images so that min(w, h) == 286
            #img_src = img_src.resize((int(r*w), int(r*h)), Image.BILINEAR)
            img_src = img_src.resize((w_in, h_in), Image.BILINEAR)

            img_src = np.asarray(img_src).astype("f").transpose(2,0,1)/128.0-1.0

            out_image(img_src,man,files, enc, dec, 1, 1, 0, iter_num, args.gpu)
        elapsed_time = time.time() - start
        print("complete: "+files)
        print("time: " +str(elapsed_time))
finish_time = time.time()-all_time
print("total: "+str(finish_time))