#!/usr/bin/env python
# coding: utf-8

# # Nile data

# In[331]:


from __future__ import print_function
get_ipython().run_line_magic('matplotlib', 'inline')
import numpy as np
import pandas as pd
import statsmodels.api as sm
from statsmodels.tsa.statespace import structural
import matplotlib.pyplot as plt


# In[276]:


# Get dataset
nile = sm.datasets.nile.data.load_pandas().data
nile.index = pd.date_range('1871', '1970', freq='AS')


# Create exogenous array for intervention in 1899
# and for pulse in 1913
nile['intervention'] = 0
nile.ix['1899':, 'intervention'] = 1
nile['pulse'] = 0
nile.ix['1913', 'pulse'] = 1


# In[327]:


nile['volume'].plot(figsize=(10,3));


# ## Local level model

# In[330]:


mod_ll = structural.UnobservedComponents(
    nile['volume'], exog=nile[['intervention', 'pulse']],
    level=True, stochastic_level=True,
)
res_ll = mod_ll.fit()
print(res_ll.summary())
res_ll.plot_diagnostics(figsize=(11,6));


# In[329]:


fig, axes = plt.subplots(2, 1, figsize=(10,6))

dates = nile.index._mpl_repr()

# Filtered level
ax = axes[0]
ax.plot(dates, nile['volume'], 'k.')
ax.plot(dates, res_ll.filtered_state[0] + res_ll.obs_intercept[0])
ax.set_title('Filtered level');

# Residuals
ax = axes[1]
ax.plot(dates, nile['volume'] - res_ll.filtered_state[0]- res_lc.obs_intercept[0])
ax.set_title('Filtered residuals');


# ## Deterministic level + stochastic cycle

# In[302]:


# Create the model
mod_lc = structural.UnobservedComponents(
    nile['volume'], exog=nile[['intervention','pulse']],
    irregular=False,
    level=True, stochastic_level=False,
    cycle=True, stochastic_cycle=True, damped_cycle=True
)
# Fit the model
res_lc = mod_lc.fit()
print(res_lc.summary())
res_lc.plot_diagnostics(figsize=(11,6));


# In[323]:


fig, axes = plt.subplots(3, 1, figsize=(10,9))

dates = nile.index._mpl_repr()

# Plot the level
ax = axes[0]
ax.plot(dates, nile['volume'], 'k.')
ax.plot(dates, res_lc.filtered_state[0] + res_lc.obs_intercept[0])
ax.set_title('Filtered level');

# Plot the cycle
ax = axes[1]
ax.plot(dates, res_lc.filtered_state[1])
ax.set_title('Filtered cycle');

# Plot the filtered residuals
ax = axes[2]
ax.plot(dates, nile['volume'] - res_lc.filtered_state[0] - res_lc.filtered_state[1] - res_lc.obs_intercept[0])
ax.set_title('Filtered residuals')
ax.set_ylim(-300,400);