#!/usr/bin/env python3
"""
Feature Engineering for Economic Recession Prediction
======================================================
Transforms raw economic indicators into predictive features including:
- Year-over-Year (YoY) and Month-over-Month (MoM) changes
- Rolling statistics (averages and volatility)
- Temporal lags (1, 3, 6, 12 months)
- Stationarity analysis via Augmented Dickey-Fuller tests

CRITICAL METHODOLOGICAL REQUIREMENTS:
- No look-ahead bias: rolling windows use ONLY past data
- Strict temporal ordering preserved
- Target alignment verified
"""

import pandas as pd
import numpy as np
from statsmodels.tsa.stattools import adfuller
from pathlib import Path
import warnings
warnings.filterwarnings('ignore')

# Set random seed for reproducibility
np.random.seed(42)

# Paths
INPUT_PATH = '/app/sandbox/session_20251228_135128_cd3e46eedd35/workflow/data/consolidated_macro_data.csv'
OUTPUT_DATA_PATH = '/app/sandbox/session_20251228_135128_cd3e46eedd35/workflow/data/feature_engineered_data.csv'
OUTPUT_REPORT_PATH = '/app/sandbox/session_20251228_135128_cd3e46eedd35/results/feature_engineering_report.txt'

print("=" * 80)
print("FEATURE ENGINEERING PIPELINE")
print("=" * 80)

# ============================================================================
# STEP 1: LOAD DATA
# ============================================================================
print("\n[1/7] Loading harmonized dataset...")
df = pd.read_csv(INPUT_PATH)
print(f"  ✓ Loaded: {df.shape[0]} rows × {df.shape[1]} columns")

# Check for date column (try both 'DATE' and 'date')
date_col = 'DATE' if 'DATE' in df.columns else 'date'
print(f"  ✓ Date range: {df[date_col].min()} to {df[date_col].max()}")

# Convert date to datetime and set as index
df[date_col] = pd.to_datetime(df[date_col])
df = df.sort_values(date_col).set_index(date_col)
print(f"  ✓ Index set to datetime")

# Identify target columns
target_cols = [col for col in df.columns if col.startswith('target_')]
feature_cols = [col for col in df.columns if col not in target_cols]
print(f"  ✓ Identified {len(feature_cols)} features and {len(target_cols)} targets")

# ============================================================================
# STEP 2: FEATURE TRANSFORMATIONS
# ============================================================================
print("\n[2/7] Calculating feature transformations...")

# Define variable groups
yoy_vars = ['INDPRO', 'RSAFS', 'PCE', 'PERMIT', 'HOUST', 'CSUSHPISA',
            'M2SL', 'BUSLOANS', 'MARKET_INDEX']
mom_vars = ['UNRATE', 'CIVPART', 'TCU', 'VIXCLS', 'BAA10Y', 'T10Y2Y', 'T10Y3M']

# YoY % changes (12-month lag)
yoy_features = {}
for var in yoy_vars:
    if var in df.columns:
        yoy_features[f'{var}_yoy_pct'] = ((df[var] - df[var].shift(12)) / df[var].shift(12)) * 100
print(f"  ✓ Created {len(yoy_features)} YoY % change features")

# MoM changes
mom_features = {}
for var in mom_vars:
    if var in df.columns:
        mom_features[f'{var}_mom'] = df[var] - df[var].shift(1)
print(f"  ✓ Created {len(mom_features)} MoM change features")

# Add transformed features to dataframe
for name, series in {**yoy_features, **mom_features}.items():
    df[name] = series

# ============================================================================
# STEP 3: ROLLING STATISTICS
# ============================================================================
print("\n[3/7] Calculating rolling statistics...")

rolling_features = {}

# Rolling averages for volatile indicators
for var, windows in [('VIXCLS', [3, 6]), ('ICSA', [3, 6])]:
    if var in df.columns:
        for window in windows:
            # Use min_periods=window to ensure full window before calculating
            rolling_features[f'{var}_ma{window}'] = df[var].rolling(
                window=window, min_periods=window
            ).mean()

# Rolling volatility for market indicators
for var, window in [('MARKET_INDEX', 3), ('BAA10Y', 3)]:
    if var in df.columns:
        rolling_features[f'{var}_vol{window}'] = df[var].rolling(
            window=window, min_periods=window
        ).std()

print(f"  ✓ Created {len(rolling_features)} rolling statistic features")

# Add rolling features
for name, series in rolling_features.items():
    df[name] = series

# ============================================================================
# STEP 4: LAG GENERATION
# ============================================================================
print("\n[4/7] Generating temporal lags...")

