# =============================================================================
# STEP 0: LOAD ESSENTIAL LIBRARIES   
# =============================================================================
#Core data manipulation
import pandas as pd
import numpy as np
import os
# =============================================================================
# STEP 1: Create helper function for extracting date for multiple time series 
# =============================================================================
def fetch_bql_series(assets, bql_service, date_range):
    """
    Fetches time series data for given assets using Bloomberg Query Language (BQL).
    
    This function provides a streamlined interface to Bloomberg's data,
    handling multiple asset queries efficiently and returning clean pandas DataFrames.

    Parameters:
    -----------
    assets : dict
        Dictionary mapping asset names to Bloomberg tickers
        Example: {'Brazil': 'MXBR Index', 'USD_Index': 'DXY Curncy'}
    bql_service : bql.Service
        Authenticated Bloomberg Query Language service instance
    date_range : bql.func.range
        Bloomberg date range object (e.g., bq.func.range('-10Y', '0D'))

    Returns:
    --------
    pd.DataFrame
        DataFrame with datetime index and columns for each asset
        All data is automatically forward-filled by Bloomberg for quality
        
    Notes:
    ------
    - Uses PX_LAST field for consistent daily closing values
    - Bloomberg handles missing values with native forward-fill
    """
    data = {}
    
    # Iterate through each asset and fetch time series
    for asset_name, ticker in assets.items():
        # Execute BQL query for last price over specified date range
        query = bql_service.time_series(ticker, "PX_LAST", date_range[0], date_range[1])
        
        # Extract value column and store with descriptive asset name
        data[asset_name] = query.to_frame()["value"]
    
    # Combine all series into single DataFrame with datetime alignment
    return pd.DataFrame(data)

# =============================================================================
#  STEP 2: DATA EXPORT CONFIGURATION & ASSET UNIVERSE DEFINITION
# =============================================================================

# Define output path for the combined dataset
output_path = '../data/combined_em_macro_data.csv'

print("🔧 Configuration Summary:")
print(f"   📂 Output Path: {output_path}")
print(f"   🌍 Setting up comprehensive EM and macro factor universe...")

# =============================================================================
# STEP 3: DEFINE EMERGING MARKETS EQUITY INDICES (MSCI)
# =============================================================================
# Using MSCI indices for consistency and institutional benchmarking
em_assets = {
    # LATIN AMERICA - Commodity-driven economies with US trade linkages
    'Brazil': 'MXBR Index',        # MSCI Brazil - LatAm largest economy
    'Mexico': 'MXMX Index',        # MSCI Mexico - USMCA integration
    
    # ASIA PACIFIC - Technology and manufacturing powerhouses
    'India': 'MXIN Index',         # MSCI India - South Asian growth market
    'China': 'MXCN Index',         # MSCI China - East Asian manufacturing hub
    'Taiwan': 'TAMSCI Index',      # Taiwan MSCI - Technology manufacturing
    'Korea': 'MXKR Index',         # MSCI Korea - Advanced EM market
    'Indonesia': 'MXID Index',     # MSCI Indonesia - SE Asian commodity exporter
    
    # AFRICA & MIDDLE EAST - Resource-rich markets
    'SouthAfrica': 'MXZA Index',   # MSCI South Africa - African gateway
    
    # DEVELOPED MARKET BENCHMARK - For relative performance analysis
    'US': 'MXUS Index'             # MSCI USA - Developed market benchmark
}
# =============================================================================
# STEP 4: DEFINE MACROECONOMIC FACTORS 
# =============================================================================
# Comprehensive factor set covering monetary policy, risk sentiment, 
# currency dynamics, and commodity cycles
macro_assets = {
    # MONETARY POLICY & RATES
    'USD_Index': 'DXY Curncy',         # US Dollar Index - currency strength
    'US_10Y_Yield': 'USGG10YR Index', # 10Y Treasury - long-term risk-free rate
    'US_2Y_Yield': 'USGG2YR Index',   # 2Y Treasury - Fed policy proxy
    
    # RISK SENTIMENT & VOLATILITY  
    'VIX': 'VIX Index',                # CBOE Volatility Index - fear gauge
    'BAA_spread': 'CSI BB Index',      # Corporate credit spreads - risk premium
    
    # COMMODITIES & REAL ASSETS
    'Oil_Brent': 'CO1 Comdty',        # Brent crude oil - energy prices
    'Copper': 'LMCADY Comdty'         # LME copper - industrial demand proxy
}

