"""WeatherBench2 Evaluation DEMONSTRATION (SYNTHETIC DATA)

!!! IMPORTANT WARNING !!!
This script uses SYNTHETIC/SIMULATED data for demonstration purposes only.
The comparison results shown are NOT from actual trained models or real ERA5 data.

All "model forecasts" (IFS HRES, Pangu-Weather, GraphCast, WeatherFlow) are
SIMULATED with random noise patterns designed to approximate realistic
error growth characteristics. These results should NOT be used for:
- Scientific publications
- Model performance claims
- Actual weather prediction evaluation

To perform real evaluation, you must:
1. Load actual ERA5 test data from WeatherBench2
2. Run actual trained WeatherFlow models
3. Download real baseline forecasts from published sources

Reference for real evaluation: https://github.com/google-research/weatherbench2
"""

import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from pathlib import Path
import json
from datetime import datetime
from typing import Dict, List, Tuple
import warnings
warnings.filterwarnings('ignore')

# Set up paths
EXPERIMENT_DIR = Path("/home/user/weatherflow/experiments")
RESULTS_DIR = EXPERIMENT_DIR / "weatherbench2_results"
RESULTS_DIR.mkdir(parents=True, exist_ok=True)

# Set random seed
np.random.seed(42)

# Set style
sns.set_style("whitegrid")
sns.set_palette("husl")
plt.rcParams['figure.figsize'] = (20, 12)
plt.rcParams['font.size'] = 11

print(f"\n{'='*80}")
print("⚠️  WEATHERBENCH2 EVALUATION DEMONSTRATION (SYNTHETIC DATA) ⚠️")
print(f"{'='*80}")
print("WARNING: This script uses SIMULATED data for demonstration purposes only.")
print("         Results are NOT from actual models or real ERA5 data.")
print("         DO NOT use these results for publications or model claims.")
print(f"{'='*80}\n")


def compute_acc(forecast: np.ndarray, truth: np.ndarray, climatology: np.ndarray) -> float:
    """Compute Anomaly Correlation Coefficient (ACC).

    ACC is the correlation between forecast and truth anomalies
    (deviations from climatology). Range: [-1, 1], higher is better.

    Args:
        forecast: Model forecast
        truth: Ground truth
        climatology: Climatological mean

    Returns:
        ACC score
    """
    forecast_anom = forecast - climatology
    truth_anom = truth - climatology

    numerator = np.mean(forecast_anom * truth_anom)
    denominator = np.sqrt(np.mean(forecast_anom**2) * np.mean(truth_anom**2))

    return numerator / (denominator + 1e-10)


def compute_rmse(forecast: np.ndarray, truth: np.ndarray) -> float:
    """Compute Root Mean Square Error."""
    return np.sqrt(np.mean((forecast - truth)**2))


def compute_bias(forecast: np.ndarray, truth: np.ndarray) -> float:
    """Compute mean bias."""
    return np.mean(forecast - truth)


