//! Sample data and configurations for WeatherCast demo.
//!
//! Provides pre-configured scenarios for different types of weather
//! predictions including global forecasts, regional high-resolution
//! forecasts, and ensemble predictions.

use weathercast_shared::{
    AtmosphericState, GridPoint, MeshConfig, PredictionConfig, TrainingConfig,
};

// ============================================================================
// Global Forecast Configurations
// ============================================================================

/// Create a global medium-range forecast configuration.
#[must_use]
pub fn global_forecast() -> (MeshConfig, PredictionConfig) {
    let mesh_config = MeshConfig::from_refinement(5); // ~40km resolution

    let pred_config = PredictionConfig {
        lead_time: 240.0, // 10 days
        ensemble_size: 1,
        resolution: 0.25,
        time_step: 6.0,
        num_steps: 40,
        multi_scale: true,
        variables: vec![
            "temperature".to_string(),
            "geopotential".to_string(),
            "humidity".to_string(),
            "wind_u".to_string(),
            "wind_v".to_string(),
        ],
        pressure_levels: vec![1000.0, 925.0, 850.0, 700.0, 500.0, 300.0, 200.0, 50.0],
    };

    (mesh_config, pred_config)
}

/// Create a short-range global forecast (0-72h).
#[must_use]
pub fn global_short_range() -> (MeshConfig, PredictionConfig) {
    let mesh_config = MeshConfig::from_refinement(5);

    let pred_config = PredictionConfig {
        lead_time: 72.0,
        time_step: 3.0,
        num_steps: 24,
        ..PredictionConfig::default()
    };

    (mesh_config, pred_config)
}

// ============================================================================
// Regional Forecast Configurations
// ============================================================================

/// Create a regional high-resolution forecast for North America.
#[must_use]
pub fn regional_forecast() -> (MeshConfig, PredictionConfig) {
    let mesh_config = MeshConfig {
        refinement_levels: 6,
        num_nodes: 40962,
        num_edges: 122880,
        num_faces: 81920,
        avg_edge_length: 20.0, // ~20km
        multi_scale: true,
        regional_focus: Some((40.0, -100.0, 2000.0)), // Central US, 2000km radius
    };

    let pred_config = PredictionConfig {
        lead_time: 48.0, // 2 days
        time_step: 1.0,  // Hourly
        num_steps: 48,
        resolution: 0.1, // ~10km effective
        ensemble_size: 1,
        multi_scale: true,
        ..PredictionConfig::default()
    };

    (mesh_config, pred_config)
}

/// Create a European regional forecast.
#[must_use]
pub fn european_regional() -> (MeshConfig, PredictionConfig) {
    let mesh_config = MeshConfig {
        refinement_levels: 6,
        num_nodes: 40962,
        num_edges: 122880,
        num_faces: 81920,
        avg_edge_length: 15.0,
        multi_scale: true,
        regional_focus: Some((50.0, 10.0, 1500.0)), // Central Europe
    };

    let pred_config = PredictionConfig {
        lead_time: 72.0,
        time_step: 1.0,
        num_steps: 72,
        resolution: 0.1,
        ..PredictionConfig::default()
    };

    (mesh_config, pred_config)
}

/// Create an Asian regional forecast.
#[must_use]
pub fn asian_regional() -> (MeshConfig, PredictionConfig) {
    let mesh_config = MeshConfig::regional(35.0, 135.0, 2000.0); // Japan/East Asia

    let pred_config = PredictionConfig {
        lead_time: 72.0,
        time_step: 3.0,
        num_steps: 24,
        ..PredictionConfig::default()
    };

    (mesh_config, pred_config)
}

// ============================================================================
// Ensemble Forecast Configurations
// ============================================================================

/// Create ensemble forecast configuration (50 members).
#[must_use]
pub fn ensemble_forecast() -> (MeshConfig, PredictionConfig) {
    let mesh_config = MeshConfig::from_refinement(4); // Coarser for ensemble

    let pred_config = PredictionConfig {
        lead_time: 360.0, // 15 days
        ensemble_size: 50,
        time_step: 6.0,
        num_steps: 60,
        resolution: 0.5,
        multi_scale: false,
        ..PredictionConfig::default()
    };

    (mesh_config, pred_config)
}

/// Create large ensemble (100 members) for probabilistic forecasting.
#[must_use]
pub fn large_ensemble() -> (MeshConfig, PredictionConfig) {
    let mesh_config = MeshConfig::from_refinement(3); // Even coarser

    let pred_config = PredictionConfig {
        lead_time: 240.0,
        ensemble_size: 100,
        time_step: 12.0,
        num_steps: 20,
        resolution: 1.0,
        ..PredictionConfig::default()
    };

    (mesh_config, pred_config)
}

