"""
================================================================================
IMPROVED MAPE/SMAPE METRICS FOR CASH FLOW DATA
================================================================================
AstraZeneca Cash Flow Forecasting - Datathon 2025

Problem: Standard MAPE/SMAPE explode when:
  1. Actual values are close to zero
  2. Actual and predicted have opposite signs (inflow vs outflow)

Solutions:
  1. Modified MAPE with threshold (only calculate when |actual| > threshold)
  2. SMAPE with sign-aware calculation
  3. Weighted MAPE (weight by absolute value)
  4. Direction-aware MAPE (only when same direction)
  5. Scaled MAPE (scale by average magnitude)
"""

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')

from sklearn.metrics import mean_absolute_error, mean_squared_error
from prophet import Prophet
import os

# 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("IMPROVED MAPE/SMAPE METRICS FOR CASH FLOW DATA")
print("=" * 80)

# =============================================================================
# HELPER FUNCTIONS - IMPROVED METRICS
# =============================================================================

def standard_mape(y_true, y_pred):
    """Standard MAPE (unreliable for cash flow)"""
    mask = y_true != 0
    if mask.sum() > 0:
        return np.mean(np.abs((y_true[mask] - y_pred[mask]) / y_true[mask])) * 100
    return np.nan

def standard_smape(y_true, y_pred):
    """Standard SMAPE (better than MAPE but still problematic)"""
    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

def threshold_mape(y_true, y_pred, threshold=100000):
    """
    Modified MAPE: Only calculate when |actual| > threshold
    Avoids explosion near zero
    """
    mask = np.abs(y_true) > threshold
    if mask.sum() > 0:
        return np.mean(np.abs((y_true[mask] - y_pred[mask]) / y_true[mask])) * 100
    return np.nan

def sign_aware_smape(y_true, y_pred):
    """
    SMAPE that handles opposite signs better
    When signs differ, use absolute error relative to average magnitude
    """
    same_sign = (np.sign(y_true) == np.sign(y_pred))
    
    # For same sign: standard SMAPE
    if same_sign.sum() > 0:
        numerator_same = np.abs(y_true[same_sign] - y_pred[same_sign])
        denominator_same = (np.abs(y_true[same_sign]) + np.abs(y_pred[same_sign])) / 2
        smape_same = np.mean(numerator_same / denominator_same) * 100
    else:
        smape_same = 0
    
    # For opposite sign: use absolute error / average magnitude
    opposite_sign = ~same_sign
    if opposite_sign.sum() > 0:
        avg_magnitude = (np.abs(y_true).mean() + np.abs(y_pred).mean()) / 2
        if avg_magnitude > 0:
            smape_opposite = np.mean(np.abs(y_true[opposite_sign] - y_pred[opposite_sign]) / avg_magnitude) * 100
        else:
            smape_opposite = 0
    else:
        smape_opposite = 0
    
    # Weighted average
    total = len(y_true)
    if total > 0:
        return (smape_same * same_sign.sum() + smape_opposite * opposite_sign.sum()) / total
    return np.nan

def weighted_mape(y_true, y_pred):
    """
    Weighted MAPE: Weight each error by absolute value
    More important weeks (larger cash flows) have more weight
    """
    abs_errors = np.abs(y_true - y_pred)
    abs_actuals = np.abs(y_true)
    
    # Weight by absolute value
    weights = abs_actuals / abs_actuals.sum() if abs_actuals.sum() > 0 else np.ones_like(abs_actuals) / len(abs_actuals)
    
    mask = abs_actuals > 0
    if mask.sum() > 0:
        weighted_errors = (abs_errors[mask] / abs_actuals[mask]) * weights[mask]
        return np.sum(weighted_errors) * 100
    return np.nan

def direction_aware_mape(y_true, y_pred):
    """
    MAPE only calculated when actual and predicted have same direction
    More meaningful for cash flow (inflow vs outflow)
    """
    same_direction = (np.sign(y_true) == np.sign(y_pred))
    
    if same_direction.sum() > 0:
        mask = (same_direction) & (y_true != 0)
        if mask.sum() > 0:
            return np.mean(np.abs((y_true[mask] - y_pred[mask]) / y_true[mask])) * 100
    return np.nan

