import numpy as np
import pandas as pd

from statsforecast.core import StatsForecast
from statsforecast.models import AutoETS

from hierarchicalforecast.core import HierarchicalReconciliation
from hierarchicalforecast.evaluation import HierarchicalEvaluation
from hierarchicalforecast.methods import TopDown

from hierarchicalforecast import utils

# Parameters for dataset creation
n_rows = 1000
date_range = pd.date_range(start='2020-01-01', periods=n_rows, freq='MS')
group_col_one_values = ['A', 'B', 'C']
group_col_two_values = ['X', 'Y', 'Z']

# Create the dataset
data = pd.DataFrame({
    'group_col_one': np.random.choice(group_col_one_values, size=n_rows),
    'group_col_two': np.random.choice(group_col_two_values, size=n_rows),
    'ds': date_range,
    'y': np.random.randint(1, 100, size=n_rows)
})

# Create Top Level to Generate Forecasts
top_data = data.groupby(by=['group_col_one', 'ds'])['y'].sum().reset_index()
Y_top_df, S_top_df, tags_top = utils.aggregate(top_data, [['group_col_one']])

# Create Historical Values of Both 'Top' & 'Bottom'
Y_hist_df, S_hist_df, tags_hist = utils.aggregate(data, [['group_col_one'], ['group_col_one', 'group_col_two']])

# Produce Top Level Forecasts to use for Disaggregation
fcst = StatsForecast(models=[AutoETS(season_length=12)],
                     freq='MS')
Y_hat_df = fcst.forecast(h=12, df=Y_top_df)
reconcilers = [
            TopDown(method='proportion_averages'),
        ]
hrec = HierarchicalReconciliation(reconcilers=reconcilers)
Y_rec_df = hrec.reconcile(Y_hat_df, S_hist_df, tags_hist, Y_hist_df)