import numpy as np
import pandas as pd

from statsforecast.core import StatsForecast
from statsforecast.models import ETS, Naive

from hierarchicalforecast.utils import aggregate
from hierarchicalforecast.core import HierarchicalReconciliation
from hierarchicalforecast.methods import BottomUp, MinTrace

# Load TourismSmall dataset
df = pd.read_csv('<https://raw.githubusercontent.com/Nixtla/transfer-learning-time-series/main/datasets/tourism.csv>')
df = df.rename({'Trips': 'y', 'Quarter': 'ds'}, axis=1)
df.insert(0, 'Country', 'Australia')

# Create hierarchical seires based on geographic levels and purpose
# And Convert quarterly ds string to pd.datetime format
hierarchy_levels = [['Country'],
                    ['Country', 'State'], 
                    ['Country', 'Purpose'], 
                    ['Country', 'State', 'Region'], 
                    ['Country', 'State', 'Purpose'], 
                    ['Country', 'State', 'Region', 'Purpose']]

Y_df, S, tags = aggregate(df=df, spec=hierarchy_levels)
qs = Y_df['ds'].str.replace(r'(\d+) (Q\d)', r'\1-\2', regex=True)
Y_df['ds'] = pd.PeriodIndex(qs, freq='Q').to_timestamp()
Y_df = Y_df.reset_index()

# Split train/test sets
Y_test_df  = Y_df.groupby('unique_id').tail(4)
Y_train_df = Y_df.drop(Y_test_df.index)

# Compute base auto-ETS predictions
# Careful identifying correct data freq, this data quarterly 'Q'
fcst = StatsForecast(df=Y_train_df,
                     #models=[ETS(season_length=12), Naive()],
                     models=[Naive()],
                     freq='Q', n_jobs=-1) 
Y_hat_df = fcst.forecast(h=4)

# Reconcile the base predictions
Y_train_df = Y_train_df.reset_index().set_index('unique_id')
Y_hat_df = Y_hat_df.reset_index().set_index('unique_id')
reconcilers = [BottomUp(),
               MinTrace(method='ols')]
hrec = HierarchicalReconciliation(reconcilers=reconcilers)
Y_rec_df = hrec.reconcile(Y_hat_df=Y_hat_df, Y_df=Y_train_df,
                          S=S, tags=tags)
Y_rec_df.groupby('unique_id').head(2)