# Imports
import sys
sys.path.append('../')   # This is where all the python files are!
from importlib import reload
import emos_network_theano; reload(emos_network_theano)
from  emos_network_theano import EMOS_Network
from losses import crps_cost_function
import utils; reload(utils)
from utils import *
import keras_models; reload(keras_models)
from keras_models import *
from collections import OrderedDict

# Basic setup
DATA_DIR = '/Volumes/STICK/data/ppnn_data/'  # Mac
# DATA_DIR = '/project/meteo/w2w/C7/ppnn_data/'   # LMU
results_dir = '../results/'
window_size = 25   # Days in rolling window
fclt = 48   # Forecast lead time in hours
train_dates = ['2015-01-01', '2016-01-01']
test_dates =  ['2016-01-01', '2017-01-01']

"""
Explanation: <h1>Table of Contents<span class="tocSkip"></span></h1>
<div class="toc"><ul class="toc-item"><li><span><a href="#Fully-connected-linear-network" data-toc-modified-id="Fully-connected-linear-network-1"><span class="toc-item-num">1&nbsp;&nbsp;</span>Fully connected linear network</a></span><ul class="toc-item"><li><span><a href="#Get-temperature-data" data-toc-modified-id="Get-temperature-data-1.1"><span class="toc-item-num">1.1&nbsp;&nbsp;</span>Get temperature data</a></span></li><li><span><a href="#Build-fully-connected-model" data-toc-modified-id="Build-fully-connected-model-1.2"><span class="toc-item-num">1.2&nbsp;&nbsp;</span>Build fully connected model</a></span></li><li><span><a href="#Predict-for-one-day" data-toc-modified-id="Predict-for-one-day-1.3"><span class="toc-item-num">1.3&nbsp;&nbsp;</span>Predict for one day</a></span></li><li><span><a href="#Post-processing-with-rolling-window-for-2016" data-toc-modified-id="Post-processing-with-rolling-window-for-2016-1.4"><span class="toc-item-num">1.4&nbsp;&nbsp;</span>Post processing with rolling window for 2016</a></span></li><li><span><a href="#Train-2015,-predict-2016" data-toc-modified-id="Train-2015,-predict-2016-1.5"><span class="toc-item-num">1.5&nbsp;&nbsp;</span>Train 2015, predict 2016</a></span></li></ul></li><li><span><a href="#Neural-network-with-one-hidden-layer" data-toc-modified-id="Neural-network-with-one-hidden-layer-2"><span class="toc-item-num">2&nbsp;&nbsp;</span>Neural network with one hidden layer</a></span><ul class="toc-item"><li><span><a href="#Build-network" data-toc-modified-id="Build-network-2.1"><span class="toc-item-num">2.1&nbsp;&nbsp;</span>Build network</a></span></li><li><span><a href="#Train-2015,-predict-2016" data-toc-modified-id="Train-2015,-predict-2016-2.2"><span class="toc-item-num">2.2&nbsp;&nbsp;</span>Train 2015, predict 2016</a></span></li><li><span><a href="#Making-the-hidden-model-more-complex" data-toc-modified-id="Making-the-hidden-model-more-complex-2.3"><span class="toc-item-num">2.3&nbsp;&nbsp;</span>Making the hidden model more complex</a></span></li></ul></li><li><span><a href="#Add-station-embeddings" data-toc-modified-id="Add-station-embeddings-3"><span class="toc-item-num">3&nbsp;&nbsp;</span>Add station embeddings</a></span><ul class="toc-item"><li><span><a href="#Build-linear-embedding-model" data-toc-modified-id="Build-linear-embedding-model-3.1"><span class="toc-item-num">3.1&nbsp;&nbsp;</span>Build linear embedding model</a></span></li><li><span><a href="#Train-2015,-predict-2016" data-toc-modified-id="Train-2015,-predict-2016-3.2"><span class="toc-item-num">3.2&nbsp;&nbsp;</span>Train 2015, predict 2016</a></span></li><li><span><a href="#Embedding-size-hyper-parameter-tuning" data-toc-modified-id="Embedding-size-hyper-parameter-tuning-3.3"><span class="toc-item-num">3.3&nbsp;&nbsp;</span>Embedding size hyper-parameter tuning</a></span></li></ul></li><li><span><a href="#Adding-auxiliary-variables" data-toc-modified-id="Adding-auxiliary-variables-4"><span class="toc-item-num">4&nbsp;&nbsp;</span>Adding auxiliary variables</a></span><ul class="toc-item"><li><span><a href="#Load-extended-dataset" data-toc-modified-id="Load-extended-dataset-4.1"><span class="toc-item-num">4.1&nbsp;&nbsp;</span>Load extended dataset</a></span></li><li><span><a href="#Linear-model" data-toc-modified-id="Linear-model-4.2"><span class="toc-item-num">4.2&nbsp;&nbsp;</span>Linear model</a></span></li><li><span><a href="#Hidden-model" data-toc-modified-id="Hidden-model-4.3"><span class="toc-item-num">4.3&nbsp;&nbsp;</span>Hidden model</a></span></li><li><span><a href="#Even-more-variables" data-toc-modified-id="Even-more-variables-4.4"><span class="toc-item-num">4.4&nbsp;&nbsp;</span>Even more variables</a></span></li><li><span><a href="#Save-pickled-datasets" data-toc-modified-id="Save-pickled-datasets-4.5"><span class="toc-item-num">4.5&nbsp;&nbsp;</span>Save pickled datasets</a></span></li><li><span><a href="#Train-model" data-toc-modified-id="Train-model-4.6"><span class="toc-item-num">4.6&nbsp;&nbsp;</span>Train model</a></span></li></ul></li><li><span><a href="#Additional-variables-with-the-embedding-model" data-toc-modified-id="Additional-variables-with-the-embedding-model-5"><span class="toc-item-num">5&nbsp;&nbsp;</span>Additional variables with the embedding model</a></span><ul class="toc-item"><li><span><a href="#Linear-model-with-embeddings" data-toc-modified-id="Linear-model-with-embeddings-5.1"><span class="toc-item-num">5.1&nbsp;&nbsp;</span>Linear model with embeddings</a></span></li><li><span><a href="#Hidden-model-with-embeddings" data-toc-modified-id="Hidden-model-with-embeddings-5.2"><span class="toc-item-num">5.2&nbsp;&nbsp;</span>Hidden model with embeddings</a></span></li><li><span><a href="#A-longer-training-period" data-toc-modified-id="A-longer-training-period-5.3"><span class="toc-item-num">5.3&nbsp;&nbsp;</span>A longer training period</a></span></li><li><span><a href="#A-longer-training-period-with-more-data" data-toc-modified-id="A-longer-training-period-with-more-data-5.4"><span class="toc-item-num">5.4&nbsp;&nbsp;</span>A longer training period with more data</a></span></li></ul></li><li><span><a href="#Data-augmentation" data-toc-modified-id="Data-augmentation-6"><span class="toc-item-num">6&nbsp;&nbsp;</span>Data augmentation</a></span></li></ul></div>

Fully connected and neural networks
Now that we have established that we can get equivalent EMOS results using a network architecture with SGD, we can now extend this approach to fully connected networks. 
Here we will try several approaches:
- Simple linear fully connected networks
- Neural networks with hidden layers
- Adding station embeddings
- Adding additional data
End of explanation
"""