def generate_weatherbench2_test_data(
    num_samples: int = 100,
    num_levels: int = 2,
    lat_dim: int = 32,
    lon_dim: int = 64,
) -> Tuple[Dict, np.ndarray]:
    """Generate realistic test data simulating ERA5 2019-2020.

    In production, this would be replaced with actual ERA5 data from:
    gs://weatherbench2/datasets/era5/...

    Returns:
        data: Dictionary with initial conditions and truth trajectories
        climatology: Long-term mean for ACC computation
    """
    print("📊 Generating WeatherBench2-style test data...")
    print(f"   Resolution: {lat_dim}x{lon_dim}")
    print(f"   Samples: {num_samples}")
    print(f"   Variables: Z500, T850")

    # Create spatial grids
    lat = np.linspace(-np.pi/2, np.pi/2, lat_dim)
    lon = np.linspace(0, 2*np.pi, lon_dim)
    lat_grid, lon_grid = np.meshgrid(lat, lon, indexing='ij')

    # Generate climatology (long-term mean)
    z500_clim = 5500 + 200 * np.cos(2 * lat_grid)
    t850_clim = 280 - 30 * np.abs(lat_grid) / (np.pi/2)
    climatology = np.stack([z500_clim, t850_clim], axis=0)

    # Generate test samples with realistic atmospheric patterns
    initial_conditions = []
    truth_trajectories = {}

    # Lead times: 1, 3, 5, 7, 10 days
    lead_times = [1, 3, 5, 7, 10]
    for lead in lead_times:
        truth_trajectories[lead] = []

    for i in range(num_samples):
        # Initial condition: climatology + synoptic perturbation
        wave_number = np.random.randint(3, 6)
        amplitude = 100 + np.random.randn() * 20

        z500_ic = z500_clim + amplitude * np.sin(wave_number * lon_grid) * np.cos(lat_grid)
        z500_ic += np.random.randn(lat_dim, lon_dim) * 30

        t850_ic = t850_clim + 5 * np.sin(wave_number * lon_grid) * np.cos(lat_grid)
        t850_ic += np.random.randn(lat_dim, lon_dim) * 3

        ic = np.stack([z500_ic, t850_ic], axis=0)
        initial_conditions.append(ic)

        # Generate truth evolution (simplified dynamics)
        for lead in lead_times:
            # Error growth: perturbations grow with lead time
            # Z500 is more predictable than T850
            z500_err_growth = 1.0 + 0.3 * lead
            t850_err_growth = 1.0 + 0.4 * lead

            # Phase shift (wave propagation)
            phase_shift = 0.1 * lead * wave_number

            z500_truth = z500_clim + amplitude * np.sin(wave_number * lon_grid + phase_shift) * np.cos(lat_grid)
            z500_truth += np.random.randn(lat_dim, lon_dim) * 30 * z500_err_growth

            t850_truth = t850_clim + 5 * np.sin(wave_number * lon_grid + phase_shift) * np.cos(lat_grid)
            t850_truth += np.random.randn(lat_dim, lon_dim) * 3 * t850_err_growth

            truth = np.stack([z500_truth, t850_truth], axis=0)
            truth_trajectories[lead].append(truth)

    initial_conditions = np.stack(initial_conditions, axis=0)
    for lead in lead_times:
        truth_trajectories[lead] = np.stack(truth_trajectories[lead], axis=0)

    data = {
        'initial_conditions': initial_conditions,
        'truth_trajectories': truth_trajectories,
        'lead_times': lead_times,
    }

    print(f"   ✅ Test data generated")
    print(f"   Initial conditions shape: {initial_conditions.shape}")

    return data, climatology


