from __future__ import print_function
import tensorflow as tf
import argparse
import os
from six.moves import cPickle
from model import Model

sess = None
with open(os.path.join('save/save', 'config.pkl'), 'rb') as f:
    saved_args = cPickle.load(f)
with open(os.path.join('save/save', 'words_vocab.pkl'), 'rb') as f:
    words, vocab = cPickle.load(f)
model = Model(saved_args, True)
sess = tf.Session()
with sess.as_default():
    tf.initialize_all_variables().run()
    saver = tf.train.Saver(tf.all_variables())
    ckpt = tf.train.get_checkpoint_state('save/save')


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--save_dir', type=str, default='save',
                       help='model directory to store checkpointed models')
    parser.add_argument('-n', type=int, default=200,
                       help='number of words to sample')
    parser.add_argument('--prime', type=str, default=' ',
                       help='prime text')
    parser.add_argument('--sample', type=int, default=1,
                       help='0 to use max at each timestep, 1 to sample at each timestep, 2 to sample on spaces')

    args = parser.parse_args()
    print(sample(args))


def sample(args):
    if ckpt and ckpt.model_checkpoint_path:
        saver.restore(sess, ckpt.model_checkpoint_path)
    return model.sample(sess, words, vocab, args.n, args.prime, args.sample)


if __name__ == '__main__':
    main()
