"""
Data Ingestion & Harmonization Script (FIXED VERSION)
Fetches macroeconomic indicators from FRED and harmonizes to monthly frequency.

FIXES APPLIED:
- Replaced SP500 with SPASTT01USM661N (longer history proxy)
- Use .last() resampling for financial data instead of .mean()
- Removed limit_direction='both' to prevent backfilling artifacts
"""

import pandas as pd
import numpy as np
from pandas_datareader import data as pdr
from datetime import datetime
import warnings
warnings.filterwarnings('ignore')

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

print("="*80)
print("STEP 1: DATA INGESTION & HARMONIZATION (FIXED)")
print("="*80)

# Define data series to fetch from FRED
FRED_SERIES = {
    # Yield Curve
    'T10Y2Y': 'Yield Curve: 10Y-2Y Spread',
    'T10Y3M': 'Yield Curve: 10Y-3M Spread',

    # Labor Market
    'ICSA': 'Initial Jobless Claims',
    'UNRATE': 'Unemployment Rate',
    'CIVPART': 'Labor Force Participation Rate',

    # Production
    'INDPRO': 'Industrial Production Index',
    'TCU': 'Capacity Utilization',

    # Consumer
    'RSAFS': 'Retail Sales',
    'PCE': 'Personal Consumption Expenditures',

    # Financial Markets
    'BAA10Y': 'Credit Spread: BAA-10Y Treasury',
    'VIXCLS': 'VIX Volatility Index',

    # Housing
    'PERMIT': 'Building Permits',
    'HOUST': 'Housing Starts',
    'CSUSHPISA': 'Case-Shiller Home Price Index',

    # Money & Credit
    'M2SL': 'M2 Money Supply',
    'BUSLOANS': 'Commercial & Industrial Loans',

    # Target Variable
    'USREC': 'NBER Recession Indicator'
}

# Alternative market indices with longer history
MARKET_PROXIES = [
    ('SPASTT01USM661N', 'Total Share Prices for All Shares (OECD)'),
    ('SP500', 'S&P 500 Index'),
]

# Start date for data collection (going back far enough for comprehensive history)
START_DATE = '1980-01-01'
END_DATE = datetime.now().strftime('%Y-%m-%d')

print(f"\n1. Fetching {len(FRED_SERIES)} economic indicators from FRED")
print(f"   Date Range: {START_DATE} to {END_DATE}")
print("-"*80)

# Fetch data from FRED
data_dict = {}
failed_series = []

for i, (ticker, description) in enumerate(FRED_SERIES.items(), 1):
    try:
        print(f"   [{i}/{len(FRED_SERIES)}] Fetching {ticker}: {description}...", end=' ')
        df = pdr.DataReader(ticker, 'fred', START_DATE, END_DATE)
        data_dict[ticker] = df
        print(f"✓ ({len(df)} records, {df.index.min().strftime('%Y-%m')} to {df.index.max().strftime('%Y-%m')})")
    except Exception as e:
        print(f"✗ Failed: {str(e)[:50]}")
        failed_series.append((ticker, description, str(e)))

print(f"\n   Successfully fetched: {len(data_dict)}/{len(FRED_SERIES)} series")

if failed_series:
    print(f"   Failed series: {len(failed_series)}")
    for ticker, desc, error in failed_series:
        print(f"      - {ticker} ({desc}): {error[:60]}")

# Try to fetch a market index with long history
print("\n   Attempting to fetch stock market index with long history...")
market_index_fetched = False
for ticker, description in MARKET_PROXIES:
    try:
        print(f"      Trying {ticker}: {description}...", end=' ')
        df = pdr.DataReader(ticker, 'fred', START_DATE, END_DATE)
        if len(df) > 100:  # Require at least 100 data points
            data_dict['MARKET_INDEX'] = df
            print(f"✓ ({len(df)} records, {df.index.min().strftime('%Y-%m')} to {df.index.max().strftime('%Y-%m')})")
            print(f"         Using {ticker} as market proxy")
            market_index_fetched = True
            break
        else:
            print(f"✗ Insufficient data ({len(df)} records)")
    except Exception as e:
        print(f"✗ Failed: {str(e)[:50]}")

