import os
import tensorflow

import click
import tensorflow as tf

# import keras.backend.tensorflow_backend as KTF
from model.utils.data_generator import DataGenerator
from model.seqGAN import SeqGAN
from model.utils.lr_schedule import LRSchedule
from model.utils.general import Config
from model.utils.text import Vocab
from model.utils.image import greyscale


@click.command()
# @click.option('--data', default="configs/data_small.json",
#               help='Path to data json config')
# @click.option('--vocab', default="configs/vocab_small.json",
#               help='Path to vocab json config')
# @click.option('--training', default="configs/training_small.json",
#               help='Path to training json config')
# @click.option('--model', default="configs/model.json",
#               help='Path to model json config')
# @click.option('--output', default="results/small/",
#               help='Dir for results and model weights')

@click.option('--data', default="configs/data_full.json",
              help='Path to data json config')
@click.option('--vocab', default="configs/vocab.json",
              help='Path to vocab json config')
@click.option('--training', default="configs/training.json",
              help='Path to training json config')
@click.option('--model', default="configs/model.json",
              help='Path to model json config')
@click.option('--output', default="results/full/",
              help='Dir for results and model weights')
def main(data, vocab, training, model, output):
    # Load configs
    # os.environ["CUDA_VISIBLE_DEVICES"] = "0"

    # config = tf.compat.v1.ConfigProto()

    # config.gpu_options.per_process_gpu_memory_fraction = 0.9
    # config.gpu_options.allow_growth = True
    # session = tf.compat.v1.Session(config=config)  # 设置session KTF.set_session(session )
    dir_output = output
    config = Config([data, vocab, training, model])
    config.save(dir_output)
    vocab = Vocab(config)

    # Load datasets
    train_set = DataGenerator(path_formulas=config.path_formulas_train,
                              dir_images=config.dir_images_train,
                              img_prepro=greyscale,
                              max_iter=config.max_iter,
                              #   iter_mode='formulas_only',
                              bucket=config.bucket_train,
                              path_matching=config.path_matching_train,
                              max_len=config.max_length_formula,
                              form_prepro=vocab.form_prepro)
    val_set = DataGenerator(path_formulas=config.path_formulas_val,
                            dir_images=config.dir_images_val,
                            img_prepro=greyscale,
                            max_iter=config.max_iter,
                            # iter_mode='formulas_only',
                            bucket=config.bucket_val,
                            path_matching=config.path_matching_val,
                            max_len=config.max_length_formula,
                            form_prepro=vocab.form_prepro)

    # Define learning rate schedule
    n_batches_epoch = ((len(train_set) + config.batch_size - 1) //
                       config.batch_size)
    lr_schedule = LRSchedule(lr_init=config.lr_init,
                             start_decay=config.start_decay*n_batches_epoch,
                             end_decay=config.end_decay*n_batches_epoch,
                             end_warm=config.end_warm*n_batches_epoch,
                             lr_warm=config.lr_warm,
                             lr_min=config.lr_min)
    # Build model and train
    model = SeqGAN(config, dir_output, vocab)
    model.build_train(config)
    model.train(config, train_set, val_set, lr_schedule)


if __name__ == "__main__":
    main()