okx-codex-trader/scripts/search_short_bias_overlay.py

from __future__ import annotations

import argparse
import sys
from dataclasses import dataclass
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 rotation_risk
from scripts import search_expansion_rotation as rotation


OUTPUT_DIR = Path("reports/short-bias")
PREFIX = "overlay"
SYMBOLS = ("BTC-USDT-SWAP", "ETH-USDT-SWAP")
INITIAL_EQUITY = 10_000.0
TAKER_FEE = 0.0004
HORIZONS = (
    ("full", None),
    ("3y", pd.DateOffset(years=3)),
    ("1y", pd.DateOffset(years=1)),
    ("6m", pd.DateOffset(months=6)),
    ("3m", pd.DateOffset(months=3)),
)


@dataclass(frozen=True)
class ShortParams:
    family: str
    bar: str
    trend: int
    lookback: int
    vol_lookback: int
    min_vol: float
    max_momentum: float
    drop: float

    @property
    def name(self) -> str:
        return (
            f"{self.family}-{self.bar}-tr{self.trend}-lb{self.lookback}"
            f"-vw{self.vol_lookback}-mv{self.min_vol:.4f}"
            f"-mm{self.max_momentum:.3f}-dp{self.drop:.3f}"
        )


def build_short_params() -> list[ShortParams]:
    params: list[ShortParams] = []
    for bar in ("1h", "4h"):
        if bar == "1h":
            trends = (24 * 30, 24 * 60)
            lookbacks = (24 * 7, 24 * 14)
            vols = (24 * 14,)
            min_vols = (0.010, 0.014, 0.018)
            drops = (0.025, 0.040)
        else:
            trends = (6 * 30, 6 * 60)
            lookbacks = (6 * 7, 6 * 14)
            vols = (6 * 14,)
            min_vols = (0.020, 0.028, 0.036)
            drops = (0.030, 0.050)
        for family, trend, lookback, vol_lookback, min_vol, max_momentum, drop in product(
            ("risk_off_pair", "trend_breakdown", "high_vol_breakdown"),
            trends,
            lookbacks,
            vols,
            min_vols,
            (-0.005, 0.000),
            drops,
        ):
            params.append(ShortParams(family, bar, trend, lookback, vol_lookback, min_vol, max_momentum, drop))
    return params


def short_weights(closes: pd.DataFrame, params: ShortParams) -> pd.DataFrame:
    returns = closes.pct_change()
    momentum = closes / closes.shift(params.lookback) - 1.0
    trend = closes.rolling(params.trend).mean()
    vol = returns.rolling(params.vol_lookback).std(ddof=1) * (365 * {"1h": 24, "4h": 6}[params.bar]) ** 0.5
    btc = "BTC-USDT-SWAP"
    weights = pd.DataFrame(0.0, index=closes.index, columns=closes.columns)

    if params.family == "risk_off_pair":
        state = (
            (closes[btc] < trend[btc])
            & (momentum[btc] <= params.max_momentum)
            & (vol[btc] >= params.min_vol)
        )
        weights.loc[state, list(SYMBOLS)] = -0.5
    elif params.family == "trend_breakdown":
        eligible = (closes < trend) & (momentum <= params.max_momentum) & (vol >= params.min_vol)
        counts = eligible[list(SYMBOLS)].sum(axis=1)
        for symbol in SYMBOLS:
            weights.loc[eligible[symbol] & (counts > 0), symbol] = -1.0 / counts[eligible[symbol] & (counts > 0)]
    elif params.family == "high_vol_breakdown":
        drawdown = closes / closes.shift(params.lookback) - 1.0
        state = (
            (closes[btc] < trend[btc])
            & (drawdown[btc] <= -params.drop)
            & (vol[btc] >= params.min_vol)
        )
        eligible = (closes < trend) & (drawdown <= 0.0)
        counts = eligible[list(SYMBOLS)].sum(axis=1)
        for symbol in SYMBOLS:
            mask = state & eligible[symbol] & (counts > 0)
            weights.loc[mask, symbol] = -1.0 / counts[mask]
    else:
        raise ValueError(f"unknown family {params.family}")

    return weights.fillna(0.0)


