import pandas as pd
import matplotlib.pyplot as plt
from statsmodels.tsa.holtwinters import ExponentialSmoothing
from datetime import timedelta

df = pd.read_csv("norway_covid_2020_2024.csv")
df['date'] = pd.to_datetime(df['date'])

df_full = df[['date', 'new_cases_smoothed']].dropna().sort_values('date')

#  training data: 2020-01-01 to 2021-01-01
df_train = df_full[(df_full['date'] >= '2020-01-01') & (df_full['date'] <= '2021-01-01')].copy()
df_train.set_index('date', inplace=True)

# HWES 
model = ExponentialSmoothing(
    df_train['new_cases_smoothed'],
    trend='add',
    seasonal='add',
    seasonal_periods=7  # assuming weekly seasonality
)
model_fit = model.fit()

forecast_horizon = 365
forecast = model_fit.forecast(steps=forecast_horizon)
last_date = df_train.index[-1]
forecast_index = [last_date + timedelta(days=i) for i in range(1, forecast_horizon + 1)]

start_forecast = forecast_index[0]
end_forecast = forecast_index[-1]
df_actual = df_full[(df_full['date'] >= start_forecast) & (df_full['date'] <= end_forecast)].copy()
df_actual.set_index('date', inplace=True)

plt.figure(figsize=(12, 6))
plt.plot(df_train.index, df_train['new_cases_smoothed'], label='Training data (2020)', color='blue')
plt.plot(forecast_index, forecast, label='HWES forecast (2021)', linestyle='--', color='orange')
plt.plot(df_actual.index, df_actual['new_cases_smoothed'], label='Actual values (2021)', linestyle=':', color='green')
plt.xlabel('Date')
plt.ylabel('New COVID cases (smoothed)')
plt.title('COVID-19 Cases in Norway: HWES Forecast for 2021 Based on 2020 Data')
plt.legend()
plt.grid(True)
plt.tight_layout()
plt.savefig('hwes_forecast_2021.png', dpi=400)
plt.show()
