import argparse
import numpy as np
import io, math
import os, sys, time
from PIL import Image
import cPickle as pickle
import thread
import chainer
from chainer import cuda, Variable, optimizers, serializers
import chainer.functions as F
import chainer.links as L
import net, rnn_net
import copy
import signal

parser = argparse.ArgumentParser(description='VAE and DCGAN trainer')
parser.add_argument('--gpu', '-g', default=0, type=int,
                    help='GPU ID (negative value indicates CPU)')
parser.add_argument('--input', '-i', default='../model/', type=str,
                    help='input model file path without extension')
parser.add_argument('--output', '-o', default='../model/', type=str,
                    help='output model file path without extension')
parser.add_argument('--iter', default=1000, type=int,
                    help='number of iteration')
parser.add_argument('--out_image_dir', default='../image', type=str,
                    help='output directory to output images')
parser.add_argument('--size', '-s', default=128, type=int, choices=[48, 64, 80, 96, 112, 128],
                    help='image size')
args = parser.parse_args()

cost_mode = 'GAN' # VAE , GAN, BEGAN ** BEGAN not working!
latent_size = 256
BATCH_SIZE = 100
seq_length = 1
gpus_to_use = [0] # You can use multiple GPUs by putting their indices in an array. For instance: [0,1,2,3] for four GPUs
num_gpus = len(gpus_to_use)
main_gpu = gpus_to_use[0]
max_seq_length = 5
kt = 0
lambda_k = 0.001
gamma = .5
train_LSTM = True
train_LSTM_using_cached_features = False
train_lstm_prob = .5
train_dis = True
image_save_interval = 200000
model_save_interval = image_save_interval
out_image_row_num = 7
out_image_col_num = 14
if train_LSTM:
    BATCH_SIZE /= seq_length
normer = args.size * args.size * 3 * 60
image_path = ['../../../trajectories/al5d']
np.random.seed( 1241 )
image_size = args.size
enc_model = [net.Encoder(density=8, size=image_size, latent_size=latent_size)]
gen_model = [net.Generator(density=8, size=image_size, latent_size=latent_size)]
dis_model = [net.Discriminator(density=8, size=image_size)]
for i in range(num_gpus-1):
    enc_model.append(copy.deepcopy(enc_model[0]))
    gen_model.append(copy.deepcopy(gen_model[0]))
    dis_model.append(copy.deepcopy(dis_model[0]))

enc_dis_model = net.Encoder(density=8, size=image_size, latent_size=latent_size)
gen_dis_model = net.Generator(density=8, size=image_size, latent_size=latent_size)
rnn_model = rnn_net.MDN_RNN(IN_DIM=latent_size+5, HIDDEN_DIM=300, OUT_DIM=6, NUM_MIXTURE=40)

optimizer_enc = optimizers.Adam(alpha=0.0001, beta1=0.5)
optimizer_enc.setup(enc_model[0])
optimizer_enc.add_hook(chainer.optimizer.WeightDecay(0.00001))
optimizer_gen = optimizers.Adam(alpha=0.0001, beta1=0.5)
optimizer_gen.setup(gen_model[0])
optimizer_gen.add_hook(chainer.optimizer.WeightDecay(0.00001))

optimizer_enc_dis = optimizers.Adam(alpha=0.0001, beta1=0.5)
optimizer_enc_dis.setup(enc_dis_model)
optimizer_enc_dis.add_hook(chainer.optimizer.WeightDecay(0.00001))
optimizer_gen_dis = optimizers.Adam(alpha=0.0001, beta1=0.5)
optimizer_gen_dis.setup(gen_dis_model)
optimizer_gen_dis.add_hook(chainer.optimizer.WeightDecay(0.00001))

optimizer_dis = optimizers.Adam(alpha=0.0001, beta1=0.5)
optimizer_dis.setup(dis_model[0])
optimizer_dis.add_hook(chainer.optimizer.WeightDecay(0.00001))
optimizer_rnn = optimizers.Adam(alpha=0.0001, beta1=0.5)
optimizer_rnn.setup(rnn_model)
optimizer_rnn.add_hook(chainer.optimizer.WeightDecay(0.00001))
xp = cuda.cupy

def signal_handler(signal, frame):
    sys.exit(0)
signal.signal(signal.SIGINT, signal_handler)

for i in range(num_gpus):
    enc_model[i].to_gpu(gpus_to_use[i])
    gen_model[i].to_gpu(gpus_to_use[i])
    dis_model[i].to_gpu(gpus_to_use[i])
if cost_mode == 'BEGAN':
    enc_dis_model.to_gpu(main_gpu)
    gen_dis_model.to_gpu(main_gpu)
rnn_model.to_gpu(main_gpu)