date_str = '2011-02-14'
train_set, test_set = get_train_test_sets(DATA_DIR, predict_date=date_str,
                                          fclt=fclt, window_size=window_size)

"""
Explanation: Fully connected linear network
As a first step, we can build a linear model which also connects the means and standard deviations.
Get temperature data
This follows the steps in the EMOS Network data
End of explanation
"""


fc_model = build_fc_model(2, 2, compile=True)

fc_model.summary()

"""
Explanation: Build fully connected model
End of explanation
"""


# Define some parameters
early_stopping_delta = 1e-4   # How much the CRPS must improve before stopping
steps_max = 1000   # How many steps to fit at max
batch_size = train_set.features.shape[0]

fc_model.fit(train_set.features, train_set.targets, epochs=steps_max, 
             batch_size=batch_size,
             validation_data=[test_set.features, test_set.targets], 
             verbose=0,
             callbacks=[EarlyStopping(monitor='loss', 
                                      min_delta=early_stopping_delta,
                                      patience=2)]);

# Get train and test CRPS
(fc_model.evaluate(train_set.features, train_set.targets, batch_size, verbose=0), 
 fc_model.evaluate(test_set.features, test_set.targets, batch_size, verbose=0))

"""
Explanation: Now we have 6 parameters instead of 4 with the standard EMOS Network.
Predict for one day
End of explanation
"""


date_str_start = '2016-01-01'
date_str_stop = '2017-01-01'

