"""
Feature Engineering Pipeline for Recession Prediction
======================================================

This script performs data cleaning, feature transformation, and generates
lagged targets for recession prediction modeling.

Input: workflow/data_raw.csv
Outputs:
  - workflow/data_processed.csv
  - results/feature_engineering_report.txt
"""

import pandas as pd
import numpy as np
from statsmodels.tsa.stattools import adfuller
from pathlib import Path
import sys

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

# Define paths
BASE_DIR = Path('/app/sandbox/session_20251228_135128_cd3e46eedd35')
INPUT_FILE = BASE_DIR / 'workflow' / 'data_raw.csv'
OUTPUT_FILE = BASE_DIR / 'workflow' / 'data_processed.csv'
REPORT_FILE = BASE_DIR / 'results' / 'feature_engineering_report.txt'

# Define feature groups
GROWTH_METRICS = ['INDPRO', 'RSXFS', 'PCE', 'HOUST', 'PERMIT', 'CSUSHPISA', 'M2SL']
RATES_AND_SPREADS = ['T10Y2Y', 'T10Y3M', 'UNRATE', 'CIVPART', 'TCU', 'VIXCLS', 'CORP_SPREAD']
TARGET_COLUMN = 'USREC'
DATE_COLUMN = 'DATE'

# Lag windows (in months)
LAG_WINDOWS = [3, 6, 12]


def print_progress(message):
    """Print progress with timestamp."""
    print(f"[INFO] {message}")
    sys.stdout.flush()


def load_and_clean_data():
    """
    Load raw data and perform initial cleaning.

    Steps:
    1. Load data with DATE parsing
    2. Drop SP500 column (insufficient historical data)
    3. Filter to start from 1982-01 (ensures T10Y3M availability)
    4. Sort by date
    """
    print_progress("Loading raw data...")
    df = pd.read_csv(INPUT_FILE, parse_dates=[DATE_COLUMN])
    print_progress(f"Loaded {len(df)} rows, {len(df.columns)} columns")
    print_progress(f"Date range: {df[DATE_COLUMN].min()} to {df[DATE_COLUMN].max()}")

    # Drop SP500 (data only from 2015)
    if 'SP500' in df.columns:
        print_progress("Dropping SP500 column (insufficient historical data)")
        df = df.drop(columns=['SP500'])

    # Filter to start from 1982-01
    print_progress("Filtering data to start from 1982-01...")
    df = df[df[DATE_COLUMN] >= '1982-01-01'].copy()
    print_progress(f"After filtering: {len(df)} rows from {df[DATE_COLUMN].min()} to {df[DATE_COLUMN].max()}")

    # Sort by date
    df = df.sort_values(DATE_COLUMN).reset_index(drop=True)

    # Drop ICSA (initial claims) - not in primary feature set
    if 'ICSA' in df.columns:
        print_progress("Dropping ICSA column (not in primary feature set)")
        df = df.drop(columns=['ICSA'])

    return df


def create_forward_targets(df):
    """
    Create forward-looking binary recession targets.

    Creates targets for 3, 6, and 12-month prediction horizons.
    USREC_lead3 = 1 if recession occurs in next 3 months
    """
    print_progress("Creating forward-looking targets...")

    for horizon in [3, 6, 12]:
        target_col = f'{TARGET_COLUMN}_lead{horizon}'
        # Shift target backward (lead forward)
        df[target_col] = df[TARGET_COLUMN].shift(-horizon)

        # Count positive cases
        positive_count = df[target_col].sum()
        prevalence = positive_count / df[target_col].notna().sum() * 100
        print_progress(f"  {target_col}: {positive_count} positive cases ({prevalence:.1f}%)")

    return df


def calculate_yoy_changes(df):
    """
    Calculate Year-over-Year percentage changes for growth metrics.

    YoY change = (Value_t - Value_{t-12}) / Value_{t-12} * 100
    """
    print_progress("Calculating Year-over-Year changes for growth metrics...")

    for col in GROWTH_METRICS:
        if col in df.columns:
            yoy_col = f'{col}_yoy'
            df[yoy_col] = df[col].pct_change(periods=12) * 100
            print_progress(f"  Created {yoy_col}")

    return df


