# In[2]:


import acoustic_model.generators_wav as generators
import acoustic_model.utils as utils
import acoustic_model.deepspeech2_model_conv2d as deepspeech2conv2d
import pandas as pd


# In[3]:


import warnings
warnings.filterwarnings("ignore")


# In[4]:


from keras.backend import set_session
import tensorflow as tf
import os
# # Only for server
# set settings
os.environ["CUDA_VISIBLE_DEVICES"] = '3'

config = tf.ConfigProto()
config.gpu_options.allow_growth = True
set_session(tf.Session(config=config))


# In[14]:


BATCH_PATH_TRAIN = '/media/sdb-dados/users/cps/batches/batches_train/'
BATCH_PATH_VALID = '/media/sdb-dados/users/cps/batches/batches_valid/'


# In[15]:


noise_df = pd.read_csv('/media/sdb-dados/users/cps/noise_audio_csv/noise_audio_mel_wav.csv')


# In[31]:


generator_train = generators.DeepSpeechBatchGeneratorConv2dTf(batch_path=BATCH_PATH_TRAIN, noise_data=noise_df,
                                                            batch_size=64, no_aug=False, nfilt=80)
generator_valid = generators.DeepSpeechBatchGeneratorConv2dTf(batch_path=BATCH_PATH_VALID, noise_data=noise_df,
                                                            batch_size=64, no_aug=False, nfilt=80)

generator_train.TIME_CONV_KERNELS_SZ = [5, 5, 5]
generator_train.CONV_STRIDES = [2, 1, 1]

generator_valid.TIME_CONV_KERNELS_SZ = [5, 5, 5]
generator_valid.CONV_STRIDES = [2, 1, 1]

# In[8]:


model = deepspeech2conv2d.DeepSpeechModel2Conv2D(model_save_dir='./notebooks/conv2d_80_big_5/', input_dim=80)


# In[9]:


model.build_model(gru_layers=4, gru_size=384, conv_filters_sz=[32, 32, 64], conv_kernels=[(5,11), (5,7), (5,5)], conv_strides=[(2,2), (1,1), (1,1)], tf_features=True)



# In[22]:
callbacks = model.build_callbacks()


model.fit_generator(generator_train,
                    steps_per_epoch=len(generator_train.batchs),
                    epochs=100,
                    verbose=1,
                    validation_data=generator_valid,
                    validation_steps=len(generator_valid.batchs),
                    callbacks=callbacks,
                    workers=16,
                    max_queue_size=8,
                    use_multiprocessing=False)