def scaled_mape(y_true, y_pred):
    """
    Scaled MAPE: Scale by average absolute value of actuals
    More stable than standard MAPE
    """
    abs_errors = np.abs(y_true - y_pred)
    avg_magnitude = np.abs(y_true).mean()
    
    if avg_magnitude > 0:
        return np.mean(abs_errors / avg_magnitude) * 100
    return np.nan

def robust_smape(y_true, y_pred, epsilon=1e-6):
    """
    Robust SMAPE: Add small epsilon to denominator to avoid division by zero
    """
    numerator = np.abs(y_true - y_pred)
    denominator = (np.abs(y_true) + np.abs(y_pred)) / 2 + epsilon
    return np.mean(numerator / denominator) * 100

# =============================================================================
# LOAD DATA AND RUN VALIDATION
# =============================================================================
print("\n[1/4] Loading data and running validation...")

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)

# Simple holdout validation
train_size = 39
train_df = df.iloc[:train_size].copy()
test_df = df.iloc[train_size:].copy()

# Train Prophet
prophet_df = pd.DataFrame({
    'ds': train_df['Week_Start'],
    'y': train_df['Net_Cash_Flow_USD']
})

model = Prophet(
    yearly_seasonality=False,
    weekly_seasonality=False,
    daily_seasonality=False,
    seasonality_mode='additive',
    interval_width=0.95,
    changepoint_prior_scale=0.05
)
model.add_seasonality(name='monthly', period=30.5, fourier_order=5)
model.fit(prophet_df)

# Predict
future = pd.DataFrame({'ds': test_df['Week_Start']})
forecast = model.predict(future)

y_test = test_df['Net_Cash_Flow_USD'].values
y_pred = forecast['yhat'].values

print(f"   Test period: {len(y_test)} weeks")
print(f"   Actual range: ${y_test.min()/1e6:.2f}M to ${y_test.max()/1e6:.2f}M")
print(f"   Predicted range: ${y_pred.min()/1e6:.2f}M to ${y_pred.max()/1e6:.2f}M")

# =============================================================================
# CALCULATE ALL METRICS
# =============================================================================
print("\n[2/4] Calculating improved metrics...")

metrics = {}

# Standard metrics (for comparison)
metrics['Standard MAPE'] = standard_mape(y_test, y_pred)
metrics['Standard SMAPE'] = standard_smape(y_test, y_pred)

# Improved metrics
metrics['Threshold MAPE ($100K)'] = threshold_mape(y_test, y_pred, threshold=100000)
metrics['Threshold MAPE ($500K)'] = threshold_mape(y_test, y_pred, threshold=500000)
metrics['Threshold MAPE ($1M)'] = threshold_mape(y_test, y_pred, threshold=1000000)

metrics['Sign-Aware SMAPE'] = sign_aware_smape(y_test, y_pred)
metrics['Weighted MAPE'] = weighted_mape(y_test, y_pred)
metrics['Direction-Aware MAPE'] = direction_aware_mape(y_test, y_pred)
metrics['Scaled MAPE'] = scaled_mape(y_test, y_pred)
metrics['Robust SMAPE'] = robust_smape(y_test, y_pred)

# Also calculate MAE for reference
mae = mean_absolute_error(y_test, y_pred)
rmse = np.sqrt(mean_squared_error(y_test, y_pred))

print(f"\n   METRIC COMPARISON:")
print(f"   +---------------------------+------------------+------------------+")
print(f"   | Metric                   | Value            | Interpretation   |")
print(f"   +---------------------------+------------------+------------------+")

for metric_name, value in metrics.items():
    if not np.isnan(value):
        if 'Standard' in metric_name:
            interp = "Unreliable (explodes)"
        elif value < 20:
            interp = "Excellent"
        elif value < 40:
            interp = "Good"
        elif value < 60:
            interp = "Acceptable"
        else:
            interp = "High error"
        print(f"   | {metric_name:<25} | {value:>12.1f}% | {interp:<17} |")
    else:
        print(f"   | {metric_name:<25} | {'N/A':>12} | Cannot calculate  |")

print(f"   +---------------------------+------------------+------------------+")
print(f"   | MAE (reference)          | ${mae/1e6:>12.2f}M | Absolute error   |")
print(f"   | RMSE (reference)         | ${rmse/1e6:>12.2f}M | Penalizes large  |")
print(f"   +---------------------------+------------------+------------------+")