def create_lagged_features(df):
    """
    Create lagged features for key economic indicators.

    Captures momentum by including 3, 6, and 12-month lagged values.
    """
    print_progress("Creating lagged features...")

    # Determine which features to lag (YoY growth + rates/spreads)
    yoy_features = [f'{col}_yoy' for col in GROWTH_METRICS if col in df.columns]
    features_to_lag = yoy_features + RATES_AND_SPREADS

    lagged_count = 0
    for col in features_to_lag:
        if col in df.columns:
            for lag in LAG_WINDOWS:
                lag_col = f'{col}_lag{lag}'
                df[lag_col] = df[col].shift(lag)
                lagged_count += 1

                if lagged_count % 10 == 0:
                    print_progress(f"  Created {lagged_count} lagged features...")

    print_progress(f"  Total lagged features created: {lagged_count}")
    return df


def run_stationarity_tests(df):
    """
    Run Augmented Dickey-Fuller tests on all features.

    Returns a dictionary with test results for each feature.
    """
    print_progress("Running ADF stationarity tests...")

    # Get all numeric columns except DATE and targets
    feature_cols = [col for col in df.columns
                    if col not in [DATE_COLUMN, TARGET_COLUMN] +
                    [f'{TARGET_COLUMN}_lead{h}' for h in [3, 6, 12]] and
                    df[col].dtype in ['float64', 'int64']]

    adf_results = {}
    tested_count = 0

    for col in feature_cols:
        # Drop NaN values for testing
        data = df[col].dropna()

        if len(data) > 20:  # Need sufficient data for ADF test
            try:
                result = adfuller(data, autolag='AIC')
                adf_results[col] = {
                    'adf_stat': result[0],
                    'p_value': result[1],
                    'n_lags': result[2],
                    'n_obs': result[3],
                    'stationary': result[1] < 0.05  # p-value < 0.05 implies stationary
                }
                tested_count += 1

                if tested_count % 10 == 0:
                    print_progress(f"  Tested {tested_count}/{len(feature_cols)} features...")
            except Exception as e:
                print_progress(f"  Warning: ADF test failed for {col}: {str(e)}")
                adf_results[col] = {'error': str(e)}

    print_progress(f"  Completed ADF tests on {tested_count} features")
    return adf_results


