# Load modules
from __future__ import print_function
import os
import pandas as pd
import numpy as np
from matplotlib import pyplot as plt

# read the data from into a pandas.DataFrame
wisc_emp = pd.read_csv("datasets/wisconsin-employment-time-series.csv")

# Let's find out the shape of the DataFrame
print("Shape of the DataFrame:", wisc_emp.shape)

# Let's see first 10 rows of it
wisc_emp.head()

# plot the wisconsin employment dataset
wisc_emp.plot()

# Capture seasonality component
def initialize_T(x, seasonLength):
    total = 0.0
    for i in range(seasonLength):
        total += float(x[i + seasonLength] - x[i]) / seasonLength
    return total


initialize_T(wisc_emp["Employment"], 12)

# Initialize seasonal trend
def initialize_seasonalilty(x, seasonLength):
    seasons = {}
    seasonsMean = []
    num_season = int(len(x) / seasonLength)
    # Compute season average
    for i in range(num_season):
        seasonsMean.append(
            sum(x[seasonLength * i : seasonLength * i + seasonLength])
            / float(seasonLength)
        )

    # compute season intial values
    for i in range(seasonLength):
        tot = 0.0
        for j in range(num_season):
            tot += x[seasonLength * j + i] - seasonsMean[j]
        seasons[i] = tot / num_season
    return seasons


initialize_seasonalilty(wisc_emp["Employment"], 12)


# Triple Exponential Smoothing Forecast
def triple_exp_smoothing(x, seasonLength, alpha, beta, gamma, h):
    yhat = []
    S = initialize_seasonalilty(x, seasonLength)
    for i in range(len(x) + h):
        if i == 0:
            F = x[0]
            T = initialize_T(x, seasonLength)
            yhat.append(x[0])
            continue
        if i >= len(x):
            m = i - len(x) + 1
            yhat.append((F + m * T) + S[i % seasonLength])
        else:
            obsval = x[i]
            F_last, F = (
                F,
                alpha * (obsval - S[i % seasonLength]) + (1 - alpha) * (F + T),
            )
            T = beta * (F - F_last) + (1 - beta) * T
            S[i % seasonLength] = (
                gamma * (obsval - F) + (1 - gamma) * S[i % seasonLength]
            )
            yhat.append(F + T + S[i % seasonLength])
    return yhat


# Triple exponential smoothing
wisc_emp["TES"] = triple_exp_smoothing(wisc_emp["Employment"], 12, 0.4, 0.6, 0.2, 0)

### Plot Single Exponential Smoothing forecasted value
fig = plt.figure(figsize=(5.5, 5.5))
ax = fig.add_subplot(2, 1, 1)
wisc_emp["Employment"].plot(ax=ax)
ax.set_title("Beer Production")
ax = fig.add_subplot(2, 1, 2)
wisc_emp["TES"].plot(ax=ax, color="r")
ax.set_title("Triple Smoothing Forecast")
plt.savefig("plots/ch2/B07887_03_14.png", format="png", dpi=300)
