prototype_pipeline.py

"""
Speculorix Pipeline
===================
ML-driven stock selection using XGBoost on financial fundamentals.

Main workflow:
1. load Compustat-CRSP data and handle duplicates
2. create financial ratio features (profitability, leverage, valuation, etc.)
3. train XGBoost to predict monthly returns
4. backtest top-k portfolio selection strategy
5. evaluate via Information Coefficient and quintile analysis
"""

import pandas as pd
import numpy as np
import warnings
from typing import Tuple, Dict, List 
import matplotlib.pyplot as plt
from scipy.stats import spearmanr
import xgboost as xgb

warnings.filterwarnings('ignore')
np.random.seed(42)


class Config:
    DATA_PATH = 'data.csv'
    
    # investability filters
    MIN_PRICE = 1.0
    MIN_MARKET_CAP_M = 50
    
    # time periods (YYYYMM format)
    TRAIN_START = 201001
    TRAIN_END = 201912
    VAL_START = 202001
    VAL_END = 202112
    TEST_START = 202201
    TEST_END = 202412
    
    # outlier treatment (clip at 1st and 99th percentile)
    WINSORIZE_LOWER = 0.01
    WINSORIZE_UPPER = 0.99
    
    # xgboost hyperparameters (tuned for financial data)
    XGB_PARAMS = {
        'objective': 'reg:squarederror',
        'max_depth': 4,
        'learning_rate': 0.05,
        'n_estimators': 2000,
        'subsample': 0.8,
        'colsample_bytree': 0.8,
        'reg_alpha': 1.0,
        'reg_lambda': 1.0,
        'random_state': 42,
        'verbosity': 0,
        'early_stopping_rounds': 50
    }
    
    # portfolio construction
    TOP_K = 30
    BASE_BUDGET = 1000
    MIN_STOCKS_FOR_IC = 30


def load_data(filepath: str) -> pd.DataFrame:
    """Load and preprocess data"""
    print("Loading data...")
    
    df = pd.read_csv(filepath, low_memory=False)
    print(f"Loaded {len(df):,} rows, {len(df.columns)} columns")
    
    # handle CRSP negative price convention (bid-ask average indicator)
    if 'MthPrc' in df.columns:
        neg_prices = (df['MthPrc'] < 0).sum()
        print(f"Negative MthPrc (bid-ask avg): {neg_prices:,}")
        df['MthPrc_abs'] = df['MthPrc'].abs()
    
    # extract calendar month for time-based grouping
    if 'YYYYMM' not in df.columns:
        raise ValueError("YYYYMM column required")
    
    before = len(df)
    df = df[df['YYYYMM'].notna()].copy()
    df['calendar_month'] = df['YYYYMM'].astype(int)
    print(f"Created calendar_month (dropped {before - len(df):,} rows with missing YYYYMM)")
    
    months = sorted(df['calendar_month'].unique())
    print(f"Calendar months: {len(months)} unique ({min(months)} to {max(months)})")
    print(f"Avg companies per month: {df.groupby('calendar_month').size().mean():.0f}")
    
    return df


def handle_duplicates(df: pd.DataFrame) -> pd.DataFrame:
    """Remove duplicates with deterministic priority"""
    print("\nHandling duplicates...")
    
    dups = df.duplicated(subset=['gvkey', 'calendar_month']).sum()
    print(f"Duplicates (gvkey, calendar_month): {dups:,}")
    
    if dups > 0:
        # calculate market cap for sorting
        df['mcap_crsp'] = df['MthPrc_abs'] * df['ShrOut'] / 1000
        
        # prioritize link quality: LC (primary) > LU (unresearched) > LS (secondary)
        link_priority = {'LC': 1, 'LU': 2, 'LS': 3}
        df['link_priority'] = df['LINKTYPE'].map(link_priority).fillna(99)
        
        df = df.sort_values(
            ['gvkey', 'calendar_month', 'link_priority', 'mcap_crsp'],
            ascending=[True, True, True, False]
        )
        df = df.drop_duplicates(subset=['gvkey', 'calendar_month'], keep='first')
        df = df.drop(columns=['link_priority'])
        
        print(f"Resolved: {len(df):,} rows remaining")
    
    return df


