feature_engineering.py

"""
Comprehensive Feature Engineering for Battery SOH Prediction

Creates 50+ engineered features per cycle including:
- Raw measurements (normalized)
- Statistical features
- Temporal features
- Degradation features
- Physics-informed features
"""

import pandas as pd
import numpy as np
from sklearn.preprocessing import MinMaxScaler
import warnings
warnings.filterwarnings('ignore')


def normalize_features(df, columns):
    """Normalize features to 0-1 range."""
    df_norm = df.copy()
    scaler = MinMaxScaler()
    
    for col in columns:
        if col in df_norm.columns:
            df_norm[col] = scaler.fit_transform(df_norm[[col]])
    
    return df_norm


def calculate_energy_throughput(df):
    """
    Calculate energy throughput: ∫(voltage * current) dt per cycle.
    
    Parameters:
    -----------
    df : DataFrame
        DataFrame with cycle_index, voltage, current, time columns
    
    Returns:
    --------
    Series: Energy throughput per cycle
    """
    if 'time' not in df.columns:
        # Estimate time from cycle progression
        df = df.copy()
        df['time'] = df.groupby('cycle_index').cumcount()
    
    # Calculate power (voltage * current)
    df['power'] = df['voltage'] * df['current']
    
    # Calculate energy throughput per cycle
    energy_per_cycle = df.groupby('cycle_index').apply(
        lambda x: np.trapz(x['power'], x['time']) if len(x) > 1 else 0
    )
    
    return energy_per_cycle