# =============================================================================
# STEP 5: FINAL CONFIGURATION SUMMARY
# =============================================================================
print("✅ Asset Universe Configuration Complete:")
print(f"   📂 Output Path: {output_path}")
print(f"   🌍 EM Markets: {len(em_assets)} total ({len(em_assets)-1} EM + 1 DM benchmark)")
print(f"   📊 Macro Factors: {len(macro_assets)} base factors")
print(f"   🔍 Derived Factors: 1 (Term Spread = 10Y - 2Y)")
print(f"   📈 Total Variables: {len(em_assets) + len(macro_assets) + 1}")
print(f"   🎯 Ready for Bloomberg data collection pipeline...")

# =============================================================================
# STEP 6: Bloomberg Authentication & Connection Configuration
# =============================================================================
print("🔗 Establishing Bloomberg BQL Connection...")
print("   • Authenticating with Bloomberg Professional services")
print("   • Initializing BQL query interface")
bq = bql.Service()  # Create authenticated Bloomberg service instance

# Configure 10-year lookback period for robust statistical analysis
print("\n📅 Setting Data Collection Parameters...")
date_range = bq.func.range('-10Y', '0D')  # 10 years back from today
print(f"   • Time Horizon: 10-year lookback for statistical robustness")
print(f"   • Frequency: Daily closing prices (business days)")
print(f"   • Data Quality: Bloomberg native forward-fill for missing values")

# =============================================================================
# STEP 7: EMERGING MARKETS EQUITY DATA COLLECTION
# =============================================================================
print("\n🌍 Step 1: Fetching EM Equity Index Data...")
print(f"   • Markets: {list(em_assets.keys())}")
print(f"   • Data Field: PX_LAST (daily closing levels)")
em_df = fetch_bql_series(em_assets, bq, date_range)

# Validate EM data collection
print(f"\n✅ EM Data Collection Complete:")
print(f"   📊 Shape: {em_df.shape[0]} observations × {em_df.shape[1]} indices")
print(f"   📅 Period: {em_df.index.min().strftime('%Y-%m-%d')} to {em_df.index.max().strftime('%Y-%m-%d')}")
print(f"   🕰️ Trading Days: {len(em_df)} business days")

# =============================================================================
# STEP 8: MACROECONOMIC FACTORS DATA COLLECTION
# =============================================================================
print(f"\n📈 Step 2: Fetching Macroeconomic Factor Data...")
print(f"   • Base Factors: {list(macro_assets.keys())}")
print(f"   • Asset Classes: Rates, FX, Commodities, Volatility, Credit")
macro_df = fetch_bql_series(macro_assets, bq, date_range)

# =============================================================================
# STEP 9: FEATURE ENGINEERING - TERM SPREAD CALCULATION
# =============================================================================
print(f"\n🔧 Step 3: Engineering Derived Factors...")
# Calculate yield curve term spread as growth/policy indicator
macro_df['Term_Spread'] = macro_df['US_10Y_Yield'] - macro_df['US_2Y_Yield']
print(f"   • Term Spread: US 10Y - 2Y yields (yield curve slope)")
print(f"   • Economic Significance: Growth expectations vs. policy stance")

# Validate macro data collection
print(f"\n✅ Macro Data Collection Complete:")
print(f"   📊 Shape: {macro_df.shape[0]} observations × {macro_df.shape[1]} factors")
print(f"   📈 Total Factors: {len(macro_assets)} base + 1 derived = {macro_df.shape[1]}")