def apply_filters(df: pd.DataFrame) -> pd.DataFrame:
    """Apply investability filters"""
    print("\nApplying filters...")
    
    init = len(df)
    
    # remove penny stocks
    df = df[df['MthPrc_abs'] >= Config.MIN_PRICE]
    print(f"Price >= ${Config.MIN_PRICE}: {init - len(df):,} removed")
    
    # remove micro caps (liquidity requirement)
    df['mcap_musd'] = df['MthPrc_abs'] * df['ShrOut'] / 1000
    before = len(df)
    df = df[df['mcap_musd'] >= Config.MIN_MARKET_CAP_M]
    print(f"Mcap >= ${Config.MIN_MARKET_CAP_M}M: {before - len(df):,} removed")
    
    # require valid return data
    before = len(df)
    df = df[df['MthRet'].notna()]
    print(f"MthRet not null: {before - len(df):,} removed")
    
    print(f"Final: {len(df):,} observations")
    
    return df


def create_labels(df: pd.DataFrame) -> pd.DataFrame:
    """Create return labels"""
    print("\nCreating labels...")
    
    df['label'] = df['MthRet']
    
    print(f"Label statistics:")
    print(f"  Mean: {df['label'].mean():.4f}")
    print(f"  Std:  {df['label'].std():.4f}")
    
    return df


def create_features(df: pd.DataFrame) -> pd.DataFrame:
    """Create financial features"""
    print("\nCreating features...")
    
    # profitability ratios
    df['roa'] = df['ni'] / df['at'].replace(0, np.nan)
    df['roe'] = df['ni'] / df['ceq'].replace(0, np.nan)
    df['profit_margin'] = df['ni'] / df['sale'].replace(0, np.nan)
    df['ebitda_margin'] = df['ebitda'] / df['sale'].replace(0, np.nan)
    df['gross_margin'] = (df['sale'] - df['cogs']) / df['sale'].replace(0, np.nan)
    
    # leverage and solvency
    df['total_debt'] = df['dlc'].fillna(0) + df['dltt'].fillna(0)
    df['leverage'] = df['total_debt'] / df['at'].replace(0, np.nan)
    df['debt_to_equity'] = df['total_debt'] / df['ceq'].replace(0, np.nan)
    df['interest_coverage'] = df['ebit'] / df['xint'].replace(0, np.nan)
    
    # liquidity ratios
    df['current_ratio'] = df['act'] / df['lct'].replace(0, np.nan)
    df['quick_ratio'] = (df['act'] - df.get('invt', 0)) / df['lct'].replace(0, np.nan)
    df['cash_ratio'] = df['che'] / df['lct'].replace(0, np.nan)
    df['cash_to_assets'] = df['che'] / df['at'].replace(0, np.nan)
    
    # cash flow metrics
    df['cfo_to_assets'] = df['oancf'] / df['at'].replace(0, np.nan)
    df['capex_to_assets'] = df['capx'] / df['at'].replace(0, np.nan)
    df['fcf_to_assets'] = (df['oancf'] - df['capx']) / df['at'].replace(0, np.nan)
    
    # valuation multiples
    df['book_to_market'] = df['ceq'] / df['mcap_musd'].replace(0, np.nan)
    df['earnings_yield'] = df['ni'] / df['mcap_musd'].replace(0, np.nan)
    df['sales_to_price'] = df['sale'] / df['mcap_musd'].replace(0, np.nan)
    df['ebitda_to_ev'] = df['ebitda'] / (df['mcap_musd'] + df['total_debt']).replace(0, np.nan)
    
    # size factors (log transform for skewness)
    df['log_assets'] = np.log1p(df['at'])
    df['log_mcap'] = np.log1p(df['mcap_musd'])
    df['log_sales'] = np.log1p(df['sale'])
    
    # dividend policy
    df['dividend_yield'] = df['dvpsx_f'] / df['MthPrc_abs'].replace(0, np.nan)
    df['payout_ratio'] = df['dvt'] / df['ni'].replace(0, np.nan)
    
    print("created profitability, leverage, liquidity, cash flow, valuation features")
    
    return df


