predictor.py

"""
Real prediction pipeline using EAGLE-I historical data, weather data, and trained XGBoost model.
"""

import io
import os
import pickle
import numpy as np
import pandas as pd
import boto3
import httpx
import asyncio
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Tuple
import logging

logger = logging.getLogger(__name__)

# Major city coordinates for weather sampling (representative of states)
STATE_WEATHER_POINTS = {
    'FL': [(27.9, -82.5), (25.8, -80.2), (30.3, -81.7)],  # Tampa, Miami, Jacksonville
    'TX': [(29.8, -95.4), (32.8, -96.8), (29.4, -98.5)],  # Houston, Dallas, San Antonio
    'LA': [(30.0, -90.1), (30.2, -93.2)],  # New Orleans, Lake Charles
    'GA': [(33.7, -84.4), (32.1, -81.1)],  # Atlanta, Savannah
    'NC': [(35.8, -78.6), (35.2, -80.8), (34.2, -77.9)],  # Raleigh, Charlotte, Wilmington
    'SC': [(34.0, -81.0), (32.8, -79.9)],  # Columbia, Charleston
    'AL': [(33.5, -86.8), (30.7, -88.1)],  # Birmingham, Mobile
    'MS': [(32.3, -90.2), (30.4, -89.1)],  # Jackson, Gulfport
    'TN': [(36.2, -86.8), (35.1, -90.0)],  # Nashville, Memphis
    'PA': [(40.0, -75.1), (40.4, -80.0)],  # Philadelphia, Pittsburgh
    'NY': [(40.7, -74.0), (42.9, -78.9)],  # NYC, Buffalo
    'OH': [(39.1, -84.5), (41.5, -81.7)],  # Cincinnati, Cleveland
    'MI': [(42.3, -83.0), (42.7, -84.6)],  # Detroit, Lansing
    'CA': [(34.1, -118.2), (37.8, -122.4)],  # LA, SF
}

# S3 Configuration
S3_BUCKET = os.environ.get('S3_BUCKET_NAME', 'lightgrid-eaglei-data')
AWS_REGION = os.environ.get('AWS_REGION', 'us-east-2')

# Feature lists (must match training)
LAG_FEATURES = [
    'customers_out_lag_1h',
    'customers_out_lag_6h',
    'customers_out_lag_12h',
    'customers_out_lag_24h',
    'customers_out_rolling_mean_6h',
    'customers_out_rolling_mean_12h',
    'customers_out_rolling_mean_24h',
    'customers_out_rolling_max_6h',
    'customers_out_rolling_max_12h',
    'customers_out_rolling_max_24h',
    'customers_out_rolling_std_6h',
    'customers_out_rolling_std_12h',
    'customers_out_rolling_std_24h',
    'had_outage_past_6h',
    'had_outage_past_12h',
    'had_outage_past_24h',
    'hours_since_outage',
]

TEMPORAL_FEATURES = [
    'hour', 'hour_sin', 'hour_cos',
    'day_of_week', 'dow_sin', 'dow_cos',
    'month', 'month_sin', 'month_cos',
    'day_of_year', 'is_weekend',
]

COUNTY_FEATURES = [
    'county_outage_rate',
    'county_mean_outage',
    'county_mean_outage_log',
    'county_max_outage',
    'county_std_outage',
    'county_mean_severity',
    'county_mean_severity_log',
    'county_median_severity',
    'state_outage_rate',
    'state_mean_outage',
]

# All features in order
ALL_FEATURES = TEMPORAL_FEATURES + LAG_FEATURES + COUNTY_FEATURES

# EAGLE-I data cache
_eaglei_cache = {}
_s3_client = None
_county_stats = None
_trained_model = None
_model_feature_names = None
_socioeconomic_data = None

# Features used by trained model v2 (must match training)
# Includes socioeconomic features: density, poverty, income, mobile homes, housing age, metro, age 65+
MODEL_FEATURES = [
    'customers_out_lag_1h',
    'customers_out_lag_3h',
    'customers_out_lag_6h',
    'customers_out_lag_12h',
    'customers_out_lag_24h',
    'rolling_mean_6h',
    'rolling_mean_12h',
    'rolling_mean_24h',
    'rolling_max_6h',
    'rolling_max_24h',
    'rolling_std_6h',
    'rolling_std_24h',
    'pct_change_1h',
    'pct_change_6h',
    'hour',
    'hour_sin',
    'hour_cos',
    'day_of_week',
    'dow_sin',
    'dow_cos',
    'month',
    'month_sin',
    'month_cos',
    'is_weekend',
    'is_peak_hour',
    'county_mean_outage',
    'county_max_outage',
    'county_std_outage',
    'county_outage_rate',
    'state_mean_outage',
    'temp_c',
    'temp_max_6h',
    'wind_speed',
    'wind_gust_max_6h',
    'precip_sum_6h',
    'storm_severity',
    'temp_stress',
    # Socioeconomic features (v2)
    'density',
    'poverty_rate',
    'median_income',
    'mobile_home_pct',
    'housing_age_pct',
    'is_metro',
    'pop_over_65_pct',
]


