ea_impact_events.py

#!/usr/bin/env python3
"""
East Africa Impact Events - EM-DAT Data Processing and Visualization

This script processes EM-DAT disaster data for East Africa (ICPAC region),
assigns admin codes to events, and generates frequency maps as PNG files.

Data sources:
- EM-DAT: public_emdat_custom_request_2026-01-21.xlsx
- Boundaries: icpac_adm1v3.geojson

Output:
- Frequency maps (PNG)
- Yearly extent maps (PNG)
- Updated Excel files with admin codes
"""

import os
import pandas as pd
import geopandas as gpd
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
from matplotlib.colors import LinearSegmentedColormap
import numpy as np
from pathlib import Path

# =============================================================================
# Configuration
# =============================================================================

# File paths
EMDAT_FILE = "public_emdat_custom_request_2026-01-21.xlsx"
ICPAC_GEOJSON = "icpac_adm1v3.geojson"
OUTPUT_DIR = "output"

# East Africa (ICPAC) countries
EA_COUNTRIES = [
    'Djibouti', 'Eritrea', 'Ethiopia', 'Kenya', 'Rwanda',
    'Somalia', 'South Sudan', 'Sudan', 'United Republic of Tanzania',
    'Uganda', 'Burundi'
]

# ISO code mapping
COUNTRY_ISO_MAP = {
    'Djibouti': 'DJI',
    'Eritrea': 'ERI',
    'Ethiopia': 'ETH',
    'Kenya': 'KEN',
    'Rwanda': 'RWA',
    'Somalia': 'SOM',
    'South Sudan': 'SSD',
    'Sudan': 'SDN',
    'United Republic of Tanzania': 'TZA',
    'Uganda': 'UGA',
    'Burundi': 'BDI'
}

# Map extent for East Africa
MAP_EXTENT = {
    'x_min': 21.84,
    'x_max': 51.42,
    'y_min': -11.75,
    'y_max': 23.15
}

# Year range
YEAR_START = 1990
YEAR_END = 2025

# =============================================================================
# Data Loading Functions
# =============================================================================

def load_emdat_data(filepath):
    """
    Load EM-DAT data from Excel file.

    Parameters
    ----------
    filepath : str
        Path to the EM-DAT Excel file

    Returns
    -------
    pd.DataFrame
        Loaded EM-DAT data
    """
    print(f"Loading EM-DAT data from {filepath}...")
    df = pd.read_excel(filepath, engine='openpyxl', sheet_name='EM-DAT Data')
    print(f"  Loaded {len(df)} records")
    return df


def filter_ea_data(df, countries=EA_COUNTRIES):
    """
    Filter data for East Africa countries.

    Parameters
    ----------
    df : pd.DataFrame
        Full EM-DAT dataset
    countries : list
        List of country names to filter

    Returns
    -------
    pd.DataFrame
        Filtered dataset for East Africa
    """
    df_ea = df[df['Country'].isin(countries)].copy()
    print(f"  Filtered to {len(df_ea)} East Africa records")
    return df_ea


def filter_by_disaster_type(df, disaster_type):
    """
    Filter data by disaster type.

    Parameters
    ----------
    df : pd.DataFrame
        EM-DAT dataset
    disaster_type : str
        'Drought' or 'Flood'

    Returns
    -------
    pd.DataFrame
        Filtered dataset
    """
    df_filtered = df[df['Disaster Type'] == disaster_type].copy()
    print(f"  {disaster_type}: {len(df_filtered)} events")
    return df_filtered


def filter_by_year_range(df, start_year, end_year):
    """
    Filter data by year range.

    Parameters
    ----------
    df : pd.DataFrame
        EM-DAT dataset
    start_year : int
        Start year (inclusive)
    end_year : int
        End year (inclusive)

    Returns
    -------
    pd.DataFrame
        Filtered dataset
    """
    df_filtered = df[(df['Start Year'] >= start_year) &
                     (df['Start Year'] <= end_year)].copy()
    print(f"  Year range {start_year}-{end_year}: {len(df_filtered)} events")
    return df_filtered


def load_icpac_boundaries(filepath):
    """
    Load ICPAC admin1 boundaries from GeoJSON.

    Parameters
    ----------
    filepath : str
        Path to GeoJSON file

    Returns
    -------
    gpd.GeoDataFrame
        GeoDataFrame with admin1 boundaries
    """
    print(f"Loading ICPAC boundaries from {filepath}...")
    gdf = gpd.read_file(filepath)

    # Extract ISO code from GID_1 (e.g., "ETH.1_1" -> "ETH")
    gdf['ISO'] = gdf['GID_1'].str.split('.').str[0]

    print(f"  Loaded {len(gdf)} admin1 regions")
    print(f"  Countries: {sorted(gdf['ISO'].unique())}")
    return gdf