def generate_report(df, adf_results):
    """
    Generate comprehensive feature engineering report.
    """
    print_progress("Generating feature engineering report...")

    report_lines = []
    report_lines.append("=" * 80)
    report_lines.append("FEATURE ENGINEERING REPORT")
    report_lines.append("=" * 80)
    report_lines.append("")

    # Dataset summary
    report_lines.append("1. DATASET SUMMARY")
    report_lines.append("-" * 80)
    report_lines.append(f"Date range: {df[DATE_COLUMN].min()} to {df[DATE_COLUMN].max()}")
    report_lines.append(f"Total observations: {len(df)}")
    report_lines.append(f"Total features: {len(df.columns) - 1}")  # Excluding DATE
    report_lines.append("")

    # Target distribution
    report_lines.append("2. TARGET DISTRIBUTION")
    report_lines.append("-" * 80)
    for target in [TARGET_COLUMN, f'{TARGET_COLUMN}_lead3', f'{TARGET_COLUMN}_lead6', f'{TARGET_COLUMN}_lead12']:
        if target in df.columns:
            positive = df[target].sum()
            total = df[target].notna().sum()
            prevalence = positive / total * 100 if total > 0 else 0
            report_lines.append(f"{target}: {positive}/{total} positive ({prevalence:.1f}%)")
    report_lines.append("")

    # Feature categories
    report_lines.append("3. FEATURE CATEGORIES")
    report_lines.append("-" * 80)

    yoy_features = [col for col in df.columns if col.endswith('_yoy')]
    lagged_features = [col for col in df.columns if '_lag' in col]
    level_features = [col for col in RATES_AND_SPREADS if col in df.columns]

    report_lines.append(f"YoY Growth Features: {len(yoy_features)}")
    for feat in yoy_features:
        report_lines.append(f"  - {feat}")
    report_lines.append("")

    report_lines.append(f"Level Features (Rates/Spreads): {len(level_features)}")
    for feat in level_features:
        report_lines.append(f"  - {feat}")
    report_lines.append("")

    report_lines.append(f"Lagged Features: {len(lagged_features)}")
    report_lines.append(f"  (3, 6, 12-month lags for {len(lagged_features) // 3} base features)")
    report_lines.append("")

    # Missing data summary
    report_lines.append("4. MISSING DATA SUMMARY")
    report_lines.append("-" * 80)
    missing_summary = df.isnull().sum()
    missing_pct = (missing_summary / len(df) * 100).round(1)

    features_with_missing = missing_summary[missing_summary > 0].sort_values(ascending=False)
    report_lines.append(f"Features with missing values: {len(features_with_missing)}/{len(df.columns)-1}")
    report_lines.append("")
    report_lines.append("Top 10 features by missingness:")
    for feat, count in features_with_missing.head(10).items():
        pct = missing_pct[feat]
        report_lines.append(f"  {feat}: {count} ({pct}%)")
    report_lines.append("")

    # Stationarity test results
    report_lines.append("5. STATIONARITY TEST RESULTS (ADF)")
    report_lines.append("-" * 80)

    stationary = [k for k, v in adf_results.items() if 'stationary' in v and v['stationary']]
    non_stationary = [k for k, v in adf_results.items() if 'stationary' in v and not v['stationary']]

    report_lines.append(f"Stationary features (p < 0.05): {len(stationary)}/{len(adf_results)}")
    report_lines.append(f"Non-stationary features (p >= 0.05): {len(non_stationary)}/{len(adf_results)}")
    report_lines.append("")

    if non_stationary:
        report_lines.append("Non-stationary features requiring attention:")
        for feat in non_stationary[:20]:  # Show top 20
            p_val = adf_results[feat]['p_value']
            report_lines.append(f"  {feat}: p-value = {p_val:.4f}")
        if len(non_stationary) > 20:
            report_lines.append(f"  ... and {len(non_stationary) - 20} more")
    report_lines.append("")

    # Success criteria verification
    report_lines.append("6. SUCCESS CRITERIA VERIFICATION")
    report_lines.append("-" * 80)
    criteria = {
        'Starts from 1982': df[DATE_COLUMN].min().year >= 1982,
        'SP500 removed': 'SP500' not in df.columns,
        'USREC_lead3 exists': 'USREC_lead3' in df.columns,
        'USREC_lead6 exists': 'USREC_lead6' in df.columns,
        'USREC_lead12 exists': 'USREC_lead12' in df.columns,
        'YoY features created': len(yoy_features) > 0,
        'Lagged features created': len(lagged_features) > 0,
    }

    all_passed = all(criteria.values())
    for criterion, passed in criteria.items():
        status = "✓ PASS" if passed else "✗ FAIL"
        report_lines.append(f"{status}: {criterion}")

    report_lines.append("")
    report_lines.append("=" * 80)
    report_lines.append(f"OVERALL STATUS: {'SUCCESS' if all_passed else 'INCOMPLETE'}")
    report_lines.append("=" * 80)

    # Write report
    with open(REPORT_FILE, 'w') as f:
        f.write('\n'.join(report_lines))

    print_progress(f"Report saved to {REPORT_FILE}")

    # Print summary to console
    print("\n" + "\n".join(report_lines[-10:]))


def main():
    """Main execution pipeline."""
    print_progress("=" * 80)
    print_progress("FEATURE ENGINEERING PIPELINE")
    print_progress("=" * 80)

    # Step 1: Load and clean data
    df = load_and_clean_data()

    # Step 2: Create forward-looking targets
    df = create_forward_targets(df)

    # Step 3: Calculate YoY changes for growth metrics
    df = calculate_yoy_changes(df)

    # Step 4: Create lagged features
    df = create_lagged_features(df)

    # Step 5: Run stationarity tests
    adf_results = run_stationarity_tests(df)

    # Step 6: Save processed data
    print_progress(f"Saving processed data to {OUTPUT_FILE}...")
    df.to_csv(OUTPUT_FILE, index=False)
    print_progress(f"Saved {len(df)} rows, {len(df.columns)} columns")

    # Step 7: Generate report
    generate_report(df, adf_results)

    print_progress("=" * 80)
    print_progress("FEATURE ENGINEERING COMPLETE")
    print_progress("=" * 80)
    print_progress(f"Processed data: {OUTPUT_FILE}")
    print_progress(f"Report: {REPORT_FILE}")

    return 0


if __name__ == '__main__':
    sys.exit(main())