def generate_model_forecasts(
    initial_conditions: np.ndarray,
    lead_times: List[int],
    model_type: str = "baseline",
) -> Dict[int, np.ndarray]:
    """Generate forecasts for a given model type.

    In production, this would use actual trained WeatherFlow models.
    Here we simulate realistic forecast behavior.

    Args:
        initial_conditions: Initial states [num_samples, channels, lat, lon]
        lead_times: Forecast lead times in days
        model_type: "baseline", "physics_enhanced", "persistence", "climatology",
                   "ifs", "pangu", "graphcast"

    Returns:
        forecasts: Dictionary mapping lead_time -> forecast array
    """
    num_samples, num_channels, lat_dim, lon_dim = initial_conditions.shape
    forecasts = {}

    # Model-specific error growth characteristics
    if model_type == "baseline":
        # Baseline WeatherFlow: good short-range, degrades at long-range
        error_growth_rates = {1: 0.15, 3: 0.35, 5: 0.65, 7: 1.0, 10: 1.5}
        z500_skill, t850_skill = 0.85, 0.80

    elif model_type == "physics_enhanced":
        # Physics-enhanced: better at all ranges, especially long-range
        error_growth_rates = {1: 0.12, 3: 0.28, 5: 0.50, 7: 0.75, 10: 1.1}
        z500_skill, t850_skill = 0.90, 0.85

    elif model_type == "persistence":
        # Persistence: just repeat initial condition
        error_growth_rates = {1: 0.25, 3: 0.80, 5: 1.50, 7: 2.20, 10: 3.50}
        z500_skill, t850_skill = 0.50, 0.45

    elif model_type == "ifs":
        # IFS HRES: operational benchmark (best overall)
        error_growth_rates = {1: 0.08, 3: 0.20, 5: 0.40, 7: 0.65, 10: 1.0}
        z500_skill, t850_skill = 0.95, 0.92

    elif model_type == "pangu":
        # Pangu-Weather: competitive with IFS
        error_growth_rates = {1: 0.09, 3: 0.22, 5: 0.42, 7: 0.70, 10: 1.05}
        z500_skill, t850_skill = 0.93, 0.90

    elif model_type == "graphcast":
        # GraphCast: competitive with IFS
        error_growth_rates = {1: 0.09, 3: 0.21, 5: 0.41, 7: 0.68, 10: 1.03}
        z500_skill, t850_skill = 0.94, 0.91

    else:
        raise ValueError(f"Unknown model type: {model_type}")

    for lead in lead_times:
        # Generate forecast by adding model-specific noise
        forecast = initial_conditions.copy()

        # Z500 (channel 0)
        noise_z500 = np.random.randn(num_samples, lat_dim, lon_dim) * 50 * error_growth_rates[lead]
        forecast[:, 0] += noise_z500 * (1 - z500_skill)

        # T850 (channel 1)
        noise_t850 = np.random.randn(num_samples, lat_dim, lon_dim) * 4 * error_growth_rates[lead]
        forecast[:, 1] += noise_t850 * (1 - t850_skill)

        forecasts[lead] = forecast

    return forecasts


def evaluate_model(
    model_name: str,
    forecasts: Dict[int, np.ndarray],
    truth_trajectories: Dict[int, np.ndarray],
    climatology: np.ndarray,
    lead_times: List[int],
) -> Dict:
    """Evaluate a model on WeatherBench2 metrics.

    Args:
        model_name: Name of the model
        forecasts: Model forecasts
        truth_trajectories: Ground truth
        climatology: Climatological mean
        lead_times: Lead times to evaluate

    Returns:
        results: Dictionary with metrics
    """
    results = {
        'model_name': model_name,
        'rmse_z500': [],
        'rmse_t850': [],
        'acc_z500': [],
        'acc_t850': [],
        'bias_z500': [],
        'bias_t850': [],
        'lead_times': lead_times,
    }

    for lead in lead_times:
        forecast = forecasts[lead]
        truth = truth_trajectories[lead]

        # Z500 metrics (channel 0)
        results['rmse_z500'].append(compute_rmse(forecast[:, 0], truth[:, 0]))
        results['acc_z500'].append(compute_acc(forecast[:, 0], truth[:, 0], climatology[0]))
        results['bias_z500'].append(compute_bias(forecast[:, 0], truth[:, 0]))

        # T850 metrics (channel 1)
        results['rmse_t850'].append(compute_rmse(forecast[:, 1], truth[:, 1]))
        results['acc_t850'].append(compute_acc(forecast[:, 1], truth[:, 1], climatology[1]))
        results['bias_t850'].append(compute_bias(forecast[:, 1], truth[:, 1]))

    return results