def get_feature_columns(df: pd.DataFrame) -> List[str]:
    """Select feature columns with proper exclusions"""
    print("\nSelecting features...")
    
    # exclude identifiers, dates, prices, returns (prevent leakage)
    exclude = [
        'gvkey', 'fyear', 'fyr', 'datadate', 'tic', 'conm', 'costat',
        'permno', 'permco', 'linktype', 'cusip',
        'yyyymm', 'mthcaldt', 'mthprc', 'mthret', 'mthretx', 'shrout',
        'calendar_month', 'fyr_month', 'mcap', 'market_cap',
        'label', 'pred', 'score', 'weight', '_abs',
        'indfmt', 'consol', 'popsrc', 'datafmt', 'curcd',
        'prcc_', 'prc', '_ret', 'ret'
    ]
    
    feature_cols = []
    for col in df.columns:
        col_lower = col.lower()
        should_exclude = any(x in col_lower for x in exclude)
        
        if not should_exclude and df[col].dtype in [np.float64, np.int64, np.float32, np.int32]:
            if df[col].notna().sum() > 100:
                feature_cols.append(col)
    
    print(f"Selected {len(feature_cols)} features")
    
    return feature_cols


def fit_winsor_bounds(df: pd.DataFrame, cols: List[str]) -> Dict:
    """Fit winsorization bounds on training data only"""
    bounds = {}
    for col in cols:
        if col in df.columns and df[col].notna().sum() > 0:
            lower = df[col].quantile(Config.WINSORIZE_LOWER)
            upper = df[col].quantile(Config.WINSORIZE_UPPER)
            bounds[col] = (lower, upper)
    return bounds


def apply_winsor(df: pd.DataFrame, bounds: Dict) -> pd.DataFrame:
    """Clip extreme values to reduce outlier impact"""
    df = df.copy()
    for col, (lower, upper) in bounds.items():
        if col in df.columns:
            df[col] = df[col].clip(lower, upper)
    return df


def create_ranks(df: pd.DataFrame, cols: List[str]) -> pd.DataFrame:
    """Create cross-sectional ranks per month"""
    print("\nCreating cross-sectional ranks...")
    
    df = df.copy()
    # rank stocks within each month (0 to 1 scale)
    for col in cols:
        if col in df.columns:
            df[f'{col}_rank'] = df.groupby('calendar_month')[col].rank(pct=True)
    
    rank_cols = [c for c in df.columns if c.endswith('_rank')]
    print(f"created {len(rank_cols)} rank features")
    
    return df


def split_data(df: pd.DataFrame) -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
    """Time-based split (no shuffling to prevent lookahead bias)"""
    print("\nsplitting data...")
    
    actual_max = df['calendar_month'].max()
    test_end = min(Config.TEST_END, actual_max)
    
    # strict chronological split: train → validate → test
    train = df[(df['calendar_month'] >= Config.TRAIN_START) & (df['calendar_month'] <= Config.TRAIN_END)]
    val = df[(df['calendar_month'] >= Config.VAL_START) & (df['calendar_month'] <= Config.VAL_END)]
    test = df[(df['calendar_month'] >= Config.TEST_START) & (df['calendar_month'] <= test_end)]
    
    print(f"Train: {Config.TRAIN_START}-{Config.TRAIN_END} ({len(train):,} obs, {train['calendar_month'].nunique()} months)")
    print(f"Val: {Config.VAL_START}-{Config.VAL_END} ({len(val):,} obs, {val['calendar_month'].nunique()} months)")
    print(f"Test: {Config.TEST_START}-{test_end} ({len(test):,} obs, {test['calendar_month'].nunique()} months)")
    
    return train, val, test


def train_model(X_train, y_train, X_val, y_val):
    """Train XGBoost model with early stopping"""
    print("\nTraining model...")
    
    params = Config.XGB_PARAMS.copy()
    model = xgb.XGBRegressor(**params)
    
    try:
        # monitor validation performance to prevent overfitting
        model.fit(
            X_train, y_train,
            eval_set=[(X_val, y_val)],
            verbose=100
        )
        best = getattr(model, 'best_iteration', params['n_estimators'])
        print(f"stopped at iteration: {best}")
    except TypeError:
        # fallback for older xgboost versions
        print("fallback: training without early stopping")
        params.pop('early_stopping_rounds', None)
        model = xgb.XGBRegressor(**params)
        model.fit(X_train, y_train)
    
    return model


