fithia2/apps/tools/peer_relative_idle_alpha_probe.py

"""Probe additivity from peer-relative surprise features on top of idle-alpha.

This is a standalone research tool. It does not modify shared runtime code.

The probe compares a baseline experiment config against variants that add
peer-relative surprise filters to the existing post-allocation idle-alpha
earnings engines.

Example:
    python -m apps.tools.peer_relative_idle_alpha_probe \
        --config configs/experiments/return_max_long_v17.1_idle_alpha_extended.json \
        --start 2022-01-01 \
        --end 2026-03-31 \
        --period-start y2026=2026-01-01
"""
from __future__ import annotations

import argparse
import datetime as dt
import json
from collections import Counter, defaultdict
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from apps.backtester.run import BacktestRunner, _build_merged_snapshot_store, load_manifest, resolve_config
from libs.backtest.artifacts import _resolve_trade_sleeve
from libs.common.logging import configure_logging


DEFAULT_CONFIG = "configs/experiments/return_max_long_v17.1_idle_alpha_extended.json"
TARGET_EVENT_TYPE = "earnings_release"
TARGET_SCOPE_BULLISH_ENGINE_ID = "next_open_long_bullish_raised_strong"


@dataclass(frozen=True)
class VariantSpec:
    name: str
    scope: str
    peer_sector_event_count_365d_min: float | None = None
    peer_relative_surprise_pct_365d_min: float | None = None
    peer_relative_sue_hist_mean_4q_365d_min: float | None = None


