okx-codex-trader/scripts/search_short_overlay_mix.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
from scripts.search_short_bias_overlay import markdown_table


OUTPUT_DIR = Path("reports/short-bias")
PREFIX = "overlay-mix"
INITIAL_EQUITY = 10_000.0
TAKER_FEE = 0.0004
SYMBOLS = ("BTC-USDT-SWAP", "ETH-USDT-SWAP", "SOL-USDT-SWAP")
SHORTABLE = ("ETH-USDT-SWAP", "SOL-USDT-SWAP")
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 BtcRiskShort:
    family: str
    bar: str
    btc_trend: int
    btc_lookback: int
    symbol_trend: int
    vol_lookback: int
    btc_max_momentum: float
    btc_min_drop: float
    min_btc_vol: float
    symbol_max_momentum: float
    short_symbols: tuple[str, ...]

    @property
    def name(self) -> str:
        suffix = "ethsol" if self.short_symbols == SHORTABLE else "eth"
        return (
            f"{self.family}-{self.bar}-{suffix}-bt{self.btc_trend}"
            f"-bl{self.btc_lookback}-st{self.symbol_trend}-vw{self.vol_lookback}"
            f"-bm{self.btc_max_momentum:.3f}-bd{self.btc_min_drop:.3f}"
            f"-bv{self.min_btc_vol:.4f}-sm{self.symbol_max_momentum:.3f}"
        )


def build_short_params() -> list[BtcRiskShort]:
    params: list[BtcRiskShort] = []
    for bar in ("1h", "4h"):
        if bar == "1h":
            btc_trends = (24 * 60, 24 * 120)
            btc_lookbacks = (24 * 7, 24 * 14)
            symbol_trends = (24 * 30, 24 * 60)
            vol_lookbacks = (24 * 14,)
            min_vols = (0.012, 0.018, 0.024)
            drops = (0.025, 0.040, 0.060)
        else:
            btc_trends = (6 * 60, 6 * 120)
            btc_lookbacks = (6 * 7, 6 * 14)
            symbol_trends = (6 * 30, 6 * 60)
            vol_lookbacks = (6 * 14,)
            min_vols = (0.020, 0.030, 0.040)
            drops = (0.030, 0.050, 0.070)
        for family, btc_trend, btc_lookback, symbol_trend, vol_lookback, btc_max_momentum, drop, min_vol, symbol_max_momentum, short_symbols in product(
            ("btc_risk_pair", "btc_breakdown_symbol"),
            btc_trends,
            btc_lookbacks,
            symbol_trends,
            vol_lookbacks,
            (-0.005, 0.000),
            drops,
            min_vols,
            (-0.010, 0.000, 0.020),
            (("ETH-USDT-SWAP",), SHORTABLE),
        ):
            params.append(
                BtcRiskShort(
                    family=family,
                    bar=bar,
                    btc_trend=btc_trend,
                    btc_lookback=btc_lookback,
                    symbol_trend=symbol_trend,
                    vol_lookback=vol_lookback,
                    btc_max_momentum=btc_max_momentum,
                    btc_min_drop=drop,
                    min_btc_vol=min_vol,
                    symbol_max_momentum=symbol_max_momentum,
                    short_symbols=short_symbols,
                )
            )
    return params


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 daily(series: pd.Series) -> pd.Series:
    index = pd.date_range(series.index[0].normalize(), series.index[-1].normalize(), freq="1D", tz="UTC")
    return series.resample("1D").last().reindex(index).ffill()


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):
        rolled = (1.0 + returns).rolling(window).apply(lambda values: float(values.prod() - 1.0), raw=True)
        rolled = rolled.dropna()
        if rolled.empty:
            continue
        end_month = rolled.idxmin()
        rows.append(
            {
                "name": name,
                "kind": kind,
                "window_months": window,
                "end_month": end_month,
                "return": float(rolled.loc[end_month]),
            }
        )
    return rows


def rotation_base(years: float) -> tuple[str, pd.Series, dict[str, float | int]]:
    row = pd.read_csv(rotation_risk.OUTPUT_DIR / "rotation-risk-top.csv").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"]),
    )
    rotation.SYMBOLS = ("BTC-USDT-SWAP", "ETH-USDT-SWAP")
    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)
    equity = rotation_risk.equity_curve(closes, weights, params)
    stats = rotation_risk.trade_stats(weights, closes, params.leverage)
    return str(row["strategy"]), equity, stats