# =============================================================================
# DETAILED WEEK-BY-WEEK ANALYSIS
# =============================================================================
print("\n[3/4] Week-by-week metric analysis...")

week_analysis = []
for i, (actual, pred) in enumerate(zip(y_test, y_pred)):
    abs_error = abs(actual - pred)
    same_sign = (actual > 0) == (pred > 0)
    
    # Calculate per-week metrics
    if actual != 0:
        mape_week = abs((actual - pred) / actual) * 100
    else:
        mape_week = np.nan
    
    avg_mag = (abs(actual) + abs(pred)) / 2
    if avg_mag > 0:
        smape_week = (abs_error / avg_mag) * 100
    else:
        smape_week = np.nan
    
    week_analysis.append({
        'Week': test_df['Week_Number'].iloc[i],
        'Actual': actual,
        'Predicted': pred,
        'Error': actual - pred,
        'Abs_Error': abs_error,
        'MAPE': mape_week,
        'SMAPE': smape_week,
        'Same_Sign': same_sign
    })

week_df = pd.DataFrame(week_analysis)

print(f"\n   WEEK-BY-WEEK METRICS:")
print(f"   +------+-------------+-------------+------------+-----------+-----------+")
print(f"   | Week | Actual      | Predicted   | Abs Error  | MAPE      | SMAPE     |")
print(f"   +------+-------------+-------------+------------+-----------+-----------+")
for _, row in week_df.iterrows():
    mape_str = f"{row['MAPE']:.1f}%" if not np.isnan(row['MAPE']) else "N/A"
    smape_str = f"{row['SMAPE']:.1f}%" if not np.isnan(row['SMAPE']) else "N/A"
    sign_marker = "[OK]" if row['Same_Sign'] else "[X]"
    print(f"   | {row['Week']:>4} | ${row['Actual']/1e6:>8.2f}M | ${row['Predicted']/1e6:>8.2f}M | ${row['Abs_Error']/1e6:>7.2f}M | {mape_str:>9} | {smape_str:>9} | {sign_marker}")
print(f"   +------+-------------+-------------+------------+-----------+-----------+")

# =============================================================================
# CREATE VISUALIZATION
# =============================================================================
print("\n[4/4] Creating visualization...")

os.makedirs('dashboard', exist_ok=True)

fig = plt.figure(figsize=(20, 14))
fig.suptitle('IMPROVED MAPE/SMAPE METRICS FOR CASH FLOW DATA\nAstraZeneca Cash Flow Forecasting', 
             fontsize=18, fontweight='bold', color=AZ_COLORS['navy'], y=0.98)

# ------ Plot 1: Metric Comparison Bar Chart ------
ax1 = fig.add_subplot(2, 3, 1)

metric_names = list(metrics.keys())
metric_values = [v for v in metrics.values() if not np.isnan(v)]
metric_names_clean = [n for n, v in zip(metric_names, metrics.values()) if not np.isnan(v)]

# Color code: green for good, yellow for moderate, red for high
colors = [AZ_COLORS['lime_green'] if v < 40 else AZ_COLORS['gold'] if v < 80 else AZ_COLORS['magenta'] 
          for v in metric_values]

bars = ax1.barh(metric_names_clean, metric_values, color=colors, edgecolor=AZ_COLORS['navy'], linewidth=1.5)
ax1.set_xlabel('Percentage Error (%)', fontweight='bold')
ax1.set_title('Metric Comparison\nLower is Better', fontweight='bold', fontsize=11)
ax1.grid(True, alpha=0.3, axis='x')

for bar, val in zip(bars, metric_values):
    ax1.text(val + 5, bar.get_y() + bar.get_height()/2, f'{val:.1f}%', 
             ha='left', va='center', fontweight='bold', fontsize=9)

# ------ Plot 2: Week-by-Week MAPE ------
ax2 = fig.add_subplot(2, 3, 2)

weeks = week_df['Week'].values
mape_vals = week_df['MAPE'].values
valid_mape = ~np.isnan(mape_vals)