def calculate_monthly_ic(df: pd.DataFrame, pred_col: str = 'pred') -> pd.DataFrame:
    """Calculate information coefficient (spearman correlation) per month"""
    results = []
    
    for month in sorted(df['calendar_month'].unique()):
        month_df = df[df['calendar_month'] == month]
        
        # need sufficient stocks for meaningful cross-sectional IC
        if len(month_df) >= Config.MIN_STOCKS_FOR_IC:
            ic, pval = spearmanr(month_df[pred_col], month_df['label'])
            results.append({
                'calendar_month': month,
                'ic': ic if not np.isnan(ic) else 0,
                'pval': pval if not np.isnan(pval) else 1,
                'n_stocks': len(month_df)
            })
    
    return pd.DataFrame(results)


def evaluate_ic(df: pd.DataFrame, predictions: np.ndarray, name: str) -> Dict:
    """Evaluate Information Coefficient"""
    print(f"\n{name} IC Evaluation")
    
    df_eval = df.copy()
    df_eval['pred'] = predictions
    
    overall_ic, overall_pval = spearmanr(predictions, df_eval['label'])
    print(f"Overall IC: {overall_ic:+.4f} (p={overall_pval:.4f})")
    
    monthly_ic = calculate_monthly_ic(df_eval)
    
    if len(monthly_ic) > 0:
        mean_ic = monthly_ic['ic'].mean()
        std_ic = monthly_ic['ic'].std()
        ir = mean_ic / std_ic if std_ic > 0 else 0
        
        print(f"Monthly IC: mean={mean_ic:+.4f}, std={std_ic:.4f}, IR={ir:.2f}")
        print(f"Months with IC > 0: {(monthly_ic['ic'] > 0).sum()}/{len(monthly_ic)}")
    
    return {
        'overall_ic': overall_ic,
        'monthly_ic_df': monthly_ic,
        'mean_monthly_ic': mean_ic if len(monthly_ic) > 0 else np.nan
    }


def calculate_quintile_returns(df: pd.DataFrame, pred_col: str = 'pred') -> pd.DataFrame:
    """Sort stocks into quintiles by prediction and compute average returns"""
    results = []
    
    for month in sorted(df['calendar_month'].unique()):
        month_df = df[df['calendar_month'] == month].copy()
        
        if len(month_df) >= 50:
            try:
                # split into 5 equal groups (Q1=worst, Q5=best predicted)
                month_df['quintile'] = pd.qcut(month_df[pred_col], q=5, labels=[1,2,3,4,5], duplicates='drop')
                
                for q in range(1, 6):
                    q_df = month_df[month_df['quintile'] == q]
                    if len(q_df) > 0:
                        results.append({
                            'calendar_month': month,
                            'quintile': q,
                            'return': q_df['label'].mean(),
                            'n_stocks': len(q_df)
                        })
            except:
                pass
    
    return pd.DataFrame(results)