# Get all feature columns (exclude targets and date)
current_features = [col for col in df.columns if col not in target_cols]

lag_windows = [1, 3, 6, 12]
lag_features = {}

for lag in lag_windows:
    for feature in current_features:
        lag_name = f'{feature}_lag{lag}'
        lag_features[lag_name] = df[feature].shift(lag)

print(f"  ✓ Created {len(lag_features)} lagged features")
print(f"  ✓ Lag windows: {lag_windows} months")

# Add lagged features
for name, series in lag_features.items():
    df[name] = series

# ============================================================================
# STEP 5: STATIONARITY ANALYSIS
# ============================================================================
print("\n[5/7] Performing stationarity analysis (ADF tests)...")

# Get all engineered features (exclude original data columns and targets)
engineered_features = [col for col in df.columns if col not in target_cols]

adf_results = []
stationary_count = 0
non_stationary_count = 0

print(f"  → Testing {len(engineered_features)} features...")

for i, feature in enumerate(engineered_features):
    # Print progress every 50 features
    if i % 50 == 0 and i > 0:
        print(f"    Progress: {i}/{len(engineered_features)} features tested...")

    # Drop NaN values for the test
    feature_data = df[feature].dropna()

    if len(feature_data) > 12:  # Minimum observations for ADF test
        try:
            result = adfuller(feature_data, autolag='AIC')
            adf_stat = result[0]
            p_value = result[1]
            is_stationary = p_value < 0.05

            adf_results.append({
                'feature': feature,
                'adf_statistic': adf_stat,
                'p_value': p_value,
                'stationary': is_stationary
            })

            if is_stationary:
                stationary_count += 1
            else:
                non_stationary_count += 1
        except Exception as e:
            adf_results.append({
                'feature': feature,
                'adf_statistic': np.nan,
                'p_value': np.nan,
                'stationary': False,
                'error': str(e)
            })
            non_stationary_count += 1

print(f"  ✓ Completed ADF tests")
print(f"  ✓ Stationary features: {stationary_count}")
print(f"  ✓ Non-stationary features: {non_stationary_count}")

# ============================================================================
# STEP 6: DATA CLEANING
# ============================================================================
print("\n[6/7] Cleaning data (removing NaN values)...")

print(f"  → Shape before cleaning: {df.shape}")
print(f"  → Missing values: {df.isnull().sum().sum()}")

# Drop rows with ANY missing values
df_clean = df.dropna()

print(f"  ✓ Shape after cleaning: {df_clean.shape}")
print(f"  ✓ Missing values: {df_clean.isnull().sum().sum()}")
print(f"  ✓ Date range after cleaning: {df_clean.index.min()} to {df_clean.index.max()}")

# Verify target variables are preserved
targets_preserved = all(target in df_clean.columns for target in target_cols)
print(f"  ✓ Target variables preserved: {targets_preserved}")

# ============================================================================
# STEP 7: OUTPUT GENERATION
# ============================================================================
print("\n[7/7] Generating outputs...")

# Save feature-engineered dataset
df_clean.to_csv(OUTPUT_DATA_PATH)
print(f"  ✓ Saved dataset: {OUTPUT_DATA_PATH}")
print(f"    → Dimensions: {df_clean.shape[0]} rows × {df_clean.shape[1]} columns")

# Generate detailed report
report_lines = []
report_lines.append("=" * 80)
report_lines.append("FEATURE ENGINEERING REPORT")
report_lines.append("=" * 80)
report_lines.append("")

# Dataset info
report_lines.append("DATASET SUMMARY")
report_lines.append("-" * 80)
report_lines.append(f"Input file: {INPUT_PATH}")
report_lines.append(f"Output file: {OUTPUT_DATA_PATH}")
report_lines.append(f"Original dimensions: {len(df)} rows × {len(feature_cols)} features")
report_lines.append(f"Final dimensions: {df_clean.shape[0]} rows × {df_clean.shape[1]} columns")
report_lines.append(f"Date range: {df_clean.index.min().strftime('%Y-%m-%d')} to {df_clean.index.max().strftime('%Y-%m-%d')}")
report_lines.append(f"Rows trimmed: {len(df) - len(df_clean)} ({100*(len(df)-len(df_clean))/len(df):.1f}%)")
report_lines.append("")

