from statsforecast import StatsForecast
from statsforecast.models import HoltWinters
from statsforecast.utils import ConformalIntervals
import numpy as np
import polars as pl

# Polars option to display all rows
pl.Config.set_tbl_rows(None)

# Initialize the models
models = [
    HoltWinters(season_length=4)
]

# Initialize the StatsForecast model
sf = StatsForecast(models=models, freq='1w', n_jobs=-1)

# Perform cross-validation with a step size of 1 to mimic an expanding window
crossvalidation_df = sf.cross_validation(df=y_cl4_over_19, h=4, step_size=1, n_windows=5, sort_df=True)

# Calculate the ensemble mean
ensemble = crossvalidation_df[['HoltWinters']].mean(axis=1)

# Create a Series for the ensemble mean
ensemble_series = pl.Series('Ensemble', ensemble)

# Add the ensemble mean as a new column to the DataFrame
crossvalidation_df = crossvalidation_df.with_columns(ensemble_series)

def wmape(y_true, y_pred):
    return np.abs(y_true - y_pred).sum() / np.abs(y_true).sum()

# Calculate the WMAPE for the ensemble model
wmape_value = wmape(crossvalidation_df['y'], crossvalidation_df['Ensemble'])
print('Average WMAPE for Ensemble: ', round(wmape_value, 4))

# Calculate the errors for the ensemble model
errors = crossvalidation_df['y'] - crossvalidation_df['Ensemble']

# For an individual forecast
individual_accuracy = 1 - (abs(crossvalidation_df['y'] - crossvalidation_df['Ensemble']) / crossvalidation_df['y'])
individual_bias = (crossvalidation_df['Ensemble'] / crossvalidation_df['y']) - 1

# Add these calculations as new columns to DataFrame
crossvalidation_df = crossvalidation_df.with_columns([
    individual_accuracy.alias("individual_accuracy"),
    individual_bias.alias("individual_bias")
])

# Print the individual accuracy and bias for each week
for row in crossvalidation_df.to_dicts():
    id = row['unique_id']
    date = row['ds'] 
    accuracy = row['individual_accuracy']
    bias = row['individual_bias']
    print(f"{id}, {date}, Individual Accuracy: {accuracy:.4f}, Individual Bias: {bias:.4f}")

# For groups of forecasts
group_accuracy = 1 - (errors.abs().sum() / crossvalidation_df['y'].sum())
group_bias = (crossvalidation_df['Ensemble'].sum() / crossvalidation_df['y'].sum()) - 1

# Print the average group accuracy and group bias over all folds for the ensemble model
print('Average Group Accuracy: ', round(group_accuracy, 4))
print('Average Group Bias: ', round(group_bias, 4))
 12519007_US01_8_1108349, 2023-11-13 00:00:00, Individual Accuracy: -0.4972, Individual Bias: 1.4972
12519007_US01_8_1108349, 2023-11-20 00:00:00, Individual Accuracy: 0.4678, Individual Bias: -0.5322
12519007_US01_8_1108349, 2023-11-06 00:00:00, Individual Accuracy: 0.9100, Individual Bias: -0.0900
12519007_US01_8_1108349, 2023-11-13 00:00:00, Individual Accuracy: -1.7148, Individual Bias: 2.7148
12519007_US01_8_1108349, 2023-11-20 00:00:00, Individual Accuracy: 0.6192, Individual Bias: -0.3808
12519007_US01_8_1108349, 2023-11-27 00:00:00, Individual Accuracy: -5.4835, Individual Bias: 6.4835
12519007_US01_8_5338737, 2023-07-31 00:00:00, Individual Accuracy: -97.3718, Individual Bias: 98.3718
12519007_US01_8_5338737, 2023-08-14 00:00:00, Individual Accuracy: -8.2802, Individual Bias: 9.2802
12519007_US01_8_5338737, 2023-08-28 00:00:00, Individual Accuracy: -14.6154, Individual Bias: 15.6154
12519007_US01_8_5338737, 2023-09-04 00:00:00, Individual Accuracy: 0.3744, Individual Bias: -0.6256
12519007_US01_8_5338737, 2023-08-14 00:00:00, Individual Accuracy: -4.6083, Individual Bias: 5.6083
12519007_US01_8_5338737, 2023-08-28 00:00:00, Individual Accuracy: -10.1394, Individual Bias: 11.1394 there are negative values in here leading to big spread on inidividual accuracy and bias calculations, is there a way to put a constraint on group accuracy and bias calculation to ignore individual accurray and individual bias with values greater than -15 or 15? , only at the group level # For groups of forecasts
group_accuracy = 1 - (errors.abs().sum() / crossvalidation_df['y'].sum())
group_bias = (crossvalidation_df['Ensemble'].sum() / crossvalidation_df['y'].sum()) - 1