// ============================================================================
// High-Resolution Configurations
// ============================================================================

/// Create ultra high-resolution forecast (convection-permitting).
#[must_use]
pub fn high_resolution_forecast() -> (MeshConfig, PredictionConfig) {
    let mesh_config = MeshConfig {
        refinement_levels: 7,
        num_nodes: 163842,
        num_edges: 491520,
        num_faces: 327680,
        avg_edge_length: 5.0, // ~5km (convection-permitting)
        multi_scale: true,
        regional_focus: Some((40.0, -100.0, 500.0)), // Small regional domain
    };

    let pred_config = PredictionConfig {
        lead_time: 24.0, // 1 day
        time_step: 0.5,  // 30-minute steps
        num_steps: 48,
        resolution: 0.05,
        ensemble_size: 1,
        multi_scale: true,
        variables: vec![
            "temperature".to_string(),
            "geopotential".to_string(),
            "humidity".to_string(),
            "wind_u".to_string(),
            "wind_v".to_string(),
            "precipitation".to_string(),
            "cloud_cover".to_string(),
        ],
        pressure_levels: vec![
            1000.0, 925.0, 850.0, 700.0, 600.0, 500.0, 400.0, 300.0, 250.0, 200.0, 150.0, 100.0,
            50.0,
        ],
    };

    (mesh_config, pred_config)
}

// ============================================================================
// Sample Atmospheric States
// ============================================================================

/// Create initial state for mid-latitude cyclone.
#[must_use]
pub fn midlatitude_cyclone() -> AtmosphericState {
    AtmosphericState {
        temperature: vec![280.0, 278.0, 275.0, 265.0, 248.0, 225.0, 215.0, 200.0],
        geopotential: vec![0.0, 800.0, 1500.0, 3100.0, 5600.0, 9200.0, 11800.0, 20500.0],
        humidity: vec![0.015, 0.012, 0.008, 0.003, 0.001, 0.0002, 0.00005, 0.000001],
        wind_u: vec![-5.0, -8.0, -15.0, -25.0, -40.0, -50.0, -45.0, -25.0],
        wind_v: vec![10.0, 15.0, 20.0, 25.0, 20.0, 10.0, 5.0, 0.0],
        pressure_levels: vec![1000.0, 925.0, 850.0, 700.0, 500.0, 300.0, 200.0, 50.0],
        surface_pressure: 99000.0, // Low pressure center
        temperature_2m: 283.0,
        wind_u_10m: -5.0,
        wind_v_10m: 8.0,
        precipitation: 5.0,
        timestamp: 0.0,
    }
}

/// Create initial state for tropical cyclone.
#[must_use]
pub fn tropical_cyclone() -> AtmosphericState {
    AtmosphericState {
        temperature: vec![299.0, 297.0, 292.0, 280.0, 260.0, 228.0, 217.0, 195.0],
        geopotential: vec![0.0, 750.0, 1500.0, 3000.0, 5800.0, 9500.0, 12200.0, 20800.0],
        humidity: vec![
            0.020, 0.018, 0.015, 0.008, 0.002, 0.0001, 0.00001, 0.0000001,
        ],
        wind_u: vec![-30.0, -40.0, -50.0, -55.0, -45.0, -20.0, -5.0, 5.0],
        wind_v: vec![30.0, 40.0, 50.0, 55.0, 45.0, 20.0, 5.0, -5.0],
        pressure_levels: vec![1000.0, 925.0, 850.0, 700.0, 500.0, 300.0, 200.0, 50.0],
        surface_pressure: 94000.0, // Very low for hurricane
        temperature_2m: 300.0,
        wind_u_10m: -35.0,
        wind_v_10m: 35.0,
        precipitation: 50.0,
        timestamp: 0.0,
    }
}

