from qrnn import get_model, qloss
from keras.callbacks import *
import numpy as np

# Generate the synthetic data
x1 = np.sin(np.arange(0, 9, 0.01))
x2 = np.cos(np.arange(0, 9, 0.01))
x3 = x1**2
x4 = (x1+x2)/2

Xtrain = np.vstack((x2, x3, x4)).T #(900, 3)
Ytrain = np.array([x1]*1).T #(900, 1)

# Parameters
input_dim = 3
num_hidden_layers = 2
num_units = [200, 200]
act = ['relu', 'relu']
dropout = [0.1, 0.1]
gauss_std = [0.3, 0.3]
num_quantiles = 9

# Get model
model = get_model(input_dim, num_units, act, dropout, gauss_std, num_hidden_layers, num_quantiles)
print(model.summary())

# Train
early_stopping = EarlyStopping(monitor='val_loss', patience=5)
model.compile(loss=lambda y_t, y_p: qloss(y_true=y_t, y_pred=y_p, n_q=num_quantiles), 
              optimizer='adam')
model.fit(x=Xtrain, y=Ytrain, 
          epochs=10, 
          validation_split=0.2, 
          batch_size=64, 
          shuffle=True, 
          callbacks=[early_stopping]
         )