# =============================================================================
# STEP 10: DATA INTEGRATION & QUALITY ASSURANCE
# =============================================================================
print(f"\n🔄 Step 4: Data Integration & Quality Control...")
print(f"   • Aligning EM and macro datasets on trading calendar")
print(f"   • Performing inner join to ensure data completeness")

# Merge datasets on datetime index with inner join for data integrity
combined_df = pd.merge(em_df, macro_df, left_index=True, right_index=True)
combined_df = combined_df.sort_index().dropna()  # Remove any remaining missing values

# =============================================================================
# STEP 11: DATA VALIDATION
# =============================================================================
print(f"\n📊 Final Dataset Validation:")
print(f"   📈 Combined Shape: {combined_df.shape[0]} observations × {combined_df.shape[1]} variables")
print(f"   📅 Final Period: {combined_df.index.min().strftime('%Y-%m-%d')} to {combined_df.index.max().strftime('%Y-%m-%d')}")
print(f"   🌍 EM Markets: {len(em_assets)} indices")
print(f"   📊 Macro Factors: {len(macro_assets) + 1} total (inc. term spread)")
print(f"   🔍 Data Completeness: {((1 - combined_df.isnull().sum().sum() / (combined_df.shape[0] * combined_df.shape[1])) * 100):.2f}%")
print(f"   ❌ Missing Values: {combined_df.isnull().sum().sum()} total")

# =============================================================================
# STEP 12: DATA EXPORT 
# =============================================================================
print(f"\n💾 Step 6: Data Export & Preservation...")
# Ensure output directory exists
os.makedirs(os.path.dirname(output_path), exist_ok=True)

# Export to CSV with datetime index preservation
combined_df.to_csv(output_path)
print(f"✅ Dataset successfully exported to: {output_path}")
print(f"   📁 File Format: CSV with datetime index")
print(f"   🔗 Ready for: Factor modeling, PCA analysis, regression modeling")

# Final summary statistics
print(f"\n📋 Dataset Summary Statistics:")
print(f"   • Memory Usage: {combined_df.memory_usage(deep=True).sum() / 1024**2:.2f} MB")
print(f"   • Index Type: {type(combined_df.index).__name__}")
print(f"   • Data Types: {dict(combined_df.dtypes.value_counts())}")
print(f"\n🎯 Data Pipeline Complete - Ready for Factor Analysis!")

# =============================================================================
# STEP 0: LOAD ESSENTIAL LIBRARIES   
# =============================================================================
# Core data manipulation
import pandas as pd
import numpy as np
import os

# Machine learning components
from sklearn.decomposition import PCA
from sklearn.linear_model import LinearRegression
from sklearn.preprocessing import StandardScaler

# Visualization libraries
import matplotlib.pyplot as plt
import seaborn as sns

# Set plotting style
plt.style.use('default')
sns.set_palette("husl")

# Define EM market prefixes
# =============================================================================
# STEP 1: DATA LOADING
# =============================================================================
df = pd.read_csv('../data/combined_em_macro_data.csv', parse_dates=['date'], index_col='date')

print(f"📊 Dataset loaded: {df.shape[0]} rows × {df.shape[1]} columns")
print(f"📅 Date range: {df.index.min()} to {df.index.max()}")
# =============================================================================
# STEP 2: ENGINEERING TERM SPREAD TO REMOVE NEGATIVE VALUES
# =============================================================================
shift_value = df['Term_Spread'].min()
if shift_value < 0:
    print("Shift value is negative, adding a buffer to avoid negative values")
    shift_value = df['Term_Spread'].min() - 0.1 # add a buffer to avoid negative values
    print(f"Shift value: {shift_value}")
else:
    print("Shift value is positive, no buffer needed")

df['Term_Spread_Engineered'] = df['Term_Spread'] - shift_value

