okx-codex-trader/scripts/stress_expansion_rotation_risk.py

from __future__ import annotations

import argparse
import sys
from dataclasses import dataclass, replace
from itertools import product
from pathlib import Path

import pandas as pd

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from scripts import refine_expansion_rotation_risk as refine
from scripts import search_expansion_rotation as rotation


OUTPUT_DIR = Path("reports/strategy-expansion")
PREFIX = "rotation-risk-stress"
FEE_RATES = (0.0004, 0.0006, 0.0008, 0.0010)
SLIPPAGE_RATES = (0.0, 0.0002, 0.0005, 0.0010)
FUNDING_RATES_8H = (-0.00005, 0.0, 0.00005)
EXPOSURES = (0.50, 0.65, 0.80)
VOL_TARGETS = (0.0, 0.20, 0.30)

rotation.SYMBOLS = ("BTC-USDT-SWAP", "ETH-USDT-SWAP")


@dataclass(frozen=True)
class StressParams:
    risk: refine.RiskParams
    fee_rate: float
    slippage_rate: float
    funding_rate_8h: float

    @property
    def name(self) -> str:
        return (
            f"{self.risk.name}-fee{self.fee_rate:.4f}"
            f"-slip{self.slippage_rate:.4f}-fund{self.funding_rate_8h:.5f}"
        )


def params_from_risk_row(row: pd.Series) -> refine.RiskParams:
    base = refine.params_from_row(row)
    return refine.RiskParams(
        base=base,
        leverage=float(row["leverage"]),
        exposure=float(row["exposure"]),
        vol_target=float(row["vol_target"]),
    )


def bar_hours(bar: str) -> float:
    return {"1h": 1.0, "4h": 4.0, "1d": 24.0}[bar]


def source_candidates(limit: int) -> list[refine.RiskParams]:
    frame = pd.read_csv(OUTPUT_DIR / "rotation-risk-top.csv").head(limit)
    candidates: dict[str, refine.RiskParams] = {}
    for _, row in frame.iterrows():
        risk = params_from_risk_row(row)
        candidates[risk.name] = risk
    return list(candidates.values())


def robustness_variants(candidate: refine.RiskParams) -> list[refine.RiskParams]:
    variants: dict[str, refine.RiskParams] = {}
    for exposure, vol_target in product(EXPOSURES, VOL_TARGETS):
        risk = replace(candidate, exposure=exposure, vol_target=vol_target)
        variants[risk.name] = risk
    return list(variants.values())


def equity_curve(closes: pd.DataFrame, weights: pd.DataFrame, params: StressParams) -> pd.Series:
    returns = closes.pct_change().fillna(0.0)
    executed = weights.shift(1).fillna(0.0)
    turnover = executed.diff().abs().sum(axis=1).fillna(executed.abs().sum(axis=1))
    gross_returns = (executed * returns * params.risk.leverage).sum(axis=1)
    trade_cost = turnover * (params.fee_rate + params.slippage_rate) * params.risk.leverage
    funding = executed.abs().sum(axis=1) * params.risk.leverage * params.funding_rate_8h * bar_hours(params.risk.base.bar) / 8.0
    net_returns = gross_returns - trade_cost - funding
    equity = refine.INITIAL_EQUITY * (1.0 + net_returns).cumprod()
    equity.name = "equity"
    return equity


def row_for_result(
    candidate_name: str,
    closes: pd.DataFrame,
    weights: pd.DataFrame,
    params: StressParams,
    equity: pd.Series,
) -> dict[str, object]:
    years = (equity.index[-1] - equity.index[0]).total_seconds() / 86_400 / 365
    turnover_per_year = float(weights.shift(1).fillna(0.0).diff().abs().sum(axis=1).sum() / years)
    horizon_rows = rotation.horizon_rows(params.name, equity)
    horizons_by_label = {row["horizon"]: row for row in horizon_rows}
    return {
        "candidate": candidate_name,
        "strategy": params.name,
        "family": params.risk.base.family,
        "bar": params.risk.base.bar,
        "universe": ",".join(closes.columns),
        "lookback": params.risk.base.lookback,
        "trend": params.risk.base.trend,
        "btc_trend": params.risk.base.btc_trend,
        "rebalance": params.risk.base.rebalance,
        "top_n": params.risk.base.top_n,
        "min_momentum": params.risk.base.min_momentum,
        "btc_min_momentum": params.risk.base.btc_min_momentum,
        "vol_lookback": params.risk.base.vol_lookback,
        "max_vol": params.risk.base.max_vol,
        "leverage": params.risk.leverage,
        "exposure": params.risk.exposure,
        "vol_target": params.risk.vol_target,
        "fee_rate": params.fee_rate,
        "slippage_rate": params.slippage_rate,
        "funding_rate_8h": params.funding_rate_8h,
        "first_candle": equity.index[0].strftime("%Y-%m-%d %H:%M"),
        "last_candle": equity.index[-1].strftime("%Y-%m-%d %H:%M"),
        "years": years,
        "turnover_per_year": turnover_per_year,
        "h3y_return": horizons_by_label["3y"]["total_return"],
        "h1y_return": horizons_by_label["1y"]["total_return"],
        "h6m_return": horizons_by_label["6m"]["total_return"],
        "h3m_return": horizons_by_label["3m"]["total_return"],
        **rotation.metrics(equity),
    }


def markdown_table(frame: pd.DataFrame) -> str:
    rows = [list(frame.columns), ["---" for _ in frame.columns]]
    rows.extend(frame.astype(object).where(pd.notna(frame), "").values.tolist())
    return "\n".join("| " + " | ".join(rotation.format_cell(value) for value in row) + " |" for row in rows)


