feature_generation.py

# ---------------------------------------------------------------------------
#  Feature generation functions for NYC taxi trip data
# ---------------------------------------------------------------------------
# Provides functions to add various feature groups to taxi trip dataframes:
# - Time-based features (time of day, day of week, month, etc.)
# - Weather features (merged from weather data)
# - Holiday features (US holidays)
# - Heat zone features (location frequency-based)
# - Duration-related features (copies and predictions for fare modeling)
#
# All functions take a dataframe and return a modified copy with new columns added.
# Used by FE_DA_helper_functions.py for batch processing multiple datasets.
# ---------------------------------------------------------------------------

import pandas as pd
import numpy as np
import holidays


# ---------------------------------------------------------------------------
# Time-based feature generation
# ---------------------------------------------------------------------------

def add_time_features(df):
    """
    Add time-based features extracted from pickup_datetime.
    
    Creates four time-related features:
    - pickup_tod: Time of day in minutes (0-1439)
    - pickup_day_of_month: Day of month (1-31)
    - pickup_dow: Day of week (0=Monday, 6=Sunday)
    - pickup_month: Month (1-12)
    
    Args:
        df: DataFrame with 'pickup_datetime' column (datetime type)
    
    Returns:
        DataFrame with added time feature columns
    """
    df_out = df.copy()
    
    # Time of day as minutes since midnight (0–1439)
    df_out["pickup_tod"] = df_out['pickup_datetime'].dt.hour * 60 + df_out['pickup_datetime'].dt.minute
    df_out['pickup_day_of_month'] = df_out['pickup_datetime'].dt.day
    df_out["pickup_dow"] = df_out["pickup_datetime"].dt.dayofweek
    df_out["pickup_month"] = df_out["pickup_datetime"].dt.month
    return df_out


# ---------------------------------------------------------------------------
# Weather feature generation
# ---------------------------------------------------------------------------

def add_weather_features(df, weather_df):
    """
    Merge weather data into the trip dataframe based on date and hour.
    
    Matches weather records by date (YYYY-MM-DD) and hour (0-23), then merges
    weather_code (and any other weather columns) into the trip data. Temporary
    merge columns (date, hour) are dropped after merging.
    
    Args:
        df: DataFrame with 'pickup_datetime' column
        weather_df: DataFrame with 'date', 'hour', and 'weather_code' columns
    
    Returns:
        DataFrame with weather features merged (weather_code column added)
    """
    df_out = df.copy()
    # Create date and hour columns for merging with weather data
    df_out['date'] = df_out['pickup_datetime'].dt.strftime('%Y-%m-%d')
    df_out['hour'] = df_out['pickup_datetime'].dt.hour
    
    # Prepare weather dataframe: ensure date is string format
    weather_df_copy = weather_df.copy()
    weather_df_copy['date'] = pd.to_datetime(weather_df_copy['date']).dt.strftime('%Y-%m-%d')
    
    # Merge on date and hour
    df_out = pd.merge(df_out, weather_df_copy, on=['date', 'hour'], how='left')
    
    # Drop temporary merge columns
    df_out = df_out.drop(columns=['date', 'hour'])
    return df_out


# ---------------------------------------------------------------------------
# Holiday feature generation
# ---------------------------------------------------------------------------

def add_holiday_features(df):
    """
    Add a binary holiday indicator feature.
    
    Checks if the pickup date is a US federal holiday in 2022 and creates
    an 'is_holiday' boolean column (True if holiday, False otherwise).
    
    Args:
        df: DataFrame with 'pickup_datetime' column
    
    Returns:
        DataFrame with 'is_holiday' column added
    """
    df_out = df.copy()
    us_holidays = holidays.US(years=[2022])
    
    def is_holiday(date):
        """Check if date is a US holiday."""
        return date.date() in us_holidays
    
    # Apply holiday check to each pickup date
    df_out['is_holiday'] = df_out['pickup_datetime'].apply(is_holiday)
    return df_out


