HEADER = ['time_index','value']
TIME_INDEX_FEATURE_NAME = 'time_index'
VALUE_FEATURE_NAMES = ['value']

"""
Explanation: Steps to use the ARRegressor Estimator API

Define the metadata
Define a data (Numpy) input function
Define a create Estimator function
Instantiate an esimator with the run_config and required params
Train the estimaor
Evaluate the estimator
Predict using the estimator

1. Define Metadata
End of explanation
"""


def generate_input_fn(file_name, mode, header_lines=1, batch_size = None, windows_size = None, tail_count=None):
    
    dataframe = pd.read_csv(file_name, names=HEADER, skiprows=header_lines)
    

    if tail_count is not None:
        dataframe = dataframe.tail(tail_count)
    
    print("Dataset Size: {}".format(len(dataframe)))
    print("")
    
    
    data = {
        ts.TrainEvalFeatures.TIMES: dataframe[TIME_INDEX_FEATURE_NAME],
        ts.TrainEvalFeatures.VALUES: dataframe[VALUE_FEATURE_NAMES],
    }
    
    reader = ts.NumpyReader(data)
    
    num_threads = multiprocessing.cpu_count() if MULTI_THREADING else 1
    
    if mode == tf.estimator.ModeKeys.TRAIN:
        input_fn = tf.contrib.timeseries.RandomWindowInputFn(
            reader, 
            batch_size=batch_size, 
            window_size=windows_size,
            num_threads= num_threads
        )
        
    elif mode == tf.estimator.ModeKeys.EVAL:
        input_fn = tf.contrib.timeseries.WholeDatasetInputFn(reader)
    
    return input_fn

"""
Explanation: 2. Define a Data Input Function
End of explanation
"""


def create_estimator(run_config, hparams):

    estimator = ts.ARRegressor(
        periodicities= hparams.periodicities, 
        input_window_size= hparams.input_window_size, 
        output_window_size= hparams.output_window_size,
        num_features=len(VALUE_FEATURE_NAMES),
        loss=hparams.loss,
        hidden_layer_sizes = hparams.hidden_units,
        optimizer = tf.train.AdagradOptimizer(learning_rate=hparams.learning_rate),
        config=run_config
    )
    
    print("")
    print("Estimator Type: {}".format(type(estimator)))
    print("")

    return estimator

"""
Explanation: 3. Define a Create Estimator Function
End of explanation
"""


CHECKPOINT_STEPS=1000

hparams  = tf.contrib.training.HParams(
    training_steps = 10000,
    periodicities = [200],
    input_window_size = 40,
    output_window_size=10,
    batch_size = 15,
    loss = tf.contrib.timeseries.ARModel.NORMAL_LIKELIHOOD_LOSS, # NORMAL_LIKELIHOOD_LOSS | SQUARED_LOSS
    hidden_units = None,
    learning_rate = 0.1
    
)


model_dir = 'trained_models/{}'.format(MODEL_NAME)

run_config = tf.estimator.RunConfig().replace(
    save_checkpoints_steps=CHECKPOINT_STEPS,
    tf_random_seed=19830610,
    model_dir=model_dir
)
                                             
print("Model directory: {}".format(run_config.model_dir))
print("Hyper-parameters: {}".format(hparams))
print("")

train_input_fn = generate_input_fn(
    file_name=TRAIN_DATA_FILE,
    mode = tf.estimator.ModeKeys.TRAIN,
    batch_size=hparams.batch_size,
    windows_size = hparams.input_window_size + hparams.output_window_size
)

estimator = create_estimator(run_config, hparams)

"""
Explanation: 4. Instantiate and Estimator with the Required Parameters
End of explanation
"""


if not RESUME_TRAINING:
    shutil.rmtree(model_dir, ignore_errors=True)
    
tf.logging.set_verbosity(tf.logging.INFO)

time_start = datetime.utcnow() 
print("Estimator training started at {}".format(time_start.strftime("%H:%M:%S")))
print(".......................................")

estimator.train(input_fn=train_input_fn, steps=hparams.training_steps)

time_end = datetime.utcnow() 
print(".......................................")
print("Estimator training finished at {}".format(time_end.strftime("%H:%M:%S")))
print("")
time_elapsed = time_end - time_start
print("Estimator training elapsed time: {} seconds".format(time_elapsed.total_seconds()))

