quant/factor_robustness.py

"""
Robustness checks for winning factor strategies.

Tests:
1. Rolling 2-year window performance (stability)
2. Top-N sensitivity (5, 10, 15, 20)
3. Rebalance frequency sensitivity (5d, 10d, 21d, 42d)
4. Transaction cost sensitivity (0, 10bps, 20bps, 50bps)
5. Drawdown analysis
"""

from __future__ import annotations

import argparse
import warnings

import numpy as np
import pandas as pd

import data_manager
from universe import UNIVERSES
from factor_real_backtest import (
    f_recovery_mom,
    f_momentum_12_1,
    f_recovery,
    f_recovery_deep,
    f_up_volume_proxy,
    f_gap_up_freq,
    f_earnings_drift_proxy,
    f_reversal_vol_cn,
    f_consistent_winner,
    combo_signal,
    make_strategy,
    run_backtest,
    compute_stats,
)

warnings.filterwarnings("ignore")


def rolling_window_performance(equity: pd.Series, window_years: int = 2):
    """Compute rolling window returns."""
    daily_ret = equity.pct_change().dropna()
    window = 252 * window_years
    results = []
    for end_idx in range(window, len(daily_ret), 63):  # step 3 months
        start_idx = end_idx - window
        chunk = daily_ret.iloc[start_idx:end_idx]
        total = (1 + chunk).prod() - 1
        ann = (1 + total) ** (252 / len(chunk)) - 1
        sharpe = chunk.mean() / chunk.std() * np.sqrt(252) if chunk.std() > 0 else 0
        results.append({
            "end_date": chunk.index[-1].date(),
            "ann_return": ann,
            "sharpe": sharpe,
        })
    return pd.DataFrame(results)


def drawdown_analysis(equity: pd.Series) -> pd.DataFrame:
    """Find top 5 drawdown episodes."""
    running_max = equity.cummax()
    drawdown = (equity - running_max) / running_max

    # Find drawdown episodes
    episodes = []
    in_dd = False
    start = None
    for i in range(len(drawdown)):
        if drawdown.iloc[i] < -0.05 and not in_dd:
            in_dd = True
            start = i
        elif drawdown.iloc[i] >= 0 and in_dd:
            in_dd = False
            trough_idx = drawdown.iloc[start:i].idxmin()
            episodes.append({
                "start": drawdown.index[start].date(),
                "trough": trough_idx.date(),
                "end": drawdown.index[i].date(),
                "depth": drawdown.loc[trough_idx],
                "duration_days": i - start,
            })
    # Handle ongoing drawdown
    if in_dd:
        trough_idx = drawdown.iloc[start:].idxmin()
        episodes.append({
            "start": drawdown.index[start].date(),
            "trough": trough_idx.date(),
            "end": "ongoing",
            "depth": drawdown.loc[trough_idx],
            "duration_days": len(drawdown) - start,
        })

    df = pd.DataFrame(episodes)
    if df.empty:
        return df
    return df.nsmallest(5, "depth")