# ---------------------------------------------------------------------------
# Heat zone feature generation
# ---------------------------------------------------------------------------

def add_heat_features(df):
    """
    Add heat zone features based on pickup/dropoff location frequency.
    
    Calculates normalized frequency (proportion) of trips for each pickup and
    dropoff location ID. Higher frequency indicates a "hot" (popular) zone.
    Creates two features:
    - pickup_location_id_heat: Normalized frequency of pickup location
    - dropoff_location_id_heat: Normalized frequency of dropoff location
    
    Args:
        df: DataFrame with 'pickup_location_id' and 'dropoff_location_id' columns
    
    Returns:
        DataFrame with heat zone features added (if location IDs exist)
    """
    df_out = df.copy()
    
    # Add heat zone features if location IDs exist
    if 'pickup_location_id' in df_out.columns and 'dropoff_location_id' in df_out.columns:
        # Calculate normalized frequency (proportion) for each location
        heat_pickup = df_out["pickup_location_id"].value_counts(normalize=True)
        heat_dropoff = df_out["dropoff_location_id"].value_counts(normalize=True)
        
        # Map frequencies to each row
        df_out["pickup_location_id_heat"] = df_out["pickup_location_id"].map(heat_pickup)
        df_out["dropoff_location_id_heat"] = df_out["dropoff_location_id"].map(heat_dropoff)
    
    return df_out


# ---------------------------------------------------------------------------
# Duration-related feature generation (for fare modeling)
# ---------------------------------------------------------------------------

def add_timeduration_copy(df):
    """
    Add a copy of duration_seconds for fare_amount modeling.
    
    Creates 'duration_seconds_copy' column as a duplicate of 'duration_seconds'.
    Used in fare_amount pipeline: the original 'duration_seconds' is dropped to
    prevent leakage, while 'duration_seconds_copy' is kept as a feature to
    measure the impact of actual duration on fare prediction.
    
    Args:
        df: DataFrame with 'duration_seconds' column
    
    Returns:
        DataFrame with 'duration_seconds_copy' column added
    """
    df_out = df.copy()
    if 'duration_seconds' in df_out.columns:
        df_out['duration_seconds_copy'] = df_out['duration_seconds'].copy()
    return df_out