def run_backtest(df: pd.DataFrame, model, feature_cols: List[str]) -> pd.DataFrame:
    """Run walk-forward backtest with top-k portfolio"""
    print("\nRunning backtest...")
    
    results = []
    months = sorted(df['calendar_month'].unique())
    
    print(f"backtesting {len(months)} months")
    
    for month in months:
        month_df = df[df['calendar_month'] == month].copy()
        
        if len(month_df) < Config.TOP_K:
            continue
        
        # predict and select top k stocks
        X_month = month_df[feature_cols].fillna(0)
        month_df['pred'] = model.predict(X_month)
        
        top_k = month_df.nlargest(Config.TOP_K, 'pred')
        
        # equal-weighted portfolio vs equal-weighted benchmark
        port_ret = top_k['label'].mean()
        bench_ret = month_df['label'].mean()
        
        results.append({
            'calendar_month': month,
            'portfolio_return': port_ret,
            'benchmark_return': bench_ret,
            'excess_return': port_ret - bench_ret,
            'n_stocks': len(top_k)
        })
    
    df_results = pd.DataFrame(results)
    
    if len(df_results) > 0:
        mean_monthly = df_results['portfolio_return'].mean()
        std_monthly = df_results['portfolio_return'].std()
        mean_excess = df_results['excess_return'].mean()
        
        annual_ret = mean_monthly * 12
        annual_alpha = mean_excess * 12
        sharpe = (mean_monthly / std_monthly * np.sqrt(12)) if std_monthly > 0 else 0
        
        print(f"\nBacktest Summary:")
        print(f"  Monthly return: {mean_monthly:+.4f} ({mean_monthly*100:+.2f}%)")
        print(f"  Monthly alpha: {mean_excess:+.4f} ({mean_excess*100:+.2f}%)")
        print(f"  Annualized return: {annual_ret:+.4f} ({annual_ret*100:+.2f}%)")
        print(f"  Annualized alpha: {annual_alpha:+.4f} ({annual_alpha*100:+.2f}%)")
        print(f"  Sharpe ratio: {sharpe:.3f}")
        print(f"  Win rate: {(df_results['excess_return'] > 0).sum()}/{len(df_results)}")
        
        if len(df_results) < 24:
            print(f"\nWARNING: Only {len(df_results)} test months")
            print(f"Results are NOT statistically reliable!")
    
    return df_results


def create_plots(df_backtest: pd.DataFrame, df_ic: pd.DataFrame):
    """Generate cumulative return and IC time series charts"""
    print("\ncreating visualizations...")
    
    try:
        import matplotlib
        matplotlib.use('Agg')  # non-interactive backend for saving to file
        
        fig, axes = plt.subplots(2, 1, figsize=(12, 10))
        
        # cumulative returns over time
        df_backtest = df_backtest.copy()
        df_backtest['cum_port'] = (1 + df_backtest['portfolio_return']).cumprod()
        df_backtest['cum_bench'] = (1 + df_backtest['benchmark_return']).cumprod()
        
        axes[0].plot(range(len(df_backtest)), df_backtest['cum_port'], 
                     label='Portfolio', linewidth=2, color='blue')
        axes[0].plot(range(len(df_backtest)), df_backtest['cum_bench'], 
                     label='Benchmark', linewidth=2, color='gray', alpha=0.7)
        axes[0].axhline(y=1, color='k', linestyle='--', alpha=0.3)
        axes[0].set_xlabel('Month')
        axes[0].set_ylabel('Cumulative Return')
        axes[0].set_title('Cumulative Returns')
        axes[0].legend()
        axes[0].grid(True, alpha=0.3)
        
        # monthly IC bar chart
        if len(df_ic) > 0:
            colors = ['green' if x > 0 else 'red' for x in df_ic['ic']]
            axes[1].bar(range(len(df_ic)), df_ic['ic'], color=colors, alpha=0.7)
            axes[1].axhline(y=0, color='k', linestyle='-', linewidth=0.5)
            axes[1].axhline(y=df_ic['ic'].mean(), color='blue', linestyle='--', 
                            label=f"Mean IC: {df_ic['ic'].mean():.3f}")
            axes[1].set_xlabel('Month')
            axes[1].set_ylabel('IC')
            axes[1].set_title('Information Coefficient by Month')
            axes[1].legend()
            axes[1].grid(True, alpha=0.3)
        
        plt.tight_layout()
        plt.savefig('performance_charts.png', dpi=150, bbox_inches='tight')
        print("Saved performance_charts.png")
        plt.close()
    except Exception as e:
        print(f"Could not create plots: {e}")