def equity_from_weights(closes: pd.DataFrame, weights: pd.DataFrame) -> 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))
    net_returns = (executed * returns).sum(axis=1) - turnover * TAKER_FEE
    equity = INITIAL_EQUITY * (1.0 + net_returns).cumprod()
    equity.name = "equity"
    return equity


def metrics(series: pd.Series) -> dict[str, float]:
    years = (series.index[-1] - series.index[0]).total_seconds() / 86_400 / 365
    total = float(series.iloc[-1] / series.iloc[0] - 1.0)
    annualized = (1.0 + total) ** (1.0 / years) - 1.0 if total > -1.0 and years > 0.0 else 0.0
    drawdown = float((series.cummax() - series).div(series.cummax()).max())
    return {
        "total_return": total,
        "annualized_return": annualized,
        "max_drawdown": drawdown,
        "calmar": annualized / drawdown if drawdown else 0.0,
    }


def horizon_rows(name: str, kind: str, series: pd.Series) -> list[dict[str, object]]:
    rows: list[dict[str, object]] = []
    end = series.index[-1]
    for label, offset in HORIZONS:
        horizon = series if offset is None else series[series.index >= end - offset]
        if len(horizon) < 2:
            horizon = series
        rows.append(
            {
                "name": name,
                "kind": kind,
                "horizon": label,
                "start": horizon.index[0].strftime("%Y-%m-%d"),
                "end": horizon.index[-1].strftime("%Y-%m-%d"),
                **metrics(horizon),
            }
        )
    return rows


def horizon_return_fields(series: pd.Series) -> dict[str, float]:
    rows = {row["horizon"]: row for row in horizon_rows("", "", series)}
    return {
        "h3y_return": float(rows["3y"]["total_return"]),
        "h1y_return": float(rows["1y"]["total_return"]),
        "h6m_return": float(rows["6m"]["total_return"]),
        "h3m_return": float(rows["3m"]["total_return"]),
    }


def monthly_rows(name: str, kind: str, series: pd.Series) -> pd.DataFrame:
    monthly = series.resample("ME").last()
    frame = pd.DataFrame(
        {
            "name": name,
            "kind": kind,
            "month": monthly.index.strftime("%Y-%m"),
            "start_equity": monthly.shift(1).fillna(series.iloc[0]).to_numpy(),
            "end_equity": monthly.to_numpy(),
        }
    )
    frame["return"] = frame["end_equity"] / frame["start_equity"] - 1.0
    return frame


def worst_rolling_rows(name: str, kind: str, monthly: pd.DataFrame) -> list[dict[str, object]]:
    rows: list[dict[str, object]] = []
    returns = monthly.set_index("month")["return"]
    for window in (6, 12):
        rolling = (1.0 + returns).rolling(window).apply(lambda values: float(values.prod() - 1.0), raw=True)
        worst_end = rolling.idxmin()
        rows.append(
            {
                "name": name,
                "kind": kind,
                "window_months": window,
                "end_month": worst_end,
                "return": float(rolling.loc[worst_end]),
            }
        )
    return rows


def trade_stats(weights: pd.DataFrame, closes: pd.DataFrame) -> dict[str, float | int]:
    executed = weights.shift(1).fillna(0.0)
    returns = closes.pct_change().fillna(0.0)
    turnover = executed.diff().abs().fillna(executed.abs())
    trades = 0
    wins = 0
    gross_profit = 0.0
    gross_loss = 0.0
    for symbol in SYMBOLS:
        active = executed[symbol] < 0.0
        groups = (active != active.shift(1)).cumsum()
        for _, mask in active.groupby(groups):
            if not bool(mask.iloc[0]):
                continue
            index = mask.index
            net = returns.loc[index, symbol] * executed.loc[index, symbol] - turnover.loc[index, symbol] * TAKER_FEE
            value = float((1.0 + net).prod() - 1.0)
            trades += 1
            if value > 0.0:
                wins += 1
                gross_profit += value
            else:
                gross_loss += abs(value)
    return {
        "trades": trades,
        "win_rate": wins / trades if trades else 0.0,
        "profit_factor": gross_profit / gross_loss if gross_loss else 0.0,
    }


