"""ORB Per-Bar State Dataset Builder.

For each entry in a V49.91 ORB intraday backtest run, expand to one row per
5-min bar from entry through 16:00 ET (regardless of when V49.91's rule
exited — the ML model may decide to hold longer or exit earlier).

Each row is a "decision point" with state features available at that time.
The continuation-value labels are computed in Phase C; this script only
prepares features.

Usage:
    python scripts/orb_per_bar_state.py \
        --run tmp/v49_91_baseline_200_20260506/intraday_20260506_081008_effaac09.json \
        --out tmp/orb_perbar_v49_91_200d
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
import pandas as pd

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


def _infer_risk_per_share(t: dict) -> float:
    r = t.get("r_multiple_at_exit")
    entry = t["entry_price"]
    exit_p = t["exit_price"]
    direction = t.get("orb_direction", "long")
    if r not in (None, 0) and abs(r) > 1e-6:
        if direction == "long":
            return (exit_p - entry) / r
        return (entry - exit_p) / r
    atr = t.get("atr_at_entry") or 0.0
    return max(atr * DEFAULT_ATR_STOP_MULT, 1e-6)


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
    df = pd.read_parquet(p).sort_values("timestamp").reset_index(drop=True)
    df["ts"] = pd.to_datetime(df["timestamp"], utc=True)
    return df


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:
    """Compute running VWAP from market open onward; bars before open get NaN."""
    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)
    vwap = np.where(in_session, vwap, np.nan)
    return vwap


def _bars_since_peak(running_peak: np.ndarray) -> np.ndarray:
    """For each i, how many bars ago was the most recent strict new high in running_peak."""
    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]:
    """Align market index closes to target bar timestamps via asof; compute returns since entry / open."""
    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-time close & open-time close from the same series
    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


# Trade-level passthrough columns (broadcast across all bars of that trade).
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_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",
    "is_liquid_largecap",
    "is_moderate_gap_liquid",
    "entry_market_guard_active",
    "entry_market_guard_return_pct",
    "trigger_type",
]


def build_for_trade(
    trade_id: int,
    t: dict,
    intraday_root: Path,
    spy_by_date: dict[str, pd.DataFrame],
    qqq_by_date: dict[str, pd.DataFrame],
) -> pd.DataFrame | None:
    ticker = t["ticker"]
    date = t["date"]
    bars = _load_intraday(ticker, date, intraday_root)
    if bars is None or bars.empty:
        return None

    entry_time = pd.Timestamp(t["entry_time"]).tz_convert("UTC")
    exit_time = pd.Timestamp(t["exit_time"]).tz_convert("UTC")
    entry_price = float(t["entry_price"])
    direction = t.get("orb_direction", "long")
    risk = _infer_risk_per_share(t)
    realized_r = float(t.get("r_multiple_at_exit") or 0.0)

    market_open_utc, market_close_utc = _market_open_close_utc(date)

    # Compute running VWAP on full-day bars (so it accumulates correctly), then slice.
    bars["running_vwap"] = _running_vwap(bars, market_open_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"]
    minutes_since_entry = ((ts - entry_time).dt.total_seconds() / 60.0).to_numpy()
    minutes_to_close = ((market_close_utc - ts).dt.total_seconds() / 60.0).to_numpy()
    minutes_since_open = ((ts - 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()

    # Favorable-direction R-multiples per bar (long: price-entry; short: entry-price).
    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-bar open R: actionable execution price if you decide at this bar's close.
    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

    # Running MFE / MAE / peak (vectorized over bars).
    mfe_so_far = np.maximum.accumulate(fav_high)
    mae_so_far = np.minimum.accumulate(fav_low)
    giveback_from_peak = mfe_so_far - fav_close
    bars_since_peak = _bars_since_peak(mfe_so_far)

    # Bar dynamics.
    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  # 0 = closed at low, 1 = closed at high

    # Liquidity / VWAP.
    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)

    # Bar volume vs avg-daily expectation: avg_daily_dollar_volume / entry_price / 78 bars.
    avg_dvol_30d = t.get("avg_dollar_vol_30d") or 0.0
    expected_bar_vol = (avg_dvol_30d / entry_price / 78.0) if entry_price > 0 else np.nan
    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 (SPY, QQQ).
    spy_close, spy_ret_entry, spy_ret_open = _market_context_aligned(
        spy_by_date.get(date), ts, entry_time, market_open_utc
    )
    qqq_close, qqq_ret_entry, qqq_ret_open = _market_context_aligned(
        qqq_by_date.get(date), ts, entry_time, market_open_utc
    )

    # Baseline benchmark anchors (reference, not features).
    is_held_by_baseline = (ts <= exit_time).to_numpy()
    baseline_exit_at_or_before = (ts >= exit_time).to_numpy()
    # Index of bar nearest to (>=) baseline exit_time.
    diffs = (ts - 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

    # Build the per-bar DataFrame.
    base = {
        # ID
        "trade_id": trade_id,
        "ticker": ticker,
        "date": date,
        "direction": direction,
        "entry_time": entry_time,
        "bar_ts": ts.to_numpy(),
        "bar_idx": bar_idx,
        # Time
        "minutes_since_entry": minutes_since_entry,
        "minutes_to_close": minutes_to_close,
        "minutes_since_open": minutes_since_open,
        # Trade-level static
        "entry_price": entry_price,
        "risk_per_share": risk,
        # Bar OHLCV
        "bar_open": opens,
        "bar_high": highs,
        "bar_low": lows,
        "bar_close": closes,
        "bar_volume": vols,
        # Bar dynamics
        "bar_return_pct": bar_return_pct,
        "bar_close_loc": bar_close_loc,
        # Trade R state
        "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_from_peak,
        "bars_since_peak": bars_since_peak,
        # Liquidity
        "vwap": vwap,
        "vwap_dev_pct": vwap_dev_pct,
        "vol_vs_first_bar": vol_vs_first,
        "vol_vs_avg_dvol30d": vol_vs_avg,
        # Market context
        "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,
        # Baseline benchmark anchors (NOT features for training; used in eval)
        "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,
        "baseline_exit_reason": t.get("exit_reason"),
    }

    out = pd.DataFrame(base)
    # Broadcast trade-level passthrough columns (entry-time-known features).
    for col in PASSTHROUGH_COLS:
        out[col] = t.get(col)
    return out


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--run", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--intraday-cache", default=str(DEFAULT_INTRADAY_CACHE))
    args = ap.parse_args()

    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)
    intraday_root = Path(args.intraday_cache)

    payload = json.loads(Path(args.run).read_text())
    trades = payload["trades"]
    print(f"Loaded {len(trades)} trades")

    # Pre-load SPY / QQQ once per unique trade date.
    unique_dates = sorted({t["date"] for t in trades})
    spy_by_date: dict[str, pd.DataFrame] = {}
    qqq_by_date: dict[str, pd.DataFrame] = {}
    for d in unique_dates:
        s = _load_intraday("SPY", d, intraday_root)
        q = _load_intraday("QQQ", d, intraday_root)
        if s is not None:
            spy_by_date[d] = s
        if q is not None:
            qqq_by_date[d] = q
    print(f"Loaded SPY for {len(spy_by_date)}/{len(unique_dates)} dates, QQQ for {len(qqq_by_date)}")

    frames = []
    skipped = 0
    for i, t in enumerate(trades):
        try:
            df = build_for_trade(i, t, intraday_root, spy_by_date, qqq_by_date)
            if df is None or df.empty:
                skipped += 1
                continue
            frames.append(df)
        except Exception as e:
            print(f"  trade {i} {t.get('ticker')} {t.get('date')}: {e}")
            skipped += 1
        if (i + 1) % 25 == 0:
            print(f"  processed {i+1}/{len(trades)}")

    if not frames:
        raise SystemExit("No trades produced rows")

    big = pd.concat(frames, ignore_index=True)
    parquet = out_dir / "per_bar_states.parquet"
    big.to_parquet(parquet, index=False)
    print(
        f"\nWrote {parquet} | trades: {big['trade_id'].nunique()} | "
        f"rows: {len(big):,} | columns: {big.shape[1]} | skipped: {skipped}"
    )

    # Quick sanity report.
    rep_path = out_dir / "summary.md"
    rep = build_summary(big, payload.get("metrics", {}))
    rep_path.write_text(rep)
    print(f"Wrote {rep_path}")


def build_summary(df: pd.DataFrame, run_meta: dict) -> str:
    n_trades = df["trade_id"].nunique()
    n_rows = len(df)
    bars_per_trade = df.groupby("trade_id").size()
    cov = {
        "spy_return_since_entry": df["spy_return_since_entry"].notna().mean(),
        "qqq_return_since_entry": df["qqq_return_since_entry"].notna().mean(),
        "vwap_dev_pct": df["vwap_dev_pct"].notna().mean(),
        "vol_vs_avg_dvol30d": df["vol_vs_avg_dvol30d"].notna().mean(),
    }
    lines = [
        "# Per-Bar State Dataset",
        "",
        f"**Run:** `{run_meta.get('run_id')}` "
        f"({run_meta.get('start_date')} → {run_meta.get('end_date')}, "
        f"{run_meta.get('total_trades')} entries)",
        "",
        f"- Trades expanded: **{n_trades}**",
        f"- Total decision-point rows: **{n_rows:,}**",
        f"- Bars per trade: median {int(bars_per_trade.median())}, "
        f"min {int(bars_per_trade.min())}, max {int(bars_per_trade.max())}",
        f"- Schema: **{df.shape[1]}** columns",
        "",
        "## Feature coverage (non-null fraction)",
        "",
        "| feature | coverage |",
        "|---|---|",
    ]
    for k, v in cov.items():
        lines.append(f"| {k} | {v:.1%} |")
    lines += [
        "",
        "## Sanity distributions (random sample of decision points)",
        "",
        f"- current_close_r quartiles: {df['current_close_r'].quantile([0.1,0.25,0.5,0.75,0.9]).round(3).to_dict()}",
        f"- mfe_so_far_r quartiles:    {df['mfe_so_far_r'].quantile([0.1,0.25,0.5,0.75,0.9]).round(3).to_dict()}",
        f"- mae_so_far_r quartiles:    {df['mae_so_far_r'].quantile([0.1,0.25,0.5,0.75,0.9]).round(3).to_dict()}",
        f"- giveback_from_peak_r:      {df['giveback_from_peak_r'].quantile([0.1,0.25,0.5,0.75,0.9]).round(3).to_dict()}",
        f"- vwap_dev_pct:              {df['vwap_dev_pct'].quantile([0.1,0.25,0.5,0.75,0.9]).round(4).to_dict()}",
        f"- spy_return_since_entry:    {df['spy_return_since_entry'].quantile([0.1,0.25,0.5,0.75,0.9]).round(4).to_dict()}",
        "",
        "## Held vs post-baseline-exit bars",
        "",
        f"- Bars while baseline still held: **{int(df['is_held_by_baseline'].sum()):,}** "
        f"({df['is_held_by_baseline'].mean():.1%})",
        f"- Bars after baseline exited (extension data, ML may use to learn 'hold longer'): "
        f"**{int((~df['is_held_by_baseline']).sum()):,}**",
        "",
    ]
    return "\n".join(lines)


if __name__ == "__main__":
    main()