# Display the results
print(f"🔧 Term_Spread Engineering:")
print(f"   • Original min: {df['Term_Spread'].min():.4f}")
print(f"   • Original max: {df['Term_Spread'].max():.4f}")
print(f"   • Shift applied: +{shift_value:.4f}")
print(f"   • Engineered min: {df['Term_Spread_Engineered'].min():.4f}")
print(f"   • Engineered max: {df['Term_Spread_Engineered'].max():.4f}")
df.drop(columns=['Term_Spread'], inplace=True)
df.rename(columns={'Term_Spread_Engineered': 'Term_Spread'}, inplace=True)

# =============================================================================
# STEP 3: CALCULATE LOG RETURNS
# =============================================================================
log_returns = np.log(df / df.shift(1)).dropna()

print(f"📈 Log returns calculated: {log_returns.shape[0]} observations")
print(f"🧹 Data cleaning: {df.shape[0] - log_returns.shape[0]} rows dropped (missing/infinite values)")

# Display basic statistics
print(f"\n📋 Log Returns Summary:")
log_returns.describe().round(4)

# =============================================================================
# STEP 0: DEFINE EM AND US COLUMNS 
# =============================================================================
em_prefixes = (
    "Brazil", "India", "China", "SouthAfrica",
    "Mexico", "Indonesia", "Taiwan", "Korea", "US"
)

# the macro columns you want to avoid
exclude = {"USD_Index", "US_10Y_Yield", "US_2Y_Yield"}

# =============================================================================
# STEP 1: EXTRACT EM COLUMNS AND TAKE THE LOG OF THE RETURNS
# =============================================================================
em_columns = [
    col for col in df.columns
    if col.startswith(em_prefixes)
    and col not in exclude
]
Y = log_returns[em_columns]

# =============================================================================
# STEP 2: EXTRACT MACRO COLUMNS AND TAKE THE LOG OF THE RETURNS 
# =============================================================================
macro_columns = [col for col in df.columns if col.startswith(('USD_Index', 'Oil_Brent', 'US_10Y_Yield', 'US_2Y_Yield', 'VIX', 'BAA_spread', 'Term_Spread'))]

log_returns = np.log(df / df.shift(1)).dropna()


X = log_returns[macro_columns]

print(f"\n📊 Model Setup:")
print(f"   • Y matrix (EM returns): {Y.shape}")
print(f"   • X matrix (Macro factors): {X.shape}")

# =============================================================================
# STEP 3: SCALE INPUT VARIABLES FOR PCA 
# =============================================================================
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)
print(f"📏 Standardization completed:")
print(f"   • Original X shape: {X.shape}")
print(f"   • Scaled X shape: {X_scaled.shape}")
print(f"   • Mean: {X_scaled.mean():.6f}")
print(f"   • Std: {X_scaled.std():.6f}")

# =============================================================================
# STEP 0: EXECUTE PCA ANALYSIS WITH 3 COMPONENTS 
# =============================================================================
n_components = 3
pca = PCA(n_components=n_components)
X_pca = pca.fit_transform(X_scaled)

# =============================================================================
# STEP 1: EXTRACT EXPLAINED VARIANCE AND CUMULATIVE VARIANCE
# =============================================================================
explained_var = pca.explained_variance_ratio_
cumulative_var = np.cumsum(explained_var)

print(f"\n🔍 PCA Results:")
for i in range(n_components):
    print(f"   • PC{i+1}: {explained_var[i]:.1%} variance explained")
print(f"   • Total: {cumulative_var[-1]:.1%} variance captured")

print(f"\n📊 Principal Components Matrix: {X_pca.shape}")

# 🔍 DETAILED VARIANCE CAPTURE ANALYSIS
print(f"\n🔍 DETAILED VARIANCE CAPTURE ANALYSIS:")
print("="*60)
   
# Full PCA to see all components
pca_full = PCA()
X_pca_full = pca_full.fit_transform(X_scaled)
   
# Calculate cumulative variance for all components
explained_var_full = pca_full.explained_variance_ratio_
cumulative_var_full = np.cumsum(explained_var_full)
   
