fithia2/scripts/recompute_reset_common_window_batch.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import datetime as dt
import json
from pathlib import Path

from apps.backtester.run import BacktestRunner, _build_merged_snapshot_store, _extend_store_to_requested_window
from libs.backtest.domain import CommonWindowSummary, ResetCommonWindowSummary
from libs.backtest.manifests import load_manifest, resolve_config
from libs.backtest.tracker import (
    compute_reset_common_window_score,
    journal_lock,
    load_journal,
    rebuild_registry,
    refresh_public_scores,
)


def _add_years(value: dt.date, years: int) -> dt.date:
    try:
        return value.replace(year=value.year + years)
    except ValueError:
        return value.replace(month=2, day=28, year=value.year + years)


def _iter_segments(start_date: dt.date, end_date: dt.date, *, segment_years: int) -> list[tuple[dt.date, dt.date]]:
    segments: list[tuple[dt.date, dt.date]] = []
    cursor = start_date
    while cursor <= end_date:
        next_cursor = _add_years(cursor, segment_years)
        segment_end = min(end_date, next_cursor - dt.timedelta(days=1))
        segments.append((cursor, segment_end))
        cursor = segment_end + dt.timedelta(days=1)
    return segments


def _sort_key(entry: dict) -> tuple:
    return (
        entry.get("sqs_score") is None,
        -(entry.get("sqs_score") or 0.0),
        entry.get("valid_total_return_pct") is None,
        -(entry.get("valid_total_return_pct") or 0.0),
        entry.get("total_return_pct") is None,
        -(entry.get("total_return_pct") or 0.0),
        entry.get("annualized_return_pct") is None,
        -(entry.get("annualized_return_pct") or 0.0),
        entry.get("max_drawdown_pct") is None,
        +(entry.get("max_drawdown_pct") or 100.0),
    )


def _pick_targets(registry_path: Path, *, top_n: int, include_experiments: list[str]) -> list[str]:
    payload = json.loads(registry_path.read_text())
    entries = [entry for entry in payload.get("entries", []) if not entry.get("is_retired")]
    entries.sort(key=_sort_key)
    selected: list[str] = []
    seen: set[str] = set()
    for entry in entries[:top_n]:
        experiment_name = entry["experiment_name"]
        if experiment_name in seen:
            continue
        selected.append(experiment_name)
        seen.add(experiment_name)
    for experiment_name in include_experiments:
        if experiment_name and experiment_name not in selected:
            selected.append(experiment_name)
    return selected


def _pick_common_window_range(
    *,
    experiment_name: str,
    journal_entries_by_name: dict[str, list],
    registry_entries_by_name: dict[str, list[dict]],
) -> tuple[dt.date, dt.date] | None:
    for entry in journal_entries_by_name.get(experiment_name, []):
        if entry.common_window_summary is not None:
            return (entry.common_window_summary.start_date, entry.common_window_summary.end_date)
    for entry in registry_entries_by_name.get(experiment_name, []):
        summary = entry.get("common_window_summary")
        if summary:
            parsed = CommonWindowSummary.model_validate(summary)
            return (parsed.start_date, parsed.end_date)
    return None


def _build_reset_summary(
    *,
    experiment_name: str,
    start_date: dt.date,
    end_date: dt.date,
    segment_years: int,
    initial_equity: float,
) -> ResetCommonWindowSummary:
    config_path = Path("configs/experiments") / f"{experiment_name}.json"
    manifest = load_manifest(config_path)
    config = resolve_config(manifest)
    base_store = _build_merged_snapshot_store(manifest, config, snapshot_dir_override=None)
    snapshot_id = (
        config.canonical_snapshot_id
        or config.requested_snapshot_id
        or config.dataset_snapshot_id
        or ""
    )

    segment_summaries: list[CommonWindowSummary] = []
    for segment_start, segment_end in _iter_segments(
        start_date,
        end_date,
        segment_years=segment_years,
    ):
        segment_store = base_store.slice_by_date_range(segment_start, segment_end)
        segment_store = _extend_store_to_requested_window(
            store=segment_store,
            config=config,
            start_date=segment_start,
            end_date=segment_end,
            snapshot_dir_override=None,
        )
        runner = BacktestRunner(
            manifest=manifest,
            config=config,
            store=segment_store,
            initial_equity=initial_equity,
            split_name="reset_common_window",
        )
        result = runner.run(output_root=None)
        segment_summaries.append(
            CommonWindowSummary(
                snapshot_id=snapshot_id,
                start_date=segment_start,
                end_date=segment_end,
                initial_equity=initial_equity,
                run_id=result.run_id,
                metrics=result.metrics,
            )
        )

    return ResetCommonWindowSummary(
        snapshot_id=snapshot_id,
        start_date=start_date,
        end_date=end_date,
        reset_initial_equity=initial_equity,
        segment_days=365 * segment_years,
        segment_summaries=segment_summaries,
    )


