CA/scripts/alert_engine.py

#!/usr/bin/env python3
"""
Alert Engine for Wuhan Respiratory Disease Risk Prediction.
Dual-path alert logic: Monitoring (medical z-scores) + Warning (model predictions)
"""

import os
import sys
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import warnings
warnings.filterwarnings('ignore')

import numpy as np
import pandas as pd
from pathlib import Path
from datetime import datetime, timedelta
import json

PROCESSED_DIR = Path('processed')
OUTPUT_DIR = Path('outputs/daily')
OUTPUT_DIR.mkdir(exist_ok=True)


class AlertLevel:
    """Alert level enumeration with comparison support."""
    GREEN = 0
    YELLOW = 1
    ORANGE = 2
    RED = 3

    @classmethod
    def from_str(cls, s):
        return {'Green': cls.GREEN, 'Yellow': cls.YELLOW,
                'Orange': cls.ORANGE, 'Red': cls.RED}[s]

    @classmethod
    def to_str(cls, level):
        return {0: 'Green', 1: 'Yellow', 2: 'Orange', 3: 'Red'}[level]


def compute_zscore(value, historical_mean, historical_std):
    """Compute z-score; return 0 if std is 0."""
    if historical_std == 0 or np.isnan(historical_std):
        return 0.0
    return (value - historical_mean) / historical_std


def evaluate_monitoring_alert(outpatient_cases, inpatient_cases,
                               out_hist_mean, out_hist_std,
                               inp_hist_mean, inp_hist_std):
    """
    Evaluate monitoring alert based on medical data z-scores.

    Thresholds per PRD:
    - Yellow: outpatient z > 2.0
    - Orange: inpatient z > 2.5
    - Red: combined z > 3.0

    Returns:
        tuple: (AlertLevel, dict with z-scores)
    """
    out_z = compute_zscore(outpatient_cases, out_hist_mean, out_hist_std)
    inp_z = compute_zscore(inpatient_cases, inp_hist_mean, inp_hist_std)
    combined_z = np.sqrt(out_z**2 + inp_z**2)

    if combined_z > 3.0:
        return AlertLevel.RED, {'out_z': out_z, 'inp_z': inp_z, 'combined_z': combined_z}
    elif inp_z > 2.5:
        return AlertLevel.ORANGE, {'out_z': out_z, 'inp_z': inp_z, 'combined_z': combined_z}
    elif out_z > 2.0:
        return AlertLevel.YELLOW, {'out_z': out_z, 'inp_z': inp_z, 'combined_z': combined_z}
    else:
        return AlertLevel.GREEN, {'out_z': out_z, 'inp_z': inp_z, 'combined_z': combined_z}


def evaluate_warning_alert(risk_3d, risk_7d):
    """
    Evaluate warning alert based on model predictions.

    Thresholds per PRD:
    - Orange: risk_3d > 0.6
    - Red: risk_7d > 0.7

    Returns:
        tuple: (AlertLevel, dict with risk values)
    """
    if risk_7d > 0.7:
        return AlertLevel.RED, {'risk_3d': risk_3d, 'risk_7d': risk_7d}
    elif risk_3d > 0.6:
        return AlertLevel.ORANGE, {'risk_3d': risk_3d, 'risk_7d': risk_7d}
    else:
        return AlertLevel.GREEN, {'risk_3d': risk_3d, 'risk_7d': risk_7d}


def resolve_alert(monitoring_level, warning_level):
    """
    Conflict resolution: risk_level = GREATEST(monitoring, warning)
    Where Red > Orange > Yellow > Green
    """
    return max(monitoring_level, warning_level)


def generate_alerts(predictions_df, medical_df=None, date=None):
    """
    Generate alerts with dual-path logic.

    Args:
        predictions_df: DataFrame with risk predictions (node_id, risk_1d, risk_3d, risk_7d, district)
        medical_df: Optional DataFrame with medical data (district, outpatient, inpatient)
        date: Date for alert generation

    Returns:
        list: Alert dictionaries
    """
    if date is None:
        date = datetime.now().date()
    if isinstance(date, str):
        date = datetime.fromisoformat(date).date()

    alerts = []
    districts = predictions_df['district'].unique() if 'district' in predictions_df.columns else []

    for district in districts:
        district_preds = predictions_df[predictions_df['district'] == district]
        risk_1d = district_preds['risk_1d'].mean()
        risk_3d = district_preds['risk_3d'].mean()
        risk_7d = district_preds['risk_7d'].mean()

        # Warning path
        warn_level, warn_info = evaluate_warning_alert(risk_3d, risk_7d)

        # Monitoring path (if medical data provided)
        if medical_df is not None and district in medical_df['district'].values:
            med_row = medical_df[medical_df['district'] == district].iloc[0]
            mon_level, mon_info = evaluate_monitoring_alert(
                med_row.get('outpatient', 0),
                med_row.get('inpatient', 0),
                med_row.get('out_hist_mean', 0),
                med_row.get('out_hist_std', 1),
                med_row.get('inp_hist_mean', 0),
                med_row.get('inp_hist_std', 1)
            )
        else:
            mon_level = AlertLevel.GREEN
            mon_info = {'out_z': 0, 'inp_z': 0, 'combined_z': 0}

        # Resolve final level
        final_level = resolve_alert(mon_level, warn_level)

        # Determine alert type
        if mon_level > AlertLevel.GREEN and warn_level > AlertLevel.GREEN:
            alert_type = 'combined'
        elif mon_level > AlertLevel.GREEN:
            alert_type = 'monitoring'
        elif warn_level > AlertLevel.GREEN:
            alert_type = 'warning'
        else:
            continue  # Skip green alerts

        # Build trigger description
        triggers = []
        if mon_level == AlertLevel.RED:
            triggers.append(f"combined z={mon_info['combined_z']:.2f}")
        elif mon_level == AlertLevel.ORANGE:
            triggers.append(f"inpatient z={mon_info['inp_z']:.2f}")
        elif mon_level == AlertLevel.YELLOW:
            triggers.append(f"outpatient z={mon_info['out_z']:.2f}")

        if warn_level == AlertLevel.RED:
            triggers.append(f"7d risk={risk_7d:.2f}")
        elif warn_level == AlertLevel.ORANGE:
            triggers.append(f"3d risk={risk_3d:.2f}")

        alert = {
            'alert_id': f"ALERT_{date.strftime('%Y%m%d')}_{datetime.now().strftime('%H%M%S')}",
            'alert_type': alert_type,
            'district': district,
            'risk_level': AlertLevel.to_str(final_level),
            'risk_1d': round(float(risk_1d), 4),
            'risk_3d': round(float(risk_3d), 4),
            'risk_7d': round(float(risk_7d), 4),
            'trigger': ' | '.join(triggers),
            'timestamp': datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        }
        alerts.append(alert)

    return alerts