print(f"📈 PCA Variance Analysis:")
print(f"   Total variance to explain: 100%")
print(f"   Variance captured by 3 PCs: {cumulative_var_full[2]:.1%}")
print(f"   Variance captured by 4 PCs: {cumulative_var_full[3]:.1%}")
print(f"   Variance captured by 5 PCs: {cumulative_var_full[4]:.1%}")
print(f"   Variance captured by 6 PCs: {cumulative_var_full[5]:.1%}")
print(f"   Variance captured by 7 PCs: {cumulative_var_full[6]:.1%}")

results = {}
for n_components_test in [1, 2, 3, 4, 5, 6, 7]:
    # Apply PCA with n_components
    pca_test = PCA(n_components=n_components_test)
    X_pca_test = pca_test.fit_transform(X_scaled)
       
    # Calculate R² for each market
    r2_scores_test = []
    for col in Y.columns:
        model_test = LinearRegression().fit(X_pca_test, Y[col])
        r2_test = model_test.score(X_pca_test, Y[col])
        r2_scores_test.append(r2_test)
       
    avg_r2_test = np.mean(r2_scores_test)
    variance_captured_test = sum(pca_test.explained_variance_ratio_)
       
    results[n_components_test] = {
        'avg_r2': avg_r2_test,
        'variance_captured': variance_captured_test,
        'r2_scores': r2_scores_test
    }
       
    print(f"   {n_components_test} PCs: Variance={variance_captured_test:.1%}, Avg R²={avg_r2_test:.3f}")

# =============================================================================
# STEP 0: RUN REGRESSION ANALYSIS FOR EACH EM MARKET 
# =============================================================================
betas = {}       # Store factor loadings (sensitivities)
r2_scores = {}   # Store model fit statistics

print(f"🔄 Fitting {len(Y.columns)} regression models...\n")

for col in Y.columns:
    # Fit regression: EM_return = α + β₁×PC1 + β₂×PC2 + β₃×PC3 + ε
    model = LinearRegression().fit(X_pca, Y[col])
    
    # Store results
    betas[col] = model.coef_
    r2_scores[col] = model.score(X_pca, Y[col])
    
    print(f"✅ {col}:")
    print(f"   • R² Score: {r2_scores[col]:.3f}")
    print(f"   • Factor Loadings: [{betas[col][0]:.3f}, {betas[col][1]:.3f}, {betas[col][2]:.3f}]")

# Create DataFrame for factor loadings (betas)
beta_df = pd.DataFrame(betas, index=['PC1', 'PC2', 'PC3']).T
beta_df.index.name = 'EM Index'

print(f"\n📊 Factor Loadings Matrix:")
print(beta_df.round(3))

# =============================================================================
# STEP 0: DEFINE PERIODS FOR ANALYSIS 
# =============================================================================
yearly_periods = {
    '2022/2023': ('2022-01-01', '2023-12-31'),  # Post-COVID recovery, inflation concerns
    '2023/2024': ('2023-01-01', '2024-12-31'),  # Central bank policy normalization
    '2024/2025': ('2024-01-01', '2025-08-06')   # Current market regime
}
# =============================================================================
# STEP 1: RUN REGRESSION ANALYSIS FOR EACH PERIOD    
# =============================================================================
for period_name, (start_date, end_date) in yearly_periods.items():
    print(f"🗓️  Analyzing {period_name} Period")
    
    # Filter and validate data for the specific period
    period_mask = (log_returns.index >= start_date) & (log_returns.index <= end_date)
    Y_period = log_returns[em_columns][period_mask]
    X_period = log_returns[macro_columns][period_mask]
    
    # Apply period-specific PCA and regression analysis
    scaler_period = StandardScaler()
    X_scaled_period = scaler_period.fit_transform(X_period.fillna(method='ffill'))
    
    pca_period = PCA(n_components=3)
    X_pca_period = pca_period.fit_transform(X_scaled_period)
    
    # Calculate R² for each EM market in this period with enhanced validation
    for em_market in em_columns:
        # Robust regression with data quality checks
        model_period = LinearRegression()
        model_period.fit(X_pca_aligned, y_aligned)
        r2_period = model_period.score(X_pca_aligned, y_aligned)