def main() -> None:
    parser = argparse.ArgumentParser(description="Recompute reset common-window summaries for leaderboard leaders")
    parser.add_argument("--journal-dir", default="journal")
    parser.add_argument("--top", type=int, default=30)
    parser.add_argument("--segment-years", type=int, default=1)
    parser.add_argument("--initial-equity", type=float, default=10_000.0)
    parser.add_argument(
        "--include-experiment",
        action="append",
        default=[],
        help="Extra experiment names to include even if outside the top-N",
    )
    parser.add_argument("--force", action="store_true", help="Recompute even if reset summary already exists")
    parser.add_argument("--default-start-date", type=dt.date.fromisoformat)
    parser.add_argument("--default-end-date", type=dt.date.fromisoformat)
    args = parser.parse_args()

    journal_dir = Path(args.journal_dir)
    journal_path = journal_dir / "improvement_journal.jsonl"
    registry_path = journal_dir / "experiment_registry.json"
    leaderboard_path = journal_dir / "LEADERBOARD.md"
    summary_dir = journal_dir / "reset_common_window"
    summary_dir.mkdir(parents=True, exist_ok=True)

    targets = _pick_targets(
        registry_path,
        top_n=args.top,
        include_experiments=args.include_experiment,
    )
    journal_entries_by_name: dict[str, list] = {}
    for entry in load_journal(journal_path):
        journal_entries_by_name.setdefault(entry.experiment_name, []).append(entry)
    registry_payload = json.loads(registry_path.read_text())
    registry_entries_by_name: dict[str, list[dict]] = {}
    for entry in registry_payload.get("entries", []):
        registry_entries_by_name.setdefault(entry["experiment_name"], []).append(entry)

    built: list[tuple[str, Path]] = []
    for experiment_name in targets:
        common_window_range = _pick_common_window_range(
            experiment_name=experiment_name,
            journal_entries_by_name=journal_entries_by_name,
            registry_entries_by_name=registry_entries_by_name,
        )
        if common_window_range is None and args.default_start_date and args.default_end_date:
            common_window_range = (args.default_start_date, args.default_end_date)
        if common_window_range is None:
            print(f"SKIP {experiment_name}: missing journal/common_window_summary")
            continue
        has_existing_reset = any(
            entry.reset_common_window_summary is not None
            for entry in journal_entries_by_name.get(experiment_name, [])
        )
        if has_existing_reset and not args.force:
            print(f"SKIP {experiment_name}: reset_common_window_summary already exists")
            continue

        try:
            summary = _build_reset_summary(
                experiment_name=experiment_name,
                start_date=common_window_range[0],
                end_date=common_window_range[1],
                segment_years=args.segment_years,
                initial_equity=args.initial_equity,
            )
        except FileNotFoundError:
            print(f"SKIP {experiment_name}: missing manifest/config file")
            continue
        summary_path = summary_dir / f"{experiment_name}.json"
        summary_path.write_text(summary.model_dump_json(indent=2), encoding="utf-8")
        built.append((experiment_name, summary_path))
        print(
            f"BUILT {experiment_name}: segments={len(summary.segment_summaries)} "
            f"range={summary.start_date}..{summary.end_date}"
        )

    if not built:
        print("No reset common-window summaries built.")
        return

    summary_by_name = {
        experiment_name: ResetCommonWindowSummary.model_validate_json(summary_path.read_text())
        for experiment_name, summary_path in built
    }
    target_set = set(summary_by_name)
    with journal_lock(journal_path):
        entries = load_journal(journal_path)
        rewritten = []
        for entry in entries:
            summary = summary_by_name.get(entry.experiment_name)
            if summary is None:
                rewritten.append(entry)
                continue
            score, breakdown = compute_reset_common_window_score(summary)
            rewritten.append(
                entry.model_copy(
                    update={
                        "reset_common_window_summary": summary,
                        "reset_common_window_score": score,
                        "reset_common_window_breakdown": breakdown,
                    }
                )
            )
        journal_path.write_text(
            "".join(item.model_dump_json() + "\n" for item in rewritten),
            encoding="utf-8",
        )
        refresh_public_scores(journal_path, selector=lambda entry: entry.experiment_name in target_set)
        rebuild_registry(journal_path, registry_path, leaderboard_path)

    updated_registry = json.loads(registry_path.read_text())
    updated_entries = {
        entry["experiment_name"]: entry
        for entry in updated_registry.get("entries", [])
        if entry["experiment_name"] in target_set
    }
    for experiment_name in targets:
        entry = updated_entries.get(experiment_name)
        if entry is None:
            continue
        print(
            f"UPDATED {experiment_name}: "
            f"SQS={entry.get('sqs_score')} "
            f"RCW={entry.get('reset_common_window_score')} "
            f"CW={entry.get('common_window_score')}"
        )


if __name__ == "__main__":
    main()