/// Create initial state for high pressure system (anticyclone).
#[must_use]
pub fn anticyclone() -> AtmosphericState {
    AtmosphericState {
        temperature: vec![295.0, 292.0, 287.0, 275.0, 255.0, 228.0, 218.0, 205.0],
        geopotential: vec![0.0, 820.0, 1550.0, 3150.0, 5750.0, 9400.0, 12100.0, 20600.0],
        humidity: vec![
            0.008, 0.006, 0.004, 0.002, 0.0005, 0.0001, 0.00002, 0.000001,
        ],
        wind_u: vec![2.0, 3.0, 5.0, 8.0, 12.0, 15.0, 10.0, 5.0],
        wind_v: vec![-2.0, -3.0, -5.0, -8.0, -10.0, -8.0, -5.0, -2.0],
        pressure_levels: vec![1000.0, 925.0, 850.0, 700.0, 500.0, 300.0, 200.0, 50.0],
        surface_pressure: 103000.0, // High pressure
        temperature_2m: 298.0,
        wind_u_10m: 2.0,
        wind_v_10m: -1.0,
        precipitation: 0.0,
        timestamp: 0.0,
    }
}

/// Create initial state for winter storm.
#[must_use]
pub fn winter_storm() -> AtmosphericState {
    AtmosphericState {
        temperature: vec![268.0, 265.0, 260.0, 250.0, 235.0, 215.0, 208.0, 195.0],
        geopotential: vec![0.0, 700.0, 1350.0, 2800.0, 5200.0, 8800.0, 11500.0, 20000.0],
        humidity: vec![
            0.003, 0.003, 0.002, 0.001, 0.0005, 0.0001, 0.00002, 0.000001,
        ],
        wind_u: vec![-15.0, -20.0, -30.0, -45.0, -55.0, -60.0, -55.0, -30.0],
        wind_v: vec![5.0, 8.0, 12.0, 15.0, 12.0, 5.0, 0.0, -5.0],
        pressure_levels: vec![1000.0, 925.0, 850.0, 700.0, 500.0, 300.0, 200.0, 50.0],
        surface_pressure: 98500.0,
        temperature_2m: 268.0,
        wind_u_10m: -12.0,
        wind_v_10m: 5.0,
        precipitation: 15.0, // Snow equivalent
        timestamp: 0.0,
    }
}

// ============================================================================
// Sample Locations
// ============================================================================

/// Create sample verification locations.
#[must_use]
pub fn verification_locations() -> Vec<GridPoint> {
    vec![
        GridPoint {
            lat: 40.7128,
            lon: -74.0060,
            elevation: 10.0,
            land_mask: 1.0,
        }, // New York
        GridPoint {
            lat: 34.0522,
            lon: -118.2437,
            elevation: 71.0,
            land_mask: 1.0,
        }, // Los Angeles
        GridPoint {
            lat: 51.5074,
            lon: -0.1278,
            elevation: 11.0,
            land_mask: 1.0,
        }, // London
        GridPoint {
            lat: 35.6762,
            lon: 139.6503,
            elevation: 40.0,
            land_mask: 1.0,
        }, // Tokyo
        GridPoint {
            lat: -33.8688,
            lon: 151.2093,
            elevation: 58.0,
            land_mask: 1.0,
        }, // Sydney
        GridPoint {
            lat: 55.7558,
            lon: 37.6173,
            elevation: 156.0,
            land_mask: 1.0,
        }, // Moscow
        GridPoint {
            lat: -22.9068,
            lon: -43.1729,
            elevation: 11.0,
            land_mask: 1.0,
        }, // Rio de Janeiro
        GridPoint {
            lat: 1.3521,
            lon: 103.8198,
            elevation: 15.0,
            land_mask: 1.0,
        }, // Singapore
    ]
}

/// Create sample ocean locations for marine forecasting.
#[must_use]
pub fn marine_locations() -> Vec<GridPoint> {
    vec![
        GridPoint {
            lat: 30.0,
            lon: -60.0,
            elevation: 0.0,
            land_mask: 0.0,
        }, // North Atlantic
        GridPoint {
            lat: 0.0,
            lon: -160.0,
            elevation: 0.0,
            land_mask: 0.0,
        }, // Central Pacific
        GridPoint {
            lat: -45.0,
            lon: 20.0,
            elevation: 0.0,
            land_mask: 0.0,
        }, // Southern Ocean
        GridPoint {
            lat: 20.0,
            lon: 120.0,
            elevation: 0.0,
            land_mask: 0.0,
        }, // South China Sea
    ]
}

// ============================================================================
// Training Configurations
// ============================================================================

/// Create quick training configuration for testing.
#[must_use]
pub fn quick_training() -> TrainingConfig {
    TrainingConfig {
        epochs: 10,
        batch_size: 8,
        learning_rate: 1e-3,
        weight_decay: 1e-5,
        num_ar_steps: 2,
        gradient_accumulation: 1,
        curriculum: false,
        variable_weights: vec![
            ("temperature".to_string(), 1.0),
            ("geopotential".to_string(), 1.0),
        ],
        noise_level: 0.01,
    }
}