# =============================================================================
# Admin Code Assignment Functions
# =============================================================================

def assign_admin1_from_country(df, gdf_boundaries):
    """
    Assign admin1 regions to events based on country.
    For events without specific location, assigns all admin1 regions of that country.

    Parameters
    ----------
    df : pd.DataFrame
        EM-DAT disaster data
    gdf_boundaries : gpd.GeoDataFrame
        Admin1 boundaries

    Returns
    -------
    pd.DataFrame
        Data with admin1_list column
    """
    df = df.copy()

    # Create country to admin1 mapping
    country_admin1_map = {}
    for iso in gdf_boundaries['ISO'].unique():
        admin1_list = gdf_boundaries[gdf_boundaries['ISO'] == iso]['GID_1'].tolist()
        country_admin1_map[iso] = admin1_list

    # Assign admin1 codes based on country ISO
    def get_admin1_list(row):
        iso = COUNTRY_ISO_MAP.get(row['Country'])
        if iso and iso in country_admin1_map:
            return country_admin1_map[iso]
        return []

    df['admin1_list'] = df.apply(get_admin1_list, axis=1)

    return df


def count_events_by_admin1(df, gdf_boundaries):
    """
    Count disaster events per admin1 region.

    Parameters
    ----------
    df : pd.DataFrame
        Disaster data with admin1_list column
    gdf_boundaries : gpd.GeoDataFrame
        Admin1 boundaries

    Returns
    -------
    gpd.GeoDataFrame
        Boundaries with event_count column
    """
    # Explode admin1_list to get one row per admin1 per event
    df_exploded = df.explode('admin1_list')

    # Count events per admin1
    counts = df_exploded.groupby('admin1_list').size().reset_index(name='event_count')
    counts.columns = ['GID_1', 'event_count']

    # Merge with boundaries
    gdf = gdf_boundaries.merge(counts, on='GID_1', how='left')
    gdf['event_count'] = gdf['event_count'].fillna(0).astype(int)

    return gdf


def get_affected_admin1_by_year(df, year):
    """
    Get list of admin1 regions affected in a specific year.

    Parameters
    ----------
    df : pd.DataFrame
        Disaster data with admin1_list column
    year : int
        Year to filter

    Returns
    -------
    list
        List of affected admin1 GID codes
    """
    df_year = df[df['Start Year'] == year]

    # Collect all admin1 codes
    affected = []
    for admin_list in df_year['admin1_list']:
        if isinstance(admin_list, list):
            affected.extend(admin_list)

    return list(set(affected))

# =============================================================================
# Visualization Functions
# =============================================================================

def get_colormap():
    """
    Create custom colormap for frequency maps.

    Returns
    -------
    LinearSegmentedColormap
        Custom colormap (yellow to red)
    """
    colors = ['#ffffcc', '#ffeda0', '#fed976', '#feb24c',
              '#fd8d3c', '#fc4e2a', '#e31a1c', '#b10026']
    return LinearSegmentedColormap.from_list('disaster_cmap', colors)


def create_frequency_map(gdf, title, output_path, disaster_type='Disaster'):
    """
    Create choropleth frequency map.

    Parameters
    ----------
    gdf : gpd.GeoDataFrame
        Boundaries with event_count column
    title : str
        Map title
    output_path : str
        Output PNG file path
    disaster_type : str
        Type of disaster for labeling
    """
    fig, ax = plt.subplots(1, 1, figsize=(12, 14))

    # Set map extent
    ax.set_xlim(MAP_EXTENT['x_min'], MAP_EXTENT['x_max'])
    ax.set_ylim(MAP_EXTENT['y_min'], MAP_EXTENT['y_max'])

    # Plot background (all regions in light gray)
    gdf.plot(ax=ax, color='lightgray', edgecolor='white', linewidth=0.5)

    # Plot regions with events
    gdf_with_events = gdf[gdf['event_count'] > 0]
    if len(gdf_with_events) > 0:
        max_count = gdf['event_count'].max()
        gdf_with_events.plot(
            ax=ax,
            column='event_count',
            cmap=get_colormap(),
            edgecolor='black',
            linewidth=0.3,
            legend=True,
            legend_kwds={
                'label': f'Number of {disaster_type} Events',
                'orientation': 'horizontal',
                'shrink': 0.6,
                'pad': 0.05
            },
            vmin=0,
            vmax=max_count
        )

    # Add title
    ax.set_title(title, fontsize=16, fontweight='bold', pad=20)

    # Add source annotation
    ax.annotate(
        'Data Source: EM-DAT, https://www.emdat.be/',
        xy=(0.99, 0.01), xycoords='axes fraction',
        ha='right', va='bottom', fontsize=8, color='gray'
    )

    # Remove axis labels
    ax.set_xlabel('')
    ax.set_ylabel('')
    ax.tick_params(labelsize=8)

    # Save
    plt.tight_layout()
    plt.savefig(output_path, dpi=200, bbox_inches='tight', facecolor='white')
    plt.close()

    print(f"  Saved: {output_path}")