def main():
    """Main pipeline execution"""
    print("\n" + "=" * 60)
    print(" " * 15 + "SPECULORIX PIPELINE")
    print("=" * 60)
    
    # load and clean data
    df = load_data(Config.DATA_PATH)
    df = handle_duplicates(df)
    df = apply_filters(df)
    df = create_labels(df)
    df = create_features(df)
    
    feature_cols = get_feature_columns(df)
    
    # time-based split to prevent lookahead bias
    train_df, val_df, test_df = split_data(df)
    
    # fit winsorization on train only, then apply to all
    print("\nwinsorizing features...")
    ratio_cols = [c for c in feature_cols if df[c].dtype == np.float64]
    winsor_bounds = fit_winsor_bounds(train_df, ratio_cols)
    print(f"fitted bounds on {len(winsor_bounds)} columns")
    
    train_df = apply_winsor(train_df, winsor_bounds)
    val_df = apply_winsor(val_df, winsor_bounds)
    test_df = apply_winsor(test_df, winsor_bounds)
    
    # add cross-sectional rank features
    train_df = create_ranks(train_df, feature_cols)
    val_df = create_ranks(val_df, feature_cols)
    test_df = create_ranks(test_df, feature_cols)
    
    all_features = get_feature_columns(train_df)
    print(f"final features: {len(all_features)}")
    
    # prepare matrices
    X_train = train_df[all_features].fillna(0)
    y_train = train_df['label']
    X_val = val_df[all_features].fillna(0)
    y_val = val_df['label']
    X_test = test_df[all_features].fillna(0)
    y_test = test_df['label']
    
    print(f"\ndata shapes: Train={X_train.shape}, Val={X_val.shape}, Test={X_test.shape}")
    
    # train model
    model = train_model(X_train, y_train, X_val, y_val)
    
    # evaluate predictive power
    val_pred = model.predict(X_val)
    test_pred = model.predict(X_test)
    
    val_ic = evaluate_ic(val_df, val_pred, "VALIDATION")
    test_ic = evaluate_ic(test_df, test_pred, "TEST")
    
    # check if model ranks stocks correctly (Q5 should beat Q1)
    print("\nquintile analysis")
    test_df_eval = test_df.copy()
    test_df_eval['pred'] = test_pred
    
    quintile_df = calculate_quintile_returns(test_df_eval)
    if len(quintile_df) > 0:
        avg_by_q = quintile_df.groupby('quintile')['return'].mean()
        print("\naverage return by quintile:")
        for q in range(1, 6):
            if q in avg_by_q.index:
                print(f"  Q{q}: {avg_by_q[q]:+.4f} ({avg_by_q[q]*100:+.2f}%)")
        
        if 5 in avg_by_q.index and 1 in avg_by_q.index:
            spread = avg_by_q[5] - avg_by_q[1]
            print(f"\nQ5-Q1 spread: {spread:+.4f} monthly, {spread*12:+.4f} annualized")
    
    # simulate actual portfolio performance
    df_backtest = run_backtest(test_df, model, all_features)
    
    # generate visualizations
    if len(df_backtest) > 0 and len(test_ic.get('monthly_ic_df', [])) > 0:
        create_plots(df_backtest, test_ic['monthly_ic_df'])
    
    # save results to csv files
    print("\nsaving results...")
    df_backtest.to_csv('backtest_results_monthly.csv', index=False)
    
    if 'monthly_ic_df' in test_ic and len(test_ic['monthly_ic_df']) > 0:
        test_ic['monthly_ic_df'].to_csv('ic_by_month.csv', index=False)
    
    importance = pd.DataFrame({
        'feature': all_features,
        'importance': model.feature_importances_
    }).sort_values('importance', ascending=False)
    importance.to_csv('feature_importance.csv', index=False)
    
    print("\n" + "=" * 60)
    print("RESULTS")
    print("=" * 60)
    print(f"\nValidation IC: {val_ic['overall_ic']:+.4f}")
    print(f"Test IC: {test_ic['overall_ic']:+.4f}")
    
    if len(df_backtest) > 0:
        print(f"\nMonthly Alpha: {df_backtest['excess_return'].mean():+.4f}")
        print(f"Annualized Alpha: {df_backtest['excess_return'].mean()*12:+.4f}")
        print(f"Test Months: {len(df_backtest)}")
    
    print("\n" + "=" * 60)
    print("COMPLETE")
    print("=" * 60)
    
    return {
        'model': model,
        'features': all_features,
        'backtest': df_backtest,
        'val_ic': val_ic,
        'test_ic': test_ic
    }


if __name__ == "__main__":
    results = main()