""" Issues: "ImportError: cannot import name '_maybe_get_pandas_wrapper_freq' from 'statsmodels.tsa.filters._utils' " install an old version of statsmodel works "pip install statsmodels==0.10.2" Conference: 1. Cleveland et al. 1990 [ https://www.wessa.net/download/stl.pdf ] 2. https://github.com/jrmontag/STLDecompose """

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import sys
sys.path.append('../lib')
import statsmodels.api as sm
from stldecompose.stl import decompose, forecast
from stldecompose.forecast_funcs import (naive,drift,mean,seasonal_naive)



np.random.seed(42)
x = np.arange(0, 300)
idx = pd.date_range('2018-01-01', periods = len(x))
season = pd.Series(((x % 7) - 3) * 1.5, idx)
trend = pd.Series(x * 0.15, idx) - pd.Series(np.power(x * 0.022, 2), idx)
error = pd.Series(np.random.random(len(x))*7, idx)
obs = season + trend + error

print("len(obs):", len(obs))

decomp = decompose(obs, period= 30)

decomp.plot()
plt.show()

short_obs = obs[ : int(len(obs) // 3 * 2)]
print("len(short_obs):", len(short_obs))

# apply the decomp to the truncated observation
short_decomp = decompose(short_obs, period=30)

fcast = forecast(short_decomp, steps=100, fc_func=drift)

plt.figure(figsize=(12, 6))
plt.plot(obs, '--', label='truth')
plt.plot(short_obs, '--', label='obs')
plt.plot(short_decomp.trend, ':', label='decomp.trend')
plt.plot(fcast, '-', label=fcast.columns[0])
# plt.xlim('1970','2004')
# plt.ylim(330,380)
plt.legend()
plt.show()

plt.figure(figsize=(12,6))
fcast = forecast(short_decomp, steps=100, fc_func=drift, seasonal=True)
plt.plot(obs, '--', label='truth')
plt.plot(short_obs, '--', label='obs')
plt.plot(short_decomp.trend, ':', label='decomp.trend')
plt.plot(fcast, '-', label=fcast.columns[0])
# plt.xlim('1970','2004')
# plt.ylim(330,380)
plt.legend()
plt.show()