def get_s3_client():
    """Get or create S3 client."""
    global _s3_client
    if _s3_client is None:
        _s3_client = boto3.client('s3', region_name=AWS_REGION)
    return _s3_client


def load_trained_model():
    """Load trained XGBoost model v2 from S3 (includes socioeconomic features)."""
    global _trained_model, _model_feature_names

    if _trained_model is not None:
        return _trained_model, _model_feature_names

    s3 = get_s3_client()
    model_key = 'models/outage_model_v2.pkl'  # v2 with socioeconomic features

    try:
        logger.info(f"Loading trained model from S3: {model_key}")
        response = s3.get_object(Bucket=S3_BUCKET, Key=model_key)
        model_data = pickle.load(io.BytesIO(response['Body'].read()))

        _trained_model = model_data['model']
        _model_feature_names = model_data['feature_names']

        logger.info(f"Model loaded successfully. Features: {len(_model_feature_names)}")
        logger.info(f"Model metrics: R²={model_data['metrics']['test_r2']:.4f}, corr={model_data['metrics']['test_corr']:.4f}")

        return _trained_model, _model_feature_names
    except Exception as e:
        logger.error(f"Failed to load trained model: {e}")
        return None, None


def load_socioeconomic_data() -> pd.DataFrame:
    """Load county-level socioeconomic data from S3."""
    global _socioeconomic_data

    if _socioeconomic_data is not None:
        return _socioeconomic_data

    s3 = get_s3_client()

    try:
        logger.info("Loading socioeconomic data from S3...")
        response = s3.get_object(Bucket=S3_BUCKET, Key='reference/county_complete.txt')
        df = pd.read_csv(io.BytesIO(response['Body'].read()), sep='\t', dtype={'fips': str})

        # Normalize FIPS codes to 5 digits
        df['fips_code'] = df['fips'].astype(str).str.zfill(5)

        # Select and process key features
        socio = pd.DataFrame({
            'fips_code': df['fips_code'],
            'density': np.log1p(df.get('density_2010', pd.Series([0]*len(df))).fillna(0)),
            'poverty_rate': df.get('poverty_2019', df.get('poverty_2017', pd.Series([10]*len(df)))).fillna(10) / 100,
            'median_income': df.get('median_household_income_2019', df.get('median_household_income_2017', pd.Series([50000]*len(df)))).fillna(50000) / 100000,
            'mobile_home_pct': df.get('housing_mobile_homes_2019', pd.Series([10]*len(df))).fillna(10) / 100,
            'is_metro': (df.get('metro_2013', pd.Series([0]*len(df))) == 1).astype(int),
            'pop_over_65_pct': df.get('age_over_65_2019', df.get('age_over_65_2017', pd.Series([15]*len(df)))).fillna(15) / 100,
        })

        # Housing age proxy from home value
        median_val = df.get('median_val_owner_occupied_2010', pd.Series([150000]*len(df))).fillna(150000)
        socio['housing_age_pct'] = np.clip(100 - (median_val / 5000), 0, 100) / 100

        _socioeconomic_data = socio
        logger.info(f"Loaded socioeconomic data for {len(socio)} counties")
        return socio

    except Exception as e:
        logger.warning(f"Failed to load socioeconomic data: {e}")
        return None


def load_eaglei_data(year: int) -> pd.DataFrame:
    """Load EAGLE-I data for a year from S3."""
    if year in _eaglei_cache:
        return _eaglei_cache[year]

    s3 = get_s3_client()
    s3_key = f"parquet/eaglei_{year}.parquet"

    logger.info(f"Loading EAGLE-I {year} from S3...")
    response = s3.get_object(Bucket=S3_BUCKET, Key=s3_key)
    df = pd.read_parquet(io.BytesIO(response['Body'].read()))

    # Ensure proper types
    df['run_start_time'] = pd.to_datetime(df['run_start_time'])
    df['fips_code'] = df['fips_code'].astype(str).str.zfill(5)

    # Rename columns if needed
    if 'sum' in df.columns and 'customers_out' not in df.columns:
        df = df.rename(columns={'sum': 'customers_out'})

    _eaglei_cache[year] = df
    logger.info(f"Loaded {len(df):,} rows for {year}")
    return df