fc_model = build_fc_model(2, 2, compile=True, optimizer='sgd')

# Use the loop function in utils
train_crps_list, valid_crps_list, results_df = loop_over_days(
    DATA_DIR,
    fc_model,
    date_str_start, date_str_stop, 
    window_size=window_size,
    fclt=fclt,     
    epochs_max=steps_max, 
    early_stopping_delta=early_stopping_delta, 
    lr=0.1,   
    verbose=0)

np.mean(train_crps_list), np.mean(valid_crps_list)

"""
Explanation: For this particular day we get a score that is slightly better than the standard EMOS network.
Post processing with rolling window for 2016
As with the EMOS models let's do a rolling window global post-processing for 2016.
End of explanation
"""


results_df.to_csv(results_dir + 'fc_network_rolling_window.csv')

"""
Explanation: So we get a slightly better training score and a slightly worse test score. This is a sign of overfitting. But the differences are small.
End of explanation
"""


train_dates = ['2015-01-01', '2016-01-01']
test_dates =  ['2016-01-01', '2017-01-01']
train_set, test_set = get_train_test_sets(DATA_DIR, train_dates, test_dates)

train_set.features.shape

fc_model = build_fc_model(2, 2, compile=True)

from IPython.display import SVG
from keras.utils.vis_utils import model_to_dot

SVG(model_to_dot(fc_model).create(prog='dot', format='svg'))

fc_model.compile(keras.optimizers.Adam(0.001), loss=crps_cost_function)

# Note: I am running this cell several times (40)
fc_model.fit(train_set.features, train_set.targets, epochs=10, batch_size=1024,
             validation_data=[test_set.features, test_set.targets])

fc_model.evaluate(test_set.features, test_set.targets, 4096)

"""
Explanation: Train 2015, predict 2016
End of explanation
"""


preds = fc_model.predict(test_set.features)
results_df = create_results_df(test_set.date_strs, test_set.station_ids,
                               preds[:, 0], preds[:, 1])
results_df.to_csv(results_dir + 'fc_network_train_2015_pred_2016.csv')

"""
Explanation: Very similar to the standard EMOS Network. This indicates that there is not much additional information in the two extra connections we added.
End of explanation
"""


hidden_model = build_hidden_model(2, 2, hidden_nodes=10, compile=True)

hidden_model.summary()

"""
Explanation: Neural network with one hidden layer
Now we will build the first neural network with a hidden layer and a non-linear activation function. We will restrict ourselves to testing the 2015 training, 2016 prediction case.
Build network
End of explanation
"""


hidden_model.compile(keras.optimizers.Adam(0.0001), loss=crps_cost_function)

# We can use the same data from above!
# Note I am running this cell several times
hidden_model.fit(train_set.features, train_set.targets, epochs=10, batch_size=1024,
                 validation_data=[test_set.features, test_set.targets])

"""
Explanation: Train 2015, predict 2016
End of explanation
"""


preds = hidden_model.predict(test_set.features)
results_df = create_results_df(test_set.date_strs, test_set.station_ids,
                               preds[:, 0], preds[:, 1])
results_df.to_csv(results_dir + 'hidden_nn_train_2015_pred_2016.csv')

"""
Explanation: Again, the results are pretty similar. This indicates that for the given data, the added nonlinearity is not important.
End of explanation
"""


hidden_model = build_hidden_model(2, 2, hidden_nodes=[100, 100, 100], compile=True)

hidden_model.summary()

hidden_model.compile(keras.optimizers.Adam(0.0001), loss=crps_cost_function)

hidden_model.fit(train_set.features, train_set.targets, epochs=10, batch_size=4096,
                 validation_data=[test_set.features, test_set.targets])

"""
Explanation: Making the hidden model more complex
Let's see what happens if we make the model more complex
End of explanation
"""


emb_size = 3
max_id = int(np.max([train_set.cont_ids.max(), test_set.cont_ids.max()]))
max_id

emb_model = build_emb_model(2, 2, [], emb_size, max_id, compile=True,
                            lr=0.01)

emb_model.summary()

"""
Explanation: So we can see that even for a model with 20,000 parameters then training score only goes down a few percent. For a simple bias and spread correction, a linear model seems fully sufficient.
Add station embeddings
Next we will add a station embedding. Here we are giving every station additional parameters which the model can learn.
Build linear embedding model
Let's build a linear embedding model. I tried out hidden layers, but they seem to make the validation score worse!
End of explanation
"""


