"""
Vector Autogregression using statsmodels
http://statsmodels.sourceforge.net/devel/vector_ar.html
"""

import numpy as np
import pandas as pd
from statsmodels.tsa.api import *
from statsmodels.tsa.vector_ar.var_model import VAR
from datetime import datetime
import matplotlib.pyplot as plt
from infodenguepredict.data.infodengue import get_alerta_table, build_multicity_dataset


def build_model(data):
    data.index = pd.DatetimeIndex(data.index)
    model = VAR(data)
    return model


if __name__ == "__main__":
    prediction_window = 5  # weeks
    scenario = 'global'
    if scenario == 'local':
        data = get_alerta_table(3303500)  # Nova Iguaçu: 3303500
        data = data[['casos', 'nivel']]
    else:
        data = build_multicity_dataset('RJ')
        data = data[[col for col in data.columns if col.startswith('casos') and not col.startswith('casos_est')]][:5]
    print(data.info())
    # data.casos_est.plot(title="Series")
    model = build_model(data)
    fit = model.fit(maxlags=11, ic='aic') # 4 lags
    print(fit.summary())
    fit.plot()
    fit.plot_acorr()

    plt.figure()
    lag_order = fit.k_ar
    forecast = fit.forecast(data.values[-lag_order:], prediction_window)
    print(forecast)
    fit.plot_forecast(prediction_window)

    plt.show()