def compute_county_stats(df: pd.DataFrame, threshold: int = 100) -> pd.DataFrame:
    """Compute county-level baseline statistics from historical data."""
    # Aggregate by county
    stats = df.groupby('fips_code').agg({
        'customers_out': ['mean', 'std', 'max', 'median']
    }).reset_index()
    stats.columns = ['fips_code', 'county_mean_outage', 'county_std_outage',
                     'county_max_outage', 'county_median_outage']

    # Outage rate (fraction of hours with outage >= threshold)
    outage_counts = df.groupby('fips_code').apply(
        lambda x: (x['customers_out'] >= threshold).mean()
    ).reset_index(name='county_outage_rate')

    stats = stats.merge(outage_counts, on='fips_code')

    # Log transforms
    stats['county_mean_outage_log'] = np.log1p(stats['county_mean_outage'])

    # Severity stats (use outage as proxy)
    stats['county_mean_severity'] = stats['county_mean_outage'] / 1000  # Scale
    stats['county_mean_severity_log'] = np.log1p(stats['county_mean_severity'])
    stats['county_median_severity'] = stats['county_median_outage'] / 1000

    return stats


def get_county_stats() -> pd.DataFrame:
    """Get or compute county baseline statistics."""
    global _county_stats
    if _county_stats is not None:
        return _county_stats

    # Load a sample year to compute stats
    logger.info("Computing county baseline statistics...")
    df = load_eaglei_data(2023)  # Use recent full year
    _county_stats = compute_county_stats(df)
    logger.info(f"Computed stats for {len(_county_stats)} counties")
    return _county_stats


async def fetch_historical_weather(target_time: datetime, states: List[str] = None) -> Dict[str, Dict]:
    """
    Fetch historical weather data from Open-Meteo archive API.
    Returns weather conditions by state for the given time.
    """
    if states is None:
        states = list(STATE_WEATHER_POINTS.keys())

    date_str = target_time.strftime('%Y-%m-%d')
    weather_by_state = {}

    async with httpx.AsyncClient(timeout=30.0) as client:
        for state in states:
            if state not in STATE_WEATHER_POINTS:
                continue

            coords = STATE_WEATHER_POINTS[state]
            lat, lon = coords[0]  # Use first point

            try:
                # Open-Meteo Archive API for historical weather
                url = (
                    f"https://archive-api.open-meteo.com/v1/archive"
                    f"?latitude={lat}&longitude={lon}"
                    f"&start_date={date_str}&end_date={date_str}"
                    f"&hourly=temperature_2m,precipitation,wind_speed_10m,wind_gusts_10m"
                )

                response = await client.get(url)
                if response.status_code == 200:
                    data = response.json()

                    # Get hourly data for target hour
                    hour_idx = target_time.hour
                    hourly = data.get('hourly', {})

                    temp = hourly.get('temperature_2m', [0] * 24)[hour_idx]
                    precip = hourly.get('precipitation', [0] * 24)[hour_idx]
                    wind = hourly.get('wind_speed_10m', [0] * 24)[hour_idx]
                    gusts = hourly.get('wind_gusts_10m', [0] * 24)[hour_idx]

                    # Compute max in surrounding hours (storm indicator)
                    start_idx = max(0, hour_idx - 6)
                    end_idx = min(24, hour_idx + 1)
                    wind_max_6h = max(hourly.get('wind_speed_10m', [0])[start_idx:end_idx] or [0])
                    precip_sum_6h = sum(hourly.get('precipitation', [0])[start_idx:end_idx] or [0])
                    gust_max_6h = max(hourly.get('wind_gusts_10m', [0])[start_idx:end_idx] or [0])

                    # Storm flags (Open-Meteo returns km/h, not m/s)
                    # Tropical storm: 63-118 km/h, Hurricane: 119+ km/h
                    high_wind = wind_max_6h > 50 or gust_max_6h > 80  # km/h thresholds
                    heavy_precip = precip_sum_6h > 10  # mm

                    # Storm severity: normalized 0-1 scale
                    # Gusts > 100 km/h = severe, > 150 km/h = extreme
                    wind_severity = np.clip(gust_max_6h / 150, 0, 1)
                    precip_severity = np.clip(precip_sum_6h / 50, 0, 1)
                    combined_severity = min(1.0, wind_severity * 0.7 + precip_severity * 0.3)

                    # Temperature stress: extreme heat (>30C/86F) or cold (<0C/32F)
                    # Causes grid stress from AC/heating demand
                    temp_val = temp or 20
                    temp_max_6h = max(hourly.get('temperature_2m', [20])[start_idx:end_idx] or [20])
                    temp_min_6h = min(hourly.get('temperature_2m', [20])[start_idx:end_idx] or [20])

                    heat_stress = np.clip((temp_max_6h - 30) / 10, 0, 1)  # Ramps up from 30C to 40C
                    cold_stress = np.clip((0 - temp_min_6h) / 15, 0, 1)  # Ramps up from 0C to -15C
                    temp_stress = max(heat_stress, cold_stress)

                    weather_by_state[state] = {
                        'temp_c': temp_val,
                        'temp_max_6h': temp_max_6h,
                        'temp_min_6h': temp_min_6h,
                        'temp_stress': temp_stress,
                        'precip_mm': precip or 0,
                        'wind_speed_kmh': wind or 0,
                        'wind_gust_kmh': gusts or 0,
                        'wind_max_6h': wind_max_6h,
                        'precip_sum_6h': precip_sum_6h,
                        'gust_max_6h': gust_max_6h,
                        'high_wind_flag': int(high_wind),
                        'storm_flag': int(high_wind or heavy_precip),
                        'storm_severity': combined_severity,
                    }
                else:
                    logger.warning(f"Weather API returned {response.status_code} for {state}")

            except Exception as e:
                logger.warning(f"Failed to fetch weather for {state}: {e}")
                weather_by_state[state] = {
                    'temp_c': 20, 'temp_max_6h': 20, 'temp_min_6h': 20, 'temp_stress': 0,
                    'precip_mm': 0, 'wind_speed_kmh': 5, 'wind_gust_kmh': 10,
                    'wind_max_6h': 5, 'precip_sum_6h': 0, 'gust_max_6h': 10,
                    'high_wind_flag': 0, 'storm_flag': 0, 'storm_severity': 0,
                }

    return weather_by_state


