evaluate_model.py

"""
Comprehensive model evaluation across multiple dates and conditions.
Tests model performance on normal days, storms, and seasonal patterns.
"""

import asyncio
import numpy as np
import pandas as pd
from datetime import datetime, timedelta
from predictor import predict_outages, load_eaglei_data, get_s3_client, S3_BUCKET
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Test dates covering different conditions
TEST_DATES = {
    # Normal days (2024) - random sampling across months
    'normal': [
        ('2024-01-15', 14, 'Winter weekday'),
        ('2024-02-20', 14, 'Late winter'),
        ('2024-03-12', 14, 'Early spring'),
        ('2024-04-18', 14, 'Spring'),
        ('2024-05-22', 14, 'Late spring'),
        ('2024-06-10', 14, 'Early summer'),
        ('2024-06-25', 14, 'Summer'),
        ('2024-07-15', 14, 'Mid summer'),
        ('2024-08-05', 14, 'August'),
        ('2024-08-20', 14, 'Late summer'),
        ('2024-10-08', 14, 'Fall'),
        ('2024-11-12', 14, 'Late fall'),
    ],
    # Storm events (known major outages)
    'storms': [
        ('2024-09-27', 18, 'Hurricane Helene'),
        ('2024-07-08', 18, 'Hurricane Beryl TX'),
        ('2024-05-27', 18, 'TX storms'),
        ('2024-04-27', 14, 'Midwest tornadoes'),
        ('2023-12-17', 12, 'Winter storm'),
        ('2023-08-30', 18, 'Hurricane Idalia'),
        ('2023-07-15', 14, 'Summer storms'),
        ('2023-06-29', 14, 'Heat wave'),
    ],
    # Weekend vs weekday comparison
    'weekends': [
        ('2024-06-08', 14, 'Saturday summer'),
        ('2024-06-09', 14, 'Sunday summer'),
        ('2024-03-16', 14, 'Saturday spring'),
        ('2024-03-17', 14, 'Sunday spring'),
    ],
}


def get_actual_outages(target_time: datetime) -> dict:
    """Get actual EAGLE-I outages for a specific time."""
    year = target_time.year
    df = load_eaglei_data(year)

    # Find closest timestamp
    df['time_diff'] = abs((df['run_start_time'] - target_time).dt.total_seconds())
    mask = df['time_diff'] <= 900  # Within 15 minutes
    snapshot = df[mask].copy()

    if len(snapshot) == 0:
        nearest_idx = df['time_diff'].idxmin()
        nearest_time = df.loc[nearest_idx, 'run_start_time']
        snapshot = df[df['run_start_time'] == nearest_time].copy()

    # Aggregate by state
    state_totals = snapshot.groupby('state')['customers_out'].sum().to_dict()
    total = snapshot['customers_out'].sum()

    return {
        'total': int(total),
        'by_state': state_totals,
        'timestamp': snapshot['run_start_time'].iloc[0] if len(snapshot) > 0 else target_time
    }


async def evaluate_single_date(date_str: str, hour: int, description: str) -> dict:
    """Evaluate model for a single date."""
    try:
        target_time = datetime.strptime(f"{date_str} {hour:02d}:00:00", "%Y-%m-%d %H:%M:%S")

        # Get predictions
        predictions = await predict_outages(target_time, horizon=6)
        if predictions is None or 'error' in predictions:
            return {'date': date_str, 'error': predictions.get('error', 'No predictions') if predictions else 'None returned'}

        # Get actuals
        actuals = get_actual_outages(target_time)
        if actuals is None:
            return {'date': date_str, 'error': 'Failed to get actuals'}

        pred_total = predictions.get('total_predicted', 0)
        actual_total = actuals['total']

        # State-level correlation
        pred_by_state = {p['state']: p['predicted_outages'] for p in predictions.get('by_state', [])}
        actual_by_state = actuals['by_state']

        common_states = set(pred_by_state.keys()) & set(actual_by_state.keys())
        if len(common_states) > 5:
            pred_vals = [pred_by_state[s] for s in common_states]
            actual_vals = [actual_by_state[s] for s in common_states]
            correlation = np.corrcoef(pred_vals, actual_vals)[0, 1]
        else:
            correlation = 0

        return {
            'date': date_str,
            'hour': hour,
            'description': description,
            'predicted': pred_total,
            'actual': actual_total,
            'ratio': pred_total / actual_total if actual_total > 0 else None,
            'error': pred_total - actual_total,
            'abs_error': abs(pred_total - actual_total),
            'pct_error': (pred_total - actual_total) / actual_total * 100 if actual_total > 0 else None,
            'correlation': correlation if not np.isnan(correlation) else 0,
            'states_matched': len(common_states),
        }
    except Exception as e:
        logger.error(f"Error evaluating {date_str}: {e}")
        return {'date': date_str, 'error': str(e)}