def load_frames(years: float) -> dict[tuple[str, str], pd.DataFrame]:
    rotation.SYMBOLS = SYMBOLS
    return rotation.load_symbol_bar_frames(years)


def short_weights(closes: pd.DataFrame, params: BtcRiskShort) -> pd.DataFrame:
    returns = closes.pct_change()
    btc = closes["BTC-USDT-SWAP"]
    btc_trend = btc.rolling(params.btc_trend).mean()
    btc_momentum = btc / btc.shift(params.btc_lookback) - 1.0
    btc_vol = returns["BTC-USDT-SWAP"].rolling(params.vol_lookback).std(ddof=1) * (365 * {"1h": 24, "4h": 6}[params.bar]) ** 0.5
    risk_state = (
        (btc < btc_trend)
        & (btc_momentum <= params.btc_max_momentum)
        & (btc_momentum <= -params.btc_min_drop)
        & (btc_vol >= params.min_btc_vol)
    )
    symbol_trend = closes.rolling(params.symbol_trend).mean()
    symbol_momentum = closes / closes.shift(params.btc_lookback) - 1.0
    eligible = (closes < symbol_trend) & (symbol_momentum <= params.symbol_max_momentum)
    weights = pd.DataFrame(0.0, index=closes.index, columns=closes.columns)

    if params.family == "btc_risk_pair":
        active_symbols = [symbol for symbol in params.short_symbols if symbol in closes.columns]
        if active_symbols:
            weights.loc[risk_state, active_symbols] = -1.0 / len(active_symbols)
    elif params.family == "btc_breakdown_symbol":
        active_symbols = [symbol for symbol in params.short_symbols if symbol in closes.columns]
        active = eligible[active_symbols].where(risk_state, False)
        counts = active.sum(axis=1)
        weights.loc[:, active_symbols] = active.astype(float).div(counts.where(counts > 0.0), axis=0).fillna(0.0) * -1.0
    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 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())
    wins = 0
    trades = 0
    gross_profit = 0.0
    gross_loss = 0.0
    for symbol in SHORTABLE:
        if symbol not in closes.columns:
            continue
        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 combine(base: pd.Series, short: pd.Series, allocation: float) -> pd.Series:
    returns = pd.DataFrame({"base": base.pct_change().fillna(0.0), "short": short.pct_change().fillna(0.0)}).dropna()
    combined = (1.0 - allocation) * returns["base"] + allocation * returns["short"]
    equity = INITIAL_EQUITY * (1.0 + combined).cumprod()
    equity.name = "equity"
    return equity


def recent_not_worse(combo: pd.Series, base: pd.Series) -> bool:
    combo_returns = horizon_return_fields(combo)
    base_returns = horizon_return_fields(base)
    return all(combo_returns[key] >= base_returns[key] for key in ("h1y_return", "h6m_return", "h3m_return"))