def collect_duration_predictions_for_fare(splits_by_dataset, models_by_dataset, source_name='base', save_path='data/duration_predictions_base.csv'):
    """
    Collect duration model predictions for train/test splits to use in fare modeling.
    
    Uses a trained duration model (from the duration pipeline) to generate predictions
    for X_train and X_test. These predictions can then be merged with the base dataset
    as a feature ('duration_estimated') for fare_amount modeling.
    
    The function selects the first compatible model (whose feature set matches the
    available data) and generates predictions aligned to the original dataframe indices.
    
    Args:
        splits_by_dataset: Dict mapping dataset name to (X_train, X_val, X_test, y_train, y_val, y_test)
                          from the duration pipeline
        models_by_dataset: Dict mapping dataset name to list of (model_name, fitted_model) tuples
                          from the duration pipeline
        source_name: Which dataset's model to use (default: 'base')
        save_path: Optional path to save predictions CSV (default: 'data/duration_predictions_base.csv')
    
    Returns:
        Dictionary with keys:
            - 'full': DataFrame with all predictions (train + test) and split indicator
            - 'train': Series with train predictions, indexed by X_train.index
            - 'test': Series with test predictions, indexed by X_test.index
            - 'model_name': Name of the model used for predictions
        Returns None if source not found or no compatible model available
    """
    if source_name not in splits_by_dataset:
        return None
    
    X_train, X_val, X_test, y_train, y_val, y_test = splits_by_dataset[source_name]
    fitted_list = models_by_dataset.get(source_name, [])
    if not fitted_list:
        return None
    
    # Use first model whose fitted feature set matches X_train
    model_to_use = None
    for model_name, model in fitted_list:
        try:
            feature_names = getattr(model, 'feature_names_in_', None)
            if feature_names is not None:
                missing = set(feature_names) - set(X_train.columns)
                if missing:
                    continue
                model_to_use = (model_name, model)
                break
            else:
                model_to_use = (model_name, model)
                break
        except Exception:
            continue
    
    if model_to_use is None:
        return None
    
    model_name, model = model_to_use
    print(f"[Duration predictions] Using {model_name} model from {source_name} dataset")
    
    # Get predictions for train and test sets
    feature_names = getattr(model, 'feature_names_in_', None)
    if feature_names is not None:
        X_train_use = X_train[list(feature_names)]
        X_test_use = X_test[list(feature_names)]
    else:
        X_train_use = X_train
        X_test_use = X_test
    
    pred_train = model.predict(X_train_use)
    pred_test = model.predict(X_test_use)
    
    # Create dataframe with predictions aligned to original indices
    # Note: train_test_split shuffles, so we need to track indices
    # For now, create a dataframe indexed by position (0..len-1)
    # The fare pipeline will need to merge on the same split indices
    df_predictions = pd.DataFrame({
        'duration_estimated': np.concatenate([pred_train, pred_test]),
        'split': ['train'] * len(pred_train) + ['test'] * len(pred_test)
    })
    
    # Also create separate series for train and test for easier merging
    pred_train_series = pd.Series(pred_train, index=X_train.index, name='duration_estimated')
    pred_test_series = pd.Series(pred_test, index=X_test.index, name='duration_estimated')
    
    if save_path:
        df_predictions.to_csv(save_path, index=False)
        print(f"[Duration predictions] Saved to {save_path}")
    
    return {
        'full': df_predictions,
        'train': pred_train_series,
        'test': pred_test_series,
        'model_name': model_name
    }


def add_duration_estimated_dataset(encoded_datasets, models_by_dataset, source_name='base', new_name='base+timeduration_estimation'):
    """
    Create a dataset with duration predictions merged into the base encoded data.
    
    Uses a trained duration model to predict duration for all rows in the base encoded
    dataset, then adds these predictions as a 'duration_estimated' column. This creates
    a new dataset (e.g., 'base+timeduration_estimation') for fare_amount modeling to
    compare using predicted vs actual duration.
    
    The function selects the first compatible model whose feature set matches the
    available data, allowing the fare pipeline to work even if re-run separately
    from the duration pipeline.
    
    Args:
        encoded_datasets: List of (name, dataframe) tuples from the duration pipeline
        models_by_dataset: Dict mapping dataset name to list of (model_name, fitted_model) tuples
                          from the duration pipeline
        source_name: Which encoded dataset to use (default: 'base')
        new_name: Name for the new dataset (default: 'base+timeduration_estimation')
    
    Returns:
        Tuple (new_name, dataframe) with 'duration_estimated' column added, or None if
        source not found or no compatible model available
    """
    for name, df in encoded_datasets:
        if name != source_name:
            continue
        cols_to_drop = [c for c in ['duration_seconds', 'fare_amount'] if c in df.columns]
        X = df.drop(columns=cols_to_drop)
        fitted_list = models_by_dataset.get(source_name, [])
        if not fitted_list:
            return None
            
        # Use first model whose fitted feature set is contained in X (handles re-run fare-only or mismatched state)
        for model_name, model in fitted_list:
            try:
                feature_names = getattr(model, 'feature_names_in_', None)
                if feature_names is not None:
                    missing = set(feature_names) - set(X.columns)
                    if missing:
                        continue
                    X_use = X[list(feature_names)]
                else:
                    X_use = X
                pred = model.predict(X_use)
                df_est = df.copy()
                df_est['duration_estimated'] = pred
                return (new_name, df_est)
            except Exception:
                continue
        return None
    return None