'''
Script for training the model
'''
import tensorflow as tf
import numpy as np
from input import BatchGenerator
from model import MultiRnn
import time
from datetime import datetime
import os
import matplotlib as mpl
mpl.use('Agg')
from matplotlib import pyplot as plt


sum_dir = 'sum'  # dir to write summary
train_dir = 'ckpt'  # dir to store the model
data_dir = 'train.pkl'  # dir of the data set
NEFF = 129  # effective FFT points
batch_size = 128
num_steps = 20
epochs = 2000
cell_type = 'NL_LSTM'
state_size = 256
output_size = 129
num_layer = 3
learning_rate = 0.0001

# build the model
rnn_model = MultiRnn(
    cell_type, state_size, output_size,
    batch_size, num_layer, learning_rate, num_steps)

# input data and referene data placeholder
in_data = tf.placeholder(
    tf.float32, [batch_size, num_steps, 2 * NEFF])
ref_data = tf.placeholder(
    tf.float32, [batch_size, num_steps, NEFF])

# make inference
init_state, final_state, inf_data = rnn_model.inference(in_data)

# compute loss
loss = rnn_model.loss(inf_data, ref_data)

saver = tf.train.Saver(tf.all_variables())

summary_op = tf.merge_all_summaries()

train_op = rnn_model.train(loss)

batch_gen = BatchGenerator(data_dir, batch_size, num_steps, epochs)

with tf.Session() as sess:
    summary_writer = tf.train.SummaryWriter(
        sum_dir, sess.graph)
    sess.run(tf.initialize_all_variables())
    steps = 0
    # generator for epoch data
    for idx, epoch in enumerate(batch_gen.gen_epochs()):
        training_state = None
        # generator for batch data
        for f_data, b_data, r_data, v_data in epoch:
            start_time = time.time()
            steps += 1
            in_data_np = np.concatenate((f_data, b_data), axis=2)
            if steps % 100 == 0:
                feed_dict = {in_data: in_data_np, ref_data: r_data}
                if training_state is not None:
                    feed_dict[init_state] = training_state
                # training the net
                loss_value, training_state, _, summary_str, test_inf = sess.run(
                    [loss, final_state, train_op, summary_op, inf_data], feed_dict)
                duration = time.time() - start_time
                sec_per_batch = float(duration)
                examples_per_sec = batch_size / duration
                format_str = (
                    '%s: step %d, loss = %.2f (%.1f examples/sec; %.3f '
                    'sec/batch, epoch %d)')
                print (format_str % (datetime.now(), steps, loss_value,
                                     examples_per_sec, sec_per_batch,
                                     idx))
                summary_writer.add_summary(summary_str, steps)
            else:
                feed_dict = {in_data: in_data_np, ref_data: r_data}
                if training_state is not None:
                    feed_dict[init_state] = training_state

                loss_value, training_state, _ = sess.run(
                    [loss, final_state, train_op], feed_dict)
            if steps % 10000 == 0:
                checkpoint_path = os.path.join(train_dir, 'model.ckpt')
                saver.save(sess, checkpoint_path, global_step=steps)

[((12, 0, 12, 14), 'matplotlib.use', 'mpl.use', ({(12, 8, 12, 13): '"""Agg"""'}, {}), "('Agg')", True, 'import matplotlib as mpl\n'), ((30, 12, 32, 52), 'model.MultiRnn', 'MultiRnn', ({(31, 4, 31, 13): 'cell_type', (31, 15, 31, 25): 'state_size', (31, 27, 31, 38): 'output_size', (32, 4, 32, 14): 'batch_size', (32, 16, 32, 25): 'num_layer', (32, 27, 32, 40): 'learning_rate', (32, 42, 32, 51): 'num_steps'}, {}), '(cell_type, state_size, output_size, batch_size, num_layer,\n    learning_rate, num_steps)', False, 'from model import MultiRnn\n'), ((35, 10, 36, 50), 'tensorflow.placeholder', 'tf.placeholder', ({(36, 4, 36, 14): 'tf.float32', (36, 16, 36, 49): '[batch_size, num_steps, 2 * NEFF]'}, {}), '(tf.float32, [batch_size, num_steps, 2 * NEFF])', True, 'import tensorflow as tf\n'), ((37, 11, 38, 46), 'tensorflow.placeholder', 'tf.placeholder', ({(38, 4, 38, 14): 'tf.float32', (38, 16, 38, 45): '[batch_size, num_steps, NEFF]'}, {}), '(tf.float32, [batch_size, num_steps, NEFF])', True, 'import tensorflow as tf\n'), ((48, 13, 48, 37), 'tensorflow.merge_all_summaries', 'tf.merge_all_summaries', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((52, 12, 52, 67), 'input.BatchGenerator', 'BatchGenerator', ({(52, 27, 52, 35): 'data_dir', (52, 37, 52, 47): 'batch_size', (52, 49, 52, 58): 'num_steps', (52, 60, 52, 66): 'epochs'}, {}), '(data_dir, batch_size, num_steps, epochs)', False, 'from input import BatchGenerator\n'), ((46, 23, 46, 41), 'tensorflow.all_variables', 'tf.all_variables', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((54, 5, 54, 17), 'tensorflow.Session', 'tf.Session', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((55, 21, 56, 28), 'tensorflow.train.SummaryWriter', 'tf.train.SummaryWriter', ({(56, 8, 56, 15): 'sum_dir', (56, 17, 56, 27): 'sess.graph'}, {}), '(sum_dir, sess.graph)', True, 'import tensorflow as tf\n'), ((57, 13, 57, 42), 'tensorflow.initialize_all_variables', 'tf.initialize_all_variables', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((64, 25, 64, 36), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((66, 25, 66, 65), 'numpy.concatenate', 'np.concatenate', (), '', True, 'import numpy as np\n'), ((92, 34, 92, 71), 'os.path.join', 'os.path.join', ({(92, 47, 92, 56): 'train_dir', (92, 58, 92, 70): '"""model.ckpt"""'}, {}), "(train_dir, 'model.ckpt')", False, 'import os\n'), ((74, 27, 74, 38), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((80, 37, 80, 51), 'datetime.datetime.now', 'datetime.now', ({}, {}), '()', False, 'from datetime import datetime\n')]