# Ran this for 40 epochs
emb_model.fit([train_set.features, train_set.cont_ids], train_set.targets, 
              epochs=10, batch_size=1024, 
              validation_data=[[test_set.features, test_set.cont_ids], test_set.targets])

preds = emb_model.predict([test_set.features, test_set.cont_ids])
results_df = create_results_df(test_set.date_strs, test_set.station_ids,
                               preds[:, 0], preds[:, 1])
results_df.to_csv(results_dir + 'embedding_fc_train_2015_pred_2016.csv')

"""
Explanation: Train 2015, predict 2016
End of explanation
"""


def build_and_run_emb_model(emb_size):
    emb_model = build_emb_model(2, 2, [], emb_size, max_id, compile=True)
    emb_model.fit([train_set.features, train_set.cont_ids], train_set.targets, 
                  epochs=40,batch_size=1024, verbose=0,
                  validation_data=[[test_set.features, test_set.cont_ids], test_set.targets])
    print(emb_model.evaluate([train_set.features, train_set.cont_ids], train_set.targets, verbose=0),
          emb_model.evaluate([test_set.features, test_set.cont_ids], test_set.targets, verbose=0))

for emb_size in [1, 2, 3, 5, 10, 20]:
    print(emb_size)
    build_and_run_emb_model(emb_size)

"""
Explanation: Embedding size hyper-parameter tuning
Since embeddings appear to work very well, we will test the impact of the embedding size before building more complex models. Of course, a larger embedding size might be useful when adding more variables, but this should give us some feeling.
End of explanation
"""


# The prepare_data function takes an ordered dict as an input
aux_dict = OrderedDict()
aux_dict['data_aux_geo_interpolated.nc'] = ['orog', 
                                            'station_alt', 
                                            'station_lat', 
                                            'station_lon']
aux_dict['data_aux_pl500_interpolated_00UTC.nc'] = ['u_pl500_fc',
                                                    'v_pl500_fc',
                                                    'gh_pl500_fc']
aux_dict['data_aux_pl850_interpolated_00UTC.nc'] = ['u_pl850_fc',
                                                    'v_pl850_fc',
                                                    'q_pl850_fc']
aux_dict['data_aux_surface_interpolated_00UTC.nc'] = ['cape_fc',
                                                      'sp_fc',
                                                      'tcc_fc']

train_dates = ['2015-01-01', '2016-01-01']
test_dates =  ['2016-01-01', '2017-01-01']
train_set, test_set = get_train_test_sets(DATA_DIR, train_dates, test_dates,
                                         aux_dict=aux_dict)

train_set.features.shape

"""
Explanation: Note that there is some variability. In our first experiment above with an embedding size of 5 we got a better score than here. For this very simple network an embedding size of three seems sufficient.
Adding auxiliary variables
Now we can try adding additional variables.
Load extended dataset
Using the function defined in utils.
End of explanation
"""


fc_model = build_fc_model(train_set.features.shape[1], 2, compile=True, 
                          lr=0.01)

# Note that I am running this cell multiple times
fc_model.fit(train_set.features, train_set.targets, epochs=10, batch_size=1024,
             validation_data=[test_set.features, test_set.targets])

preds = fc_model.predict([test_set.features, test_set.cont_ids])
results_df = create_results_df(test_set.date_strs, test_set.station_ids,
                               preds[:, 0], preds[:, 1])
results_df.to_csv(results_dir + 'embedding_fc_train_2015_pred_2016.csv')

"""
Explanation: Linear model
End of explanation
"""


hidden_model = build_hidden_model(train_set.features.shape[1], 2, 
                                  hidden_nodes=[50], compile=True)

hidden_model.summary()

# Note that I am running this cell multiple times
hidden_model.fit(train_set.features, train_set.targets, epochs=10, batch_size=1024,
             validation_data=[test_set.features, test_set.targets])

"""
Explanation: Hidden model
End of explanation
"""


more_aux_dict = aux_dict

more_aux_dict['data_aux_surface_more_interpolated_part1_00UTC.nc']  = [
    'sshf_fc', 'slhf_fc', 'u10_fc','v10_fc'
]
more_aux_dict['data_aux_surface_more_interpolated_part2_00UTC.nc']  = [
    'ssr_fc', 'str_fc', 'd2m_fc','sm_fc'
]

