CA/scripts/etl_weather.py

#!/usr/bin/env python3
"""
Weather ETL for Wuhan Respiratory Disease Risk Prediction Platform.
Processes weather CSV files into daily Wuhan parquet.
"""

import argparse
from pathlib import Path

import pandas as pd


# Wuhan station metadata (from station list CSV)
WUHAN_STATIONS = {}

# Coordinate bounding box for Wuhan area
WUHAN_LAT_MIN, WUHAN_LAT_MAX = 29.9, 31.5
WUHAN_LON_MIN, WUHAN_LON_MAX = 113.7, 115.2

# Pollutant type mapping to output schema
POLLUTANT_MAP = {
    'AQI': 'AQI',
    'PM2.5': 'PM25',
    'PM2.5_24h': 'PM25_24h',
    'PM10': 'PM10',
    'PM10_24h': 'PM10_24h',
    'SO2': 'SO2',
    'SO2_24h': 'SO2_24h',
    'NO2': 'NO2',
    'NO2_24h': 'NO2_24h',
    'O3': 'O3',
    'O3_24h': 'O3_24h',
    'O3_8h': 'O3_8h',
    'O3_8h_24h': 'O3_8h_24h',
    'CO': 'CO',
    'CO_24h': 'CO_24h',
    'NOx': 'NOX',
    'primary_pollutant': 'PRIMARY_POLLUTANT',
    'air_quality_level': 'AIR_QUALITY_LEVEL',
}

# Core pollutants for output (7 pollutants as per plan)
OUTPUT_POLLUTANTS = ['AQI', 'PM25', 'PM10', 'SO2', 'NO2', 'O3', 'CO']


def load_station_list(station_file: str) -> dict:
    """Load station metadata from station list CSV."""
    global WUHAN_STATIONS
    stations = {}
    df = pd.read_csv(station_file, encoding='utf-8')
    for _, row in df.iterrows():
        station_id = str(row['监测点编码']).strip()
        city = str(row['城市']).strip() if pd.notna(row['城市']) else ''
        
        lat_val = row['纬度']
        lon_val = row['经度']
        
        try:
            lat = float(lat_val) if pd.notna(lat_val) and lat_val != '-' else 0
        except (ValueError, TypeError):
            lat = 0
            
        try:
            lon = float(lon_val) if pd.notna(lon_val) and lon_val != '-' else 0
        except (ValueError, TypeError):
            lon = 0
            
        district = str(row['监测点名称']).strip() if pd.notna(row['监测点名称']) else ''

        if (city == '武汉' or
            (WUHAN_LAT_MIN <= lat <= WUHAN_LAT_MAX and
             WUHAN_LON_MIN <= lon <= WUHAN_LON_MAX)):
            stations[station_id] = {
                'name': district,
                'lat': lat,
                'lon': lon,
                'district': district,
            }
    WUHAN_STATIONS = stations
    return stations


def process_daily_csv(csv_path: str, wuhan_stations: list[str]) -> pd.DataFrame:
    """Process a single daily CSV file.

    Args:
        csv_path: Path to china_sites_YYYYMMDD.csv
        wuhan_stations: List of Wuhan station IDs. If empty, process ALL stations.

    Returns:
        DataFrame with columns: datetime, station_id, pollutant, value
    """
    df = pd.read_csv(csv_path)

    if wuhan_stations:
        wuhan_cols = ['date', 'hour', 'type'] + wuhan_stations
        available_cols = [c for c in wuhan_cols if c in df.columns]
    else:
        available_cols = df.columns.tolist()
    
    df = df[available_cols]

    id_vars = ['date', 'hour', 'type']
    value_vars = [c for c in available_cols if c not in id_vars]

    if not value_vars:
        return pd.DataFrame(columns=['datetime', 'station_id', 'pollutant', 'value'])

    df_long = df.melt(
        id_vars=id_vars,
        value_vars=value_vars,
        var_name='station_id',
        value_name='value',
    )

    df_long['datetime'] = pd.to_datetime(
        df_long['date'].astype(str) + df_long['hour'].astype(str).str.zfill(2),
        format='%Y%m%d%H'
    )

    df_long['pollutant'] = df_long['type'].map(POLLUTANT_MAP)

    return df_long[['datetime', 'station_id', 'pollutant', 'value']]