async def fetch_nws_alerts(states: List[str] = None) -> Dict[str, Dict]:
    """
    Fetch active NWS alerts for each state.
    Returns alert counts and severity by state.

    Note: Only works for current/future alerts, not historical.
    """
    if states is None:
        states = list(STATE_WEATHER_POINTS.keys())

    # Alert types that cause power outages (weighted by impact)
    OUTAGE_ALERTS = {
        'Hurricane Warning': 1.0,
        'Hurricane Watch': 0.7,
        'Tornado Warning': 0.9,
        'Tornado Watch': 0.5,
        'Severe Thunderstorm Warning': 0.7,
        'Severe Thunderstorm Watch': 0.4,
        'Winter Storm Warning': 0.8,
        'Winter Storm Watch': 0.5,
        'Ice Storm Warning': 0.9,
        'Blizzard Warning': 0.8,
        'High Wind Warning': 0.6,
        'Wind Advisory': 0.3,
        'Extreme Cold Warning': 0.5,
        'Excessive Heat Warning': 0.5,
        'Heat Advisory': 0.3,
        'Flood Warning': 0.4,
        'Flash Flood Warning': 0.5,
    }

    alerts_by_state = {}

    async with httpx.AsyncClient(timeout=15.0) as client:
        for state in states:
            try:
                url = f"https://api.weather.gov/alerts/active?area={state}"
                response = await client.get(url, headers={'User-Agent': 'lightgrid-api'})

                if response.status_code == 200:
                    data = response.json()
                    features = data.get('features', [])

                    alert_count = 0
                    max_severity = 0.0
                    alert_types = []

                    for feature in features:
                        props = feature.get('properties', {})
                        event = props.get('event', '')

                        if event in OUTAGE_ALERTS:
                            alert_count += 1
                            severity = OUTAGE_ALERTS[event]
                            max_severity = max(max_severity, severity)
                            alert_types.append(event)

                    alerts_by_state[state] = {
                        'alert_count': alert_count,
                        'max_severity': max_severity,
                        'alert_types': list(set(alert_types)),
                        'has_severe_alert': max_severity >= 0.7,
                    }
                else:
                    alerts_by_state[state] = {
                        'alert_count': 0, 'max_severity': 0,
                        'alert_types': [], 'has_severe_alert': False
                    }

            except Exception as e:
                logger.warning(f"Failed to fetch NWS alerts for {state}: {e}")
                alerts_by_state[state] = {
                    'alert_count': 0, 'max_severity': 0,
                    'alert_types': [], 'has_severe_alert': False
                }

    return alerts_by_state