def rolling_r2_scores(X, Y, window=60, n_components=3):
    """
    Calculate rolling R² scores for EM indices using PCA-based factor models.
    
    Parameters:
    -----------
    X : pd.DataFrame
        Macro factor returns (features)
    Y : pd.DataFrame 
        EM index returns (targets)
    window : int
        Rolling window size in days (default: 60)
    n_components : int
        Number of PCA components to use (default: 3)
        
    Returns:
    --------
    pd.DataFrame
        Rolling R² scores for each EM index
    """
    
    # Initialize results storage
    rolling_results = {}
    
    # Get overlapping date range
    common_dates = X.index.intersection(Y.index)
    X_aligned = X.loc[common_dates]
    Y_aligned = Y.loc[common_dates]
    
    print(f"🔄 Computing rolling {window}-day R² scores...")
    print(f"   Date range: {common_dates.min()} to {common_dates.max()}")
    print(f"   Total observations: {len(common_dates)}")
    
    # Process each EM index
    for em_name in Y_aligned.columns:
        print(f"\n📊 Processing {em_name}...")
        
        r2_scores = []
        dates = []
        
        # Rolling window analysis
        for i in range(window, len(common_dates)):
            try:
                # Define current window
                start_idx = i - window
                end_idx = i
                window_dates = common_dates[start_idx:end_idx]
                
                # Extract window data
                X_window = X_aligned.loc[window_dates]
                y_window = Y_aligned.loc[window_dates, em_name]
                
                # Remove any NaN values
                valid_mask = ~(X_window.isnull().any(axis=1) | y_window.isnull())
                X_clean = X_window[valid_mask]
                y_clean = y_window[valid_mask]
                
                # Skip if insufficient data
                if len(X_clean) < 30:  # Minimum 30 observations
                    r2_scores.append(np.nan)
                    dates.append(common_dates[end_idx-1])
                    continue
                
                # Standardize and apply PCA
                scaler = StandardScaler()
                X_scaled = scaler.fit_transform(X_clean)
                
                pca = PCA(n_components=n_components)
                X_pca = pca.fit_transform(X_scaled)
                
                # Fit regression model
                model = LinearRegression()
                model.fit(X_pca, y_clean)
                
                # Calculate R²
                r2 = model.score(X_pca, y_clean)
                r2_scores.append(max(0, r2))  # Ensure non-negative
                dates.append(common_dates[end_idx-1])
                
            except Exception as e:
                print(f"   Warning: Error at window {i}: {str(e)[:50]}...")
                r2_scores.append(np.nan)
                dates.append(common_dates[end_idx-1] if end_idx <= len(common_dates) else common_dates[-1])
        
        # Store results
        rolling_results[em_name] = pd.Series(r2_scores, index=dates)
        
        # Print summary statistics
        valid_r2 = [x for x in r2_scores if not np.isnan(x)]
        if valid_r2:
            print(f"   Average R²: {np.mean(valid_r2):.3f}")
            print(f"   R² Range: {np.min(valid_r2):.3f} - {np.max(valid_r2):.3f}")
            print(f"   Valid windows: {len(valid_r2)}/{len(r2_scores)}")
    
    return pd.DataFrame(rolling_results)

# Execute rolling analysis
print("🚀 Starting comprehensive rolling window analysis...")
rolling_r2_df = rolling_r2_scores(macro_returns, em_returns, window=60, n_components=3)

# Display results summary
print(f"\n📊 Rolling Analysis Complete!")
print(f"Results shape: {rolling_r2_df.shape}")
print(f"Date range: {rolling_r2_df.index.min()} to {rolling_r2_df.index.max()}")