import tensorflow as tf
import numpy as np

# preprocessed data
from datasets.cmudict import data
import data_utils

# load data from pickle and npy files
data_ctl, idx_words, idx_phonemes = data.load_data(PATH='datasets/cmudict/')
(trainX, trainY), (testX, testY), (validX, validY) = data_utils.split_dataset(idx_phonemes, idx_words)

# parameters 
xseq_len = trainX.shape[-1]
yseq_len = trainY.shape[-1]
batch_size = 128
xvocab_size = len(data_ctl['idx2pho'].keys())  
yvocab_size = len(data_ctl['idx2alpha'].keys())
emb_dim = 128

"""
Explanation: Demonstrate Seq2Seq Wrapper with CMUDict dataset
End of explanation
"""


import seq2seq_wrapper

import importlib
importlib.reload(seq2seq_wrapper)

model = seq2seq_wrapper.Seq2Seq(xseq_len=xseq_len,
                               yseq_len=yseq_len,
                               xvocab_size=xvocab_size,
                               yvocab_size=yvocab_size,
                               ckpt_path='ckpt/cmudict/',
                               emb_dim=emb_dim,
                               num_layers=3
                               )

"""
Explanation: Create an instance of the Wrapper
End of explanation
"""


val_batch_gen = data_utils.rand_batch_gen(validX, validY, 16)
train_batch_gen = data_utils.rand_batch_gen(trainX, trainY, 128)

"""
Explanation: Create data generators
Read data_utils.py for more information
End of explanation
"""


sess = model.train(train_batch_gen, val_batch_gen, sess=sess1)

"""
Explanation: Computational graph was built when the model was instantiated
Now all we need to do is train the model using processed CMUdict dataset, via data generators
Internally a loop is run for epochs times for training
Evaluation is done periodically.

Train
End of explanation
"""


sess = model.restore_last_session()

"""
Explanation: Restore last saved session from disk
End of explanation
"""


output = model.predict(sess, val_batch_gen.__next__()[0])
print(output.shape)

output

"""
Explanation: Predict
End of explanation
"""


for oi in output:
    print(data_utils.decode(sequence=oi, lookup=data_ctl['idx2alpha'],
                           separator=''))

"""
Explanation: Let us decode and see the words
End of explanation
"""