if not market_index_fetched:
    print(f"      WARNING: Could not fetch market index with long history")

# Note about UMCSENT - if it failed, we already have PCE as backup
# We'll try to add it if available but won't require it
try:
    print("\n   Attempting to fetch UMCSENT (Consumer Sentiment)...", end=' ')
    umcsent = pdr.DataReader('UMCSENT', 'fred', START_DATE, END_DATE)
    data_dict['UMCSENT'] = umcsent
    print(f"✓ ({len(umcsent)} records)")
except Exception as e:
    print(f"✗ Not available, using PCE as consumer indicator")

print("\n2. Combining all series into a single dataframe")
print("-"*80)

# Combine all series
df_combined = pd.DataFrame()
for ticker, df in data_dict.items():
    if df_combined.empty:
        df_combined = df.copy()
        df_combined.columns = [ticker]
    else:
        df_combined = df_combined.join(df, how='outer')
        if ticker not in df_combined.columns:
            df_combined.columns = list(df_combined.columns[:-1]) + [ticker]

print(f"   Combined shape: {df_combined.shape}")
print(f"   Date range: {df_combined.index.min()} to {df_combined.index.max()}")
print(f"   Features: {list(df_combined.columns)}")

print("\n3. Resampling to monthly frequency")
print("-"*80)

# Identify financial market variables (use last value, not mean)
financial_vars = ['VIXCLS', 'MARKET_INDEX', 'BAA10Y', 'T10Y2Y', 'T10Y3M']

print("   Resampling strategy:")
print("      - Financial market data: Use LAST value of month (end-of-month)")
print("      - Economic indicators: Use MEAN for smoothing")

# Separate financial and non-financial columns
financial_cols = [col for col in df_combined.columns if col in financial_vars]
other_cols = [col for col in df_combined.columns if col not in financial_vars]

print(f"\n   Financial columns (using .last()): {financial_cols}")
print(f"   Other columns (using .mean()): {other_cols}")

# Resample separately
df_monthly_parts = []

if financial_cols:
    df_financial = df_combined[financial_cols].resample('M').last()
    df_monthly_parts.append(df_financial)

if other_cols:
    df_other = df_combined[other_cols].resample('M').mean()
    df_monthly_parts.append(df_other)

# Combine
df_monthly = pd.concat(df_monthly_parts, axis=1)

print(f"\n   Resampled shape: {df_monthly.shape}")
print(f"   New date range: {df_monthly.index.min()} to {df_monthly.index.max()}")
print(f"   Total months: {len(df_monthly)}")

print("\n4. Handling missing values (IMPROVED STRATEGY)")
print("-"*80)

# Check for missing values
missing_before = df_monthly.isnull().sum()
print(f"   Missing values per column (before handling):")
for col in df_monthly.columns:
    if missing_before[col] > 0:
        pct = 100 * missing_before[col] / len(df_monthly)
        print(f"      {col}: {missing_before[col]} ({pct:.1f}%)")

# IMPROVED Strategy to avoid backfilling artifacts:
# 1. Forward fill for minor gaps (limit to 3 months) - fills forward from known values
# 2. Drop rows at the start/end if data is insufficient
# 3. Linear interpolation ONLY forward (no backfilling to prevent flat-line artifacts)

print("\n   Applying missing value handling strategy:")
print("      - Forward fill (limit: 3 months)")
print("      - Drop initial/final rows with >50% missing values")
print("      - Linear interpolation FORWARD ONLY (no backfilling)")

# Step 1: Forward fill with limit
df_monthly = df_monthly.fillna(method='ffill', limit=3)

# Step 2: Drop rows where more than 50% of values are missing
# This handles the start of the dataset where some series don't exist yet
threshold = 0.5 * len(df_monthly.columns)
rows_before = len(df_monthly)
df_monthly = df_monthly.dropna(thresh=threshold)
rows_dropped = rows_before - len(df_monthly)
print(f"      Dropped {rows_dropped} rows with >50% missing values")

