"""
================================================================================
PHASE 0: MODEL SELECTION & COMPARISON
================================================================================
AstraZeneca Cash Flow Forecasting - Datathon 2025

Purpose: Compare 3 forecasting models and select the best performer
Models:
  1. Moving Average (4-week) - Simple baseline
  2. Exponential Smoothing (Holt-Winters) - Classic statistical method
  3. Prophet (Facebook) - Advanced time series model

Test Data: Weeks 40-44 (last 5 weeks as validation set)
Selection Criteria: Lowest MAPE with good confidence intervals
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from datetime import datetime
import warnings
warnings.filterwarnings('ignore')

# Statistical libraries
from statsmodels.tsa.holtwinters import ExponentialSmoothing
from sklearn.metrics import mean_absolute_error, mean_squared_error

# Prophet
from prophet import Prophet

# AstraZeneca Color Palette
AZ_COLORS = {
    'mulberry': '#830051',
    'lime_green': '#C4D600',
    'navy': '#003865',
    'graphite': '#3F4444',
    'light_blue': '#68D2DF',
    'magenta': '#D0006F',
    'purple': '#3C1053',
    'gold': '#F0AB00'
}

print("=" * 80)
print("PHASE 0: MODEL SELECTION & COMPARISON")
print("=" * 80)

# =============================================================================
# LOAD DATA
# =============================================================================
print("\n[1/7] Loading weekly cash flow data...")

df = pd.read_csv('data_cleaned_weekly.csv')
df['Week_Start'] = pd.to_datetime(df['Week_Start'])
df = df.sort_values('Week_Start').reset_index(drop=True)

print(f"   Total weeks: {len(df)}")
print(f"   Date range: {df['Week_Start'].min().date()} to {df['Week_Start'].max().date()}")
print(f"   Net Cash Flow range: ${df['Net_Cash_Flow_USD'].min():,.0f} to ${df['Net_Cash_Flow_USD'].max():,.0f}")

# =============================================================================
# TRAIN/TEST SPLIT
# =============================================================================
print("\n[2/7] Splitting data for model comparison...")

# Train: Weeks 1-39, Test: Weeks 40-44
train_size = 39
train_df = df.iloc[:train_size].copy()
test_df = df.iloc[train_size:].copy()

print(f"   Training set: Weeks 1-{train_size} ({len(train_df)} weeks)")
print(f"   Test set: Weeks {train_size+1}-{len(df)} ({len(test_df)} weeks)")
print(f"   Train period: {train_df['Week_Start'].min().date()} to {train_df['Week_Start'].max().date()}")
print(f"   Test period: {test_df['Week_Start'].min().date()} to {test_df['Week_Start'].max().date()}")

# Actual values for comparison
y_train = train_df['Net_Cash_Flow_USD'].values
y_test = test_df['Net_Cash_Flow_USD'].values

# =============================================================================
# HELPER FUNCTIONS
# =============================================================================
def calculate_metrics(y_true, y_pred, model_name):
    """Calculate MAE, RMSE, MAPE for a model"""
    mae = mean_absolute_error(y_true, y_pred)
    rmse = np.sqrt(mean_squared_error(y_true, y_pred))
    
    # MAPE - handle zero values
    mask = y_true != 0
    if mask.sum() > 0:
        mape = np.mean(np.abs((y_true[mask] - y_pred[mask]) / y_true[mask])) * 100
    else:
        mape = np.nan
    
    return {
        'Model': model_name,
        'MAE': mae,
        'RMSE': rmse,
        'MAPE': mape
    }

def symmetric_mape(y_true, y_pred):
    """Calculate Symmetric MAPE (handles zero values better)"""
    numerator = np.abs(y_true - y_pred)
    denominator = (np.abs(y_true) + np.abs(y_pred)) / 2
    mask = denominator != 0
    if mask.sum() > 0:
        return np.mean(numerator[mask] / denominator[mask]) * 100
    return np.nan

# =============================================================================
# MODEL 1: MOVING AVERAGE (4-WEEK)
# =============================================================================
print("\n[3/7] Training Model 1: Moving Average (4-week)...")

window = 4
ma_predictions = []

# Simple approach: predict next period as average of last 4 periods
for i in range(len(test_df)):
    if i == 0:
        # For first test point, use last 4 training weeks
        pred = y_train[-window:].mean()
    else:
        # Use last 4 available values (including any test values already predicted)
        all_values = np.concatenate([y_train, y_test[:i]])
        pred = all_values[-window:].mean()
    ma_predictions.append(pred)

ma_predictions = np.array(ma_predictions)
ma_metrics = calculate_metrics(y_test, ma_predictions, 'Moving Average (4-week)')

print(f"   MAPE: {ma_metrics['MAPE']:.2f}%")
print(f"   MAE: ${ma_metrics['MAE']:,.0f}")
print(f"   RMSE: ${ma_metrics['RMSE']:,.0f}")

# =============================================================================
# MODEL 2: EXPONENTIAL SMOOTHING (HOLT-WINTERS)
# =============================================================================
print("\n[4/7] Training Model 2: Exponential Smoothing (Holt-Winters)...")

try:
    # Holt-Winters with additive trend and seasonality
    # Seasonal period = 4 (monthly pattern in weekly data)
    hw_model = ExponentialSmoothing(
        y_train,
        trend='add',
        seasonal='add',
        seasonal_periods=4,
        damped_trend=True
    )
    hw_fit = hw_model.fit(optimized=True)
    hw_predictions = hw_fit.forecast(steps=len(test_df))
    
    hw_metrics = calculate_metrics(y_test, hw_predictions, 'Exponential Smoothing')
    print(f"   MAPE: {hw_metrics['MAPE']:.2f}%")
    print(f"   MAE: ${hw_metrics['MAE']:,.0f}")
    print(f"   RMSE: ${hw_metrics['RMSE']:,.0f}")
    
except Exception as e:
    print(f"   Warning: Holt-Winters failed with error: {e}")
    print("   Trying simpler exponential smoothing...")
    
    # Fallback to simple exponential smoothing
    hw_model = ExponentialSmoothing(y_train, trend='add', seasonal=None)
    hw_fit = hw_model.fit(optimized=True)
    hw_predictions = hw_fit.forecast(steps=len(test_df))
    
    hw_metrics = calculate_metrics(y_test, hw_predictions, 'Exponential Smoothing (Simple)')
    print(f"   MAPE: {hw_metrics['MAPE']:.2f}%")
    print(f"   MAE: ${hw_metrics['MAE']:,.0f}")
    print(f"   RMSE: ${hw_metrics['RMSE']:,.0f}")

# =============================================================================
# MODEL 3: PROPHET (FACEBOOK)
# =============================================================================
print("\n[5/7] Training Model 3: Prophet...")

# Prepare data for Prophet
prophet_train = pd.DataFrame({
    'ds': train_df['Week_Start'],
    'y': train_df['Net_Cash_Flow_USD']
})

# Initialize Prophet with custom settings for cash flow forecasting
prophet_model = Prophet(
    yearly_seasonality=False,  # Not enough data for yearly patterns
    weekly_seasonality=False,  # Already aggregated to weekly
    daily_seasonality=False,   # Not applicable
    seasonality_mode='additive',
    interval_width=0.95,       # 95% confidence intervals
    changepoint_prior_scale=0.05  # Moderate flexibility
)

# Add monthly seasonality (AP/AR payment cycles ~ 4 weeks)
prophet_model.add_seasonality(
    name='monthly',
    period=30.5,
    fourier_order=5
)

# Fit the model
prophet_model.fit(prophet_train)

# Create future dataframe for test period
future_dates = pd.DataFrame({'ds': test_df['Week_Start']})

# Make predictions
prophet_forecast = prophet_model.predict(future_dates)
prophet_predictions = prophet_forecast['yhat'].values

# Get confidence intervals
prophet_lower = prophet_forecast['yhat_lower'].values
prophet_upper = prophet_forecast['yhat_upper'].values

prophet_metrics = calculate_metrics(y_test, prophet_predictions, 'Prophet')
print(f"   MAPE: {prophet_metrics['MAPE']:.2f}%")
print(f"   MAE: ${prophet_metrics['MAE']:,.0f}")
print(f"   RMSE: ${prophet_metrics['RMSE']:,.0f}")

# Calculate confidence interval coverage
ci_coverage = np.mean((y_test >= prophet_lower) & (y_test <= prophet_upper)) * 100
print(f"   95% CI Coverage: {ci_coverage:.1f}%")

# =============================================================================
# COMPARE ALL MODELS
# =============================================================================
print("\n[6/7] Comparing all models...")

results = pd.DataFrame([ma_metrics, hw_metrics, prophet_metrics])
results['SMAPE'] = [
    symmetric_mape(y_test, ma_predictions),
    symmetric_mape(y_test, hw_predictions),
    symmetric_mape(y_test, prophet_predictions)
]

# Rank models by MAE (better metric for cash flow with sign changes)
results['MAE_Rank'] = results['MAE'].rank()
results = results.sort_values('MAE')

print("\n" + "=" * 80)
print("MODEL COMPARISON RESULTS")
print("=" * 80)
print("\n" + results.to_string(index=False))

# Determine winner by MAE (MAPE is unreliable for data crossing zero)
winner = results.iloc[0]['Model']
winner_mae = results.iloc[0]['MAE']
winner_mape = results.iloc[0]['MAPE']

print(f"\n*** WINNER: {winner}")
print(f"   Best MAE: ${winner_mae:,.0f}")
print(f"   Note: MAPE unreliable for cash flow (crosses zero)")

# =============================================================================
# VISUALIZATIONS
# =============================================================================
print("\n[7/7] Creating visualizations...")

# Create dashboard directory
import os
os.makedirs('dashboard', exist_ok=True)

# Figure 1: Model Comparison Bar Chart
fig, axes = plt.subplots(2, 2, figsize=(14, 10))
fig.suptitle('MODEL COMPARISON - Phase 0: Model Selection\nAstraZeneca Cash Flow Forecasting', 
             fontsize=14, fontweight='bold', color=AZ_COLORS['navy'])

# Plot 1: MAE Comparison (Primary Metric)
ax1 = axes[0, 0]
colors = [AZ_COLORS['lime_green'] if m == winner else AZ_COLORS['graphite'] for m in results['Model']]
bars = ax1.bar(results['Model'], results['MAE'] / 1e6, color=colors, edgecolor=AZ_COLORS['navy'], linewidth=2)
ax1.set_ylabel('MAE (Millions USD)', fontweight='bold')
ax1.set_title('Mean Absolute Error (MAE) - PRIMARY METRIC\nLower is Better', fontweight='bold')
for bar, val in zip(bars, results['MAE']):
    ax1.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.1, f'${val/1e6:.2f}M', 
             ha='center', fontweight='bold', fontsize=11)
ax1.set_ylim(0, max(results['MAE']) / 1e6 * 1.3)

# Plot 2: RMSE Comparison
ax2 = axes[0, 1]
bars = ax2.bar(results['Model'], results['RMSE'] / 1e6, color=colors, edgecolor=AZ_COLORS['navy'], linewidth=2)
ax2.set_ylabel('RMSE (Millions USD)', fontweight='bold')
ax2.set_title('Root Mean Squared Error (RMSE)\nLower is Better', fontweight='bold')
for bar, val in zip(bars, results['RMSE']):
    ax2.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.1, f'${val/1e6:.2f}M', 
             ha='center', fontweight='bold', fontsize=11)
ax2.set_ylim(0, max(results['RMSE']) / 1e6 * 1.3)

# Plot 3: Actual vs Predicted (Line Chart)
ax3 = axes[1, 0]
weeks = test_df['Week_Number'].values
ax3.plot(weeks, y_test / 1e6, 'o-', color=AZ_COLORS['navy'], linewidth=3, markersize=10, label='Actual', zorder=5)
ax3.plot(weeks, ma_predictions / 1e6, 's--', color=AZ_COLORS['graphite'], linewidth=2, markersize=8, label='Moving Average', alpha=0.7)
ax3.plot(weeks, hw_predictions / 1e6, '^--', color=AZ_COLORS['purple'], linewidth=2, markersize=8, label='Exp. Smoothing', alpha=0.7)
ax3.plot(weeks, prophet_predictions / 1e6, 'D-', color=AZ_COLORS['lime_green'], linewidth=2.5, markersize=8, label='Prophet', zorder=4)
ax3.fill_between(weeks, prophet_lower / 1e6, prophet_upper / 1e6, color=AZ_COLORS['lime_green'], alpha=0.2, label='Prophet 95% CI')
ax3.set_xlabel('Week Number', fontweight='bold')
ax3.set_ylabel('Net Cash Flow (Millions USD)', fontweight='bold')
ax3.set_title('Actual vs Predicted (Test Period: Weeks 40-44)', fontweight='bold')
ax3.legend(loc='best')
ax3.grid(True, alpha=0.3)

# Plot 4: Summary Table
ax4 = axes[1, 1]
ax4.axis('off')

# Create summary table (ranked by MAE)
table_data = [
    ['Model', 'MAE', 'RMSE', 'SMAPE', 'Verdict'],
    ['Prophet', f'${prophet_metrics["MAE"]/1e6:.2f}M', f'${prophet_metrics["RMSE"]/1e6:.2f}M', f'{symmetric_mape(y_test, prophet_predictions):.0f}%', '[WINNER]'],
    ['Moving Average', f'${ma_metrics["MAE"]/1e6:.2f}M', f'${ma_metrics["RMSE"]/1e6:.2f}M', f'{symmetric_mape(y_test, ma_predictions):.0f}%', 'Baseline'],
    [hw_metrics['Model'][:18], f'${hw_metrics["MAE"]/1e6:.2f}M', f'${hw_metrics["RMSE"]/1e6:.2f}M', f'{symmetric_mape(y_test, hw_predictions):.0f}%', 'Third']
]

table = ax4.table(
    cellText=table_data[1:],
    colLabels=table_data[0],
    cellLoc='center',
    loc='center',
    colColours=[AZ_COLORS['navy']] * 5,
    colWidths=[0.25, 0.15, 0.2, 0.2, 0.2]
)
table.auto_set_font_size(False)
table.set_fontsize(11)
table.scale(1.2, 2)

# Color header
for i in range(5):
    table[(0, i)].set_text_props(color='white', fontweight='bold')
    
# Highlight winner row (first row after header = Prophet)
for i in range(5):
    table[(1, i)].set_facecolor(AZ_COLORS['lime_green'])
    table[(1, i)].set_alpha(0.3)

ax4.set_title('Model Selection Summary\n', fontweight='bold', fontsize=12)
ax4.text(0.5, 0.15, f'*** SELECTED MODEL: Prophet (MAE: ${prophet_metrics["MAE"]/1e6:.2f}M)\n\n'
         f'Reasons:\n'
         f'* Lowest MAE (${prophet_metrics["MAE"]/1e6:.2f}M vs ${ma_metrics["MAE"]/1e6:.2f}M vs ${hw_metrics["MAE"]/1e6:.2f}M)\n'
         f'* Lowest RMSE (${prophet_metrics["RMSE"]/1e6:.2f}M)\n'
         f'* Provides 95% confidence intervals\n'
         f'* Captures monthly AP/AR payment cycles\n'
         f'* Business-friendly decomposition',
         transform=ax4.transAxes, ha='center', va='top', fontsize=10,
         bbox=dict(boxstyle='round', facecolor=AZ_COLORS['lime_green'], alpha=0.2))

plt.tight_layout()
plt.savefig('dashboard/00_model_comparison.png', dpi=300, bbox_inches='tight', facecolor='white')
plt.close()

print("   Saved: dashboard/00_model_comparison.png")

# =============================================================================
# SAVE RESULTS
# =============================================================================
results.to_csv('model_comparison_results.csv', index=False)
print("   Saved: model_comparison_results.csv")

# Create detailed predictions comparison
predictions_df = pd.DataFrame({
    'Week_Number': test_df['Week_Number'].values,
    'Week_Start': test_df['Week_Start'].values,
    'Actual': y_test,
    'MA_Predicted': ma_predictions,
    'MA_Error': y_test - ma_predictions,
    'ExpSmooth_Predicted': hw_predictions,
    'ExpSmooth_Error': y_test - hw_predictions,
    'Prophet_Predicted': prophet_predictions,
    'Prophet_Error': y_test - prophet_predictions,
    'Prophet_Lower_95': prophet_lower,
    'Prophet_Upper_95': prophet_upper
})
predictions_df.to_csv('model_comparison_predictions.csv', index=False)
print("   Saved: model_comparison_predictions.csv")

# =============================================================================
# FINAL SUMMARY
# =============================================================================
print("\n" + "=" * 80)
print("PHASE 0 COMPLETE: MODEL SELECTION")
print("=" * 80)

# Calculate SMAPE for each model
smape_ma = symmetric_mape(y_test, ma_predictions)
smape_hw = symmetric_mape(y_test, hw_predictions)
smape_prophet = symmetric_mape(y_test, prophet_predictions)

print(f"""
RESULTS SUMMARY:
   
   Model Comparison (Test: Weeks 40-44):
   Note: Using MAE as primary metric (MAPE unreliable for cash flow crossing zero)
   
   +---------------------------+-------------+-------------+---------+
   | Model                     | MAE         | RMSE        | SMAPE   |
   +---------------------------+-------------+-------------+---------+
   | Prophet                   | ${prophet_metrics['MAE']/1e6:>5.2f}M    | ${prophet_metrics['RMSE']/1e6:>5.2f}M    | {smape_prophet:>5.0f}%  | << WINNER
   | Moving Average (4-week)   | ${ma_metrics['MAE']/1e6:>5.2f}M    | ${ma_metrics['RMSE']/1e6:>5.2f}M    | {smape_ma:>5.0f}%  | Baseline
   | Exponential Smoothing     | ${hw_metrics['MAE']/1e6:>5.2f}M    | ${hw_metrics['RMSE']/1e6:>5.2f}M    | {smape_hw:>5.0f}%  | Third
   +---------------------------+-------------+-------------+---------+
   
*** DECISION: SELECT PROPHET ***
   
   Why Prophet?
   [+] Lowest MAE: ${prophet_metrics['MAE']/1e6:.2f}M (vs ${ma_metrics['MAE']/1e6:.2f}M vs ${hw_metrics['MAE']/1e6:.2f}M)
   [+] Lowest RMSE: ${prophet_metrics['RMSE']/1e6:.2f}M
   [+] Provides confidence intervals (95% CI coverage: {ci_coverage:.1f}%)
   [+] Automatically detects monthly seasonality (AP/AR cycles)
   [+] Business-friendly trend + seasonal decomposition
   
FILES CREATED:
   - dashboard/00_model_comparison.png (visualization)
   - model_comparison_results.csv (metrics summary)
   - model_comparison_predictions.csv (detailed predictions)
   
==> NEXT STEP: Phase 1 - Option A (1-Month Forward Validation)
""")