train_dates = ['2015-01-01', '2016-01-01']
test_dates =  ['2016-01-01', '2017-01-01']

more_train_set, more_test_set = get_train_test_sets(DATA_DIR, train_dates, test_dates,
                                         aux_dict=more_aux_dict)

more_train_set.features.shape

"""
Explanation: So we see a definite improvement using auxiliary variables. Again, the hidden layer does not seem to improve things a lot compared to the simple linear model.
Even more variables
End of explanation
"""


import pickle

%ls $DATA_DIR

with open(DATA_DIR + 'aux_15_16.pkl', 'wb') as f:
    pickle.dump((more_train_set, more_test_set), f)

def save_pickle(fn, train_dates=['2015-01-01', '2016-01-01'], add_current_error=False,
                current_error_len=1):
    sets = get_train_test_sets(
        DATA_DIR, train_dates, test_dates, aux_dict=more_aux_dict,
        add_current_error=add_current_error, current_error_len=current_error_len
    )
    with open(DATA_DIR + fn, 'wb') as f:
        pickle.dump(sets, f)

save_pickle('aux_15_16_current30.pkl', ['2015-01-01', '2016-01-01'], 
            add_current_error=True, current_error_len=30)

save_pickle('aux_14-15_16.pkl', ['2014-01-01', '2016-01-01'])
save_pickle('aux_13-15_16.pkl', ['2013-01-01', '2016-01-01'])
save_pickle('aux_10-15_16.pkl', ['2010-01-01', '2016-01-01'])
save_pickle('aux_08-15_16.pkl', ['2008-01-01', '2016-01-01'])

save_pickle('aux_12-15_16.pkl', ['2012-01-01', '2016-01-01'])
save_pickle('aux_11-15_16.pkl', ['2011-01-01', '2016-01-01'])

"""
Explanation: Save pickled datasets
End of explanation
"""


fc_model = build_fc_model(more_train_set.features.shape[1], 2, compile=True, 
                          lr=0.01)

# Note that I am running this cell multiple times
fc_model.fit(more_train_set.features, more_train_set.targets, epochs=10, batch_size=1024,
             validation_data=[more_test_set.features, more_test_set.targets])

"""
Explanation: Train model
End of explanation
"""


emb_model = build_emb_model(train_set.features.shape[1], 2, [], 3, max_id, 
                            compile=True, lr=0.01)

emb_model.summary()

emb_model.optimizer.lr=0.001

# Again I am running this multiple times
emb_model.fit([train_set.features, train_set.cont_ids], train_set.targets, epochs=10, 
              batch_size=1024, 
              validation_data=[[test_set.features, test_set.cont_ids], test_set.targets])

"""
Explanation: Adding these extra variables gets us another percent or so
Additional variables with the embedding model
Linear model with embeddings
End of explanation
"""


emb_model = build_emb_model(train_set.features.shape[1], 2, [50], 3, max_id, 
                            compile=True, lr=0.01)

# Again I am running this multiple times
emb_model.fit([train_set.features, train_set.cont_ids], train_set.targets, epochs=10, 
              batch_size=4096, 
              validation_data=[[test_set.features, test_set.cont_ids], test_set.targets])

"""
Explanation: Hidden model with embeddings
End of explanation
"""


max_id = int(np.max([train_set.cont_ids.max(), test_set.cont_ids.max()]))

emb_model = build_emb_model(train_set.features.shape[1], 2, [50], 3, max_id, 
                            compile=True, lr=0.01)

# Again I am running this multiple times
emb_model.fit([train_set.features, train_set.cont_ids], train_set.targets, epochs=50, 
              batch_size=4096, validation_split=0.2,
              callbacks=[EarlyStopping(monitor='val_loss', 
                                       min_delta=0,
                                       patience=2)])

emb_model.evaluate([test_set.features, test_set.cont_ids], test_set.targets, batch_size=10000)

preds = emb_model.predict([test_set.features, test_set.cont_ids])
results_df = create_results_df(test_set.date_strs, test_set.station_ids,
                               preds[:, 0], preds[:, 1])
results_df.to_csv(results_dir + 'embedding_nn_aux_train_2015_pred_2016.csv')

"""
Explanation: This is our best score so far. Here the non-linearity seems to make a difference compared to the simple linear model. But we do get some overfitting. Let's try out some techniques. Fewer or more hidden nodes does not seem to change all that much.
No we need to be very careful here. I am currently stopping when the validation score does not decrease further. THIS IS CHEATING!
Let's try doing the train valid split just with the training set, and see if we can get a good early stopping point.
End of explanation
"""


