fithia2/scripts/orb_synthetic_per_bar_state.py

"""Phase B (large dataset) — Per-Bar State Dataset from Synthetic Entries.

Reads synthetic_entries.parquet (Phase A output) and expands each entry to one
row per 5-min bar from entry through 16:00 ET. Produces the same schema as
orb_per_bar_state.py but processes ~580k entries in parallel chunks.

Output is written as per-date parquet shards that can be fed to Phase C
(orb_continuation_labels.py) after optional concatenation.

Usage:
    python scripts/orb_synthetic_per_bar_state.py \
        --entries tmp/orb_phase_a_synthetic_entries/synthetic_entries.parquet \
        --out     tmp/orb_perbar_phase_b \
        --workers 8
"""
from __future__ import annotations

import argparse
import json
import os
from multiprocessing import Pool
from pathlib import Path

import numpy as np
import pandas as pd

ET = "America/New_York"
DEFAULT_INTRADAY_CACHE = Path("data/cache/intraday")


# ── helpers shared with orb_per_bar_state.py ────────────────────────────────

def _load_intraday(ticker: str, date: str, cache_root: Path) -> pd.DataFrame | None:
    p = cache_root / ticker / f"{date}.parquet"
    if not p.exists():
        return None
    try:
        df = pd.read_parquet(p)
        if df.empty:
            return None
        df = df.sort_values("timestamp").reset_index(drop=True)
        df["ts"] = pd.to_datetime(df["timestamp"], utc=True)
        return df
    except Exception:
        return None


def _market_open_close_utc(date_str: str) -> tuple[pd.Timestamp, pd.Timestamp]:
    d = pd.Timestamp(date_str, tz=ET)
    open_et  = d + pd.Timedelta(hours=9, minutes=30)
    close_et = d + pd.Timedelta(hours=16)
    return open_et.tz_convert("UTC"), close_et.tz_convert("UTC")


def _running_vwap(df: pd.DataFrame, market_open_utc: pd.Timestamp) -> np.ndarray:
    typical = (df["high"] + df["low"] + df["close"]) / 3
    vol = df["volume"].astype(float)
    in_session = (df["ts"] >= market_open_utc)
    pv = np.where(in_session, typical * vol, 0.0)
    v  = np.where(in_session, vol, 0.0)
    cum_pv = np.cumsum(pv)
    cum_v  = np.cumsum(v)
    with np.errstate(divide="ignore", invalid="ignore"):
        vwap = np.where(cum_v > 0, cum_pv / cum_v, np.nan)
    return np.where(in_session, vwap, np.nan)


def _bars_since_peak(running_peak: np.ndarray) -> np.ndarray:
    n = len(running_peak)
    out = np.zeros(n, dtype=np.int32)
    last_peak_idx = 0
    last_peak = running_peak[0]
    for i in range(n):
        if running_peak[i] > last_peak:
            last_peak = running_peak[i]
            last_peak_idx = i
        out[i] = i - last_peak_idx
    return out