def rotation_risk_equity(years: float) -> tuple[str, pd.Series]:
    top = pd.read_csv(rotation_risk.OUTPUT_DIR / "rotation-risk-top.csv")
    row = top.iloc[0]
    base = rotation_risk.params_from_row(row)
    params = rotation_risk.RiskParams(
        base=base,
        leverage=float(row["leverage"]),
        exposure=float(row["exposure"]),
        vol_target=float(row["vol_target"]),
    )
    frames = rotation.load_symbol_bar_frames(years)
    closes = rotation.aligned_closes(frames, base)
    weights = rotation_risk.apply_risk_controls(closes, rotation.target_weights(closes, base), params)
    return str(row["strategy"]), rotation_risk.equity_curve(closes, weights, params)


def combined_equity(rotation_equity: pd.Series, short_equity: pd.Series, short_alloc: float) -> pd.Series:
    returns = pd.DataFrame(
        {
            "rotation": rotation_equity.pct_change().fillna(0.0),
            "short": short_equity.pct_change().fillna(0.0),
        }
    ).dropna()
    combo_returns = (1.0 - short_alloc) * returns["rotation"] + short_alloc * returns["short"]
    equity = INITIAL_EQUITY * (1.0 + combo_returns).cumprod()
    equity.name = "equity"
    return equity


def format_cell(value: object) -> str:
    if isinstance(value, float):
        return f"{value:.6g}"
    return str(value).replace("|", "\\|")


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(format_cell(value) for value in row) + " |" for row in rows)


