okx-codex-trader/okx_codex_trader/research_metrics.py

from __future__ import annotations

import pandas as pd

from okx_codex_trader.sampled_report import SegmentResult


DEFAULT_INITIAL_EQUITY = 10_000.0
DEFAULT_PRIMARY_COST = "maker_taker"
DEFAULT_COSTS = (
    ("maker_maker", 0.0012),
    ("maker_taker", 0.0021),
    ("taker_taker", 0.0030),
)
DEFAULT_HORIZONS = (
    ("3y", pd.DateOffset(years=3)),
    ("1y", pd.DateOffset(years=1)),
    ("6m", pd.DateOffset(months=6)),
    ("3m", pd.DateOffset(months=3)),
    ("30d", pd.DateOffset(days=30)),
)


def format_utc_ts(ts: int) -> str:
    return pd.to_datetime(ts, unit="ms", utc=True).strftime("%Y-%m-%d %H:%M")


def cost_equity_frame(result: SegmentResult, cost: float, initial_equity: float = DEFAULT_INITIAL_EQUITY) -> pd.DataFrame:
    rows = [{"ts": pd.to_datetime(result.equity_curve[0]["ts"], unit="ms", utc=True), "equity": initial_equity}]
    equity = initial_equity
    for trade in result.trades:
        equity *= 1.0 + float(trade["return_pct"]) / 100.0 - cost * float(trade.get("cost_weight", 1.0))
        rows.append({"ts": pd.to_datetime(str(trade["exit_time"]), utc=True), "equity": equity})
    return pd.DataFrame(rows)


def max_drawdown(values: list[float]) -> float:
    peak = values[0]
    dd = 0.0
    for value in values:
        peak = max(peak, value)
        dd = max(dd, (peak - value) / peak if peak else 0.0)
    return dd


def equity_metrics(frame: pd.DataFrame, first_ts: int, last_ts: int) -> dict[str, float]:
    years = (last_ts - first_ts) / 86_400_000 / 365
    total_return = float(frame["equity"].iloc[-1] / frame["equity"].iloc[0] - 1.0)
    annualized = (1.0 + total_return) ** (1.0 / years) - 1.0 if total_return > -1.0 and years > 0.0 else 0.0
    dd = max_drawdown([float(value) for value in frame["equity"]])
    return {
        "net_total_return": total_return,
        "net_annualized_return": annualized,
        "net_max_drawdown": dd,
        "net_calmar": annualized / dd if dd else 0.0,
    }


def horizon_rows(
    frame: pd.DataFrame,
    last_ts: int,
    horizons: tuple[tuple[str, pd.DateOffset], ...] = DEFAULT_HORIZONS,
) -> list[dict[str, object]]:
    rows: list[dict[str, object]] = []
    end_time = pd.to_datetime(last_ts, unit="ms", utc=True)
    for label, offset in horizons:
        cutoff = end_time - offset
        before = frame[frame["ts"] <= cutoff]
        if len(before):
            start_equity = float(before["equity"].iloc[-1])
            start_time = cutoff
            after = frame[frame["ts"] > cutoff]
            horizon_frame = pd.concat([pd.DataFrame([{"ts": cutoff, "equity": start_equity}]), after[["ts", "equity"]]], ignore_index=True)
        else:
            horizon_frame = frame[["ts", "equity"]].copy()
            start_time = pd.Timestamp(horizon_frame["ts"].iloc[0])
        rows.append(
            {
                "horizon": label,
                "horizon_start": start_time.strftime("%Y-%m-%d %H:%M"),
                "horizon_end": end_time.strftime("%Y-%m-%d %H:%M"),
                **equity_metrics(horizon_frame, int(start_time.timestamp() * 1000), last_ts),
            }
        )
    return rows


def trade_stats(trades: list[dict[str, object]]) -> dict[str, float]:
    if not trades:
        return {"avg_return_pct": 0.0, "payoff_ratio": 0.0, "profit_factor": 0.0}
    returns = [float(trade["return_pct"]) for trade in trades]
    wins = [value for value in returns if value > 0.0]
    losses = [-value for value in returns if value < 0.0]
    return {
        "avg_return_pct": sum(returns) / len(returns),
        "payoff_ratio": (sum(wins) / len(wins)) / (sum(losses) / len(losses)) if wins and losses else 0.0,
        "profit_factor": sum(wins) / sum(losses) if losses else 0.0,
    }


def worst_month(frame: pd.DataFrame) -> tuple[str, float]:
    monthly = frame.set_index("ts")["equity"].resample("ME").last().ffill().pct_change().dropna()
    if not len(monthly):
        return "", 0.0
    idx = monthly.idxmin()
    return idx.strftime("%Y-%m"), float(monthly.loc[idx])