def add_temporal_features(df: pd.DataFrame, timestamp_col: str = 'run_start_time') -> pd.DataFrame:
    """Add temporal features."""
    df = df.copy()
    ts = pd.to_datetime(df[timestamp_col])

    df['hour'] = ts.dt.hour
    df['hour_sin'] = np.sin(2 * np.pi * df['hour'] / 24)
    df['hour_cos'] = np.cos(2 * np.pi * df['hour'] / 24)

    df['day_of_week'] = ts.dt.dayofweek
    df['dow_sin'] = np.sin(2 * np.pi * df['day_of_week'] / 7)
    df['dow_cos'] = np.cos(2 * np.pi * df['day_of_week'] / 7)

    df['month'] = ts.dt.month
    df['month_sin'] = np.sin(2 * np.pi * df['month'] / 12)
    df['month_cos'] = np.cos(2 * np.pi * df['month'] / 12)

    df['day_of_year'] = ts.dt.dayofyear
    df['is_weekend'] = (df['day_of_week'] >= 5).astype(int)

    return df


def add_lag_features(df: pd.DataFrame) -> pd.DataFrame:
    """Add lag and rolling features."""
    df = df.copy()
    df = df.sort_values(['fips_code', 'run_start_time']).reset_index(drop=True)

    # Lag features
    for lag in [1, 6, 12, 24]:
        df[f'customers_out_lag_{lag}h'] = df.groupby('fips_code')['customers_out'].shift(lag)

    # Rolling features
    for window in [6, 12, 24]:
        df[f'customers_out_rolling_mean_{window}h'] = (
            df.groupby('fips_code')['customers_out']
            .transform(lambda x: x.shift(1).rolling(window, min_periods=1).mean())
        )
        df[f'customers_out_rolling_max_{window}h'] = (
            df.groupby('fips_code')['customers_out']
            .transform(lambda x: x.shift(1).rolling(window, min_periods=1).max())
        )
        df[f'customers_out_rolling_std_{window}h'] = (
            df.groupby('fips_code')['customers_out']
            .transform(lambda x: x.shift(1).rolling(window, min_periods=1).std())
        )

    # Binary lag features
    for window in [6, 12, 24]:
        df[f'had_outage_past_{window}h'] = (
            df.groupby('fips_code')['customers_out']
            .transform(lambda x: (x >= 100).shift(1).rolling(window, min_periods=1).max())
        ).fillna(0).astype(int)

    # Hours since outage
    def calc_hours_since(group):
        outage = (group.shift(1) >= 100).fillna(False)
        result = []
        hours = 168
        for val in outage:
            if val:
                hours = 1
            else:
                hours = min(hours + 1, 168)
            result.append(hours)
        return result

    df['hours_since_outage'] = df.groupby('fips_code')['customers_out'].transform(calc_hours_since)

    # Fill NaN
    lag_cols = [c for c in df.columns if 'lag' in c or 'rolling' in c or 'hours_since' in c or 'had_outage' in c]
    df[lag_cols] = df[lag_cols].fillna(0)

    return df


def get_data_window(target_time: datetime, hours_before: int = 48) -> pd.DataFrame:
    """
    Load EAGLE-I data for a time window before target_time.
    Returns data needed to compute lag features.
    """
    start_time = target_time - timedelta(hours=hours_before)
    year = target_time.year

    # Load data (may need multiple years if crossing boundary)
    df = load_eaglei_data(year)

    if start_time.year != year:
        df_prev = load_eaglei_data(start_time.year)
        df = pd.concat([df_prev, df], ignore_index=True)

    # Filter to time window
    mask = (df['run_start_time'] >= start_time) & (df['run_start_time'] <= target_time)
    df = df[mask].copy()

    return df