def engineer_features(df, initial_capacity):
    """
    Engineer comprehensive features for each cycle.
    
    Parameters:
    -----------
    df : DataFrame
        Preprocessed DataFrame with cycle_index, voltage, current, temperature, capacity, soh
    initial_capacity : float
        Initial capacity value
    
    Returns:
    --------
    DataFrame: Feature matrix with 50+ features per cycle
    """
    print("\n" + "="*60)
    print("FEATURE ENGINEERING")
    print("="*60)
    
    df_feat = df.copy()
    
    # Ensure cycle_index exists
    if 'cycle_index' not in df_feat.columns:
        print("Error: cycle_index column not found")
        return None
    
    # Sort by cycle_index
    df_feat = df_feat.sort_values('cycle_index').reset_index(drop=True)
    
    # Group by cycle to create cycle-level features
    cycle_groups = df_feat.groupby('cycle_index')
    
    # Initialize feature DataFrame
    features = pd.DataFrame()
    features['cycle_index'] = cycle_groups['cycle_index'].first()
    
    # ========================================================================
    # 1. RAW MEASUREMENTS (normalized 0-1)
    # ========================================================================
    print("\n1. Raw Measurements (normalized)...")
    
    # Normalize raw measurements
    if 'voltage' in df_feat.columns:
        features['voltage_mean'] = cycle_groups['voltage'].mean()
        features['voltage_norm'] = (features['voltage_mean'] - features['voltage_mean'].min()) / (features['voltage_mean'].max() - features['voltage_mean'].min())
    
    if 'current' in df_feat.columns:
        features['current_mean'] = cycle_groups['current'].mean()
        features['current_norm'] = (features['current_mean'] - features['current_mean'].min()) / (features['current_mean'].max() - features['current_mean'].min())
        
        # Charge/discharge indicators
        features['is_charge'] = (features['current_mean'] > 0).astype(int)
        features['is_discharge'] = (features['current_mean'] < 0).astype(int)
    
    if 'temperature' in df_feat.columns:
        features['temperature_mean'] = cycle_groups['temperature'].mean()
        features['temperature_norm'] = (features['temperature_mean'] - features['temperature_mean'].min()) / (features['temperature_mean'].max() - features['temperature_mean'].min())
    
    if 'capacity' in df_feat.columns:
        features['capacity'] = cycle_groups['capacity'].mean()
        features['capacity_norm'] = features['capacity'] / initial_capacity if initial_capacity > 0 else features['capacity']
    
    if 'soh' in df_feat.columns:
        features['soh'] = cycle_groups['soh'].mean()
    
    # ========================================================================
    # 2. STATISTICAL FEATURES (per cycle)
    # ========================================================================
    print("2. Statistical Features...")
    
    for col in ['voltage', 'current', 'temperature']:
        if col in df_feat.columns:
            features[f'{col}_mean'] = cycle_groups[col].mean()
            features[f'{col}_std'] = cycle_groups[col].std().fillna(0)
            features[f'{col}_min'] = cycle_groups[col].min()
            features[f'{col}_max'] = cycle_groups[col].max()
            features[f'{col}_range'] = features[f'{col}_max'] - features[f'{col}_min']
            features[f'{col}_median'] = cycle_groups[col].median()
    
    # Charge/discharge capacity difference
    if 'current' in df_feat.columns and 'capacity' in df_feat.columns:
        charge_capacity = df_feat[df_feat['current'] > 0].groupby('cycle_index')['capacity'].mean()
        discharge_capacity = df_feat[df_feat['current'] < 0].groupby('cycle_index')['capacity'].mean()
        features['charge_discharge_diff'] = (charge_capacity - discharge_capacity).fillna(0)
    
    # Energy throughput
    if 'voltage' in df_feat.columns and 'current' in df_feat.columns:
        try:
            energy = calculate_energy_throughput(df_feat)
            features['energy_throughput'] = features['cycle_index'].map(energy).fillna(0)
        except:
            # Fallback: simple power calculation
            features['power_mean'] = features['voltage_mean'] * features['current_mean']
            features['energy_throughput'] = features['power_mean'].abs()
    
    # ========================================================================
    # 3. TEMPORAL FEATURES
    # ========================================================================
    print("3. Temporal Features...")
    
    features['cycles_from_start'] = features['cycle_index'] - features['cycle_index'].min()
    features['cycles_to_end'] = features['cycle_index'].max() - features['cycle_index']
    
    # Cumulative operating time (if time column exists)
    if 'time' in df_feat.columns:
        features['cumulative_time'] = cycle_groups['time'].sum()
        features['time_per_cycle'] = cycle_groups['time'].mean()
    else:
        features['cumulative_time'] = features['cycles_from_start']
        features['time_per_cycle'] = 1.0
    
    # Days between cycles (estimate if not available)
    features['days_since_start'] = features['cycles_from_start'] * 0.5  # Estimate: 0.5 days per cycle
    
    # ========================================================================
    # 4. DEGRADATION FEATURES
    # ========================================================================
    print("4. Degradation Features...")
    
    if 'soh' in features.columns:
        # Capacity fade rate
        features['soh_fade_rate'] = features['soh'].diff().fillna(0)
        features['soh_fade_rate_abs'] = features['soh_fade_rate'].abs()
        
        # Rolling average of SOH (window=5 cycles)
        features['soh_rolling_mean_5'] = features['soh'].rolling(window=5, min_periods=1).mean()
        features['soh_rolling_std_5'] = features['soh'].rolling(window=5, min_periods=1).std().fillna(0)
        
        # Exponential moving average
        features['soh_ema'] = features['soh'].ewm(span=5, adjust=False).mean()
        
        # Second derivative (acceleration of degradation)
        features['soh_second_derivative'] = features['soh_fade_rate'].diff().fillna(0)
        
        # Cumulative capacity loss
        features['cumulative_capacity_loss'] = (1 - features['soh']) * initial_capacity if initial_capacity > 0 else (1 - features['soh'])
    
    if 'capacity' in features.columns:
        features['capacity_fade_rate'] = features['capacity'].diff().fillna(0)
        features['capacity_fade_rate_pct'] = features['capacity'].pct_change().fillna(0)
    
    # ========================================================================
    # 5. PHYSICS-INFORMED FEATURES
    # ========================================================================
    print("5. Physics-Informed Features...")
    
    # Charge/discharge asymmetry ratio
    if 'current' in df_feat.columns:
        charge_current = df_feat[df_feat['current'] > 0]['current'].abs().sum() if len(df_feat[df_feat['current'] > 0]) > 0 else 0
        discharge_current = df_feat[df_feat['current'] < 0]['current'].abs().sum() if len(df_feat[df_feat['current'] < 0]) > 0 else 0
        if discharge_current > 0:
            features['charge_discharge_asymmetry'] = charge_current / discharge_current
        else:
            features['charge_discharge_asymmetry'] = 1.0
    
    # Temperature stress (time spent above 40°C)
    if 'temperature' in df_feat.columns:
        temp_stress = cycle_groups['temperature'].apply(lambda x: (x > 40).sum())
        features['temperature_stress'] = temp_stress / cycle_groups['temperature'].count()
        features['max_temperature'] = cycle_groups['temperature'].max()
        features['min_temperature'] = cycle_groups['temperature'].min()
        features['temperature_range'] = features['max_temperature'] - features['min_temperature']
    
    # Voltage variance during constant current phases
    if 'voltage' in df_feat.columns and 'current' in df_feat.columns:
        # Variance during discharge (constant current phase)
        discharge_voltage_var = df_feat[df_feat['current'] < 0].groupby('cycle_index')['voltage'].var().fillna(0)
        features['voltage_variance_discharge'] = features['cycle_index'].map(discharge_voltage_var).fillna(0)
    
    # Power features
    if 'voltage_mean' in features.columns and 'current_mean' in features.columns:
        features['power_mean'] = features['voltage_mean'] * features['current_mean']
        features['power_abs'] = features['power_mean'].abs()
    
    # ========================================================================
    # 6. TARGET VARIABLES
    # ========================================================================
    print("6. Target Variables...")
    
    if 'soh' in features.columns:
        # Target: SOH 10 cycles ahead (for prediction task)
        features['soh_10_cycles_ahead'] = features['soh'].shift(-10).fillna(features['soh'].iloc[-1])
        
        # Target: Remaining cycles until 80% SOH (for RUL prediction)
        eol_threshold = 0.80
        features['remaining_cycles_to_eol'] = None
        for idx, row in features.iterrows():
            future_soh = features[features['cycle_index'] > row['cycle_index']]['soh']
            eol_cycles = future_soh[future_soh < eol_threshold]
            if len(eol_cycles) > 0:
                eol_cycle = eol_cycles.index[0]
                features.loc[idx, 'remaining_cycles_to_eol'] = features.loc[eol_cycle, 'cycle_index'] - row['cycle_index']
            else:
                features.loc[idx, 'remaining_cycles_to_eol'] = features['cycle_index'].max() - row['cycle_index']
        
        features['remaining_cycles_to_eol'] = features['remaining_cycles_to_eol'].fillna(0)
    
    # Fill any remaining NaN values
    features = features.fillna(0)
    
    print(f"\n✓ Feature engineering complete: {len(features)} cycles, {len(features.columns)} features")
    print(f"  Feature columns: {len(features.columns)}")
    
    return features