def _market_context_aligned(
    market_df: pd.DataFrame | None,
    target_ts: pd.DatetimeIndex,
    entry_ts: pd.Timestamp,
    market_open: pd.Timestamp,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    n = len(target_ts)
    if market_df is None or market_df.empty:
        nan = np.full(n, np.nan)
        return nan, nan.copy(), nan.copy()
    m = market_df.sort_values("ts").reset_index(drop=True)
    target = pd.DataFrame({"ts": target_ts})
    aligned = pd.merge_asof(target, m[["ts", "close"]], on="ts", direction="backward")
    closes = aligned["close"].to_numpy()
    entry_close = (
        m[m["ts"] <= entry_ts]["close"].iloc[-1] if (m["ts"] <= entry_ts).any() else np.nan
    )
    open_close = (
        m[m["ts"] <= market_open]["close"].iloc[-1] if (m["ts"] <= market_open).any()
        else (m["close"].iloc[0] if len(m) else np.nan)
    )
    ret_since_entry = closes / entry_close - 1.0 if pd.notna(entry_close) else np.full(n, np.nan)
    ret_since_open  = closes / open_close  - 1.0 if pd.notna(open_close)  else np.full(n, np.nan)
    return closes, ret_since_entry, ret_since_open


PASSTHROUGH_COLS = [
    "atr_at_entry", "gap_pct", "rvol", "morning_gain_pct", "entropy_20d",
    "ret_5d", "candidate_score", "score_rank_pct", "sector_confirmation_active",
    "sector_confirmation_score", "entry_market_guard_active",
    "entry_market_guard_return_pct", "is_liquid_largecap", "is_moderate_gap_liquid",
    "trigger_type", "entry_dollar_volume", "avg_dollar_vol_30d",
    "premarket_dollar_vol", "first_bar_dollar_vol", "body_ratio",
    "close_location", "gap_zscore_20d", "obv_slope_20", "obv_slope_5", "orb_return",
]


def build_for_entry(
    trade_id: int,
    row: dict,
    intraday_root: Path,
    spy_df: pd.DataFrame | None,
    qqq_df: pd.DataFrame | None,
) -> pd.DataFrame | None:
    """Expand one synthetic entry to per-bar rows (identical logic to orb_per_bar_state.py)."""
    ticker = row["ticker"]
    date   = row["date"]
    entry_price = float(row["entry_price"])
    risk  = float(row["risk_per_share"])
    direction = row.get("orb_direction", "long")

    entry_time_raw = row["entry_time"]
    exit_time_raw  = row["exit_time"]
    entry_time = pd.Timestamp(entry_time_raw).tz_convert("UTC") if "+" in str(entry_time_raw) or "Z" in str(entry_time_raw) else pd.Timestamp(entry_time_raw, tz=ET).tz_convert("UTC")
    exit_time  = pd.Timestamp(exit_time_raw ).tz_convert("UTC") if "+" in str(exit_time_raw)  or "Z" in str(exit_time_raw)  else pd.Timestamp(exit_time_raw,  tz=ET).tz_convert("UTC")

    realized_r = float(row.get("r_multiple_at_exit") or 0.0)
    is_v49 = bool(row.get("is_v49_actual_entry", False))

    bars = _load_intraday(ticker, date, intraday_root)
    if bars is None or bars.empty:
        return None

    market_open_utc, market_close_utc = _market_open_close_utc(date)
    bars["running_vwap"] = _running_vwap(bars, market_open_utc)

    # slice from entry through EOD (ts is UTC)
    win = bars[(bars["ts"] >= entry_time) &
               (bars["ts"] <= market_close_utc)].reset_index(drop=True)
    if win.empty:
        return None

    n   = len(win)
    bar_idx = np.arange(n, dtype=np.int32)
    ts  = win["ts"]      # already UTC
    ts_utc = ts          # alias for clarity

    minutes_since_entry = ((ts_utc - entry_time).dt.total_seconds() / 60.0).to_numpy()
    minutes_to_close    = ((market_close_utc - ts_utc).dt.total_seconds() / 60.0).to_numpy()
    minutes_since_open  = ((ts_utc - market_open_utc).dt.total_seconds() / 60.0).to_numpy()

    opens  = win["open"].to_numpy()
    highs  = win["high"].to_numpy()
    lows   = win["low"].to_numpy()
    closes = win["close"].to_numpy()
    vols   = win["volume"].to_numpy().astype(float)
    vwap   = win["running_vwap"].to_numpy()

    if direction == "long":
        fav_high  = (highs  - entry_price) / risk
        fav_low   = (lows   - entry_price) / risk
        fav_close = (closes - entry_price) / risk
    else:
        fav_high  = (entry_price - lows)   / risk
        fav_low   = (entry_price - highs)  / risk
        fav_close = (entry_price - closes) / risk

    next_open_price = np.append(opens[1:], closes[-1])
    if direction == "long":
        fav_next_open = (next_open_price - entry_price) / risk
    else:
        fav_next_open = (entry_price - next_open_price) / risk

    mfe_so_far      = np.maximum.accumulate(fav_high)
    mae_so_far      = np.minimum.accumulate(fav_low)
    giveback        = mfe_so_far - fav_close
    bsp             = _bars_since_peak(mfe_so_far)

    bar_return_pct  = np.where(opens > 0, (closes - opens) / opens, np.nan)
    bar_range       = np.maximum(highs - lows, 1e-9)
    bar_close_loc   = (closes - lows) / bar_range

    with np.errstate(divide="ignore", invalid="ignore"):
        vwap_dev_pct = np.where((vwap > 0) & np.isfinite(vwap), (closes - vwap) / vwap, np.nan)

    first_vol   = vols[0] if vols[0] > 0 else np.nan
    vol_vs_first = vols / first_vol if pd.notna(first_vol) and first_vol > 0 else np.full(n, np.nan)

    avg_dvol_30d = float(row.get("avg_dollar_vol_30d") or 0.0)
    expected_bar_vol = (avg_dvol_30d / entry_price / 78.0) if entry_price > 0 else None
    if expected_bar_vol and expected_bar_vol > 0 and np.isfinite(expected_bar_vol):
        vol_vs_avg = vols / expected_bar_vol
    else:
        vol_vs_avg = np.full(n, np.nan)

    # market context (ts is already UTC, passes directly)
    spy_close, spy_ret_entry, spy_ret_open = _market_context_aligned(
        spy_df, ts, entry_time, market_open_utc
    )
    qqq_close, qqq_ret_entry, qqq_ret_open = _market_context_aligned(
        qqq_df, ts, entry_time, market_open_utc
    )

    # baseline benchmark anchors
    diffs = (ts_utc - exit_time).dt.total_seconds().to_numpy()
    if (diffs >= 0).any():
        baseline_exit_bar_idx_val = int(np.argmax(diffs >= 0))
    else:
        baseline_exit_bar_idx_val = n - 1

    is_held_by_baseline = (ts_utc <= exit_time).to_numpy()

    base = {
        "trade_id":   trade_id,
        "ticker":     ticker,
        "date":       date,
        "direction":  direction,
        "entry_time": entry_time,
        "bar_ts":     ts_utc.to_numpy(),
        "bar_idx":    bar_idx,
        "minutes_since_entry":  minutes_since_entry,
        "minutes_to_close":     minutes_to_close,
        "minutes_since_open":   minutes_since_open,
        "entry_price":   entry_price,
        "risk_per_share": risk,
        "bar_open":  opens,
        "bar_high":  highs,
        "bar_low":   lows,
        "bar_close": closes,
        "bar_volume": vols,
        "bar_return_pct":  bar_return_pct,
        "bar_close_loc":   bar_close_loc,
        "current_close_r": fav_close,
        "current_high_r":  fav_high,
        "current_low_r":   fav_low,
        "next_open_r":     fav_next_open,
        "mfe_so_far_r":    mfe_so_far,
        "mae_so_far_r":    mae_so_far,
        "giveback_from_peak_r": giveback,
        "bars_since_peak":  bsp,
        "vwap":             vwap,
        "vwap_dev_pct":     vwap_dev_pct,
        "vol_vs_first_bar": vol_vs_first,
        "vol_vs_avg_dvol30d": vol_vs_avg,
        "spy_close":  spy_close,
        "spy_return_since_entry": spy_ret_entry,
        "spy_return_since_open":  spy_ret_open,
        "qqq_close":  qqq_close,
        "qqq_return_since_entry": qqq_ret_entry,
        "qqq_return_since_open":  qqq_ret_open,
        "is_held_by_baseline":    is_held_by_baseline,
        "baseline_exit_bar_idx":  baseline_exit_bar_idx_val,
        "baseline_exit_time":     exit_time,
        "baseline_realized_r":    realized_r if is_v49 else float("nan"),
        "baseline_exit_reason":   row.get("exit_reason", "eod_synthetic"),
        "is_v49_actual_entry":    is_v49,
    }

    out = pd.DataFrame(base)
    for col in PASSTHROUGH_COLS:
        out[col] = row.get(col)
    return out


# ── worker: process one trading date ────────────────────────────────────────

def _process_date(args: tuple) -> int:
    """Process all entries for one date; append rows to a parquet shard. Returns row count."""
    date, entries_for_date, intraday_root_str, out_shard_path = args
    intraday_root = Path(intraday_root_str)

    spy_df = _load_intraday("SPY", date, intraday_root)
    qqq_df = _load_intraday("QQQ", date, intraday_root)

    frames = []
    for row in entries_for_date:
        try:
            df = build_for_entry(
                int(row["trade_id"]), row, intraday_root, spy_df, qqq_df
            )
            if df is not None and not df.empty:
                frames.append(df)
        except Exception as e:
            pass  # skip individual failures silently

    if not frames:
        return 0

    shard = pd.concat(frames, ignore_index=True)
    shard.to_parquet(out_shard_path, index=False)
    return len(shard)


# ── main ─────────────────────────────────────────────────────────────────────

def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--entries", required=True, help="Phase A synthetic_entries.parquet")
    ap.add_argument("--out",     required=True)
    ap.add_argument("--intraday-cache", default=str(DEFAULT_INTRADAY_CACHE))
    ap.add_argument("--workers", type=int, default=8)
    ap.add_argument("--concat",  action="store_true",
                    help="After shards are written, concat into one big parquet")
    args = ap.parse_args()

    out_dir = Path(args.out)
    shards_dir = out_dir / "shards"
    shards_dir.mkdir(parents=True, exist_ok=True)
    intraday_root = Path(args.intraday_cache)

    entries = pd.read_parquet(args.entries)
    print(f"Loaded {len(entries):,} entries from {entries['date'].nunique()} dates")
    entries["is_v49_actual_entry"] = entries["is_v49_actual_entry"].fillna(False)

    # group by date
    dates = sorted(entries["date"].unique())
    date_groups: list[tuple] = []
    for date in dates:
        grp = entries[entries["date"] == date].to_dict("records")
        shard_path = shards_dir / f"{date}.parquet"
        if shard_path.exists():
            continue  # resume support: skip already-done dates
        date_groups.append((date, grp, str(intraday_root), str(shard_path)))

    print(f"Dates to process: {len(date_groups)} (skipping {len(dates) - len(date_groups)} cached)")

    total_rows = 0
    done = 0

    if args.workers > 1:
        with Pool(processes=args.workers) as pool:
            for n_rows in pool.imap_unordered(_process_date, date_groups, chunksize=1):
                total_rows += n_rows
                done += 1
                if done % 20 == 0 or done == len(date_groups):
                    print(f"  {done}/{len(date_groups)} dates  {total_rows:,} rows", flush=True)
    else:
        for args_tuple in date_groups:
            n_rows = _process_date(args_tuple)
            total_rows += n_rows
            done += 1
            if done % 20 == 0 or done == len(date_groups):
                print(f"  {done}/{len(date_groups)} dates  {total_rows:,} rows", flush=True)

    print(f"\nShard writing done: {total_rows:,} rows across {len(dates)} dates")

    # optionally concat all shards into one file
    if args.concat:
        print("Concatenating shards...")
        shard_files = sorted(shards_dir.glob("*.parquet"))
        all_frames = [pd.read_parquet(f) for f in shard_files]
        if all_frames:
            big = pd.concat(all_frames, ignore_index=True)
            out_path = out_dir / "per_bar_states.parquet"
            big.to_parquet(out_path, index=False)
            print(f"Wrote {out_path} ({len(big):,} rows, {big.shape[1]} columns)")


if __name__ == "__main__":
    main()