def run_us(stocks: pd.DataFrame):
    print("=" * 100)
    print("  US ROBUSTNESS — Winner: momentum_12_1 + up_volume_proxy")
    print("=" * 100)

    winner_func = combo_signal([(f_momentum_12_1, 0.5), (f_up_volume_proxy, 0.5)])
    baseline_func = f_recovery_mom

    # 1. Rolling 2-year performance
    print("\n--- 1. Rolling 2-Year Performance ---\n")
    for label, func in [("Winner: mom+upvol", winner_func),
                        ("Baseline: rec+mom", baseline_func)]:
        w = make_strategy(stocks, func, top_n=10)
        eq = run_backtest(w, stocks)
        roll = rolling_window_performance(eq)
        if roll.empty:
            continue
        win_pct = (roll["ann_return"] > 0).mean()
        print(f"  {label}:")
        print(f"    Mean 2yr ann return: {roll['ann_return'].mean():+.1%}")
        print(f"    Min 2yr ann return:  {roll['ann_return'].min():+.1%}")
        print(f"    Max 2yr ann return:  {roll['ann_return'].max():+.1%}")
        print(f"    % positive 2yr:      {win_pct:.0%}")
        print(f"    Mean 2yr Sharpe:     {roll['sharpe'].mean():.2f}")
        print()

    # 2. Top-N sensitivity
    print("--- 2. Top-N Sensitivity ---\n")
    header = f"  {'Top-N':<8}"
    for label in ["Winner: mom+upvol", "Baseline: rec+mom"]:
        header += f" | {'CAGR':>8} {'Sharpe':>8} {'MaxDD':>8}"
    print(header)
    print("  " + "-" * 70)

    for top_n in [5, 10, 15, 20, 30]:
        line = f"  {top_n:<8}"
        for func in [winner_func, baseline_func]:
            w = make_strategy(stocks, func, top_n=top_n)
            eq = run_backtest(w, stocks)
            s = compute_stats(eq, "")
            line += f" | {s['cagr']:>+7.1%} {s['sharpe']:>8.2f} {s['maxdd']:>+7.1%}"
        print(line)

    # 3. Rebalance frequency sensitivity
    print("\n--- 3. Rebalance Frequency Sensitivity ---\n")
    header = f"  {'Rebal':<8}"
    for label in ["Winner: mom+upvol", "Baseline: rec+mom"]:
        header += f" | {'CAGR':>8} {'Sharpe':>8} {'MaxDD':>8}"
    print(header)
    print("  " + "-" * 70)

    for rebal in [5, 10, 21, 42, 63]:
        line = f"  {rebal}d{'':<5}"
        for func in [winner_func, baseline_func]:
            w = make_strategy(stocks, func, top_n=10, rebal_freq=rebal)
            eq = run_backtest(w, stocks)
            s = compute_stats(eq, "")
            line += f" | {s['cagr']:>+7.1%} {s['sharpe']:>8.2f} {s['maxdd']:>+7.1%}"
        print(line)

    # 4. Transaction cost sensitivity
    print("\n--- 4. Transaction Cost Sensitivity ---\n")
    header = f"  {'Cost':<8}"
    for label in ["Winner: mom+upvol", "Baseline: rec+mom"]:
        header += f" | {'CAGR':>8} {'Sharpe':>8}"
    print(header)
    print("  " + "-" * 50)

    for cost in [0, 0.001, 0.002, 0.005]:
        line = f"  {cost*10000:.0f}bps{'':<4}"
        for func in [winner_func, baseline_func]:
            w = make_strategy(stocks, func, top_n=10)
            eq = run_backtest(w, stocks, cost=cost)
            s = compute_stats(eq, "")
            line += f" | {s['cagr']:>+7.1%} {s['sharpe']:>8.2f}"
        print(line)

    # 5. Drawdown analysis
    print("\n--- 5. Drawdown Episodes ---\n")
    for label, func in [("Winner: mom+upvol", winner_func),
                        ("Baseline: rec+mom", baseline_func)]:
        w = make_strategy(stocks, func, top_n=10)
        eq = run_backtest(w, stocks)
        dd = drawdown_analysis(eq)
        print(f"  {label}:")
        if dd.empty:
            print("    No significant drawdowns")
        else:
            for _, row in dd.iterrows():
                print(f"    {row['start']} → {row['trough']} → {row['end']}: "
                      f"{row['depth']:+.1%} ({row['duration_days']}d)")
        print()

    # 6. Also test the runner-up combos
    print("--- 6. Other Strong Combos (Top-10, 21d rebal, 10bps) ---\n")
    other_combos = [
        ("rec_deep+upvol", combo_signal([(f_recovery_deep, 0.5), (f_up_volume_proxy, 0.5)])),
        ("rec_deep+mom", combo_signal([(f_recovery_deep, 0.5), (f_momentum_12_1, 0.5)])),
        ("mom+gap_up", combo_signal([(f_momentum_12_1, 0.5), (f_gap_up_freq, 0.5)])),
        ("rec_deep+upvol+mom", combo_signal([(f_recovery_deep, 0.33), (f_up_volume_proxy, 0.33), (f_momentum_12_1, 0.34)])),
        ("mom+upvol+gap", combo_signal([(f_momentum_12_1, 0.33), (f_up_volume_proxy, 0.33), (f_gap_up_freq, 0.34)])),
    ]
    for label, func in other_combos:
        w = make_strategy(stocks, func, top_n=10)
        eq = run_backtest(w, stocks)
        s = compute_stats(eq, "")
        print(f"  {label:<25} CAGR: {s['cagr']:>+7.1%}  Sharpe: {s['sharpe']:.2f}  MaxDD: {s['maxdd']:>+7.1%}  Calmar: {s['calmar']:.2f}")