def create_yearly_extent_map(gdf, affected_admin1, year, output_path, disaster_type='Disaster'):
    """
    Create map showing affected regions for a specific year.

    Parameters
    ----------
    gdf : gpd.GeoDataFrame
        Admin1 boundaries
    affected_admin1 : list
        List of affected admin1 GID codes
    year : int
        Year
    output_path : str
        Output PNG file path
    disaster_type : str
        Type of disaster
    """
    fig, ax = plt.subplots(1, 1, figsize=(10, 12))

    # Set map extent
    ax.set_xlim(MAP_EXTENT['x_min'], MAP_EXTENT['x_max'])
    ax.set_ylim(MAP_EXTENT['y_min'], MAP_EXTENT['y_max'])

    # Plot all regions in light gray
    gdf.plot(ax=ax, color='lightgray', edgecolor='white', linewidth=0.5)

    # Plot affected regions in red
    gdf_affected = gdf[gdf['GID_1'].isin(affected_admin1)]
    if len(gdf_affected) > 0:
        gdf_affected.plot(ax=ax, color='red', edgecolor='darkred', linewidth=0.5, alpha=0.8)

    # Add year label
    ax.text(
        0.95, 0.05, str(year),
        transform=ax.transAxes,
        fontsize=36, fontweight='bold',
        ha='right', va='bottom'
    )

    # Add disaster type label
    ax.text(
        0.95, 0.12, f'{disaster_type} Events',
        transform=ax.transAxes,
        fontsize=14, fontweight='bold',
        ha='right', va='bottom'
    )

    # Remove axis labels
    ax.set_xlabel('')
    ax.set_ylabel('')
    ax.axis('off')

    # Save
    plt.tight_layout()
    plt.savefig(output_path, dpi=150, bbox_inches='tight', facecolor='white')
    plt.close()


def create_combined_yearly_grid(image_dir, output_path, disaster_type, year_start, year_end):
    """
    Create a grid of yearly maps combined into one image.

    Parameters
    ----------
    image_dir : str
        Directory containing yearly PNG files
    output_path : str
        Output combined PNG file path
    disaster_type : str
        'Drought' or 'Flood'
    year_start : int
        Start year
    year_end : int
        End year
    """
    from PIL import Image
    import glob

    # Get list of image files
    pattern = os.path.join(image_dir, f"*_{disaster_type.lower()}.png")
    files = sorted(glob.glob(pattern))

    if not files:
        print(f"  No images found for combined grid")
        return

    # Calculate grid dimensions
    n_images = len(files)
    n_cols = 6
    n_rows = (n_images + n_cols - 1) // n_cols

    # Load first image to get dimensions
    first_img = Image.open(files[0])
    img_width, img_height = first_img.size

    # Create combined image
    combined_width = n_cols * img_width
    combined_height = n_rows * img_height
    combined = Image.new('RGB', (combined_width, combined_height), 'white')

    # Paste images
    for i, filepath in enumerate(files):
        img = Image.open(filepath)
        row = i // n_cols
        col = i % n_cols
        x = col * img_width
        y = row * img_height
        combined.paste(img, (x, y))

    # Save
    combined.save(output_path, dpi=(200, 200))
    print(f"  Saved combined grid: {output_path}")

# =============================================================================
# Summary Functions
# =============================================================================

def print_summary(df, disaster_type):
    """
    Print summary statistics for disaster data.

    Parameters
    ----------
    df : pd.DataFrame
        Disaster data
    disaster_type : str
        Type of disaster
    """
    print(f"\n{disaster_type} Summary:")
    print(f"  Total events: {len(df)}")
    print(f"  Year range: {df['Start Year'].min()} - {df['Start Year'].max()}")
    print(f"  Countries affected:")
    for country in sorted(df['Country'].unique()):
        count = len(df[df['Country'] == country])
        print(f"    {country}: {count} events")