def aggregate_to_daily(df_long: pd.DataFrame, wuhan_metadata: dict) -> pd.DataFrame:
    """Aggregate hourly data to daily level per station.

    Uses mean for all pollutants.
    """
    # Filter to output pollutants only
    df_pollutants = df_long[df_long['pollutant'].isin(OUTPUT_POLLUTANTS)].copy()

    # Extract date (without time) from datetime
    df_pollutants['date'] = df_pollutants['datetime'].dt.date

    # First aggregate by (date, station_id, pollutant) to get daily mean
    df_daily_pollutant = df_pollutants.groupby(
        ['date', 'station_id', 'pollutant'], as_index=False
    )['value'].mean()

    # Pivot to wide format: one column per pollutant
    df_pivot = df_daily_pollutant.pivot_table(
        index=['date', 'station_id'],
        columns='pollutant',
        values='value',
        aggfunc='mean'
    ).reset_index()

    # Flatten column names
    df_pivot.columns.name = None

    # Add metadata
    df_pivot['district'] = df_pivot['station_id'].map(
        lambda x: wuhan_metadata.get(x, {}).get('district', '')
    )
    df_pivot['lat'] = df_pivot['station_id'].map(
        lambda x: wuhan_metadata.get(x, {}).get('lat', 0)
    )
    df_pivot['lon'] = df_pivot['station_id'].map(
        lambda x: wuhan_metadata.get(x, {}).get('lon', 0)
    )

    # Ensure output schema columns exist
    for col in OUTPUT_POLLUTANTS:
        if col not in df_pivot.columns:
            df_pivot[col] = None

    # Reorder columns
    output_cols = ['date', 'station_id', 'district', 'lat', 'lon'] + OUTPUT_POLLUTANTS
    df_pivot = df_pivot[[c for c in output_cols if c in df_pivot.columns]]

    return df_pivot


def process_year(input_dir: str, output_dir: str, year: int, station_file: str = None) -> None:
    """Process all CSV files for a given year."""
    from pathlib import Path
    import glob as glob_module
    import os
    
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)

    # Find station list file dynamically if not provided
    if station_file is None or not os.path.exists(station_file):
        base_dir = '/home/akiba/CA/Datas'
        station_files = glob_module.glob(os.path.join(base_dir, '*空气*', '*列表*.csv'))
        if station_files:
            station_file = station_files[0]
            print(f'Found station list: {station_file}')
        else:
            print(f'ERROR: No station list file found')
            return

    if station_file and os.path.exists(station_file):
        print(f'Loading station list from {station_file}...')
        wuhan_stations = load_station_list(station_file)
        station_ids = list(wuhan_stations.keys())
        print(f'  Found {len(station_ids)} Wuhan stations: {station_ids}')
    else:
        print(f'ERROR: Station file not found: {station_file}')
        return

    # Find year directory dynamically
    base_dir = '/home/akiba/CA/Datas'
    year_dirs = glob_module.glob(os.path.join(base_dir, '*空气*', f'站点_{year}*'))
    # Filter out .zip files and Zone.Identifier
    year_dirs = [d for d in year_dirs if os.path.isdir(d)]
    
    if not year_dirs:
        print(f'No directory found for year {year}')
        return
    
    year_dir = year_dirs[0]
    print(f'Using year directory: {year_dir}')
    
    csv_pattern = os.path.join(year_dir, f'china_sites_{year}*.csv')
    csv_files = sorted(glob_module.glob(csv_pattern))

    if not csv_files:
        print(f'No CSV files found for year {year}')
        return

    print(f'Processing {len(csv_files)} files for year {year}...')

    all_data = []
    for i, csv_file in enumerate(csv_files):
        try:
            df = process_daily_csv(csv_file, station_ids)
            all_data.append(df)
            if i == 0:
                print(f'  First file processed: {df.shape}')
            if (i + 1) % 50 == 0:
                print(f'  Processed {i + 1}/{len(csv_files)} files...')
        except Exception as e:
            print(f'Error processing {csv_file}: {e}')

    if not all_data:
        print('No data processed successfully.')
        return

    df_combined = pd.concat(all_data, ignore_index=True)
    print(f'Combined data shape: {df_combined.shape}')
    
    df_daily = aggregate_to_daily(df_combined, wuhan_stations)
    df_daily = df_daily.sort_values(['date', 'station_id']).reset_index(drop=True)

    output_file = output_path / f'daily_wuhan_{year}.parquet'
    df_daily.to_parquet(output_file, index=False)

    print(f'Output: {output_file}')
    print(f'Shape: {df_daily.shape}')
    print(f'Columns: {list(df_daily.columns)}')
    print(f'Date range: {df_daily["date"].min()} to {df_daily["date"].max()}')
    print(f'Stations: {df_daily["station_id"].nunique()}')


def main():
    parser = argparse.ArgumentParser(description='Process weather data for Wuhan')
    parser.add_argument('--year', type=int, required=True, help='Year to process (e.g., 2022)')
    parser.add_argument(
        '--output-dir',
        type=str,
        default='processed/weather',
        help='Output directory for parquet files'
    )
    parser.add_argument(
        '--station-file',
        type=str,
        default=None,
        help='Station list CSV file (auto-detected if not provided)'
    )
    args = parser.parse_args()

    process_year(None, args.output_dir, args.year, args.station_file)


if __name__ == '__main__':
    main()