long_train_dates = ['2008-01-01', '2016-01-01']
test_dates =  ['2016-01-01', '2017-01-01']
long_train_set, test_set = get_train_test_sets(DATA_DIR, long_train_dates, test_dates,
                                               aux_dict=aux_dict)

emb_model = build_emb_model(long_train_set.features.shape[1], 2, [50], 3, max_id, 
                            compile=True, lr=0.01)

# Again I am running this multiple times
emb_model.fit([long_train_set.features, long_train_set.cont_ids], long_train_set.targets, 
              epochs=50, 
              batch_size=4096, validation_split=0.2,
              callbacks=[EarlyStopping(monitor='val_loss', 
                                       min_delta=0,
                                       patience=5)]);

emb_model.evaluate([test_set.features, test_set.cont_ids], test_set.targets, batch_size=10000)

preds = emb_model.predict([test_set.features, test_set.cont_ids])
results_df = create_results_df(test_set.date_strs, test_set.station_ids,
                               preds[:, 0], preds[:, 1])
results_df.to_csv(results_dir + 'embedding_nn_aux_train_2008-2015_pred_2016.csv')

"""
Explanation: A longer training period
End of explanation
"""


long_more_train_set, more_test_set = get_train_test_sets(DATA_DIR, long_train_dates, test_dates,
                                                         aux_dict=more_aux_dict)

emb_model = build_emb_model(long_more_train_set.features.shape[1], 2, [100], 3, max_id, 
                            compile=True, lr=0.01)

# Again I am running this multiple times
emb_model.fit([long_more_train_set.features, long_more_train_set.cont_ids], 
              long_more_train_set.targets, 
              epochs=50, batch_size=4096, validation_split=0.2,
              callbacks=[EarlyStopping(monitor='val_loss', 
                                       min_delta=0,
                                       patience=5)])

emb_model.evaluate([more_test_set.features, more_test_set.cont_ids], more_test_set.targets, 
                   batch_size=10000)

preds = emb_model.predict([more_test_set.features, more_test_set.cont_ids])
results_df = create_results_df(test_set.date_strs, test_set.station_ids,
                               preds[:, 0], preds[:, 1])
results_df.to_csv(results_dir + 'embedding_nn_more_aux_train_2008-2015_pred_2016.csv')

"""
Explanation: A longer training period with more data
End of explanation
"""


with open(DATA_DIR + 'aux_15_16.pkl', 'rb') as f:
    train_set, test_set = pickle.load(f)

train_set.features.shape

train_set.features.std(axis=0)

train_set.features.mean(axis=0)[1::2]

train_set.features.std(axis=0)[::2]

scales = np.zeros(train_set.features.shape[1])
scales[1::2] = train_set.features.mean(axis=0)[1::2]
scales[::2] = train_set.features.std(axis=0)[::2]

features_aug = train_set.features + np.random.normal(size=train_set.features.shape) * scales * 0.05

targets_aug = train_set.targets + np.random.normal(scale=0.1, size=train_set.targets.shape)

train_set.targets[:5], targets_aug[:5]

np.concatenate([train_set.features, features_aug], axis=0).shape

emb_model = build_emb_model(train_set.features.shape[1], 2, [100], 3, max_id, 
                            compile=True, lr=0.01)

emb_model.fit([train_set.features, train_set.cont_ids], train_set.targets, 
              epochs=50, batch_size=4096, 
              validation_data=[[test_set.features, test_set.cont_ids], test_set.targets])

emb_model.fit([
    np.concatenate([train_set.features, features_aug], axis=0), 
    np.concatenate([train_set.cont_ids, train_set.cont_ids])
                   ], np.concatenate([train_set.targets, targets_aug]), 
              epochs=50, batch_size=4096, 
              validation_data=[[test_set.features, test_set.cont_ids], test_set.targets])

# Test current error
train_dates = ['2015-01-01', '2016-01-01']
test_dates =  ['2016-01-01', '2017-01-01']
train_set, test_set = get_train_test_sets(DATA_DIR, train_dates, test_dates,
                                          add_current_error=True, 
                                          current_error_len=3)

train_set.features.shape

train_set.feature_names

"""
Explanation: Data augmentation
Let's pick a setup where overfitting is a problem: emb_train_2015_aux
End of explanation
"""