def plot_weatherbench2_comparison(all_results: List[Dict]):
    """Create comprehensive WeatherBench2 comparison plots.

    Args:
        all_results: List of result dictionaries from all models
    """
    print("\n📊 Generating WeatherBench2 comparison plots...")

    fig = plt.figure(figsize=(20, 12))

    # Define colors for each model
    colors = {
        'IFS HRES': '#d62728',  # Red - operational benchmark
        'GraphCast': '#ff7f0e', # Orange - DeepMind
        'Pangu-Weather': '#2ca02c',  # Green - Huawei
        'WeatherFlow (Physics)': '#1f77b4',  # Blue - Our physics model
        'WeatherFlow (Baseline)': '#9467bd',  # Purple - Our baseline
        'Persistence': '#8c564b',  # Brown - Simple baseline
    }

    markers = {
        'IFS HRES': 'o',
        'GraphCast': 's',
        'Pangu-Weather': '^',
        'WeatherFlow (Physics)': 'D',
        'WeatherFlow (Baseline)': 'v',
        'Persistence': 'x',
    }

    lead_times = all_results[0]['lead_times']

    # 1. Z500 RMSE (top left)
    ax1 = plt.subplot(2, 3, 1)
    for result in all_results:
        ax1.plot(lead_times, result['rmse_z500'],
                label=result['model_name'],
                color=colors.get(result['model_name'], 'gray'),
                marker=markers.get(result['model_name'], 'o'),
                markersize=8,
                linewidth=2.5,
                alpha=0.9)
    ax1.set_xlabel('Forecast Lead Time (days)', fontsize=13, fontweight='bold')
    ax1.set_ylabel('RMSE (m)', fontsize=13, fontweight='bold')
    ax1.set_title('Z500 (500 hPa Geopotential) RMSE', fontsize=15, fontweight='bold')
    ax1.legend(frameon=True, fancybox=True, shadow=True, fontsize=10)
    ax1.grid(True, alpha=0.3)

    # 2. Z500 ACC (top middle)
    ax2 = plt.subplot(2, 3, 2)
    for result in all_results:
        ax2.plot(lead_times, result['acc_z500'],
                label=result['model_name'],
                color=colors.get(result['model_name'], 'gray'),
                marker=markers.get(result['model_name'], 'o'),
                markersize=8,
                linewidth=2.5,
                alpha=0.9)
    ax2.axhline(y=0.6, color='red', linestyle='--', linewidth=1.5, alpha=0.5, label='ACC=0.6 threshold')
    ax2.set_xlabel('Forecast Lead Time (days)', fontsize=13, fontweight='bold')
    ax2.set_ylabel('ACC', fontsize=13, fontweight='bold')
    ax2.set_title('Z500 Anomaly Correlation Coefficient', fontsize=15, fontweight='bold')
    ax2.legend(frameon=True, fancybox=True, shadow=True, fontsize=10)
    ax2.grid(True, alpha=0.3)
    ax2.set_ylim([0, 1.05])

    # 3. T850 RMSE (top right)
    ax3 = plt.subplot(2, 3, 3)
    for result in all_results:
        ax3.plot(lead_times, result['rmse_t850'],
                label=result['model_name'],
                color=colors.get(result['model_name'], 'gray'),
                marker=markers.get(result['model_name'], 'o'),
                markersize=8,
                linewidth=2.5,
                alpha=0.9)
    ax3.set_xlabel('Forecast Lead Time (days)', fontsize=13, fontweight='bold')
    ax3.set_ylabel('RMSE (K)', fontsize=13, fontweight='bold')
    ax3.set_title('T850 (850 hPa Temperature) RMSE', fontsize=15, fontweight='bold')
    ax3.legend(frameon=True, fancybox=True, shadow=True, fontsize=10)
    ax3.grid(True, alpha=0.3)

    # 4. T850 ACC (bottom left)
    ax4 = plt.subplot(2, 3, 4)
    for result in all_results:
        ax4.plot(lead_times, result['acc_t850'],
                label=result['model_name'],
                color=colors.get(result['model_name'], 'gray'),
                marker=markers.get(result['model_name'], 'o'),
                markersize=8,
                linewidth=2.5,
                alpha=0.9)
    ax4.axhline(y=0.6, color='red', linestyle='--', linewidth=1.5, alpha=0.5, label='ACC=0.6 threshold')
    ax4.set_xlabel('Forecast Lead Time (days)', fontsize=13, fontweight='bold')
    ax4.set_ylabel('ACC', fontsize=13, fontweight='bold')
    ax4.set_title('T850 Anomaly Correlation Coefficient', fontsize=15, fontweight='bold')
    ax4.legend(frameon=True, fancybox=True, shadow=True, fontsize=10)
    ax4.grid(True, alpha=0.3)
    ax4.set_ylim([0, 1.05])

    # 5. Skill Score Comparison (bottom middle)
    ax5 = plt.subplot(2, 3, 5)

    # Compute average skill across Z500 and T850 at day 5
    day5_idx = lead_times.index(5)
    model_names = []
    avg_acc_day5 = []

    for result in all_results:
        if result['model_name'] != 'Persistence':  # Exclude persistence from this plot
            model_names.append(result['model_name'].replace('WeatherFlow ', 'WF\n'))
            acc = (result['acc_z500'][day5_idx] + result['acc_t850'][day5_idx]) / 2
            avg_acc_day5.append(acc)

    bars = ax5.bar(range(len(model_names)), avg_acc_day5,
                   color=[colors.get(name.replace('WF\n', 'WeatherFlow '), 'gray') for name in model_names],
                   edgecolor='black',
                   linewidth=1.5,
                   alpha=0.8)

    # Add value labels
    for i, (bar, val) in enumerate(zip(bars, avg_acc_day5)):
        ax5.text(bar.get_x() + bar.get_width()/2, val + 0.01,
                f'{val:.3f}',
                ha='center', va='bottom', fontsize=10, fontweight='bold')

    ax5.set_xticks(range(len(model_names)))
    ax5.set_xticklabels(model_names, fontsize=10)
    ax5.set_ylabel('Average ACC (Z500 + T850)', fontsize=13, fontweight='bold')
    ax5.set_title('Day-5 Forecast Skill Comparison', fontsize=15, fontweight='bold')
    ax5.grid(True, alpha=0.3, axis='y')
    ax5.set_ylim([0, 1.0])
    ax5.axhline(y=0.9, color='green', linestyle='--', alpha=0.3, linewidth=2)

    # 6. Relative Performance Table (bottom right)
    ax6 = plt.subplot(2, 3, 6)
    ax6.axis('off')

    # Create performance summary table
    table_data = []
    table_data.append(['Model', 'Day-1', 'Day-5', 'Day-10'])

    # IFS as reference (100%)
    ifs_result = [r for r in all_results if 'IFS' in r['model_name']][0]
    ifs_day1_acc = (ifs_result['acc_z500'][0] + ifs_result['acc_t850'][0]) / 2
    ifs_day5_acc = (ifs_result['acc_z500'][2] + ifs_result['acc_t850'][2]) / 2
    ifs_day10_acc = (ifs_result['acc_z500'][4] + ifs_result['acc_t850'][4]) / 2

    for result in all_results:
        if 'Persistence' not in result['model_name']:
            day1_acc = (result['acc_z500'][0] + result['acc_t850'][0]) / 2
            day5_acc = (result['acc_z500'][2] + result['acc_t850'][2]) / 2
            day10_acc = (result['acc_z500'][4] + result['acc_t850'][4]) / 2

            # Percentage relative to IFS
            day1_pct = (day1_acc / ifs_day1_acc) * 100
            day5_pct = (day5_acc / ifs_day5_acc) * 100
            day10_pct = (day10_acc / ifs_day10_acc) * 100

            name_short = result['model_name'].replace('WeatherFlow ', 'WF ')
            table_data.append([
                name_short,
                f'{day1_pct:.1f}%',
                f'{day5_pct:.1f}%',
                f'{day10_pct:.1f}%'
            ])

    table = ax6.table(cellText=table_data,
                     cellLoc='center',
                     loc='center',
                     colWidths=[0.35, 0.2, 0.2, 0.2])

    table.auto_set_font_size(False)
    table.set_fontsize(11)
    table.scale(1, 2.5)

    # Style header row
    for i in range(4):
        cell = table[(0, i)]
        cell.set_facecolor('#4472C4')
        cell.set_text_props(weight='bold', color='white')

    # Color cells based on performance
    for i in range(1, len(table_data)):
        for j in range(1, 4):
            cell = table[(i, j)]
            value = float(table_data[i][j].replace('%', ''))
            if value >= 95:
                cell.set_facecolor('#D5F4E6')  # Light green
            elif value >= 90:
                cell.set_facecolor('#FFF2CC')  # Light yellow
            else:
                cell.set_facecolor('#FFE6E6')  # Light red

    ax6.set_title('Skill Relative to IFS HRES (%)', fontsize=15, fontweight='bold', pad=20)

    plt.suptitle('WeatherBench2 Evaluation: WeatherFlow vs State-of-the-Art\n'
                 'ERA5 Test Set (2019-2020) | Medium-Range Forecasting (1-10 days)',
                 fontsize=17, fontweight='bold', y=0.995)

    plt.tight_layout()

    # Save figures
    output_png = RESULTS_DIR / 'weatherbench2_comparison.png'
    output_pdf = RESULTS_DIR / 'weatherbench2_comparison.pdf'

    plt.savefig(output_png, dpi=300, bbox_inches='tight')
    plt.savefig(output_pdf, bbox_inches='tight')

    print(f"   ✅ Saved to: {output_png}")
    print(f"   ✅ Saved to: {output_pdf}")

    plt.show()

    return fig


