__doc__ = """Simple demo on toy data.

In this demo, our toy data is the set of 'valley' sequences.  These are
sequences of numbers which are first decreasing and then increasing.

We train our generator via curriculum learning.  We start by only training
it to predict the next token given a groundtruth sequence via cross-entropy
loss.  As our generator learns, we start to train it directly against the
discriminator.

The discriminator always learns at the standard rate (half real examples
and half artificial examples generated by the generator).

"""

import model
import train

import numpy as np
import tensorflow as tf
import random

NUM_EMB = 4
EMB_DIM = 5
HIDDEN_DIM = 10
SEQ_LENGTH = 5
START_TOKEN = 0

EPOCH_ITER = 1000
CURRICULUM_RATE = 0.03  # how quickly to move from supervised training to unsupervised
TRAIN_ITER = 100000  # generator/discriminator alternating
D_STEPS = 2  # how many times to train the discriminator per generator step
LEARNING_RATE = 0.01 * SEQ_LENGTH
SEED = 88


def get_trainable_model():
    return model.GRU(
        NUM_EMB, EMB_DIM, HIDDEN_DIM,
        SEQ_LENGTH, START_TOKEN,
        learning_rate=LEARNING_RATE)


def verify_sequence(seq):
    downhill = True
    prev = NUM_EMB
    for tok in seq:
        if tok == START_TOKEN:
            return False
        if downhill:
            if tok > prev:
                downhill = False
        elif tok < prev:
            return False
        prev = tok
    return True


def get_random_sequence():
    """Returns random valley sequence."""
    tokens = set(range(NUM_EMB))
    tokens.discard(START_TOKEN)
    tokens = list(tokens)

    pivot = int(random.random() * SEQ_LENGTH)
    left_of_pivot = []
    right_of_pivot = []
    for i in xrange(SEQ_LENGTH):
        tok = random.choice(tokens)
        if i <= pivot:
            left_of_pivot.append(tok)
        else:
            right_of_pivot.append(tok)

    left_of_pivot.sort(reverse=True)
    right_of_pivot.sort(reverse=False)

    return left_of_pivot + right_of_pivot


def test_sequence_definition():
    for _ in xrange(1000):
        assert verify_sequence(get_random_sequence())


def main():
    random.seed(SEED)
    np.random.seed(SEED)

    trainable_model = get_trainable_model()
    sess = tf.Session()
    sess.run(tf.initialize_all_variables())

    print 'training'
    for epoch in xrange(TRAIN_ITER // EPOCH_ITER):
        print 'epoch', epoch
        proportion_supervised = max(0.0, 1.0 - CURRICULUM_RATE * epoch)
        train.train_epoch(
            sess, trainable_model, EPOCH_ITER,
            proportion_supervised=proportion_supervised,
            g_steps=1, d_steps=D_STEPS,
            next_sequence=get_random_sequence,
            verify_sequence=verify_sequence)


if __name__ == '__main__':
    test_sequence_definition()
    main()
