"""Build continuation-focused snapshots from existing event-day snapshots.""" from __future__ import annotations import datetime as dt import json from decimal import Decimal from pathlib import Path from typing import Any import pyarrow as pa import pyarrow.parquet as pq from libs.common.config import get_settings from libs.common.time_utils import trading_days_between, utc_now from libs.export.snapshot_export import _rows_to_table, _temporal_split from libs.features.market_features import compute_market_features from libs.labeler.label_generator import _compute_labels_from_bars, _pct_return from libs.oracle_client.models import PriceBar def _parse_date(raw: Any) -> dt.date | None: if isinstance(raw, dt.date): return raw if isinstance(raw, str): try: return dt.date.fromisoformat(raw) except ValueError: return None return None def _iso_ts(date_value: dt.date) -> str: return f"{date_value.isoformat()}T21:00:00+00:00" def _to_price_bars(date_bars: dict[dt.date, dict[str, Any]]) -> list[PriceBar]: rows: list[PriceBar] = [] for date_value in sorted(date_bars): bar = date_bars[date_value] rows.append( PriceBar( date=date_value.isoformat(), open=float(bar["open"]), high=float(bar["high"]), low=float(bar["low"]), close=float(bar["close"]), volume=float(bar["volume"]), ) ) return rows def _continuation_entry_window(exec_date: dt.date, lookback_days: int) -> tuple[dt.date, dt.date] | None: trading_days = trading_days_between(exec_date, exec_date + dt.timedelta(days=14)) signal_index = lookback_days entry_index = lookback_days + 1 if len(trading_days) <= entry_index: return None return trading_days[signal_index], trading_days[entry_index] def _build_continuation_rows_from_bars( base_rows: list[dict[str, Any]], bars_by_symbol: dict[str, dict[dt.date, dict[str, Any]]], *, lookback_days: int = 3, ) -> list[dict[str, Any]]: transformed: list[dict[str, Any]] = [] price_bars_cache: dict[str, list[PriceBar]] = {} for row in base_rows: ticker = str(row.get("ticker") or row.get("symbol") or "").upper() exec_date = _parse_date(row.get("entry_date") or row.get("execution_date")) source_close_raw = row.get("event_close") or row.get("entry_price_est") if not ticker or exec_date is None or source_close_raw in (None, 0): continue source_close = float(source_close_raw) symbol_bars = bars_by_symbol.get(ticker) if not symbol_bars or exec_date not in symbol_bars: continue dates = _continuation_entry_window(exec_date, lookback_days) if dates is None: continue signal_date, entry_date = dates if signal_date not in symbol_bars or entry_date not in symbol_bars: continue signal_bar = symbol_bars[signal_date] signal_close = float(signal_bar["close"]) if source_close <= 0 or signal_close <= 0: continue drift_pct = (signal_close - source_close) / source_close price_bars = price_bars_cache.get(ticker) if price_bars is None: price_bars = _to_price_bars(symbol_bars) price_bars_cache[ticker] = price_bars signal_features = compute_market_features(price_bars, signal_date.isoformat()) sorted_dates = sorted(symbol_bars) entry_idx = sorted_dates.index(entry_date) bars_from_entry = [symbol_bars[d] for d in sorted_dates[entry_idx:]] if not bars_from_entry: continue entry_price = Decimal(str(bars_from_entry[0]["open"])) forward_bars = bars_from_entry[1:] label_status = "ok" if len(forward_bars) >= 20 else "truncated" fwd_1d = _pct_return(entry_price, Decimal(str(bars_from_entry[1]["close"]))) if len(bars_from_entry) > 1 else None lbl_3d = _compute_labels_from_bars(entry_price, forward_bars, 3) lbl_5d = _compute_labels_from_bars(entry_price, forward_bars, 5) lbl_10d = _compute_labels_from_bars(entry_price, forward_bars, 10) lbl_20d = _compute_labels_from_bars(entry_price, forward_bars, 20) new_row = dict(row) original_event_id = str(row.get("event_id") or "") new_row["event_id"] = f"{original_event_id}::cont_d{lookback_days}" new_row["original_event_id"] = original_event_id new_row["original_event_date"] = row.get("event_date") new_row["original_reaction_date"] = row.get("reaction_date") new_row["original_entry_date"] = row.get("entry_date") new_row["original_event_close"] = source_close new_row["event_date"] = signal_date.isoformat() new_row["reaction_date"] = signal_date.isoformat() new_row["entry_date"] = entry_date.isoformat() new_row["entry_convention"] = "next_open_after_continuation_signal" new_row["event_timestamp"] = _iso_ts(signal_date) new_row["event_close"] = signal_close new_row["entry_price"] = float(entry_price) new_row["reaction_day_return"] = drift_pct new_row["continuation_anchor_drift_pct"] = drift_pct new_row["continuation_anchor_day_return"] = signal_features.get("reaction_day_return") new_row["close_location"] = signal_features.get("close_location") new_row["volume_ratio_20d"] = signal_features.get("volume_ratio_20d") new_row["avg_dollar_volume_20d"] = signal_features.get("avg_dollar_volume_20d") new_row["gap_size"] = signal_features.get("gap_size") new_row["atr_14"] = signal_features.get("atr_14") new_row["reaction_day_low"] = signal_features.get("reaction_day_low") new_row["reaction_day_high"] = signal_features.get("reaction_day_high") new_row["fwd_return_1d"] = float(fwd_1d) if fwd_1d is not None else None new_row["fwd_return_3d"] = float(lbl_3d.get("fwd_return")) if lbl_3d.get("fwd_return") is not None else None new_row["fwd_return_5d"] = float(lbl_5d.get("fwd_return")) if lbl_5d.get("fwd_return") is not None else None new_row["fwd_return_10d"] = float(lbl_10d.get("fwd_return")) if lbl_10d.get("fwd_return") is not None else None new_row["fwd_return_20d"] = float(lbl_20d.get("fwd_return")) if lbl_20d.get("fwd_return") is not None else None new_row["mfe_3d"] = float(lbl_3d.get("mfe")) if lbl_3d.get("mfe") is not None else None new_row["mae_3d"] = float(lbl_3d.get("mae")) if lbl_3d.get("mae") is not None else None new_row["mfe_5d"] = float(lbl_5d.get("mfe")) if lbl_5d.get("mfe") is not None else None new_row["mae_5d"] = float(lbl_5d.get("mae")) if lbl_5d.get("mae") is not None else None new_row["mfe_10d"] = float(lbl_10d.get("mfe")) if lbl_10d.get("mfe") is not None else None new_row["mae_10d"] = float(lbl_10d.get("mae")) if lbl_10d.get("mae") is not None else None new_row["mfe_20d"] = float(lbl_20d.get("mfe")) if lbl_20d.get("mfe") is not None else None new_row["mae_20d"] = float(lbl_20d.get("mae")) if lbl_20d.get("mae") is not None else None new_row["label_status"] = label_status transformed.append(new_row) return transformed async def export_continuation_snapshot_from_base( *, base_snapshot_dir: str | Path, output_dir: str | Path, snapshot_id: str, lookback_days: int = 3, ) -> dict[str, Any]: from libs.backtest.snapshot_store import SnapshotStore base_path = Path(base_snapshot_dir) rows: list[dict[str, Any]] = [] for split in ("train", "valid", "test"): table = pq.read_table(base_path / f"{split}.parquet") rows.extend(table.to_pylist()) symbols = sorted({str(r.get("ticker") or r.get("symbol") or "").upper() for r in rows if r.get("ticker") or r.get("symbol")}) dates = [_parse_date(r.get("entry_date")) for r in rows] dates = [d for d in dates if d is not None] if not symbols or not dates: raise ValueError("Base snapshot has no symbols or dates") settings = get_settings() bars_by_symbol, _ = await SnapshotStore._fetch_price_data( symbols, (min(dates) - dt.timedelta(days=45), max(dates) + dt.timedelta(days=45)), settings.stock_oracle_url, ) transformed = _build_continuation_rows_from_bars(rows, bars_by_symbol, lookback_days=lookback_days) splits = _temporal_split(transformed, "temporal_70_15_15") out_path = Path(output_dir) / snapshot_id out_path.mkdir(parents=True, exist_ok=True) row_counts: dict[str, int] = {} for split_name, split_rows in splits.items(): pq.write_table(_rows_to_table(split_rows), out_path / f"{split_name}.parquet") row_counts[split_name] = len(split_rows) manifest = { "snapshot_id": snapshot_id, "created_at_utc": utc_now().isoformat(), "base_snapshot_id": base_path.name, "transform": f"continuation_d{lookback_days}", "split_policy": "temporal_70_15_15", "row_counts": row_counts, "total_rows": sum(row_counts.values()), "output_dir": str(out_path), } (out_path / "manifest.json").write_text(json.dumps(manifest, indent=2)) return manifest