"""
Explanation: 5. Train the Estimator
End of explanation
"""


hparams  = tf.contrib.training.HParams(
    training_steps = 10000,
    periodicities = [200],
    input_window_size = 40,
    output_window_size=10,
    batch_size = 15,
    loss = tf.contrib.timeseries.ARModel.SQUARED_LOSS, # NORMAL_LIKELIHOOD_LOSS | SQUARED_LOSS
    hidden_units = None,
    learning_rate = 0.1
    
)

estimator = create_estimator(run_config, hparams)

eval_input_fn = generate_input_fn(
    file_name=TRAIN_DATA_FILE,
    mode = tf.estimator.ModeKeys.EVAL,
)

tf.logging.set_verbosity(tf.logging.WARN)
evaluation = estimator.evaluate(input_fn=eval_input_fn, steps=1)
print("")
print(evaluation.keys())
print("")
print("Evaluation Loss ({}) : {}".format(hparams.loss, evaluation['loss']))

def compute_rmse(a, b):
    rmse =  np.sqrt(np.sum(np.square(a - b)) / len(a))
    return round(rmse,5)

def compute_mae(a, b):
    mae =  np.sqrt(np.sum(np.abs(a - b)) / len(a))
    return round(mae,5)

x_current = evaluation['times'][0]
y_current_actual = evaluation['observed'][0].reshape(-1)
y_current_estimated = evaluation['mean'][0].reshape(-1)

rmse = compute_rmse(y_current_actual, y_current_estimated)
mae = compute_mae(y_current_actual, y_current_estimated)
print("Evaluation RMSE {}".format(rmse))
print("Evaluation MAE {}".format(mae))

plt.figure(figsize=(20, 10))

plt.title("Time Series Data")
plt.plot(x_current, y_current_actual, label='actual')
plt.plot(x_current, y_current_estimated, label='estimated')
plt.xlabel("Time Index")
plt.ylabel("Value")
plt.legend(loc=2)
plt.show()

"""
Explanation: 6. Evalute the Estimator
End of explanation
"""


FORECAST_STEPS = [10,50,100,150,200,250,300]

tf.logging.set_verbosity(tf.logging.ERROR)

eval_input_fn = generate_input_fn(
    file_name=TRAIN_DATA_FILE,
    mode = tf.estimator.ModeKeys.EVAL,
    tail_count = hparams.input_window_size + hparams.output_window_size
)


evaluation = estimator.evaluate(input_fn=eval_input_fn, steps=1)

df_test = pd.read_csv(TEST_DATA_FILE, names=['time_index','value'], header=0)
print("Test Dataset Size: {}".format(len(df_test)))
print("")

for steps in FORECAST_STEPS:

    forecasts = estimator.predict(input_fn=ts.predict_continuation_input_fn(evaluation, steps=steps))
    forecasts = tuple(forecasts)[0]
    
    x_next = forecasts['times']
    
    y_next_forecast = forecasts['mean']
    y_next_actual = df_test.value[:steps].values
    
    rmse =  compute_rmse(y_next_actual, y_next_forecast)
    mae =  compute_mae(y_next_actual, y_next_forecast)

    print("Forecast Steps {}: RMSE {} - MAE {}".format(steps,rmse,mae))

print("")
print(forecasts.keys())

plt.close('all')
plt.figure(figsize=(20, 10))

plt.title("Time Series Data")
plt.plot(x_next, y_next_actual, label='actual')
plt.plot(x_next, y_next_forecast, label='forecasted')
plt.xlabel("Time Index")
plt.ylabel("Value")
plt.legend(loc=2)
plt.show()

x_all = np.concatenate( (x_current, x_next) , axis=0)
y_actual_all = np.concatenate((y_current_actual, y_next_actual), axis=0)

plt.close('all')
plt.figure(figsize=(20, 10))

plt.title("Time Series Data")
plt.plot(x_all, y_actual_all, label='actual')
plt.plot(x_current, y_current_estimated, label='estimated')
plt.plot(x_next, y_next_forecast, label='forecasted')
plt.xlabel("Time Index")
plt.ylabel("Value")
plt.legend(loc=2)
plt.show()

"""
Explanation: 7. Predict using the Estimator
End of explanation
"""