def _parse_date(value: str, *, is_end: bool = False) -> dt.date:
    parts = value.split("-")
    if len(parts) == 1 and len(value) == 4 and value.isdigit():
        year = int(value)
        return dt.date(year, 12, 31) if is_end else dt.date(year, 1, 1)
    if len(parts) == 2 and all(part.isdigit() for part in parts):
        year = int(parts[0])
        month = int(parts[1])
        if is_end:
            next_month = dt.date(year + (month // 12), (month % 12) + 1, 1)
            return next_month - dt.timedelta(days=1)
        return dt.date(year, month, 1)
    return dt.date.fromisoformat(value)


def _parse_period_start(value: str) -> tuple[str, dt.date]:
    if "=" in value:
        label, raw_date = value.split("=", 1)
    else:
        label, raw_date = value, value
    return label, _parse_date(raw_date)


def _period_stats(curve: list[Any], start_date: dt.date) -> dict[str, Any] | None:
    points = [(state.date, float(state.equity)) for state in curve if state.date >= start_date]
    if not points:
        return None
    start_equity = points[0][1]
    peak = start_equity
    max_dd_pct = 0.0
    for _, equity in points:
        peak = max(peak, equity)
        if peak > 0:
            max_dd_pct = max(max_dd_pct, (peak - equity) / peak * 100.0)
    end_equity = points[-1][1]
    return {
        "start_date": points[0][0].isoformat(),
        "end_date": points[-1][0].isoformat(),
        "return_pct": (end_equity / start_equity - 1.0) * 100.0 if start_equity else 0.0,
        "max_dd_pct": max_dd_pct,
    }


def _clear_peer_filters(engine: Any) -> None:
    engine.peer_sector_event_count_365d_min = None
    engine.peer_sector_event_count_365d_max = None
    engine.peer_relative_surprise_pct_365d_min = None
    engine.peer_relative_surprise_pct_365d_max = None
    engine.peer_relative_sue_hist_mean_4q_365d_min = None
    engine.peer_relative_sue_hist_mean_4q_365d_max = None


def _is_idle_earnings_engine(engine: Any) -> bool:
    event_types = {str(event_type) for event_type in (engine.event_types or [])}
    return bool(getattr(engine, "post_allocation_idle_only", False)) and TARGET_EVENT_TYPE in event_types


def _target_engine_ids(config: Any, scope: str) -> set[str]:
    idle_earnings = [engine.engine_id for engine in config.strategy_engines if _is_idle_earnings_engine(engine)]
    if scope == "bullish_only":
        return {engine_id for engine_id in idle_earnings if engine_id == TARGET_SCOPE_BULLISH_ENGINE_ID}
    return set(idle_earnings)


def _apply_variant(config: Any, variant: VariantSpec) -> None:
    target_engine_ids = _target_engine_ids(config, variant.scope)
    for engine in config.strategy_engines:
        _clear_peer_filters(engine)
        if engine.engine_id not in target_engine_ids:
            continue
        if variant.peer_sector_event_count_365d_min is not None:
            engine.peer_sector_event_count_365d_min = variant.peer_sector_event_count_365d_min
        if variant.peer_relative_surprise_pct_365d_min is not None:
            engine.peer_relative_surprise_pct_365d_min = variant.peer_relative_surprise_pct_365d_min
        if variant.peer_relative_sue_hist_mean_4q_365d_min is not None:
            engine.peer_relative_sue_hist_mean_4q_365d_min = variant.peer_relative_sue_hist_mean_4q_365d_min


def _build_variants() -> list[VariantSpec]:
    variants = [VariantSpec(name="baseline", scope="all_idle_earnings")]

    # Scope 1: apply to both idle earnings engines.
    variants.extend(
        [
            VariantSpec(name="all_count3", scope="all_idle_earnings", peer_sector_event_count_365d_min=3.0),
            VariantSpec(name="all_count4", scope="all_idle_earnings", peer_sector_event_count_365d_min=4.0),
            VariantSpec(name="all_surprise8", scope="all_idle_earnings", peer_relative_surprise_pct_365d_min=8.0),
            VariantSpec(name="all_surprise10", scope="all_idle_earnings", peer_relative_surprise_pct_365d_min=10.0),
            VariantSpec(name="all_surprise12", scope="all_idle_earnings", peer_relative_surprise_pct_365d_min=12.0),
            VariantSpec(
                name="all_count3_surprise8_sue0",
                scope="all_idle_earnings",
                peer_sector_event_count_365d_min=3.0,
                peer_relative_surprise_pct_365d_min=8.0,
                peer_relative_sue_hist_mean_4q_365d_min=0.0,
            ),
            VariantSpec(
                name="all_count4_surprise10_sue0",
                scope="all_idle_earnings",
                peer_sector_event_count_365d_min=4.0,
                peer_relative_surprise_pct_365d_min=10.0,
                peer_relative_sue_hist_mean_4q_365d_min=0.0,
            ),
        ]
    )

    # Scope 2: only the bullish idle earnings engine.
    variants.extend(
        [
            VariantSpec(name="bullish_count3", scope="bullish_only", peer_sector_event_count_365d_min=3.0),
            VariantSpec(name="bullish_count4", scope="bullish_only", peer_sector_event_count_365d_min=4.0),
            VariantSpec(name="bullish_surprise8", scope="bullish_only", peer_relative_surprise_pct_365d_min=8.0),
            VariantSpec(name="bullish_surprise10", scope="bullish_only", peer_relative_surprise_pct_365d_min=10.0),
            VariantSpec(name="bullish_surprise12", scope="bullish_only", peer_relative_surprise_pct_365d_min=12.0),
            VariantSpec(
                name="bullish_count3_surprise8_sue0",
                scope="bullish_only",
                peer_sector_event_count_365d_min=3.0,
                peer_relative_surprise_pct_365d_min=8.0,
                peer_relative_sue_hist_mean_4q_365d_min=0.0,
            ),
            VariantSpec(
                name="bullish_count4_surprise10_sue0",
                scope="bullish_only",
                peer_sector_event_count_365d_min=4.0,
                peer_relative_surprise_pct_365d_min=10.0,
                peer_relative_sue_hist_mean_4q_365d_min=0.0,
            ),
        ]
    )
    return variants


def _trade_sleeve_summary(trades: list[Any], candidate_map: dict[str, Any]) -> dict[str, Any]:
    by_sleeve: Counter[str] = Counter()
    by_engine: Counter[str] = Counter()
    sleeve_pnl: defaultdict[str, float] = defaultdict(float)
    idle_trades = 0
    idle_pnl = 0.0
    idle_wins = 0

    for trade in trades:
        sleeve = _resolve_trade_sleeve(trade, candidate_map)
        by_sleeve[sleeve] += 1
        by_engine[str(trade.engine_id)] += 1
        sleeve_pnl[sleeve] += float(trade.net_pnl)
        if sleeve == "idle_alpha":
            idle_trades += 1
            idle_pnl += float(trade.net_pnl)
            if float(trade.net_pnl) > 0:
                idle_wins += 1

    return {
        "by_sleeve": dict(sorted(by_sleeve.items(), key=lambda kv: (-kv[1], kv[0]))),
        "by_engine": dict(sorted(by_engine.items(), key=lambda kv: (-kv[1], kv[0]))),
        "sleeve_pnl": {k: round(v, 2) for k, v in sleeve_pnl.items()},
        "idle_alpha_trade_count": idle_trades,
        "idle_alpha_net_pnl": round(idle_pnl, 2),
        "idle_alpha_win_rate": round(idle_wins / idle_trades, 4) if idle_trades else None,
    }


def main() -> None:
    parser = argparse.ArgumentParser(description="Probe peer-relative surprise additivity for idle-alpha")
    parser.add_argument(
        "--config",
        default=DEFAULT_CONFIG,
        help=f"Experiment manifest JSON path (default: {DEFAULT_CONFIG})",
    )
    parser.add_argument("--start", required=True, help="Start date (YYYY, YYYY-MM, YYYY-MM-DD)")
    parser.add_argument("--end", required=True, help="End date (YYYY, YYYY-MM, YYYY-MM-DD)")
    parser.add_argument(
        "--period-start",
        action="append",
        default=[],
        help="Sub-period start. Format: label=YYYY-MM-DD (repeatable).",
    )
    parser.add_argument(
        "--json",
        action="store_true",
        help="Emit machine-readable JSON only.",
    )
    parser.add_argument(
        "--capital",
        type=float,
        default=10_000.0,
        help="Initial equity for the probe (default: 10000).",
    )
    args = parser.parse_args()

    configure_logging("WARNING")

    start_date = _parse_date(args.start)
    end_date = _parse_date(args.end, is_end=True)
    period_starts = [_parse_period_start(value) for value in args.period_start]
    variants = _build_variants()

    manifest = load_manifest(args.config)
    base_config = resolve_config(manifest)
    store = _build_merged_snapshot_store(
        manifest,
        base_config,
        snapshot_dir_override=None,
    ).slice_by_date_range(start_date, end_date)

    rows: list[dict[str, Any]] = []
    for variant in variants:
        config = resolve_config(manifest)
        _apply_variant(config, variant)

        runner = BacktestRunner(
            manifest=manifest,
            config=config,
            store=store,
            initial_equity=args.capital,
            split_name="peer_relative_surprise_probe",
        )
        result = runner.run(output_root=None)
        sleeve_summary = _trade_sleeve_summary(runner._closed_trades, runner._candidate_map)

        row: dict[str, Any] = {
            "variant": variant.name,
            "scope": variant.scope,
            "peer_sector_event_count_365d_min": variant.peer_sector_event_count_365d_min,
            "peer_relative_surprise_pct_365d_min": variant.peer_relative_surprise_pct_365d_min,
            "peer_relative_sue_hist_mean_4q_365d_min": variant.peer_relative_sue_hist_mean_4q_365d_min,
            "full_return_pct": round(result.metrics.total_return_pct or 0.0, 2),
            "full_dd_pct": round(result.metrics.max_drawdown_pct or 0.0, 2),
            "full_sharpe": round(result.metrics.sharpe_ratio or 0.0, 3),
            "trade_count": result.metrics.trade_count,
            **sleeve_summary,
            "last_date": runner._equity_curve[-1].date.isoformat() if runner._equity_curve else None,
        }
        for label, period_start in period_starts:
            stats = _period_stats(runner._equity_curve, period_start)
            row[label] = {
                "return_pct": round(stats["return_pct"], 2) if stats else None,
                "max_dd_pct": round(stats["max_dd_pct"], 2) if stats else None,
                "start_date": stats["start_date"] if stats else None,
                "end_date": stats["end_date"] if stats else None,
            }
        rows.append(row)

    baseline = next((row for row in rows if row["variant"] == "baseline"), None)
    if baseline is not None:
        for row in rows:
            row["delta_return_vs_baseline"] = round(
                row["full_return_pct"] - baseline["full_return_pct"],
                2,
            )
            row["delta_dd_vs_baseline"] = round(
                row["full_dd_pct"] - baseline["full_dd_pct"],
                2,
            )
            row["delta_idle_alpha_trade_count_vs_baseline"] = (
                row["idle_alpha_trade_count"] - baseline["idle_alpha_trade_count"]
            )
            row["delta_idle_alpha_net_pnl_vs_baseline"] = round(
                row["idle_alpha_net_pnl"] - baseline["idle_alpha_net_pnl"],
                2,
            )

    if args.json:
        print(json.dumps(rows, ensure_ascii=False, indent=2))
        return

    period_labels = [label for label, _ in period_starts]
    header = [
        "variant",
        "scope",
        "full_return",
        "full_dd",
        "full_sharpe",
        "trades",
        "idle_trades",
        "idle_pnl",
        "delta_return",
        "delta_dd",
        "delta_idle_trades",
        "delta_idle_pnl",
        "last_date",
    ]
    for label in period_labels:
        header.extend([f"{label}_return", f"{label}_dd"])
    print("\t".join(header))
    for row in rows:
        fields = [
            row["variant"],
            row["scope"],
            f"{row['full_return_pct']:.2f}",
            f"{row['full_dd_pct']:.2f}",
            f"{row['full_sharpe']:.3f}",
            str(row["trade_count"]),
            str(row["idle_alpha_trade_count"]),
            f"{row['idle_alpha_net_pnl']:.2f}",
            f"{row['delta_return_vs_baseline']:.2f}" if baseline else "-",
            f"{row['delta_dd_vs_baseline']:.2f}" if baseline else "-",
            str(row["delta_idle_alpha_trade_count_vs_baseline"]) if baseline else "-",
            f"{row['delta_idle_alpha_net_pnl_vs_baseline']:.2f}" if baseline else "-",
            str(row["last_date"] or "-"),
        ]
        for label in period_labels:
            stats = row.get(label) or {}
            fields.extend(
                [
                    "-" if stats.get("return_pct") is None else f"{stats['return_pct']:.2f}",
                    "-" if stats.get("max_dd_pct") is None else f"{stats['max_dd_pct']:.2f}",
                ]
            )
        print("\t".join(fields))

    if baseline is not None:
        best = max(rows, key=lambda row: (row["delta_return_vs_baseline"], -row["delta_dd_vs_baseline"]))
        best_additive = [
            row
            for row in rows
            if row["variant"] != "baseline"
            and row["delta_return_vs_baseline"] > 0
            and row["delta_dd_vs_baseline"] <= 0
        ]
        if best_additive:
            best_additive_row = max(
                best_additive,
                key=lambda row: (row["delta_return_vs_baseline"], -row["delta_dd_vs_baseline"]),
            )
            print()
            print(
                "BEST_ADDITIVE",
                best_additive_row["variant"],
                f"return_delta={best_additive_row['delta_return_vs_baseline']:.2f}",
                f"dd_delta={best_additive_row['delta_dd_vs_baseline']:.2f}",
                f"idle_delta={best_additive_row['delta_idle_alpha_net_pnl_vs_baseline']:.2f}",
            )
        else:
            print()
            print(
                "BEST_OVERALL",
                best["variant"],
                f"return_delta={best['delta_return_vs_baseline']:.2f}",
                f"dd_delta={best['delta_dd_vs_baseline']:.2f}",
                f"idle_delta={best['delta_idle_alpha_net_pnl_vs_baseline']:.2f}",
            )


if __name__ == "__main__":
    main()