# Feature counts
total_features = len([col for col in df_clean.columns if col not in target_cols])
report_lines.append("FEATURE COUNTS")
report_lines.append("-" * 80)
report_lines.append(f"Total features: {total_features}")
report_lines.append(f"  - Original features: {len(feature_cols)}")
report_lines.append(f"  - YoY % change features: {len(yoy_features)}")
report_lines.append(f"  - MoM change features: {len(mom_features)}")
report_lines.append(f"  - Rolling statistic features: {len(rolling_features)}")
report_lines.append(f"  - Lagged features: {len(lag_features)}")
report_lines.append(f"Target variables: {len(target_cols)}")
report_lines.append("")

# List features by category
report_lines.append("FEATURES BY CATEGORY")
report_lines.append("-" * 80)

report_lines.append("\n1. Original Features:")
for col in feature_cols:
    report_lines.append(f"   - {col}")

report_lines.append("\n2. YoY % Change Features:")
for col in sorted(yoy_features.keys()):
    report_lines.append(f"   - {col}")

report_lines.append("\n3. MoM Change Features:")
for col in sorted(mom_features.keys()):
    report_lines.append(f"   - {col}")

report_lines.append("\n4. Rolling Statistics:")
for col in sorted(rolling_features.keys()):
    report_lines.append(f"   - {col}")

report_lines.append(f"\n5. Lagged Features: {len(lag_features)} features (all original and transformed features with lags: {lag_windows})")

report_lines.append("\n6. Target Variables:")
for col in target_cols:
    report_lines.append(f"   - {col}")

report_lines.append("")

# Stationarity analysis summary
report_lines.append("STATIONARITY ANALYSIS (ADF TESTS)")
report_lines.append("-" * 80)
report_lines.append(f"Total features tested: {len(adf_results)}")
report_lines.append(f"Stationary features (p < 0.05): {stationary_count} ({100*stationary_count/len(adf_results):.1f}%)")
report_lines.append(f"Non-stationary features (p >= 0.05): {non_stationary_count} ({100*non_stationary_count/len(adf_results):.1f}%)")
report_lines.append("")

report_lines.append("Non-Stationary Features (p-value >= 0.05):")
adf_df = pd.DataFrame(adf_results)
non_stationary = adf_df[~adf_df['stationary']].sort_values('p_value', ascending=False)
for idx, row in non_stationary.iterrows():
    if pd.notna(row['p_value']):
        report_lines.append(f"   - {row['feature']}: p={row['p_value']:.4f}, ADF={row['adf_statistic']:.4f}")
    else:
        report_lines.append(f"   - {row['feature']}: Test failed")

report_lines.append("")

# Data quality verification
report_lines.append("DATA QUALITY VERIFICATION")
report_lines.append("-" * 80)
report_lines.append(f"Missing values in final dataset: {df_clean.isnull().sum().sum()}")
report_lines.append(f"Target variables preserved: {targets_preserved}")
report_lines.append(f"Temporal ordering maintained: {(df_clean.index == df_clean.index.sort_values()).all()}")
report_lines.append("")

# Methodological notes
report_lines.append("METHODOLOGICAL NOTES")
report_lines.append("-" * 80)
report_lines.append("1. YoY % changes calculated as: (X_t - X_{t-12}) / X_{t-12} * 100")
report_lines.append("2. MoM changes calculated as: X_t - X_{t-1}")
report_lines.append("3. Rolling statistics use min_periods=window to ensure full window coverage")
report_lines.append("4. All lags use .shift() to prevent look-ahead bias")
report_lines.append("5. Rows with any NaN values dropped to ensure clean training data")
report_lines.append("6. Target variables correctly aligned with feature timestamps")
report_lines.append("")

report_lines.append("=" * 80)
report_lines.append("END OF REPORT")
report_lines.append("=" * 80)

# Write report
report_text = "\n".join(report_lines)
with open(OUTPUT_REPORT_PATH, 'w') as f:
    f.write(report_text)

print(f"  ✓ Saved report: {OUTPUT_REPORT_PATH}")

print("\n" + "=" * 80)
print("FEATURE ENGINEERING COMPLETED SUCCESSFULLY")
print("=" * 80)
print(f"\nOutputs:")
print(f"  1. Feature-engineered dataset: {OUTPUT_DATA_PATH}")
print(f"  2. Detailed report: {OUTPUT_REPORT_PATH}")
print(f"\nKey metrics:")
print(f"  - Final dataset: {df_clean.shape[0]} rows × {df_clean.shape[1]} columns")
print(f"  - Total features: {total_features}")
print(f"  - Stationary features: {stationary_count}/{len(adf_results)} ({100*stationary_count/len(adf_results):.1f}%)")
print(f"  - No missing values: ✓")