import pandas as pd
import numpy as np
from statsmodels.tsa.holtwinters import ExponentialSmoothing
from sklearn.metrics import mean_absolute_percentage_error

print("Loading data...")
df = pd.read_csv('AstraZeneca_Cleaned_Data.csv', low_memory=False)
df['Pstng Date'] = pd.to_datetime(df['Pstng Date'], errors='coerce')
df = df.dropna(subset=['Pstng Date'])

# --- Helper Function for Forecasting ---
def generate_forecast(series, periods=26, seasonal_periods=52):
    # Try seasonal model first
    try:
        model = ExponentialSmoothing(
            series, 
            trend='add', 
            seasonal='add', 
            seasonal_periods=seasonal_periods,
            initialization_method='estimated'
        )
        fit = model.fit()
        forecast = fit.forecast(periods)
        model_type = "Exponential Smoothing (Seasonal)"
    except Exception as e:
        print(f"  [Info] Seasonal model failed ({e}), falling back to Simple ES (Trend only).")
        model = ExponentialSmoothing(series, trend='add')
        fit = model.fit()
        forecast = fit.forecast(periods)
        model_type = "Simple Exponential Smoothing"
        
    return fit, forecast, model_type

# --- 1. Baseline Forecast ---
print("\n--- Generating Baseline Forecast ---")
baseline_weekly = df.resample('W', on='Pstng Date')['Amount in USD'].sum().fillna(0) # Ensure no gaps

fit_base, fc_base, type_base = generate_forecast(baseline_weekly)

# Calculate MAPE on training data
fitted_values = fit_base.fittedvalues
mape = mean_absolute_percentage_error(baseline_weekly, fitted_values) * 100

# Approximate Confidence Intervals (using residual std dev)
residuals = baseline_weekly - fitted_values
std_resid = residuals.std()
# 95% CI is approx +/- 1.96 * std
lower_ci = fc_base - (1.96 * std_resid)
upper_ci = fc_base + (1.96 * std_resid)

baseline_df = pd.DataFrame({
    'Forecast_Week': pd.date_range(start=baseline_weekly.index.max() + pd.Timedelta(weeks=1), periods=26, freq='W'),
    'Predicted_Net_Cash': fc_base,
    'Lower_CI_95': lower_ci,
    'Upper_CI_95': upper_ci,
    'Model_Type': type_base
})

# --- 2. What-If Scenario: Delayed AR (14 Days) ---
print("\n--- Simulating What-If Scenario: Delayed AR (14 Days) ---")
# Logic: Take positive amounts (Inflows) and shift their date by +14 days
df_scenario = df.copy()
inflow_mask = df_scenario['Amount in USD'] > 0
df_scenario.loc[inflow_mask, 'Pstng Date'] = df_scenario.loc[inflow_mask, 'Pstng Date'] + pd.Timedelta(days=14)

# Re-aggregate
scenario_weekly = df_scenario.resample('W', on='Pstng Date')['Amount in USD'].sum().fillna(0)

# Forecast on new data
fit_scen, fc_scen, type_scen = generate_forecast(scenario_weekly)

scenario_df = pd.DataFrame({
    'Forecast_Week': pd.date_range(start=scenario_weekly.index.max() + pd.Timedelta(weeks=1), periods=26, freq='W'),
    'Predicted_Net_Cash_Delayed': fc_scen,
    'Model_Type': type_scen
})

# --- 3. Export to Excel ---
print("\n--- Exporting Advanced Forecast ---")
metrics_df = pd.DataFrame({
    'Metric': ['MAPE (Training Error)', 'Residual Std Dev', 'Model Used'],
    'Value': [f"{mape:.2f}%", f"{std_resid:.2f}", type_base],
    'Description': [
        'Mean Absolute Percentage Error - Lower is better.',
        'Standard deviation of prediction errors - used for Confidence Intervals.',
        'Type of forecasting model selected based on data availability.'
    ]
})

output_file = 'Advanced_Forecast.xlsx'
try:
    with pd.ExcelWriter(output_file, engine='openpyxl') as writer:
        baseline_df.to_excel(writer, sheet_name='Baseline_Forecast', index=False)
        scenario_df.to_excel(writer, sheet_name='Delayed_AR_Scenario', index=False)
        metrics_df.to_excel(writer, sheet_name='Model_Metrics', index=False)
    print(f"Success! Advanced forecast saved to '{output_file}'.")
    print("Sheets: 'Baseline_Forecast', 'Delayed_AR_Scenario', 'Model_Metrics'")
except PermissionError:
    print(f"\n[ERROR] Could not save '{output_file}'. Please close it and run the script again.")
except Exception as e:
    print(f"\n[ERROR] Unexpected error saving file: {e}")