def save_processed_data(df, output_path):
    """
    Save processed data with admin codes to Excel.

    Parameters
    ----------
    df : pd.DataFrame
        Processed data
    output_path : str
        Output Excel file path
    """
    # Convert admin1_list to string for Excel
    df_save = df.copy()
    df_save['admin1_list'] = df_save['admin1_list'].apply(
        lambda x: ','.join(x) if isinstance(x, list) else ''
    )

    df_save.to_excel(output_path, index=False)
    print(f"  Saved: {output_path}")

# =============================================================================
# Main Execution
# =============================================================================

def main():
    """Main execution function."""

    print("=" * 60)
    print("East Africa Impact Events - Data Processing")
    print("=" * 60)

    # Create output directories
    os.makedirs(os.path.join(OUTPUT_DIR, 'drought'), exist_ok=True)
    os.makedirs(os.path.join(OUTPUT_DIR, 'flood'), exist_ok=True)

    # Load data
    print("\n[1/5] Loading data...")
    df_emdat = load_emdat_data(EMDAT_FILE)
    gdf_icpac = load_icpac_boundaries(ICPAC_GEOJSON)

    # Filter for East Africa
    print("\n[2/5] Filtering for East Africa...")
    df_ea = filter_ea_data(df_emdat)
    df_ea = filter_by_year_range(df_ea, YEAR_START, YEAR_END)

    # Process Drought events
    print("\n[3/5] Processing Drought events...")
    df_drought = filter_by_disaster_type(df_ea, 'Drought')
    df_drought = assign_admin1_from_country(df_drought, gdf_icpac)
    print_summary(df_drought, 'Drought')

    # Process Flood events
    print("\n[4/5] Processing Flood events...")
    df_flood = filter_by_disaster_type(df_ea, 'Flood')
    df_flood = assign_admin1_from_country(df_flood, gdf_icpac)
    print_summary(df_flood, 'Flood')

    # Generate visualizations
    print("\n[5/5] Generating visualizations...")

    # Drought frequency map
    print("\n  Creating Drought frequency map...")
    gdf_drought_freq = count_events_by_admin1(df_drought, gdf_icpac)
    create_frequency_map(
        gdf_drought_freq,
        f'Drought Event Frequency in East Africa ({YEAR_START}-{YEAR_END})',
        os.path.join(OUTPUT_DIR, f'drought_frequency_{YEAR_START}_{YEAR_END}.png'),
        disaster_type='Drought'
    )

    # Flood frequency map
    print("\n  Creating Flood frequency map...")
    gdf_flood_freq = count_events_by_admin1(df_flood, gdf_icpac)
    create_frequency_map(
        gdf_flood_freq,
        f'Flood Event Frequency in East Africa ({YEAR_START}-{YEAR_END})',
        os.path.join(OUTPUT_DIR, f'flood_frequency_{YEAR_START}_{YEAR_END}.png'),
        disaster_type='Flood'
    )

    # Yearly drought maps
    print("\n  Creating yearly Drought maps...")
    for year in range(YEAR_START, YEAR_END + 1):
        affected = get_affected_admin1_by_year(df_drought, year)
        output_path = os.path.join(OUTPUT_DIR, 'drought', f'{year}_drought.png')
        create_yearly_extent_map(gdf_icpac, affected, year, output_path, 'Drought')
    print(f"    Generated {YEAR_END - YEAR_START + 1} yearly drought maps")

    # Yearly flood maps
    print("\n  Creating yearly Flood maps...")
    for year in range(YEAR_START, YEAR_END + 1):
        affected = get_affected_admin1_by_year(df_flood, year)
        output_path = os.path.join(OUTPUT_DIR, 'flood', f'{year}_flood.png')
        create_yearly_extent_map(gdf_icpac, affected, year, output_path, 'Flood')
    print(f"    Generated {YEAR_END - YEAR_START + 1} yearly flood maps")

    # Save processed data
    print("\n  Saving processed data...")
    save_processed_data(
        df_drought,
        os.path.join(OUTPUT_DIR, 'em_dat_ea_drought_with_admin1.xlsx')
    )
    save_processed_data(
        df_flood,
        os.path.join(OUTPUT_DIR, 'em_dat_ea_flood_with_admin1.xlsx')
    )

    print("\n" + "=" * 60)
    print("Processing complete!")
    print(f"Output files saved to: {OUTPUT_DIR}/")
    print("=" * 60)


if __name__ == "__main__":
    main()