def conservative_summary(total: pd.DataFrame) -> pd.DataFrame:
    conservative = total[
        (total["fee_rate"] == 0.0010)
        & (total["slippage_rate"] == 0.0010)
        & (total["funding_rate_8h"] == 0.00005)
    ]
    return conservative.sort_values(
        ["calmar", "annualized_return", "max_drawdown"],
        ascending=[False, False, True],
    )


def report_text(command: str, paths: list[Path], total: pd.DataFrame, best_candidate: str) -> str:
    conservative = conservative_summary(total)
    best_conservative = conservative[conservative["candidate"] == best_candidate].sort_values(
        ["calmar", "annualized_return", "max_drawdown"],
        ascending=[False, False, True],
    )
    verdict_row = best_conservative.iloc[0]
    worth = (
        verdict_row["annualized_return"] > 0.0
        and verdict_row["max_drawdown"] <= 0.35
        and verdict_row["calmar"] >= 1.0
        and verdict_row["h1y_return"] > 0.0
    )
    verdict = "worth continuing" if worth else "not worth continuing under conservative cost"
    columns = [
        "candidate",
        "strategy",
        "fee_rate",
        "slippage_rate",
        "funding_rate_8h",
        "total_return",
        "annualized_return",
        "max_drawdown",
        "calmar",
        "h3y_return",
        "h1y_return",
        "h6m_return",
        "h3m_return",
    ]
    return "\n".join(
        [
            "# Rotation risk stress test",
            "",
            f"Run command: `{command}`",
            "",
            "Output files:",
            *[f"- `{path}`" for path in paths],
            "",
            "Cost grid: one-way fee 0.04/0.06/0.08/0.10%, extra slippage 0/0.02/0.05/0.10%, funding per 8h -0.005/0/+0.005%.",
            "Funding model: positive funding is paid by long exposure; negative funding is received by long exposure.",
            "Robustness grid: exposure 0.50/0.65/0.80 and vol_target 0/0.20/0.30 on each front candidate signal.",
            "",
            f"Verdict: original best candidate `{best_candidate}` is `{verdict}`. Conservative case annualized return is {verdict_row['annualized_return']:.2%}, max drawdown is {verdict_row['max_drawdown']:.2%}, Calmar is {verdict_row['calmar']:.2f}.",
            "",
            "## Best Conservative Rows",
            "",
            markdown_table(conservative.head(10)[columns]),
            "",
            "## Best Candidate Conservative Robustness",
            "",
            markdown_table(best_conservative.head(10)[columns]),
            "",
        ]
    )


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--years", type=float, default=8.0)
    parser.add_argument("--output-dir", type=Path, default=OUTPUT_DIR)
    parser.add_argument("--candidate-limit", type=int, default=10)
    parser.add_argument("--top", type=int, default=30)
    args = parser.parse_args()

    frames = rotation.load_symbol_bar_frames(args.years)
    rows: list[dict[str, object]] = []
    horizon_rows: list[dict[str, object]] = []
    candidates = source_candidates(args.candidate_limit)
    total_runs = len(candidates) * len(EXPOSURES) * len(VOL_TARGETS) * len(FEE_RATES) * len(SLIPPAGE_RATES) * len(FUNDING_RATES_8H)
    run_index = 0

    for candidate in candidates:
        closes = rotation.aligned_closes(frames, candidate.base)
        signal_weights = rotation.target_weights(closes, candidate.base)
        for risk in robustness_variants(candidate):
            weights = refine.apply_risk_controls(closes, signal_weights, risk)
            for fee_rate, slippage_rate, funding_rate in product(FEE_RATES, SLIPPAGE_RATES, FUNDING_RATES_8H):
                run_index += 1
                stress = StressParams(risk=risk, fee_rate=fee_rate, slippage_rate=slippage_rate, funding_rate_8h=funding_rate)
                equity = equity_curve(closes, weights, stress)
                rows.append(row_for_result(candidate.name, closes, weights, stress, equity))
                horizon_rows.extend(rotation.horizon_rows(stress.name, equity))
                if run_index % 1000 == 0:
                    print(f"done {run_index}/{total_runs}")

    total = pd.DataFrame(rows)
    ranked = total.sort_values(
        ["calmar", "annualized_return", "max_drawdown"],
        ascending=[False, False, True],
    )
    top = ranked.head(args.top)
    conservative = conservative_summary(total)
    best_candidate = candidates[0].name
    horizons = pd.DataFrame(horizon_rows)
    horizons = horizons[horizons["strategy"].isin(set(top["strategy"]))]

    args.output_dir.mkdir(parents=True, exist_ok=True)
    total_path = args.output_dir / f"{PREFIX}-total.csv"
    top_path = args.output_dir / f"{PREFIX}-top.csv"
    conservative_path = args.output_dir / f"{PREFIX}-conservative.csv"
    horizon_path = args.output_dir / f"{PREFIX}-horizons.csv"
    report_path = args.output_dir / f"{PREFIX}-report.md"
    paths = [total_path, top_path, conservative_path, horizon_path, report_path]

    total.to_csv(total_path, index=False)
    top.to_csv(top_path, index=False)
    conservative.to_csv(conservative_path, index=False)
    horizons.to_csv(horizon_path, index=False)
    command = (
        "rtk .venv/bin/python scripts/stress_expansion_rotation_risk.py"
        f" --years {args.years} --candidate-limit {args.candidate_limit} --top {args.top}"
    )
    report_path.write_text(report_text(command, paths, total, best_candidate), encoding="utf-8")
    print(conservative.head(10).to_string(index=False))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())