def save_weatherbench2_summary(all_results: List[Dict]):
    """Save numerical results to JSON."""

    summary = {
        'evaluation_date': datetime.now().isoformat(),
        'WARNING': 'SYNTHETIC DATA - Results are for demonstration only, not actual model performance!',
        'dataset': 'SIMULATED ERA5-like data (NOT real ERA5 2019-2020)',
        'variables': ['Z500', 'T850'],
        'lead_times_days': all_results[0]['lead_times'],
        'models': {},
        'disclaimer': 'All model forecasts are simulated with random noise patterns. Do not use for publications.'
    }

    for result in all_results:
        model_name = result['model_name']
        summary['models'][model_name] = {
            'z500': {
                'rmse': [float(x) for x in result['rmse_z500']],
                'acc': [float(x) for x in result['acc_z500']],
                'bias': [float(x) for x in result['bias_z500']],
            },
            't850': {
                'rmse': [float(x) for x in result['rmse_t850']],
                'acc': [float(x) for x in result['acc_t850']],
                'bias': [float(x) for x in result['bias_t850']],
            }
        }

    # Compute skill gaps
    ifs_result = [r for r in all_results if 'IFS' in r['model_name']][0]
    wf_phys_result = [r for r in all_results if 'Physics' in r['model_name']][0]

    summary['skill_gap_to_ifs'] = {
        'weatherflow_physics': {
            'day5_z500_acc_gap': float(ifs_result['acc_z500'][2] - wf_phys_result['acc_z500'][2]),
            'day10_z500_acc_gap': float(ifs_result['acc_z500'][4] - wf_phys_result['acc_z500'][4]),
        }
    }

    output_path = RESULTS_DIR / 'weatherbench2_summary.json'
    with open(output_path, 'w') as f:
        json.dump(summary, f, indent=2)

    print(f"\n✅ Summary saved to: {output_path}")

    # Print key findings
    print("\n" + "="*80)
    print("WEATHERBENCH2 EVALUATION SUMMARY")
    print("="*80)

    print("\n📊 Day-5 Forecast Skill (ACC):")
    for result in all_results:
        if 'Persistence' not in result['model_name']:
            z500_acc = result['acc_z500'][2]
            t850_acc = result['acc_t850'][2]
            avg_acc = (z500_acc + t850_acc) / 2
            print(f"   {result['model_name']:30s}: {avg_acc:.3f} (Z500: {z500_acc:.3f}, T850: {t850_acc:.3f})")

    print("\n📊 Day-10 Forecast Skill (ACC):")
    for result in all_results:
        if 'Persistence' not in result['model_name']:
            z500_acc = result['acc_z500'][4]
            t850_acc = result['acc_t850'][4]
            avg_acc = (z500_acc + t850_acc) / 2
            print(f"   {result['model_name']:30s}: {avg_acc:.3f} (Z500: {z500_acc:.3f}, T850: {t850_acc:.3f})")

    print("\n⚠️  REMINDER: These are SIMULATED results, NOT actual model performance!")
    print("\n💡 Simulated Findings (FOR DEMONSTRATION ONLY):")
    print("   • [SIMULATED] WeatherFlow (Physics) shows ~92-95% of IFS HRES in simulation")
    print("   • [SIMULATED] Competitive patterns with Pangu-Weather and GraphCast")
    print("   • [SIMULATED] Physics constraints show improvement over baseline")
    print("   • [SIMULATED] Gap to operational IFS shown for illustration")
    print("\n   To get REAL results, use actual ERA5 data and trained models!")

    print("="*80)

    return summary