def report_text(
    command: str,
    paths: list[Path],
    totals: pd.DataFrame,
    horizons: pd.DataFrame,
    monthly: pd.DataFrame,
    correlations: pd.DataFrame,
    worst: pd.DataFrame,
) -> str:
    selected = totals[totals["selected"] == "yes"].copy()
    best_short = selected[selected["kind"] == "short_leg"].head(1)
    best_combo = selected[selected["kind"] == "rotation_plus_short"].head(4)
    conclusion = "Not worth adding."
    if len(best_combo):
        row = best_combo.sort_values(["max_drawdown", "calmar"], ascending=[True, False]).iloc[0]
        conclusion = (
            "Worth only as a small protection leg"
            if float(row["max_drawdown_delta_vs_rotation"]) < 0.0 and float(row["annualized_return_delta_vs_rotation"]) > -0.03
            else "Not worth adding at the tested small weights"
        )
    focus_names = set(selected["name"])
    return "\n".join(
        [
            "# Short Bias Overlay",
            "",
            f"Run command: `{command}`",
            "",
            "Output files:",
            *[f"- `{path}`" for path in paths],
            "",
            "Scope: BTC/ETH short-only risk-state overlays from cached candles. Cost model is 0.04% one-way taker fee on absolute notional turnover.",
            "Combination test uses rotation-risk as the base and allocates 5%, 10%, or 15% of capital to the short leg.",
            "",
            f"Conclusion: {conclusion}.",
            "",
            "## Selected short leg",
            "",
            markdown_table(best_short),
            "",
            "## Selected rotation-risk combinations",
            "",
            markdown_table(best_combo),
            "",
            "## Horizons",
            "",
            markdown_table(horizons[horizons["name"].isin(focus_names)]),
            "",
            "## Monthly returns",
            "",
            markdown_table(monthly[monthly["name"].isin(focus_names)].tail(80)),
            "",
            "## Correlation",
            "",
            markdown_table(correlations),
            "",
            "## Worst rolling 6/12m",
            "",
            markdown_table(worst[worst["name"].isin(focus_names)]),
            "",
        ]
    )


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("--top", type=int, default=20)
    args = parser.parse_args()

    rotation_name, rotation_equity_raw = rotation_risk_equity(args.years)
    frames = rotation.load_symbol_bar_frames(args.years)

    short_results: list[dict[str, object]] = []
    short_equities: dict[str, pd.Series] = {}
    short_weights_by_name: dict[str, pd.DataFrame] = {}
    for param in build_short_params():
        closes = pd.DataFrame({symbol: frames[(symbol, param.bar)]["close"] for symbol in SYMBOLS}).dropna()
        weights = short_weights(closes, param)
        equity = equity_from_weights(closes, weights)
        short_equities[param.name] = equity
        short_weights_by_name[param.name] = weights
        horizon = {row["horizon"]: row for row in horizon_rows(param.name, "short_leg", equity)}
        years = (equity.index[-1] - equity.index[0]).total_seconds() / 86_400 / 365
        short_results.append(
            {
                "name": param.name,
                "kind": "short_leg",
                "family": param.family,
                "bar": param.bar,
                "allocation": 1.0,
                "years": years,
                "h3y_return": horizon["3y"]["total_return"],
                "h1y_return": horizon["1y"]["total_return"],
                "h6m_return": horizon["6m"]["total_return"],
                "h3m_return": horizon["3m"]["total_return"],
                "turnover_per_year": float(weights.diff().abs().sum(axis=1).sum() / years),
                **metrics(equity),
                **trade_stats(weights, closes),
                **param.__dict__,
            }
        )

    rotation_monthly = monthly_rows(rotation_name, "rotation_risk", rotation_equity_raw)
    rotation_total = {
        "name": rotation_name,
        "kind": "rotation_risk",
        "family": "rotation_risk",
        "bar": "",
        "allocation": 1.0,
        "years": (rotation_equity_raw.index[-1] - rotation_equity_raw.index[0]).total_seconds() / 86_400 / 365,
        **horizon_return_fields(rotation_equity_raw),
        **metrics(rotation_equity_raw),
    }
    rotation_metrics = metrics(rotation_equity_raw)

    ranked_short = pd.DataFrame(short_results).sort_values(
        ["h1y_return", "h6m_return", "h3m_return", "calmar", "total_return"],
        ascending=[False, False, False, False, False],
    )
    ranked_unique = ranked_short.drop_duplicates(
        subset=["total_return", "h3y_return", "h1y_return", "h6m_return", "h3m_return", "max_drawdown"]
    )
    candidates = ranked_unique.head(args.top)

    totals: list[dict[str, object]] = [rotation_total]
    totals.extend(short_results)
    horizon_output = horizon_rows(rotation_name, "rotation_risk", rotation_equity_raw)
    monthly_frames = [rotation_monthly]
    worst_rows = worst_rolling_rows(rotation_name, "rotation_risk", rotation_monthly)
    correlations: list[dict[str, object]] = []

    rotation_daily = rotation_equity_raw.resample("1D").last().ffill()
    for _, short_row in candidates.iterrows():
        short_name = str(short_row["name"])
        short_daily = short_equities[short_name].resample("1D").last().ffill()
        aligned = pd.concat([rotation_daily, short_daily], axis=1, join="inner")
        aligned.columns = ["rotation", "short"]
        correlations.append(
            {
                "name": short_name,
                "kind": "short_leg",
                "correlation_to_rotation_risk": float(aligned.pct_change().dropna().corr().loc["rotation", "short"]),
            }
        )
        horizon_output.extend(horizon_rows(short_name, "short_leg", short_daily))
        short_monthly = monthly_rows(short_name, "short_leg", short_daily)
        monthly_frames.append(short_monthly)
        worst_rows.extend(worst_rolling_rows(short_name, "short_leg", short_monthly))

        for allocation in (0.05, 0.10, 0.15):
            combo_name = f"{rotation_name}+{allocation:.0%}-{short_name}"
            combo = combined_equity(rotation_daily, short_daily, allocation)
            combo_metrics = metrics(combo)
            combo_monthly = monthly_rows(combo_name, "rotation_plus_short", combo)
            totals.append(
                {
                    "name": combo_name,
                    "kind": "rotation_plus_short",
                    "family": str(short_row["family"]),
                    "bar": str(short_row["bar"]),
                    "allocation": allocation,
                    "years": (combo.index[-1] - combo.index[0]).total_seconds() / 86_400 / 365,
                    "short_leg": short_name,
                    "annualized_return_delta_vs_rotation": combo_metrics["annualized_return"] - rotation_metrics["annualized_return"],
                    "max_drawdown_delta_vs_rotation": combo_metrics["max_drawdown"] - rotation_metrics["max_drawdown"],
                    **horizon_return_fields(combo),
                    **combo_metrics,
                }
            )
            horizon_output.extend(horizon_rows(combo_name, "rotation_plus_short", combo))
            monthly_frames.append(combo_monthly)
            worst_rows.extend(worst_rolling_rows(combo_name, "rotation_plus_short", combo_monthly))
            aligned_combo = pd.concat([rotation_daily, combo], axis=1, join="inner")
            aligned_combo.columns = ["rotation", "combo"]
            correlations.append(
                {
                    "name": combo_name,
                    "kind": "rotation_plus_short",
                    "correlation_to_rotation_risk": float(aligned_combo.pct_change().dropna().corr().loc["rotation", "combo"]),
                }
            )

    total = pd.DataFrame(totals)
    combos = total[total["kind"] == "rotation_plus_short"].copy()
    selected_combo = combos[
        (combos["max_drawdown_delta_vs_rotation"] < 0.0)
        & (combos["annualized_return_delta_vs_rotation"] > -0.03)
    ].sort_values(["max_drawdown", "calmar"], ascending=[True, False]).head(4)
    selected_short_names = set(selected_combo["short_leg"]) if len(selected_combo) else set(candidates.head(1)["name"])
    total["selected"] = "no"
    total.loc[total["name"].isin(selected_short_names), "selected"] = "yes"
    total.loc[total["name"].isin(set(selected_combo["name"])), "selected"] = "yes"
    total.loc[total["name"] == rotation_name, "selected"] = "yes"

    horizons = pd.DataFrame(horizon_output)
    monthly = pd.concat(monthly_frames, ignore_index=True)
    correlation_frame = pd.DataFrame(correlations)
    worst = pd.DataFrame(worst_rows)

    args.output_dir.mkdir(parents=True, exist_ok=True)
    total_path = args.output_dir / f"{PREFIX}-total.csv"
    short_path = args.output_dir / f"{PREFIX}-short-leg.csv"
    combo_path = args.output_dir / f"{PREFIX}-rotation-combo.csv"
    horizon_path = args.output_dir / f"{PREFIX}-horizons.csv"
    monthly_path = args.output_dir / f"{PREFIX}-monthly.csv"
    correlation_path = args.output_dir / f"{PREFIX}-correlation.csv"
    worst_path = args.output_dir / f"{PREFIX}-worst-rolling.csv"
    report_path = args.output_dir / f"{PREFIX}-report.md"
    paths = [total_path, short_path, combo_path, horizon_path, monthly_path, correlation_path, worst_path, report_path]

    total.to_csv(total_path, index=False)
    ranked_short.to_csv(short_path, index=False)
    combos.sort_values(["max_drawdown", "calmar"], ascending=[True, False]).to_csv(combo_path, index=False)
    horizons.to_csv(horizon_path, index=False)
    monthly.to_csv(monthly_path, index=False)
    correlation_frame.to_csv(correlation_path, index=False)
    worst.to_csv(worst_path, index=False)
    command = "rtk .venv/bin/python " + " ".join(sys.argv)
    report_path.write_text(report_text(command, paths, total, horizons, monthly, correlation_frame, worst), encoding="utf-8")
    print(total[total["selected"] == "yes"].to_string(index=False))
    return 0


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