def run_alert_engine(date=None, risk_geojson_path=None, medical_csv_path=None):
    """
    Run alert engine for a specific date.

    Args:
        date: Date for alert generation
        risk_geojson_path: Path to risk GeoJSON file
        medical_csv_path: Optional path to medical data CSV
    """
    if date is None:
        date = datetime.now().date()
    if isinstance(date, str):
        date = datetime.fromisoformat(date).date()

    date_str = date.strftime('%Y%m%d')
    print(f"\n=== Alert Engine: {date_str} ===")

    # Load risk predictions from GeoJSON
    if risk_geojson_path is None:
        risk_geojson_path = OUTPUT_DIR / f'risk_{date_str}.geojson'

    if not Path(risk_geojson_path).exists():
        print(f"  Risk GeoJSON not found: {risk_geojson_path}")
        print("  Run inference_daily.py first")
        return []

    with open(risk_geojson_path) as f:
        geojson = json.load(f)

    # Convert GeoJSON to DataFrame
    predictions = []
    for feat in geojson['features']:
        props = feat['properties']
        predictions.append({
            'node_id': props['node_id'],
            'lat': props['lat'],
            'lon': props['lon'],
            'risk_1d': props['risk_1d'],
            'risk_3d': props['risk_3d'],
            'risk_7d': props['risk_7d'],
            'class_1d': props['class_1d'],
            'class_3d': props['class_3d'],
            'class_7d': props['class_7d'],
            'district': props.get('district', 'unknown')
        })

    predictions_df = pd.DataFrame(predictions)
    print(f"  Loaded predictions: {len(predictions_df)} nodes")

    # Load medical data if available
    medical_df = None
    if medical_csv_path and Path(medical_csv_path).exists():
        medical_df = pd.read_csv(medical_csv_path)
        print(f"  Loaded medical data: {len(medical_df)} districts")

    # Generate alerts
    alerts = generate_alerts(predictions_df, medical_df, date)
    print(f"  Generated alerts: {len(alerts)}")

    # Save alerts
    if len(alerts) > 0:
        out_file = OUTPUT_DIR / f'alerts_{date_str}.json'
        with open(out_file, 'w') as f:
            json.dump(alerts, f, indent=2)
        print(f"  Saved: {out_file}")

        # Print summary
        print("\n  Alert Summary:")
        for alert in alerts:
            print(f"  [{alert['risk_level']}] {alert['district']}: {alert['trigger']}")
    else:
        print("  No alerts generated")

    return alerts


# --- Unit tests ---
def test_alert_resolution():
    """Unit test: simultaneous Yellow + Orange → result Orange."""
    # Yellow monitoring + Orange warning
    result = resolve_alert(AlertLevel.YELLOW, AlertLevel.ORANGE)
    assert result == AlertLevel.ORANGE, f"Expected ORANGE, got {AlertLevel.to_str(result)}"

    # Red monitoring + Yellow warning
    result = resolve_alert(AlertLevel.RED, AlertLevel.YELLOW)
    assert result == AlertLevel.RED, f"Expected RED, got {AlertLevel.to_str(result)}"

    # Green monitoring + Red warning
    result = resolve_alert(AlertLevel.GREEN, AlertLevel.RED)
    assert result == AlertLevel.RED, f"Expected RED, got {AlertLevel.to_str(result)}"

    # Both Yellow
    result = resolve_alert(AlertLevel.YELLOW, AlertLevel.YELLOW)
    assert result == AlertLevel.YELLOW, f"Expected YELLOW, got {AlertLevel.to_str(result)}"

    # Both Green
    result = resolve_alert(AlertLevel.GREEN, AlertLevel.GREEN)
    assert result == AlertLevel.GREEN, f"Expected GREEN, got {AlertLevel.to_str(result)}"

    print("All unit tests passed!")


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(description='Alert engine for respiratory disease risk')
    parser.add_argument('--date', type=str, default=None, help='Date YYYY-MM-DD')
    parser.add_argument('--risk-geojson', type=str, default=None, help='Path to risk GeoJSON')
    parser.add_argument('--medical', type=str, default=None, help='Path to medical CSV')
    parser.add_argument('--test', action='store_true', help='Run unit tests')
    args = parser.parse_args()

    if args.test:
        test_alert_resolution()
    else:
        date = datetime.fromisoformat(args.date) if args.date else datetime.now()
        run_alert_engine(date, args.risk_geojson, args.medical)