async def run_evaluation():
    """Run full evaluation across all test dates."""
    results = {'normal': [], 'storms': [], 'weekends': []}

    for category, dates in TEST_DATES.items():
        logger.info(f"\n{'='*60}")
        logger.info(f"Evaluating {category.upper()} dates ({len(dates)} samples)")
        logger.info('='*60)

        for date_str, hour, desc in dates:
            logger.info(f"  {date_str} {hour:02d}:00 - {desc}")
            result = await evaluate_single_date(date_str, hour, desc)
            results[category].append(result)

            if 'error' not in result:
                ratio = result['ratio'] or 0
                logger.info(f"    Pred: {result['predicted']:>12,}, Actual: {result['actual']:>12,}, "
                           f"Ratio: {ratio:.2f}x, Corr: {result['correlation']:.2f}")

    return results


def compute_summary_stats(results: list) -> dict:
    """Compute summary statistics for a set of results."""
    valid = [r for r in results if 'error' not in r and r.get('actual', 0) > 0]

    if not valid:
        return {'n': 0}

    ratios = [r['ratio'] for r in valid if r['ratio'] is not None]
    correlations = [r['correlation'] for r in valid]
    abs_errors = [r['abs_error'] for r in valid]
    pct_errors = [r['pct_error'] for r in valid if r['pct_error'] is not None]

    return {
        'n': len(valid),
        'mean_ratio': np.mean(ratios) if ratios else 0,
        'median_ratio': np.median(ratios) if ratios else 0,
        'std_ratio': np.std(ratios) if ratios else 0,
        'mean_correlation': np.mean(correlations),
        'median_correlation': np.median(correlations),
        'mae': np.mean(abs_errors),
        'rmse': np.sqrt(np.mean([e**2 for e in abs_errors])),
        'mean_pct_error': np.mean(pct_errors) if pct_errors else 0,
        'median_pct_error': np.median(pct_errors) if pct_errors else 0,
    }


def print_report(results: dict):
    """Print evaluation report."""
    print("\n" + "="*80)
    print("MODEL EVALUATION REPORT")
    print("="*80)

    for category, data in results.items():
        stats = compute_summary_stats(data)

        print(f"\n{category.upper()} ({stats['n']} samples)")
        print("-"*40)

        if stats['n'] == 0:
            print("  No valid results")
            continue

        print(f"  Prediction Ratio (pred/actual):")
        print(f"    Mean:   {stats['mean_ratio']:.2f}x")
        print(f"    Median: {stats['median_ratio']:.2f}x")
        print(f"    Std:    {stats['std_ratio']:.2f}")
        print(f"  State-Level Correlation:")
        print(f"    Mean:   {stats['mean_correlation']:.3f}")
        print(f"    Median: {stats['median_correlation']:.3f}")
        print(f"  Error Metrics:")
        print(f"    MAE:    {stats['mae']:,.0f} customers")
        print(f"    RMSE:   {stats['rmse']:,.0f} customers")
        print(f"    Mean % Error: {stats['mean_pct_error']:.1f}%")

    # Overall
    all_results = results['normal'] + results['storms'] + results['weekends']
    overall = compute_summary_stats(all_results)

    print(f"\nOVERALL ({overall['n']} samples)")
    print("-"*40)
    print(f"  Mean Ratio:       {overall['mean_ratio']:.2f}x")
    print(f"  Mean Correlation: {overall['mean_correlation']:.3f}")
    print(f"  MAE:              {overall['mae']:,.0f}")
    print(f"  RMSE:             {overall['rmse']:,.0f}")

    # Detailed results table
    print("\n" + "="*80)
    print("DETAILED RESULTS")
    print("="*80)
    print(f"{'Date':<12} {'Description':<25} {'Predicted':>12} {'Actual':>12} {'Ratio':>8} {'Corr':>6}")
    print("-"*80)

    for category, data in results.items():
        for r in sorted(data, key=lambda x: x.get('actual', 0), reverse=True):
            if 'error' in r:
                print(f"{r['date']:<12} ERROR: {r['error'][:40]}")
            else:
                ratio = r['ratio'] or 0
                print(f"{r['date']:<12} {r['description']:<25} {r['predicted']:>12,} {r['actual']:>12,} "
                      f"{ratio:>7.2f}x {r['correlation']:>6.2f}")

    return overall


async def main():
    """Main evaluation pipeline."""
    logger.info("Starting comprehensive model evaluation...")

    results = await run_evaluation()
    overall = print_report(results)

    # Save results to CSV
    all_results = []
    for category, data in results.items():
        for r in data:
            r['category'] = category
            all_results.append(r)

    df = pd.DataFrame(all_results)
    df.to_csv('evaluation_results.csv', index=False)
    logger.info("\nResults saved to evaluation_results.csv")

    return results, overall


if __name__ == '__main__':
    results, overall = asyncio.run(main())