#!/usr/bin/env python3
"""
Compute Weather Lag Features for Wuhan Respiratory Disease Risk Prediction Platform.

Computes lag features (1,2,3,5,7,14 days) for weather data.
- 7 original features: PM2.5, PM10, O3, NO2, SO2, CO, temperature
- 6 lags: 1, 2, 3, 5, 7, 14 days
- CO dropped post-lag (lowest correlation with respiratory disease)
- Final: 48 features per node per day

Output: processed/weather/lag_features.parquet - 48 features per node per day
"""

import argparse
import glob
from pathlib import Path

import pandas as pd


# Wuhan station IDs from station list
WUHAN_STATIONS = [
    "1325A",  # 东湖梨园
    "1326A",  # 汉阳月湖
    "1327A",  # 汉口花桥
    "1328A",  # 武昌紫阳
    "1329A",  # 青山钢花
    "1330A",  # 沌口新区
    "1331A",  # 汉口江滩
    "1332A",  # 东湖高新
    "1333A",  # 吴家山
    "1334A",  # 沉湖七壕(对照点)
    "3153A",  # 民族大道182号
]

# Weather types to process (using _24h variants for daily averages)
# Note: Temperature may not be available in all datasets
WEATHER_TYPES = ["PM2.5", "PM10", "O3", "NO2", "SO2", "CO"]

# Lag periods in days
LAG_PERIODS = [1, 2, 3, 5, 7, 14]

# Columns to drop after lagging (CO has lowest correlation with respiratory disease)
# Per spec: CO dropped post-lag means only original CO column dropped, not its lags
DROP_COLUMNS = ["CO"]


def load_daily_weather_data(input_path: str) -> pd.DataFrame:
    """
    Load daily weather data from parquet files from US-001 processing.

    Args:
        input_path: Glob pattern for input parquet files (e.g., 'processed/weather/daily_wuhan_*.parquet')

    Returns:
        DataFrame with date, station_id, and weather features
    """
    files = glob.glob(input_path)
    if not files:
        raise FileNotFoundError(f"No files found matching pattern: {input_path}")

    print(f"Loading {len(files)} parquet files...")
    dfs = []
    for f in files:
        df = pd.read_parquet(f)
        print(f"  Loaded {f}: {df.shape}")
        dfs.append(df)

    data = pd.concat(dfs, ignore_index=True)

    # Standardize column names
    if 'PM25' in data.columns:
        data = data.rename(columns={'PM25': 'PM2.5'})

    # Select relevant weather columns (drop lat, lon, district for feature computation)
    weather_cols = ['date', 'station_id', 'AQI', 'PM2.5', 'PM10', 'SO2', 'NO2', 'O3', 'CO']
    data = data[[c for c in weather_cols if c in data.columns]]

    # Convert date to datetime
    data['date'] = pd.to_datetime(data['date'])

    print(f"Total records: {len(data)}")
    print(f"Date range: {data['date'].min()} to {data['date'].max()}")
    print(f"Weather columns: {[c for c in data.columns if c not in ['date', 'station_id']]}")

    return data


def load_weather_from_csv(csv_dir: str, year: int) -> pd.DataFrame:
    """
    Load weather data from CSV files and aggregate to daily level for Wuhan stations.

    Args:
        csv_dir: Directory containing daily CSV files
        year: Year to process

    Returns:
        DataFrame with date, station_id, and weather features
    """
    # Get all CSV files for the year
    csv_pattern = f"{csv_dir}/站点_{year}*/china_sites_*.csv"
    files = glob.glob(csv_pattern)

    if not files:
        raise FileNotFoundError(f"No weather CSV files found for year {year} in {csv_dir}")

    print(f"Processing {len(files)} CSV files for year {year}...")

    # Filter to only Wuhan stations that exist in the data
    sample_df = pd.read_csv(files[0], usecols=["date", "hour", "type"])
    available_stations = [s for s in WUHAN_STATIONS if s in pd.read_csv(files[0]).columns]
    print(f"Found {len(available_stations)} Wuhan stations in data: {available_stations}")

    if not available_stations:
        raise ValueError(f"No Wuhan stations found in data")

    # Use _24h variants for daily averages
    type_to_use = {}
    for wt in WEATHER_TYPES:
        if wt in ["PM2.5", "PM10", "SO2", "NO2", "CO"]:
            type_to_use[wt] = f"{wt}_24h"
        elif wt == "O3":
            # O3 has O3_24h variant
            type_to_use[wt] = "O3_24h"
        else:
            type_to_use[wt] = wt

    print(f"Using types: {type_to_use}")

    dfs = []
    for i, f in enumerate(files):
        if i % 50 == 0:
            print(f"  Processing file {i+1}/{len(files)}...")

        try:
            df = pd.read_csv(f)

            # Filter for hour=0 (daily values) and relevant types
            df = df[(df["hour"] == 0) & (df["type"].isin(type_to_use.values()))].copy()

            if df.empty:
                continue

            # Select only Wuhan station columns
            cols_to_keep = ["date", "type"] + available_stations
            df = df[[c for c in cols_to_keep if c in df.columns]]

            if len(df.columns) < 3:
                continue

            # Melt to long format (station_id x weather_type)
            df_melted = df.melt(
                id_vars=["date", "type"],
                var_name="station_id",
                value_name="value"
            )

            # Map back to standard type names
            reverse_map = {v: k for k, v in type_to_use.items() if k in WEATHER_TYPES}
            df_melted["type"] = df_melted["type"].map(reverse_map)

            dfs.append(df_melted)

        except Exception as e:
            print(f"Error processing {f}: {e}")
            continue

    if not dfs:
        raise ValueError(f"No valid data found for year {year}")

    data = pd.concat(dfs, ignore_index=True)
    print(f"Loaded {len(data)} records before pivot")

    # Pivot: index=(date, station_id), columns=type, values=value
    data = data.pivot_table(
        index=["date", "station_id"],
        columns="type",
        values="value"
    ).reset_index()

    data.columns.name = None

    # Convert date to datetime
    data["date"] = pd.to_datetime(data["date"], format="%Y%m%d")

    # Drop duplicate rows
    data = data.drop_duplicates(subset=["date", "station_id"])

    print(f"Loaded {len(data)} daily weather records for year {year}")
    print(f"Columns: {list(data.columns)}")
    return data