def report_text(
    command: str,
    paths: list[Path],
    totals: pd.DataFrame,
    horizons: pd.DataFrame,
    monthly: pd.DataFrame,
    worst: pd.DataFrame,
    signal_defs: pd.DataFrame,
) -> str:
    selected = totals[totals["selected"] == "yes"]
    server = selected[selected["server_readonly_candidate"] == "yes"]
    return "\n".join(
        [
            "# Short Overlay Mix",
            "",
            f"Run command: `{command}`",
            "",
            "Output files:",
            *[f"- `{path}`" for path in paths],
            "",
            "Objective: add a small short overlay to the existing rotation long strategy and keep the recent 1y/6m/3m total returns no worse than the rotation base while reducing max drawdown.",
            "Cost model: 0.04% one-way taker fee on absolute short-overlay notional turnover. Rotation-risk keeps its existing cost model.",
            "",
            "## Result",
            "",
            markdown_table(selected),
            "",
            "## Server read-only signal candidates",
            "",
            markdown_table(server),
            "",
            "## Signal definitions",
            "",
            markdown_table(signal_defs),
            "",
            "## Horizons",
            "",
            markdown_table(horizons[horizons["name"].isin(set(selected["name"]))]),
            "",
            "## Worst rolling 6/12m",
            "",
            markdown_table(worst[worst["name"].isin(set(selected["name"]))]),
            "",
            "## Recent monthly returns",
            "",
            markdown_table(monthly[monthly["name"].isin(set(selected["name"]))].tail(96)),
            "",
        ]
    )


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

    base_name, base_raw, base_stats = rotation_base(args.years)
    base_daily = daily(base_raw)
    frames = load_frames(args.years)

    totals: list[dict[str, object]] = [
        {
            "name": base_name,
            "kind": "base_rotation",
            "base": "",
            "short_leg": "",
            "allocation": 0.0,
            "server_readonly_candidate": "no",
            "selected": "yes",
            "years": (base_daily.index[-1] - base_daily.index[0]).total_seconds() / 86_400 / 365,
            **horizon_return_fields(base_daily),
            **metrics(base_daily),
            **base_stats,
        }
    ]
    horizon_output = horizon_rows(base_name, "base_rotation", base_daily)
    monthly_frames = [monthly_rows(base_name, "base_rotation", base_daily)]
    worst_rows = worst_rolling_rows(base_name, "base_rotation", monthly_frames[0])

    short_rows: list[dict[str, object]] = []
    short_equities: dict[str, pd.Series] = {}
    short_stats: dict[str, dict[str, float | int]] = {}
    params_by_name: dict[str, BtcRiskShort] = {}
    signal_defs: list[dict[str, object]] = []

    for param in build_short_params():
        params_by_name[param.name] = param
        required_symbols = ("BTC-USDT-SWAP", *param.short_symbols)
        closes = pd.DataFrame({symbol: frames[(symbol, param.bar)]["close"] for symbol in required_symbols}).dropna()
        weights = short_weights(closes, param)
        equity = equity_from_weights(closes, weights)
        equity_daily = daily(equity)
        short_equities[param.name] = equity_daily
        horizon = horizon_return_fields(equity_daily)
        years = (equity.index[-1] - equity.index[0]).total_seconds() / 86_400 / 365
        active_days = int((weights[list(set(param.short_symbols) & set(weights.columns))].abs().sum(axis=1) > 0.0).sum())
        short_rows.append(
            {
                "name": param.name,
                "kind": "short_leg",
                "family": param.family,
                "base": "",
                "short_leg": "",
                "allocation": 1.0,
                "bar": param.bar,
                "short_symbols": ",".join(param.short_symbols),
                "server_readonly_candidate": "no",
                "selected": "no",
                "years": years,
                "turnover_per_year": float(weights.diff().abs().sum(axis=1).sum() / years),
                "active_bars": active_days,
                **horizon,
                **metrics(equity_daily),
                "trades": 0,
                "win_rate": 0.0,
                "profit_factor": 0.0,
            }
        )
        signal_defs.append(
            {
                "short_leg": param.name,
                "definition": (
                    f"On {param.bar}, BTC close below SMA({param.btc_trend}), "
                    f"BTC {param.btc_lookback}-bar momentum <= {param.btc_max_momentum:.3f} "
                    f"and <= -{param.btc_min_drop:.3f}, annualized BTC vol({param.vol_lookback}) >= {param.min_btc_vol:.4f}; "
                    f"short {','.join(param.short_symbols)}"
                    + (
                        f" only when each symbol is below SMA({param.symbol_trend}) and momentum <= {param.symbol_max_momentum:.3f}."
                        if param.family == "btc_breakdown_symbol"
                        else "."
                    )
                ),
            }
        )

    short_frame = pd.DataFrame(short_rows).sort_values(
        ["h1y_return", "h6m_return", "h3m_return", "calmar"],
        ascending=[False, False, False, False],
    )
    candidates = short_frame.head(args.top)
    totals.extend(short_rows)
    for _, short_row in candidates.iterrows():
        short_name = str(short_row["name"])
        param = params_by_name[short_name]
        required_symbols = ("BTC-USDT-SWAP", *param.short_symbols)
        closes = pd.DataFrame({symbol: frames[(symbol, param.bar)]["close"] for symbol in required_symbols}).dropna()
        stats = trade_stats(short_weights(closes, param), closes)
        short_stats[short_name] = stats
        short = short_equities[short_name]
        horizon_output.extend(horizon_rows(short_name, "short_leg", short))
        short_monthly = monthly_rows(short_name, "short_leg", short)
        monthly_frames.append(short_monthly)
        worst_rows.extend(worst_rolling_rows(short_name, "short_leg", short_monthly))

        for allocation in (0.03, 0.05, 0.07, 0.10):
            combo_name = f"{base_name}+{allocation:.0%}-{short_name}"
            combo = combine(base_daily, short, allocation)
            combo_metrics = metrics(combo)
            base_metrics = metrics(base_daily.loc[combo.index])
            qualifies = combo_metrics["max_drawdown"] < base_metrics["max_drawdown"] and recent_not_worse(combo, base_daily.loc[combo.index])
            stats = short_stats[short_name]
            totals.append(
                {
                    "name": combo_name,
                    "kind": "rotation_plus_btc_risk_short",
                    "family": str(short_row["family"]),
                    "base": base_name,
                    "short_leg": short_name,
                    "allocation": allocation,
                    "bar": str(short_row["bar"]),
                    "short_symbols": str(short_row["short_symbols"]),
                    "server_readonly_candidate": "yes" if qualifies else "no",
                    "selected": "yes" if qualifies else "no",
                    "years": (combo.index[-1] - combo.index[0]).total_seconds() / 86_400 / 365,
                    "annualized_return_delta_vs_base": combo_metrics["annualized_return"] - base_metrics["annualized_return"],
                    "max_drawdown_delta_vs_base": combo_metrics["max_drawdown"] - base_metrics["max_drawdown"],
                    **horizon_return_fields(combo),
                    **combo_metrics,
                    "trades": int(base_stats["trades"]) + int(stats["trades"]),
                    "win_rate": stats["win_rate"],
                    "profit_factor": stats["profit_factor"],
                    "overlay_trades": stats["trades"],
                    "overlay_win_rate": stats["win_rate"],
                    "overlay_profit_factor": stats["profit_factor"],
                }
            )
            horizon_output.extend(horizon_rows(combo_name, "rotation_plus_btc_risk_short", combo))
            combo_monthly = monthly_rows(combo_name, "rotation_plus_btc_risk_short", combo)
            monthly_frames.append(combo_monthly)
            worst_rows.extend(worst_rolling_rows(combo_name, "rotation_plus_btc_risk_short", combo_monthly))

    total = pd.DataFrame(totals)
    combos = total[total["kind"] == "rotation_plus_btc_risk_short"].copy()
    total.loc[total["kind"] != "base_rotation", "selected"] = "no"
    qualified = combos[combos["server_readonly_candidate"] == "yes"].sort_values(
        ["max_drawdown", "calmar", "h1y_return", "h6m_return", "h3m_return"],
        ascending=[True, False, False, False, False],
    )
    selected_combos = qualified.drop_duplicates(
        subset=[
            "allocation",
            "family",
            "bar",
            "short_symbols",
            "h1y_return",
            "h6m_return",
            "h3m_return",
            "max_drawdown",
        ]
    ).head(5)
    if selected_combos.empty:
        selected_combos = combos.sort_values(
            ["max_drawdown", "h1y_return", "h6m_return", "h3m_return"],
            ascending=[True, False, False, False],
        ).head(3)
    total.loc[total["name"].isin(set(selected_combos["name"])), "selected"] = "yes"
    selected_short = set(total.loc[total["selected"] == "yes", "short_leg"]) - {""}
    total.loc[total["name"].isin(selected_short), "selected"] = "yes"
    for short_name in selected_short:
        if short_name in short_stats:
            for key, value in short_stats[short_name].items():
                total.loc[total["name"] == short_name, key] = value

    horizons = pd.DataFrame(horizon_output)
    monthly = pd.concat(monthly_frames, ignore_index=True)
    worst = pd.DataFrame(worst_rows)
    signal_defs_frame = pd.DataFrame(signal_defs)
    signal_defs_frame = signal_defs_frame[signal_defs_frame["short_leg"].isin(selected_short)]

    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}-combo.csv"
    horizon_path = args.output_dir / f"{PREFIX}-horizons.csv"
    monthly_path = args.output_dir / f"{PREFIX}-monthly.csv"
    worst_path = args.output_dir / f"{PREFIX}-worst-rolling.csv"
    signal_path = args.output_dir / f"{PREFIX}-signal-definitions.csv"
    report_path = args.output_dir / f"{PREFIX}-report.md"
    paths = [total_path, short_path, combo_path, horizon_path, monthly_path, worst_path, signal_path, report_path]

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


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