def prepare_features(df: pd.DataFrame, target_time: datetime) -> pd.DataFrame:
    """
    Prepare all features for prediction at target_time.
    """
    # Add temporal features
    df = add_temporal_features(df)

    # Add lag features
    df = add_lag_features(df)

    # Add county baseline stats
    county_stats = get_county_stats()
    df = df.merge(county_stats, on='fips_code', how='left')

    # Add state-level features
    df['state_fips'] = df['fips_code'].str[:2]
    state_stats = df.groupby('state_fips').agg({
        'county_outage_rate': 'mean',
        'county_mean_outage': 'mean'
    }).reset_index()
    state_stats.columns = ['state_fips', 'state_outage_rate', 'state_mean_outage']
    df = df.merge(state_stats, on='state_fips', how='left')

    # Get only the latest timestamp (closest to target)
    latest = df.groupby('fips_code')['run_start_time'].max().reset_index()
    latest.columns = ['fips_code', 'latest_time']
    df = df.merge(latest, on='fips_code')
    df = df[df['run_start_time'] == df['latest_time']].copy()

    # Fill missing
    for col in ALL_FEATURES:
        if col not in df.columns:
            df[col] = 0
        df[col] = df[col].fillna(0)

    return df


async def predict_outages(target_time: datetime, horizon: int = 6) -> Dict:
    """
    Generate predictions for all counties at target_time using EAGLE-I data and weather.

    Parameters
    ----------
    target_time : datetime
        Time to predict FOR (predictions would have been made horizon hours before)
    horizon : int
        Prediction horizon in hours

    Returns
    -------
    dict with predictions by county and state
    """
    logger.info(f"Generating predictions for {target_time} (t+{horizon}h horizon)")

    # Load historical data before prediction time
    prediction_time = target_time - timedelta(hours=horizon)
    df = get_data_window(prediction_time, hours_before=48)

    if len(df) == 0:
        logger.warning("No EAGLE-I data available for this time window")
        return {'error': 'No data available', 'predictions': []}

    logger.info(f"Loaded {len(df):,} rows for feature computation")

    # Fetch historical weather for the target time
    logger.info("Fetching historical weather data...")
    weather_data = await fetch_historical_weather(target_time)
    logger.info(f"Got weather for {len(weather_data)} states")

    # Prepare features
    df = prepare_features(df, prediction_time)
    logger.info(f"Prepared features for {len(df)} counties")

    # State abbreviation mapping
    STATE_ABBREV = {
        'Alabama': 'AL', 'Alaska': 'AK', 'Arizona': 'AZ', 'Arkansas': 'AR', 'California': 'CA',
        'Colorado': 'CO', 'Connecticut': 'CT', 'Delaware': 'DE', 'Florida': 'FL', 'Georgia': 'GA',
        'Hawaii': 'HI', 'Idaho': 'ID', 'Illinois': 'IL', 'Indiana': 'IN', 'Iowa': 'IA',
        'Kansas': 'KS', 'Kentucky': 'KY', 'Louisiana': 'LA', 'Maine': 'ME', 'Maryland': 'MD',
        'Massachusetts': 'MA', 'Michigan': 'MI', 'Minnesota': 'MN', 'Mississippi': 'MS',
        'Missouri': 'MO', 'Montana': 'MT', 'Nebraska': 'NE', 'Nevada': 'NV', 'New Hampshire': 'NH',
        'New Jersey': 'NJ', 'New Mexico': 'NM', 'New York': 'NY', 'North Carolina': 'NC',
        'North Dakota': 'ND', 'Ohio': 'OH', 'Oklahoma': 'OK', 'Oregon': 'OR', 'Pennsylvania': 'PA',
        'Rhode Island': 'RI', 'South Carolina': 'SC', 'South Dakota': 'SD', 'Tennessee': 'TN',
        'Texas': 'TX', 'Utah': 'UT', 'Vermont': 'VT', 'Virginia': 'VA', 'Washington': 'WA',
        'West Virginia': 'WV', 'Wisconsin': 'WI', 'Wyoming': 'WY',
    }

    # Add weather features to dataframe by state
    df['state_abbrev'] = df['state'].map(STATE_ABBREV)
    df['storm_severity'] = df['state_abbrev'].apply(
        lambda x: weather_data.get(x, {}).get('storm_severity', 0)
    )
    df['wind_max_6h'] = df['state_abbrev'].apply(
        lambda x: weather_data.get(x, {}).get('wind_max_6h', 0)
    )
    df['precip_sum_6h'] = df['state_abbrev'].apply(
        lambda x: weather_data.get(x, {}).get('precip_sum_6h', 0)
    )
    df['storm_flag'] = df['state_abbrev'].apply(
        lambda x: weather_data.get(x, {}).get('storm_flag', 0)
    )
    df['temp_stress'] = df['state_abbrev'].apply(
        lambda x: weather_data.get(x, {}).get('temp_stress', 0)
    )

    # ==========================================
    # TRAINED XGBOOST MODEL PREDICTIONS (v2 with socioeconomic features)
    # ==========================================

    # Load trained model and socioeconomic data
    model, feature_names = load_trained_model()
    socio_data = load_socioeconomic_data()

    # Merge socioeconomic features
    if socio_data is not None:
        df = df.merge(socio_data, on='fips_code', how='left')
        # Fill missing with defaults
        for col in ['density', 'poverty_rate', 'median_income', 'mobile_home_pct',
                    'housing_age_pct', 'is_metro', 'pop_over_65_pct']:
            if col in df.columns:
                df[col] = df[col].fillna(0)
        logger.info(f"Merged socioeconomic data for {df['density'].notna().sum()} counties")

    # Prepare features for XGBoost model
    # Add additional features needed by model
    df['temp_c'] = df['state_abbrev'].apply(
        lambda x: weather_data.get(x, {}).get('temp_c', 20)
    )
    df['temp_max_6h'] = df['state_abbrev'].apply(
        lambda x: weather_data.get(x, {}).get('temp_max_6h', 25)
    )
    df['wind_speed'] = df['state_abbrev'].apply(
        lambda x: weather_data.get(x, {}).get('wind_speed_kmh', 10)
    )
    df['wind_gust_max_6h'] = df['state_abbrev'].apply(
        lambda x: weather_data.get(x, {}).get('gust_max_6h', 15)
    )

    # Map feature names from our dataframe to model's expected names
    feature_mapping = {
        'customers_out_lag_1h': 'customers_out_lag_1h',
        'customers_out_lag_3h': 'customers_out_lag_1h',  # Use 1h as fallback
        'customers_out_lag_6h': 'customers_out_lag_6h',
        'customers_out_lag_12h': 'customers_out_lag_12h',
        'customers_out_lag_24h': 'customers_out_lag_24h',
        'rolling_mean_6h': 'customers_out_rolling_mean_6h',
        'rolling_mean_12h': 'customers_out_rolling_mean_12h',
        'rolling_mean_24h': 'customers_out_rolling_mean_24h',
        'rolling_max_6h': 'customers_out_rolling_max_6h',
        'rolling_max_24h': 'customers_out_rolling_max_24h',
        'rolling_std_6h': 'customers_out_rolling_std_6h',
        'rolling_std_24h': 'customers_out_rolling_std_24h',
        'is_peak_hour': 'is_peak_hour',
    }

    # Add is_peak_hour if not present
    if 'is_peak_hour' not in df.columns:
        hour = df['hour'].values if 'hour' in df.columns else 12
        df['is_peak_hour'] = ((hour >= 14) & (hour <= 19)).astype(int)

    # Add pct_change features
    if 'pct_change_1h' not in df.columns:
        df['pct_change_1h'] = (df['customers_out_lag_1h'] / (df['customers_out_lag_6h'] + 1) - 1).clip(-10, 10).fillna(0)
    if 'pct_change_6h' not in df.columns:
        df['pct_change_6h'] = (df['customers_out_lag_1h'] / (df['customers_out_lag_24h'] + 1) - 1).clip(-10, 10).fillna(0)

    # Build feature matrix for model
    if model is not None and feature_names is not None:
        logger.info(f"Using trained XGBoost model with {len(feature_names)} features")

        # Prepare feature matrix
        X = np.zeros((len(df), len(feature_names)))
        for i, feat in enumerate(feature_names):
            if feat in df.columns:
                X[:, i] = df[feat].fillna(0).values
            elif feat in feature_mapping and feature_mapping[feat] in df.columns:
                X[:, i] = df[feature_mapping[feat]].fillna(0).values
            else:
                X[:, i] = 0  # Default value

        # Make predictions (model outputs log-transformed values)
        y_pred_log = model.predict(X)
        predicted_outages = np.expm1(y_pred_log).astype(int)  # Convert from log scale

        logger.info(f"Base XGBoost predictions: min={predicted_outages.min()}, max={predicted_outages.max()}, mean={predicted_outages.mean():.0f}")

        # =====================================================
        # STORM-AWARE SCALING
        # The model underestimates storms because lag features dominate.
        # When weather indicates severe conditions, scale up predictions.
        # =====================================================
        storm_severity = df['storm_severity'].values
        wind_gust = df['wind_gust_max_6h'].values if 'wind_gust_max_6h' in df.columns else np.zeros(len(df))
        county_max = df['county_max_outage'].fillna(1000).values

        # Calculate storm multiplier per county
        # - No storm (severity < 0.2): multiplier = 1.0
        # - Moderate storm (0.2-0.5): multiplier = 1.5-3.0
        # - Severe storm (0.5-0.8): multiplier = 3.0-6.0
        # - Extreme storm (>0.8): multiplier = 6.0-10.0
        storm_multiplier = np.ones(len(df))

        # Base multiplier from storm severity
        # Tuned based on evaluation: need ~2x more scaling for major hurricanes
        storm_multiplier = np.where(
            storm_severity > 0.8,
            10.0 + (storm_severity - 0.8) * 40,  # 10-18x for extreme hurricanes
            np.where(
                storm_severity > 0.5,
                4.0 + (storm_severity - 0.5) * 20,  # 4-10x for severe
                np.where(
                    storm_severity > 0.2,
                    1.5 + (storm_severity - 0.2) * 8.3,  # 1.5-4x for moderate
                    1.0  # No scaling for normal
                )
            )
        )

        # Additional boost for very high wind gusts (>80 km/h = tropical storm+)
        wind_boost = np.clip((wind_gust - 80) / 40, 0, 3)  # Up to 3x additional
        storm_multiplier = storm_multiplier * (1 + wind_boost)

        # Apply storm multiplier
        scaled_outages = predicted_outages * storm_multiplier

        # Cap at reasonable maximum (county historical max * 10 for extreme events)
        max_cap = np.maximum(county_max * 10, 100000)  # At least 100k cap per county
        scaled_outages = np.minimum(scaled_outages, max_cap)

        predicted_outages = scaled_outages.astype(int)
        predicted_outages = np.clip(predicted_outages, 0, 10_000_000)

        # Log storm scaling info
        max_multiplier = storm_multiplier.max()
        if max_multiplier > 1.5:
            logger.info(f"Storm scaling applied: max multiplier={max_multiplier:.1f}x, max severity={storm_severity.max():.2f}")

        # Risk score based on prediction magnitude relative to county max
        risk = np.clip(predicted_outages / (county_max + 1), 0, 1)

        logger.info(f"Final predictions: min={predicted_outages.min()}, max={predicted_outages.max()}, mean={predicted_outages.mean():.0f}")
    else:
        # Fallback to simple formula if model not available
        logger.warning("Trained model not available, using fallback formula")
        lag_1h = df['customers_out_lag_1h'].fillna(0).values
        rolling_mean = df['customers_out_rolling_mean_6h'].fillna(0).values
        county_mean = df['county_mean_outage'].fillna(100).values
        storm_severity = df['storm_severity'].values

        predicted_outages = (
            0.6 * lag_1h +
            0.3 * rolling_mean +
            0.1 * county_mean * (1 + storm_severity * 2)
        ).astype(int)
        predicted_outages = np.clip(predicted_outages, 0, 10_000_000)
        risk = np.clip(predicted_outages / 10000, 0, 1)

    # Build results
    df['risk_score'] = risk
    df['predicted_outages'] = predicted_outages

    # Aggregate to state level
    state_preds = df.groupby('state').agg({
        'predicted_outages': 'sum',
        'risk_score': 'mean',
        'storm_severity': 'mean',
        'fips_code': 'count'
    }).reset_index()
    state_preds.columns = ['state', 'predicted_outages', 'avg_risk', 'storm_severity', 'county_count']

    # Sort by predicted outages
    state_preds = state_preds.sort_values('predicted_outages', ascending=False)

    # County-level results (top 100 by predicted outages)
    df_sorted = df.sort_values('predicted_outages', ascending=False)
    county_preds = df_sorted[['fips_code', 'county', 'state', 'risk_score', 'predicted_outages',
                              'customers_out_lag_1h', 'storm_severity']].head(100).to_dict('records')

    return {
        'prediction_time': prediction_time.isoformat(),
        'target_time': target_time.isoformat(),
        'horizon_hours': horizon,
        'total_predicted': int(df['predicted_outages'].sum()),
        'counties_predicted': len(df),
        'weather_data': weather_data,
        'by_state': state_preds.to_dict('records'),
        'by_county': county_preds,
    }
# Rebuild trigger Mon Jan 26 21:22:45 EST 2026