trade_pool/__init__.py

"""trade_pool — a verifiers RL environment where a coding agent writes a
causal crypto trading strategy, scored by an out-of-sample backtest.

The agent receives a task (a token's in-sample price history + the feature surface),
writes a Python `strategy(features, position) -> target_position` function, and is
rewarded by how that strategy performs on a HELD-OUT window — must beat buy-and-hold,
control drawdown, and keep sane exposure. tradewatch's MEMORY.md discipline becomes
the rubric; RL turns it into adapter weights.
"""

from __future__ import annotations

import json

import numpy as np
import verifiers as vf
from datasets import Dataset

from .backtester import benchmarks, run_backtest
from .data import build_tasks
from .executor import compile_strategy, extract_code
from .seed_principles import principles_block

FEATURE_NAMES = [
    "close", "sma_10", "sma_20", "sma_50", "ema_12", "ema_26", "rsi_14",
    "macd", "macd_signal", "macd_hist", "zscore_20", "bb_lo", "bb_mid", "bb_hi",
    "ret_1", "vol_20",
]

SYSTEM_PROMPT = """You are a quantitative trading strategy engineer. You write a Python \
function that decides a target position each bar, using ONLY causal technical features.

You must output a single Python code block defining exactly:

```python
def strategy(features, position):
    # features: dict of the current bar's indicators (see list below), all causal
    # position: your current target position, a float in [-1, 1]
    # return: new target position, float in [-1, 1]  (1=full long, -1=full short, 0=flat)
    ...
    return new_position
```

Available features (floats, computed on data up to the current bar only):
{features}

Hard rules:
- No imports, no IO, no future data. You only get the current bar's features.
- Favor capital protection: control drawdown, avoid being all-in blindly.
- Your strategy is scored on a HELD-OUT window and must beat buy-and-hold.

{seed_principles}
"""


def _task_to_prompt(task: dict) -> str:
    close = task["train_close"]
    tail = close[-30:] if len(close) > 30 else close
    return (
        f"Token: {task['symbol']}\n"
        f"In-sample bars available: {len(close)}\n"
        f"Recent closes (last {len(tail)}): {[round(c, 6) for c in tail]}\n\n"
        "Write a `strategy(features, position)` function. Use indicators like RSI "
        "(rsi_14), MACD (macd/macd_signal/macd_hist), moving averages (sma_*/ema_*), "
        "z-score (zscore_20), Bollinger bands (bb_*), and volatility (vol_20). "
        "Return a target position in [-1, 1]."
    )


# ── Rubric reward functions (each receives completion + the task info via state) ──

# Memoize per (symbol, code) so the 6 reward fns trigger ONE backtest set per rollout,
# not 6×(1 strategy + 3 benchmarks). Bounded so it can't grow unboundedly across a run.
_SCORE_CACHE: dict[tuple, dict] = {}
_PARTICIPATION_FLOOR = 0.03   # avg |exposure| below this = "did nothing", risk rewards void