def engineer_features_for_all_batteries(dataframes_dict, initial_capacities_dict):
    """
    Engineer features for all batteries.
    
    Parameters:
    -----------
    dataframes_dict : dict
        Dictionary of {battery_id: DataFrame}
    initial_capacities_dict : dict
        Dictionary of {battery_id: initial_capacity}
    
    Returns:
    --------
    dict: Dictionary of {battery_id: feature_DataFrame}
    """
    all_features = {}
    
    for battery_id, df in dataframes_dict.items():
        print(f"\n{'='*60}")
        print(f"Processing {battery_id}")
        print(f"{'='*60}")
        
        initial_capacity = initial_capacities_dict.get(battery_id, df['capacity'].max() if 'capacity' in df.columns else 1.0)
        
        features = engineer_features(df, initial_capacity)
        if features is not None:
            all_features[battery_id] = features
    
    return all_features


if __name__ == "__main__":
    # Example usage
    import sys
    sys.path.append('.')
    from unified_data_pipeline import preprocess_all_datasets
    
    # Load preprocessed data
    dataframes, summary = preprocess_all_datasets()
    
    if dataframes:
        # Get initial capacities
        initial_capacities = {}
        for battery_id, df in dataframes.items():
            initial_capacities[battery_id] = df['capacity'].max() if 'capacity' in df.columns else 1.0
        
        # Engineer features
        all_features = engineer_features_for_all_batteries(dataframes, initial_capacities)
        
        # Save features
        import os
        os.makedirs('./data/features', exist_ok=True)
        for battery_id, features in all_features.items():
            output_path = f'./data/features/{battery_id}_features.csv'
            features.to_csv(output_path, index=False)
            print(f"\nSaved features to: {output_path}")