/// Create full training configuration.
#[must_use]
pub fn full_training() -> TrainingConfig {
    TrainingConfig {
        epochs: 200,
        batch_size: 32,
        learning_rate: 1e-4,
        weight_decay: 1e-5,
        num_ar_steps: 12,
        gradient_accumulation: 4,
        curriculum: true,
        variable_weights: vec![
            ("temperature".to_string(), 1.0),
            ("geopotential".to_string(), 0.5),
            ("humidity".to_string(), 1.0),
            ("wind_u".to_string(), 1.0),
            ("wind_v".to_string(), 1.0),
        ],
        noise_level: 0.005,
    }
}

/// Create fine-tuning configuration.
#[must_use]
pub fn finetune_training() -> TrainingConfig {
    TrainingConfig {
        epochs: 50,
        batch_size: 16,
        learning_rate: 1e-5,
        weight_decay: 1e-6,
        num_ar_steps: 6,
        gradient_accumulation: 2,
        curriculum: false,
        ..TrainingConfig::default()
    }
}

// ============================================================================
// Tests
// ============================================================================

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_global_forecast_config() {
        let (mesh, pred) = global_forecast();
        assert_eq!(mesh.refinement_levels, 5);
        assert_eq!(pred.lead_time, 240.0);
    }

    #[test]
    fn test_regional_forecast_config() {
        let (mesh, pred) = regional_forecast();
        assert!(mesh.regional_focus.is_some());
        assert!(mesh.multi_scale);
        assert_eq!(pred.time_step, 1.0);
    }

    #[test]
    fn test_ensemble_config() {
        let (_, pred) = ensemble_forecast();
        assert_eq!(pred.ensemble_size, 50);
    }

    #[test]
    fn test_large_ensemble() {
        let (_, pred) = large_ensemble();
        assert_eq!(pred.ensemble_size, 100);
    }

    #[test]
    fn test_high_resolution_config() {
        let (mesh, pred) = high_resolution_forecast();
        assert!(mesh.avg_edge_length < 10.0);
        assert!(pred.variables.len() > 5);
    }

    #[test]
    fn test_midlatitude_cyclone() {
        let state = midlatitude_cyclone();
        assert!(state.surface_pressure < 101325.0); // Low pressure
        assert!(state.precipitation > 0.0);
    }

    #[test]
    fn test_tropical_cyclone() {
        let state = tropical_cyclone();
        assert!(state.surface_pressure < 95000.0); // Very low
        assert!(state.wind_speed(0) > 30.0); // Strong winds
    }

    #[test]
    fn test_anticyclone() {
        let state = anticyclone();
        assert!(state.surface_pressure > 101325.0); // High pressure
        assert!(state.precipitation == 0.0);
    }

    #[test]
    fn test_winter_storm() {
        let state = winter_storm();
        assert!(state.temperature_2m < 273.0); // Below freezing
    }

    #[test]
    fn test_verification_locations() {
        let locations = verification_locations();
        assert!(!locations.is_empty());

        for loc in &locations {
            assert!(loc.lat >= -90.0 && loc.lat <= 90.0);
            assert!(loc.lon >= -180.0 && loc.lon <= 180.0);
        }
    }

    #[test]
    fn test_marine_locations() {
        let locations = marine_locations();

        for loc in &locations {
            assert_eq!(loc.land_mask, 0.0); // All should be ocean
        }
    }

    #[test]
    fn test_training_configs() {
        let quick = quick_training();
        let full = full_training();

        assert!(quick.epochs < full.epochs);
        assert!(quick.learning_rate > full.learning_rate);
    }

    #[test]
    fn test_european_regional() {
        let (mesh, _) = european_regional();
        let (lat, lon, _) = mesh.regional_focus.unwrap();
        assert!((lat - 50.0).abs() < 1.0);
        assert!((lon - 10.0).abs() < 1.0);
    }

    #[test]
    fn test_asian_regional() {
        let (mesh, _) = asian_regional();
        assert!(mesh.regional_focus.is_some());
    }

    #[test]
    fn test_global_short_range() {
        let (_, pred) = global_short_range();
        assert_eq!(pred.lead_time, 72.0);
        assert_eq!(pred.time_step, 3.0);
    }

    #[test]
    fn test_finetune_training() {
        let config = finetune_training();
        assert!(config.learning_rate < 1e-4);
    }
}