ax2.scatter(weeks[valid_mape], mape_vals[valid_mape], 
           c=[AZ_COLORS['lime_green'] if v < 50 else AZ_COLORS['gold'] if v < 100 else AZ_COLORS['magenta'] 
              for v in mape_vals[valid_mape]], 
           s=150, edgecolors='white', linewidth=2, zorder=5)

ax2.axhline(y=50, color=AZ_COLORS['graphite'], linestyle='--', linewidth=2, label='50% threshold')
ax2.set_xlabel('Week Number', fontweight='bold')
ax2.set_ylabel('MAPE (%)', fontweight='bold')
ax2.set_title('Week-by-Week MAPE\n(Only when actual != 0)', fontweight='bold', fontsize=11)
ax2.legend()
ax2.grid(True, alpha=0.3)

# ------ Plot 3: Week-by-Week SMAPE ------
ax3 = fig.add_subplot(2, 3, 3)

smape_vals = week_df['SMAPE'].values
valid_smape = ~np.isnan(smape_vals)

ax3.scatter(weeks[valid_smape], smape_vals[valid_smape], 
           c=[AZ_COLORS['lime_green'] if v < 50 else AZ_COLORS['gold'] if v < 100 else AZ_COLORS['magenta'] 
              for v in smape_vals[valid_smape]], 
           s=150, edgecolors='white', linewidth=2, zorder=5)

ax3.axhline(y=50, color=AZ_COLORS['graphite'], linestyle='--', linewidth=2, label='50% threshold')
ax3.set_xlabel('Week Number', fontweight='bold')
ax3.set_ylabel('SMAPE (%)', fontweight='bold')
ax3.set_title('Week-by-Week SMAPE\n(More stable than MAPE)', fontweight='bold', fontsize=11)
ax3.legend()
ax3.grid(True, alpha=0.3)

# ------ Plot 4: Actual vs Predicted with Error Zones ------
ax4 = fig.add_subplot(2, 3, 4)

ax4.scatter(y_test/1e6, y_pred/1e6, 
           c=[AZ_COLORS['lime_green'] if same else AZ_COLORS['magenta'] 
              for same in week_df['Same_Sign']], 
           s=150, edgecolors='white', linewidth=2, zorder=5, alpha=0.7)

# Perfect prediction line
min_val = min(min(y_test), min(y_pred)) / 1e6
max_val = max(max(y_test), max(y_pred)) / 1e6
ax4.plot([min_val, max_val], [min_val, max_val], 'k--', linewidth=2, label='Perfect Prediction')

# Add 20% error bands
ax4.fill_between([min_val, max_val], 
                 [min_val*0.8, max_val*0.8], 
                 [min_val*1.2, max_val*1.2], 
                 alpha=0.2, color=AZ_COLORS['gold'], label='±20% Error Band')

ax4.axhline(y=0, color=AZ_COLORS['graphite'], linewidth=1, linestyle='-', alpha=0.5)
ax4.axvline(x=0, color=AZ_COLORS['graphite'], linewidth=1, linestyle='-', alpha=0.5)

ax4.set_xlabel('Actual Cash Flow (Millions USD)', fontweight='bold')
ax4.set_ylabel('Predicted Cash Flow (Millions USD)', fontweight='bold')
ax4.set_title('Actual vs Predicted\nGreen=Same Sign, Magenta=Opposite Sign', fontweight='bold', fontsize=11)
ax4.legend()
ax4.grid(True, alpha=0.3)

# ------ Plot 5: Error Distribution ------
ax5 = fig.add_subplot(2, 3, 5)

errors = week_df['Abs_Error'].values / 1e6
ax5.hist(errors, bins=10, color=AZ_COLORS['light_blue'], edgecolor=AZ_COLORS['navy'], alpha=0.7)
ax5.axvline(x=mae/1e6, color=AZ_COLORS['mulberry'], linestyle='--', linewidth=2, 
           label=f'MAE: ${mae/1e6:.2f}M')
ax5.set_xlabel('Absolute Error (Millions USD)', fontweight='bold')
ax5.set_ylabel('Frequency', fontweight='bold')
ax5.set_title('Error Distribution', fontweight='bold')
ax5.legend()
ax5.grid(True, alpha=0.3, axis='y')

# ------ Plot 6: Summary Table ------
ax6 = fig.add_subplot(2, 3, 6)
ax6.axis('off')