# Step 3: Interpolate remaining missing values (FORWARD ONLY)
# CRITICAL FIX: Use limit_direction='forward' instead of 'both' to prevent backfilling
df_monthly = df_monthly.interpolate(method='linear', limit_direction='forward', limit=6)

# Final check
missing_after = df_monthly.isnull().sum()
total_missing = missing_after.sum()

print(f"\n   Missing values after handling: {total_missing}")
if total_missing > 0:
    print("   Remaining missing values:")
    for col in df_monthly.columns:
        if missing_after[col] > 0:
            pct = 100 * missing_after[col] / len(df_monthly)
            print(f"      {col}: {missing_after[col]} ({pct:.1f}%)")

    # For any remaining NaN at the end of series (future months), drop those rows
    df_monthly = df_monthly.dropna()
    print(f"   Dropped remaining rows with NaN. Final shape: {df_monthly.shape}")

print(f"   Final clean dataset shape: {df_monthly.shape}")
print(f"   Final date range: {df_monthly.index.min()} to {df_monthly.index.max()}")

print("\n5. Creating forward-looking target variables")
print("-"*80)

# Create target variables by shifting USREC backwards
# This means target_3m at time t represents recession status at t+3
if 'USREC' in df_monthly.columns:
    print("   Creating recession prediction targets:")
    print("      - target_3m: Recession status 3 months ahead")
    print("      - target_6m: Recession status 6 months ahead")
    print("      - target_12m: Recession status 12 months ahead")

    df_monthly['target_3m'] = df_monthly['USREC'].shift(-3)
    df_monthly['target_6m'] = df_monthly['USREC'].shift(-6)
    df_monthly['target_12m'] = df_monthly['USREC'].shift(-12)

    # Drop rows where we don't have future target values
    df_monthly = df_monthly.dropna(subset=['target_3m', 'target_6m', 'target_12m'])

    print(f"\n   After creating targets and dropping NaN:")
    print(f"      Final shape: {df_monthly.shape}")
    print(f"      Usable date range: {df_monthly.index.min()} to {df_monthly.index.max()}")

    # Calculate target statistics
    for target in ['target_3m', 'target_6m', 'target_12m']:
        n_recession = df_monthly[target].sum()
        pct = 100 * n_recession / len(df_monthly)
        print(f"      {target}: {int(n_recession)} recession months ({pct:.1f}%)")
else:
    print("   WARNING: USREC not available, skipping target creation")

print("\n6. Data validation and quality checks")
print("-"*80)

# Check for flat-line artifacts (e.g., constant values indicating backfilling issues)
print("   Checking for data quality issues...")

for col in df_monthly.columns:
    if col not in ['USREC', 'target_3m', 'target_6m', 'target_12m']:
        q25 = df_monthly[col].quantile(0.25)
        q50 = df_monthly[col].quantile(0.50)
        q75 = df_monthly[col].quantile(0.75)

        # Check if quartiles are identical (indicating flat-line)
        if abs(q25 - q50) < 1e-6 and abs(q50 - q75) < 1e-6:
            print(f"      WARNING: {col} shows flat-line pattern (Q25={q25:.3f}, Q50={q50:.3f}, Q75={q75:.3f})")
            print(f"               This may indicate insufficient data or backfilling artifacts")

# Check date index is monotonic
if df_monthly.index.is_monotonic_increasing:
    print("   ✓ Date index is monotonic increasing")
else:
    print("   ✗ WARNING: Date index is not monotonic!")

print("\n7. Saving consolidated dataset")
print("-"*80)

output_path = '/app/sandbox/session_20251228_135128_cd3e46eedd35/workflow/data/consolidated_macro_data.csv'
df_monthly.to_csv(output_path)
file_size = pd.io.common.get_filepath_or_buffer(output_path)[0].__sizeof__() / 1024
print(f"   ✓ Saved to: {output_path}")
print(f"   File size: ~{file_size:.1f} KB")

print("\n" + "="*80)
print("DATA INGESTION COMPLETE")
print("="*80)
print(f"\nOutput: {output_path}")
print(f"Dataset Shape: {df_monthly.shape[0]} observations × {df_monthly.shape[1]} features")
print(f"Ready for summary generation!")