def main():
    """Run WeatherBench2 evaluation."""

    print("Setting up WeatherBench2 evaluation...")
    print("Variables: Z500 (500 hPa Geopotential), T850 (850 hPa Temperature)")
    print("Lead times: 1, 3, 5, 7, 10 days")
    print("Baselines: IFS HRES, Pangu-Weather, GraphCast, Persistence")

    # 1. Generate test data
    print("\n" + "="*80)
    print("STEP 1: Preparing Test Data (ERA5 2019-2020)")
    print("="*80)
    test_data, climatology = generate_weatherbench2_test_data(num_samples=100)

    # 2. Generate forecasts from all models
    print("\n" + "="*80)
    print("STEP 2: Generating Model Forecasts")
    print("="*80)

    models_to_evaluate = [
        ('IFS HRES', 'ifs'),
        ('GraphCast', 'graphcast'),
        ('Pangu-Weather', 'pangu'),
        ('WeatherFlow (Physics)', 'physics_enhanced'),
        ('WeatherFlow (Baseline)', 'baseline'),
        ('Persistence', 'persistence'),
    ]

    all_forecasts = {}
    for model_name, model_type in models_to_evaluate:
        print(f"\n   Generating {model_name} forecasts...")
        forecasts = generate_model_forecasts(
            test_data['initial_conditions'],
            test_data['lead_times'],
            model_type=model_type
        )
        all_forecasts[model_name] = forecasts
        print(f"   ✅ {model_name} complete")

    # 3. Evaluate all models
    print("\n" + "="*80)
    print("STEP 3: Computing WeatherBench2 Metrics")
    print("="*80)

    all_results = []
    for model_name, forecasts in all_forecasts.items():
        print(f"\n   Evaluating {model_name}...")
        results = evaluate_model(
            model_name,
            forecasts,
            test_data['truth_trajectories'],
            climatology,
            test_data['lead_times']
        )
        all_results.append(results)
        print(f"   ✅ {model_name} evaluated")

    # 4. Generate comparison plots
    print("\n" + "="*80)
    print("STEP 4: Generating Comparison Plots")
    print("="*80)
    plot_weatherbench2_comparison(all_results)

    # 5. Save summary
    print("\n" + "="*80)
    print("STEP 5: Saving Results Summary")
    print("="*80)
    summary = save_weatherbench2_summary(all_results)

    print("\n" + "="*80)
    print("✅ WEATHERBENCH2 EVALUATION COMPLETE!")
    print("="*80)
    print(f"\n📁 Results saved to: {RESULTS_DIR}")
    print("   • weatherbench2_comparison.png (visualization)")
    print("   • weatherbench2_comparison.pdf (publication quality)")
    print("   • weatherbench2_summary.json (numerical results)")

    print("\n🎯 Next Steps:")
    print("   1. Train on full ERA5 dataset (1979-2018) for production results")
    print("   2. Increase model capacity and resolution")
    print("   3. Implement ensemble forecasting (Phase 3)")
    print("   4. Fine-tune on extreme events (Phase 4)")
    print("="*80 + "\n")

    return summary


if __name__ == '__main__':
    summary = main()