def _score_strategy(completion, info) -> dict:
    """Compile + backtest the agent's strategy on the OOS window. Memoized."""
    text = completion if isinstance(completion, str) else completion[-1]["content"]
    code = extract_code(text)
    key = (info["symbol"], hash(code))
    if key in _SCORE_CACHE:
        return _SCORE_CACHE[key]

    fn, err = compile_strategy(code)
    oos = np.asarray(info["oos_close"], dtype=float)
    warmup = min(50, len(oos) // 3)
    if fn is None:
        out = {"error": err, "res": None, "bench": None}
    else:
        res = run_backtest(oos, fn, warmup=warmup)
        bench = benchmarks(oos, warmup=warmup)
        out = {"error": res.error, "res": res, "bench": bench}
    if len(_SCORE_CACHE) > 4096:
        _SCORE_CACHE.clear()
    _SCORE_CACHE[key] = out
    return out


def _participates(res) -> bool:
    return res is not None and res.avg_exposure >= _PARTICIPATION_FLOOR


def reward_sharpe(completion, info, **kwargs) -> float:
    s = _score_strategy(completion, info)
    if s["res"] is None or s["error"]:
        return 0.0
    # squash Sharpe to [0,1] via logistic; Sharpe 0 -> 0.5, 2 -> ~0.88.
    # A flat strategy gets sharpe 0 -> 0.5; gate it so do-nothing can't bank 0.5.
    if not _participates(s["res"]):
        return 0.0
    return float(1.0 / (1.0 + np.exp(-s["res"].sharpe)))


def reward_beats_benchmark(completion, info, **kwargs) -> float:
    s = _score_strategy(completion, info)
    if s["res"] is None or s["error"] or s["bench"] is None or not _participates(s["res"]):
        return 0.0
    bh = s["bench"]["buy_and_hold"].sharpe
    return 1.0 if s["res"].sharpe > bh else 0.0


def reward_drawdown(completion, info, **kwargs) -> float:
    # Drawdown reward only counts if the strategy actually traded — otherwise
    # "do nothing" banks a perfect 1.0 for taking zero risk (the inactivity exploit).
    s = _score_strategy(completion, info)
    if s["res"] is None or s["error"] or not _participates(s["res"]):
        return 0.0
    return float(max(0.0, 1.0 - s["res"].max_drawdown))


def reward_exposure(completion, info, **kwargs) -> float:
    s = _score_strategy(completion, info)
    if s["res"] is None or s["error"] or not _participates(s["res"]):
        return 0.0
    e = s["res"].avg_exposure
    return 1.0 if 0.1 <= e <= 0.85 else max(0.0, 1.0 - abs(e - 0.5))


def reward_cost(completion, info, **kwargs) -> float:
    # Likewise gated: zero turnover (no trades) must not earn the low-cost reward.
    s = _score_strategy(completion, info)
    if s["res"] is None or s["error"] or not _participates(s["res"]):
        return 0.0
    return float(max(0.0, 1.0 - s["res"].turnover / 20.0))


def reward_valid(completion, info, **kwargs) -> float:
    # Valid = compiles AND actually trades. A syntactically-valid do-nothing is not
    # a valid trading strategy for our purposes.
    s = _score_strategy(completion, info)
    if s["res"] is None or s["error"]:
        return 0.0
    return 1.0 if _participates(s["res"]) else 0.0


# Objective presets = rubric weight vectors. The recursive self-improving loop selects
# an objective (or passes explicit reward_weights) per iteration, so curriculum can shift
# emphasis (e.g. toward drawdown control if the trained model is taking too much risk)
# WITHOUT rebuilding the wheel — these arrive via the TOML [[env]] args block.
#                  [sharpe, beats_bh, drawdown, exposure, cost, valid]
OBJECTIVE_WEIGHTS = {
    "sharpe":       [0.40, 0.20, 0.15, 0.10, 0.05, 0.10],
    "min_drawdown": [0.20, 0.15, 0.35, 0.15, 0.05, 0.10],
    "balanced":     [0.30, 0.25, 0.20, 0.10, 0.05, 0.10],
}


def load_environment(
    split: str = "train",
    seed: int = 0,
    objective: str = "sharpe",
    reward_weights: list[float] | None = None,
    symbols: list[str] | None = None,
    use_seed_principles: bool = True,
    max_turns: int = 1,
) -> vf.Environment:
    """Entry point Prime calls (args come from the TOML [[env]] args block).

    Curriculum knobs for the recursive loop:
      split          train | oos | oos_symbols   (which symbol pool)
      seed           rotates the symbol shuffle  (exposes new task mixes per iteration)
      objective      sharpe | min_drawdown | balanced   (rubric weight preset)
      reward_weights explicit 6-vector, overrides objective
      symbols        restrict to these symbols (curriculum: focus on weak performers)
      use_seed_principles  inject tradewatch's 618-decision discipline block (default on);
                     set False for the "memory is the adapter" ablation (strip the prompt
                     discipline and test whether trained weights retain it)
    """
    tasks = build_tasks(split=split, seed=seed)
    if symbols:
        wanted = {s.upper() for s in symbols}
        tasks = [t for t in tasks if t["symbol"].upper() in wanted] or tasks
    seed_block = principles_block() if use_seed_principles else ""
    system = SYSTEM_PROMPT.format(features=", ".join(FEATURE_NAMES), seed_principles=seed_block)
    rows = []
    for t in tasks:
        rows.append({
            "prompt": [
                {"role": "system", "content": system},
                {"role": "user", "content": _task_to_prompt(t)},
            ],
            "info": {"symbol": t["symbol"], "oos_close": t["oos_close"]},
        })
    dataset = Dataset.from_list(rows)

    weights = reward_weights or OBJECTIVE_WEIGHTS.get(objective, OBJECTIVE_WEIGHTS["sharpe"])
    rubric = vf.Rubric(
        funcs=[reward_sharpe, reward_beats_benchmark, reward_drawdown,
               reward_exposure, reward_cost, reward_valid],
        weights=weights,
    )
    return vf.SingleTurnEnv(dataset=dataset, rubric=rubric)