def compute_lag_features(df: pd.DataFrame, lag_periods: list, drop_columns: list) -> pd.DataFrame:
    """
    Compute lag features for weather data.

    Args:
        df: DataFrame with date, station_id, and weather columns
        lag_periods: List of lag periods in days
        drop_columns: List of column names to drop after lagging (only original, not lags)

    Returns:
        DataFrame with lag features added
    """
    # Get weather columns (exclude date and station_id)
    weather_cols = [c for c in df.columns if c not in ["date", "station_id"]]

    print(f"Original weather columns: {weather_cols}")
    print(f"Number of original features: {len(weather_cols)}")

    # Sort by station and date for proper lagging
    df = df.sort_values(["station_id", "date"]).reset_index(drop=True)

    # Compute lag features for each weather column
    lag_cols_added = []
    for col in weather_cols:
        for lag in lag_periods:
            lag_col_name = f"{col}_lag{lag}"
            df[lag_col_name] = df.groupby("station_id")[col].shift(lag)
            lag_cols_added.append(lag_col_name)

    print(f"Created {len(lag_cols_added)} lag columns")

    # Drop only the ORIGINAL columns in drop_columns (not their lags)
    # Per spec: CO dropped post-lag means original CO is dropped, CO lags are kept
    for col in drop_columns:
        if col in df.columns:
            df = df.drop(columns=[col])
            print(f"Dropped original column: {col} (CO lags are kept per spec)")

    # Count final columns (excluding date and station_id)
    feature_cols = [c for c in df.columns if c not in ["date", "station_id"]]
    num_features = len(feature_cols)

    print(f"Final feature count: {num_features}")

    # Readiness gate assertion - exactly 48 columns required
    EXPECTED_FEATURES = 48
    if num_features != EXPECTED_FEATURES:
        raise ValueError(
            f"Feature count mismatch: expected {EXPECTED_FEATURES}, got {num_features}. "
            f"Features: {feature_cols}"
        )

    print(f"Readiness gate PASSED: {num_features} features per node per day")

    return df


def main():
    parser = argparse.ArgumentParser(
        description="Compute weather lag features for Wuhan Respiratory Disease Risk Prediction"
    )
    parser.add_argument(
        "--input",
        type=str,
        default="Datas/气象+空气",
        help="Input CSV directory or glob pattern for parquet files"
    )
    parser.add_argument(
        "--output",
        type=str,
        default="processed/weather/lag_features.parquet",
        help="Output parquet file path"
    )
    parser.add_argument(
        "--year",
        type=int,
        default=2022,
        help="Year to process (for CSV input)"
    )
    parser.add_argument(
        "--use-csv",
        action="store_true",
        help="Use CSV input instead of parquet"
    )

    args = parser.parse_args()

    # Create output directory
    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    # Load data
    if args.use_csv:
        print(f"Loading weather data from CSV directory: {args.input}")
        data = load_weather_from_csv(args.input, args.year)
    else:
        print(f"Loading weather data from parquet files: {args.input}")
        data = load_daily_weather_data(args.input)

    # Compute lag features
    print("Computing lag features...")
    result = compute_lag_features(data, LAG_PERIODS, DROP_COLUMNS)

    # Sort by date and station
    result = result.sort_values(["date", "station_id"]).reset_index(drop=True)

    # Save output
    print(f"Saving to {args.output}")
    result.to_parquet(args.output, index=False)

    # Verify output
    df_verify = pd.read_parquet(args.output)
    feature_cols = [c for c in df_verify.columns if c not in ["date", "station_id"]]

    print(f"\n=== Verification ===")
    print(f"Output shape: {df_verify.shape}")
    print(f"Number of features: {len(feature_cols)}")
    print(f"Date range: {df_verify['date'].min()} to {df_verify['date'].max()}")
    print(f"Stations: {df_verify['station_id'].nunique()}")
    print(f"Feature columns: {feature_cols[:10]}... (showing first 10)")

    # Final readiness gate
    assert len(feature_cols) == 48, f"Readiness gate failed: expected 48 features, got {len(feature_cols)}"
    print("\nReadiness gate PASSED: Exactly 48 features per node per day")


if __name__ == "__main__":
    main()