def run_cn(stocks: pd.DataFrame):
    print("\n" + "=" * 100)
    print("  CN ROBUSTNESS — Winners: reversal_vol + gap_up, earn_drift + reversal_vol")
    print("=" * 100)

    winner1_func = combo_signal([(f_reversal_vol_cn, 0.5), (f_gap_up_freq, 0.5)])
    winner2_func = combo_signal([(f_earnings_drift_proxy, 0.5), (f_reversal_vol_cn, 0.5)])
    baseline_func = f_recovery_mom

    # 1. Rolling 2-year performance
    print("\n--- 1. Rolling 2-Year Performance ---\n")
    for label, func in [("W1: rev_vol+gap_up", winner1_func),
                        ("W2: earn_drift+rev_vol", winner2_func),
                        ("Baseline: rec+mom", baseline_func)]:
        w = make_strategy(stocks, func, top_n=10)
        eq = run_backtest(w, stocks)
        roll = rolling_window_performance(eq)
        if roll.empty:
            continue
        win_pct = (roll["ann_return"] > 0).mean()
        print(f"  {label}:")
        print(f"    Mean 2yr ann return: {roll['ann_return'].mean():+.1%}")
        print(f"    Min 2yr ann return:  {roll['ann_return'].min():+.1%}")
        print(f"    Max 2yr ann return:  {roll['ann_return'].max():+.1%}")
        print(f"    % positive 2yr:      {win_pct:.0%}")
        print(f"    Mean 2yr Sharpe:     {roll['sharpe'].mean():.2f}")
        print()

    # 2. Top-N sensitivity
    print("--- 2. Top-N Sensitivity ---\n")
    header = f"  {'Top-N':<8}"
    for label in ["W1: rev+gap", "W2: earn+rev", "Baseline"]:
        header += f" | {'CAGR':>8} {'Sharpe':>8} {'MaxDD':>8}"
    print(header)
    print("  " + "-" * 100)

    for top_n in [5, 10, 15, 20]:
        line = f"  {top_n:<8}"
        for func in [winner1_func, winner2_func, baseline_func]:
            w = make_strategy(stocks, func, top_n=top_n)
            eq = run_backtest(w, stocks)
            s = compute_stats(eq, "")
            line += f" | {s['cagr']:>+7.1%} {s['sharpe']:>8.2f} {s['maxdd']:>+7.1%}"
        print(line)

    # 3. Rebalance frequency
    print("\n--- 3. Rebalance Frequency ---\n")
    header = f"  {'Rebal':<8}"
    for label in ["W1: rev+gap", "W2: earn+rev", "Baseline"]:
        header += f" | {'CAGR':>8} {'Sharpe':>8}"
    print(header)
    print("  " + "-" * 75)

    for rebal in [5, 10, 21, 42]:
        line = f"  {rebal}d{'':<5}"
        for func in [winner1_func, winner2_func, baseline_func]:
            w = make_strategy(stocks, func, top_n=10, rebal_freq=rebal)
            eq = run_backtest(w, stocks)
            s = compute_stats(eq, "")
            line += f" | {s['cagr']:>+7.1%} {s['sharpe']:>8.2f}"
        print(line)

    # 4. Transaction cost sensitivity
    print("\n--- 4. Transaction Cost Sensitivity ---\n")
    header = f"  {'Cost':<8}"
    for label in ["W1: rev+gap", "W2: earn+rev", "Baseline"]:
        header += f" | {'CAGR':>8} {'Sharpe':>8}"
    print(header)
    print("  " + "-" * 75)

    for cost in [0, 0.001, 0.002, 0.005]:
        line = f"  {cost*10000:.0f}bps{'':<4}"
        for func in [winner1_func, winner2_func, baseline_func]:
            w = make_strategy(stocks, func, top_n=10)
            eq = run_backtest(w, stocks, cost=cost)
            s = compute_stats(eq, "")
            line += f" | {s['cagr']:>+7.1%} {s['sharpe']:>8.2f}"
        print(line)

    # 5. Drawdown analysis
    print("\n--- 5. Drawdown Episodes ---\n")
    for label, func in [("W1: rev_vol+gap_up", winner1_func),
                        ("W2: earn_drift+rev_vol", winner2_func),
                        ("Baseline: rec+mom", baseline_func)]:
        w = make_strategy(stocks, func, top_n=10)
        eq = run_backtest(w, stocks)
        dd = drawdown_analysis(eq)
        print(f"  {label}:")
        if dd.empty:
            print("    No significant drawdowns")
        else:
            for _, row in dd.iterrows():
                print(f"    {row['start']} → {row['trough']} → {row['end']}: "
                      f"{row['depth']:+.1%} ({row['duration_days']}d)")
        print()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--market", default="both", choices=["us", "cn", "both"])
    args = parser.parse_args()

    if args.market in ("us", "both"):
        prices = data_manager.load("us")
        stocks = prices.drop(columns=["SPY"], errors="ignore")
        run_us(stocks)

    if args.market in ("cn", "both"):
        prices = data_manager.load("cn")
        stocks = prices.drop(columns=["000300.SS"], errors="ignore")
        run_cn(stocks)


if __name__ == "__main__":
    main()