# Find best metrics
best_threshold_mape = min([v for k, v in metrics.items() if 'Threshold' in k and not np.isnan(v)], default=np.nan)
best_improved = min([v for k, v in metrics.items() if k not in ['Standard MAPE', 'Standard SMAPE'] and not np.isnan(v)], default=np.nan)

summary_text = f"""
IMPROVED METRICS SUMMARY
=======================

PROBLEM WITH STANDARD METRICS:
-----------------------------
Standard MAPE:  {metrics['Standard MAPE']:.1f}% (unreliable)
Standard SMAPE: {metrics['Standard SMAPE']:.1f}% (unreliable)

These explode when:
- Actual values near zero
- Signs differ (inflow vs outflow)


RECOMMENDED METRICS:
--------------------
1. Threshold MAPE ($1M): {metrics.get('Threshold MAPE ($1M)', np.nan):.1f}%
   Only calculate when |actual| > $1M
   Avoids zero-division issues

2. Sign-Aware SMAPE: {metrics.get('Sign-Aware SMAPE', np.nan):.1f}%
   Handles opposite signs gracefully
   More stable than standard SMAPE

3. Scaled MAPE: {metrics.get('Scaled MAPE', np.nan):.1f}%
   Scales by average magnitude
   Good for volatile data

4. Direction-Aware MAPE: {metrics.get('Direction-Aware MAPE', np.nan):.1f}%
   Only when same direction
   Most meaningful for cash flow


INTERPRETATION:
---------------
< 20%:  Excellent
20-40%: Good
40-60%: Acceptable
> 60%:  High error

For cash flow crossing zero:
- Use Threshold MAPE or Scaled MAPE
- Direction-Aware MAPE for same-sign weeks
- MAE ($1.84M) as primary metric
"""

ax6.text(0.05, 0.95, summary_text, transform=ax6.transAxes, fontsize=9,
         verticalalignment='top', fontfamily='monospace',
         bbox=dict(boxstyle='round', facecolor=AZ_COLORS['lime_green'], alpha=0.2))

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

print("   Saved: dashboard/07_improved_mape_metrics.png")

# Save results
results_df = pd.DataFrame({
    'Metric': list(metrics.keys()),
    'Value_Pct': list(metrics.values()),
    'MAE_USD': [mae] * len(metrics),
    'RMSE_USD': [rmse] * len(metrics)
})
results_df.to_csv('improved_metrics_summary.csv', index=False)
print("   Saved: improved_metrics_summary.csv")

week_df.to_csv('week_by_week_metrics.csv', index=False)
print("   Saved: week_by_week_metrics.csv")

# =============================================================================
# FINAL SUMMARY
# =============================================================================
print("\n" + "=" * 80)
print("IMPROVED METRICS ANALYSIS COMPLETE")
print("=" * 80)

print(f"""
RECOMMENDATIONS FOR CASH FLOW METRICS:
======================================

1. THRESHOLD MAPE (Recommended):
   - Only calculate when |actual| > threshold (e.g., $1M)
   - Avoids explosion near zero
   - Result: {metrics.get('Threshold MAPE ($1M)', np.nan):.1f}%

2. SCALED MAPE (Recommended):
   - Scales error by average magnitude
   - More stable than standard MAPE
   - Result: {metrics.get('Scaled MAPE', np.nan):.1f}%

3. SIGN-AWARE SMAPE (Recommended):
   - Handles opposite signs gracefully
   - More meaningful for cash flow
   - Result: {metrics.get('Sign-Aware SMAPE', np.nan):.1f}%

4. DIRECTION-AWARE MAPE (For same-sign weeks):
   - Only when actual and predicted have same direction
   - Most business-meaningful
   - Result: {metrics.get('Direction-Aware MAPE', np.nan):.1f}%

5. PRIMARY METRIC: MAE
   - Always use MAE as primary metric
   - Current: ${mae/1e6:.2f}M
   - Not affected by zero-crossing issues

FILES CREATED:
   - dashboard/07_improved_mape_metrics.png
   - improved_metrics_summary.csv
   - week_by_week_metrics.csv

==> Use Threshold MAPE or Scaled MAPE for percentage-based reporting!
""")