"""ORB Exit Rule Sweeper.

Replays entry events from an ORB intraday backtest run against 5-min bars under
alternative exit rule combinations. Pure post-hoc simulation — does not change
entries, sizing, or capital allocation. Each trade gets a per-rule realized_R.

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

import argparse
import itertools
import json
from dataclasses import dataclass
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


@dataclass(frozen=True)
class Rule:
    profit_target: float | None
    trailing_activation: float | None
    trailing_dist: float | None
    partial_at: float | None  # None or e.g. 0.5
    partial_size: float = 0.5
    move_stop_to_be_after_partial: bool = True
    time_no_new_high: int | None = None  # minutes; if no new running peak in this many minutes, exit
    initial_stop: float = -1.0  # in R, baseline matches V49.91 ATR stop sizing

    def label(self) -> str:
        parts = []
        parts.append(f"pt={self.profit_target}" if self.profit_target is not None else "pt=none")
        if self.trailing_activation is not None and self.trailing_dist is not None:
            parts.append(f"tr={self.trailing_activation}/{self.trailing_dist}")
        else:
            parts.append("tr=none")
        if self.partial_at is not None:
            parts.append(f"px={self.partial_at}")
        else:
            parts.append("px=none")
        parts.append(f"tNH={self.time_no_new_high}" if self.time_no_new_high else "tNH=none")
        return "|".join(parts)


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)


@dataclass
class Bars:
    """Per-trade favorable-direction R-multiple time series."""
    ts_min: np.ndarray
    fav_high: np.ndarray
    fav_low: np.ndarray
    fav_close: np.ndarray
    fav_next_open: np.ndarray  # last bar's "next open" = last close

    @classmethod
    def from_trade(cls, t: dict, intraday_root: Path) -> "Bars | None":
        ticker = t["ticker"]
        date = t["date"]
        bar_path = intraday_root / ticker / f"{date}.parquet"
        if not bar_path.exists():
            return None
        entry_time = pd.Timestamp(t["entry_time"]).tz_convert("UTC")
        entry_price = float(t["entry_price"])
        direction = t.get("orb_direction", "long")
        risk = _infer_risk_per_share(t)

        et_close_local = pd.Timestamp(date, tz=ET) + pd.Timedelta(hours=16)
        et_close_utc = et_close_local.tz_convert("UTC")

        bars = pd.read_parquet(bar_path)
        bars["ts"] = pd.to_datetime(bars["timestamp"], utc=True)
        win = (
            bars[(bars["ts"] >= entry_time) & (bars["ts"] <= et_close_utc)]
            .sort_values("ts")
            .reset_index(drop=True)
        )
        if win.empty:
            return None

        ts_min = ((win["ts"] - entry_time).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()
        next_open = np.append(opens[1:], closes[-1])

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

        return cls(ts_min, fav_high, fav_low, fav_close, fav_next_open)


def simulate(bars: Bars, rule: Rule) -> tuple[float, str, int]:
    """Return (realized_R, exit_reason, exit_bar_index) under the rule.

    Same-bar precedence: stop > target > partial. If multiple events trigger
    inside one 5-min bar (5-min ambiguity), we always assume the adverse
    outcome wins to stay conservative.
    """
    realized_r = 0.0
    pos = 1.0
    stop_r = rule.initial_stop
    partial_done = False
    running_peak = -np.inf
    last_new_high_min = bars.ts_min[0]

    n = len(bars.ts_min)
    for i in range(n):
        fhi, flo, fcl, fno = (
            bars.fav_high[i],
            bars.fav_low[i],
            bars.fav_close[i],
            bars.fav_next_open[i],
        )
        # 1. Stop check (conservative — assume bar low touches first if it crosses stop).
        if flo <= stop_r:
            realized_r += pos * stop_r
            return realized_r, "stop", i
        # 2. Profit target.
        if rule.profit_target is not None and fhi >= rule.profit_target:
            realized_r += pos * rule.profit_target
            return realized_r, "target", i
        # 3. Partial exit (one-shot).
        if (
            not partial_done
            and rule.partial_at is not None
            and fhi >= rule.partial_at
        ):
            realized_r += rule.partial_size * rule.partial_at
            pos -= rule.partial_size
            partial_done = True
            if rule.move_stop_to_be_after_partial:
                stop_r = max(stop_r, 0.0)
        # 4. Update running peak / trailing.
        if fhi > running_peak:
            running_peak = fhi
            last_new_high_min = bars.ts_min[i]
        if (
            rule.trailing_activation is not None
            and rule.trailing_dist is not None
            and running_peak >= rule.trailing_activation
        ):
            stop_r = max(stop_r, running_peak - rule.trailing_dist)
        # 5. Time-no-new-high stop (exit at next-bar open or, on last bar, current close).
        if (
            rule.time_no_new_high is not None
            and running_peak >= 0  # only after we ever went green
            and (bars.ts_min[i] - last_new_high_min) >= rule.time_no_new_high
        ):
            exec_r = fno if i < n - 1 else fcl
            realized_r += pos * exec_r
            return realized_r, "no_new_high", i

    # EOD: exit remaining at last close.
    realized_r += pos * bars.fav_close[-1]
    return realized_r, "eod_close", n - 1


def make_grid() -> list[Rule]:
    profit_targets = [None, 0.8, 1.2]
    trail_specs = [
        (None, None),
        (0.5, 0.3),
        (1.0, 0.3),
        (1.0, 0.5),
    ]
    partial_specs = [(None, 0.0), (0.5, 0.5)]
    time_nh = [None, 30, 60]

    out = []
    for pt, (ta, td), (pa, ps), tnh in itertools.product(
        profit_targets, trail_specs, partial_specs, time_nh
    ):
        out.append(
            Rule(
                profit_target=pt,
                trailing_activation=ta,
                trailing_dist=td,
                partial_at=pa,
                partial_size=ps,
                time_no_new_high=tnh,
            )
        )
    return out


def evaluate_rule(
    rule: Rule, bars_list: list[Bars], realized_baseline: np.ndarray
) -> dict:
    rs = np.zeros(len(bars_list))
    reasons = []
    for i, b in enumerate(bars_list):
        r, why, _ = simulate(b, rule)
        rs[i] = r
        reasons.append(why)
    arr = pd.Series(rs)
    delta = rs - realized_baseline
    return {
        "rule": rule.label(),
        "profit_target": rule.profit_target,
        "trailing_activation": rule.trailing_activation,
        "trailing_dist": rule.trailing_dist,
        "partial_at": rule.partial_at,
        "time_no_new_high": rule.time_no_new_high,
        "n": len(rs),
        "total_r": float(rs.sum()),
        "mean_r": float(arr.mean()),
        "median_r": float(arr.median()),
        "win_rate": float((arr > 0).mean()),
        "p25_r": float(arr.quantile(0.25)),
        "p75_r": float(arr.quantile(0.75)),
        "worst_r": float(arr.min()),
        "best_r": float(arr.max()),
        "delta_total_r": float(delta.sum()),
        "delta_mean_r": float(delta.mean()),
        "delta_median_r": float(np.median(delta)),
        "wins_added": int(((rs > 0) & (realized_baseline <= 0)).sum()),
        "wins_lost": int(((rs <= 0) & (realized_baseline > 0)).sum()),
        "reason_counts": dict(pd.Series(reasons).value_counts().to_dict()),
    }


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")

    bars_list: list[Bars] = []
    realized_baseline = []
    trade_meta = []
    skipped = 0
    for t in trades:
        b = Bars.from_trade(t, intraday_root)
        if b is None:
            skipped += 1
            continue
        bars_list.append(b)
        realized_baseline.append(float(t.get("r_multiple_at_exit") or 0.0))
        trade_meta.append(
            {
                "ticker": t["ticker"],
                "date": t["date"],
                "entry_time": t["entry_time"],
                "rvol": t.get("rvol"),
                "gap_pct": t.get("gap_pct"),
                "entry_dollar_volume": t.get("entry_dollar_volume"),
                "candidate_score": t.get("candidate_score"),
            }
        )
    realized_baseline = np.asarray(realized_baseline)
    print(f"Prepared {len(bars_list)} trade bar series (skipped {skipped})")

    grid = make_grid()
    print(f"Evaluating {len(grid)} rule combos")

    results = []
    per_trade_long = []  # rule × trade for stratification
    for k, rule in enumerate(grid):
        row = evaluate_rule(rule, bars_list, realized_baseline)
        results.append(row)
        # also persist per-trade R under this rule for downstream stratification
        for i, b in enumerate(bars_list):
            r, why, _ = simulate(b, rule)
            per_trade_long.append(
                {
                    "rule": rule.label(),
                    **trade_meta[i],
                    "realized_r_baseline": realized_baseline[i],
                    "realized_r_rule": r,
                    "exit_reason_rule": why,
                }
            )
        if (k + 1) % 10 == 0:
            print(f"  evaluated {k+1}/{len(grid)} rules")

    df = pd.DataFrame(results).sort_values("total_r", ascending=False).reset_index(drop=True)
    long_df = pd.DataFrame(per_trade_long)

    # Baseline reference row
    baseline_total = float(realized_baseline.sum())
    baseline_mean = float(realized_baseline.mean())

    df.to_parquet(out_dir / "rule_results.parquet", index=False)
    long_df.to_parquet(out_dir / "per_trade_long.parquet", index=False)
    print(f"Wrote rule_results.parquet ({len(df)} rules) and per_trade_long.parquet ({len(long_df)} rows)")

    # Markdown summary
    md = build_markdown(df, long_df, baseline_total, baseline_mean, payload.get("metrics", {}))
    (out_dir / "summary.md").write_text(md)
    print(f"Wrote summary.md")


def build_markdown(
    df: pd.DataFrame,
    long_df: pd.DataFrame,
    baseline_total: float,
    baseline_mean: float,
    run_meta: dict,
) -> str:
    n_trades = int(df.iloc[0]["n"]) if len(df) else 0
    lines = [
        "# ORB Exit Rule Sweep — V49.91 baseline 200d",
        "",
        f"**Run:** `{run_meta.get('run_id')}` "
        f"({run_meta.get('start_date')} → {run_meta.get('end_date')}, "
        f"{n_trades} trades audited)",
        "",
        f"**Baseline (V49.91 actual):** total_R = **{baseline_total:.2f}**, "
        f"mean_R = **{baseline_mean:.3f}**",
        "",
        "## Top 15 rules by total_R",
        "",
        "| rank | rule | total_R | Δ vs base | mean_R | win | p25 | p75 | worst | wins+ | wins- |",
        "|---|---|---|---|---|---|---|---|---|---|---|",
    ]
    for i, row in df.head(15).iterrows():
        lines.append(
            f"| {i+1} | `{row['rule']}` | "
            f"{row['total_r']:.2f} | "
            f"{row['delta_total_r']:+.2f} | "
            f"{row['mean_r']:.3f} | "
            f"{row['win_rate']:.1%} | "
            f"{row['p25_r']:.2f} | {row['p75_r']:.2f} | "
            f"{row['worst_r']:.2f} | "
            f"{int(row['wins_added'])} | {int(row['wins_lost'])} |"
        )
    lines.append("")

    lines.append("## Worst 5 rules (sanity check)")
    lines.append("")
    lines.append("| rule | total_R | mean_R | win |")
    lines.append("|---|---|---|---|")
    for _, row in df.tail(5).iterrows():
        lines.append(
            f"| `{row['rule']}` | {row['total_r']:.2f} | "
            f"{row['mean_r']:.3f} | {row['win_rate']:.1%} |"
        )
    lines.append("")

    # Plateau check: median rule + neighbors by single-axis variation
    lines.append("## Plateau sensitivity (top rule and ±1 step neighbors)")
    lines.append("")
    top = df.iloc[0]
    lines.append(f"Top rule: `{top['rule']}` total_R={top['total_r']:.2f}")
    lines.append("")
    lines.append("Distribution of total_R across all rules:")
    lines.append(
        f"- min={df['total_r'].min():.2f}, "
        f"p25={df['total_r'].quantile(0.25):.2f}, "
        f"median={df['total_r'].median():.2f}, "
        f"p75={df['total_r'].quantile(0.75):.2f}, "
        f"max={df['total_r'].max():.2f}"
    )
    lines.append(
        f"- rules with total_R within 5% of best: "
        f"**{int((df['total_r'] >= top['total_r'] * 0.95).sum())}**"
    )
    lines.append(
        f"- rules with total_R better than baseline: "
        f"**{int((df['total_r'] > baseline_total).sum())}/{len(df)}**"
    )
    lines.append("")

    # Stratification of top rule vs baseline by time-to-peak bucket
    lines.append("## Top rule vs baseline by time-to-peak bucket")
    lines.append("")
    bins = [-1, 5, 15, 30, 60, 120, 240, 1000]
    labels = ["≤5m", "5-15m", "15-30m", "30-60m", "1-2h", "2-4h", ">4h"]

    audit_path = Path(run_meta.get("_audit_parquet", "")) if run_meta else None
    audit_df = None
    candidate_audit = Path("tmp/orb_exit_audit_v49_91_200d/trade_audit.parquet")
    if candidate_audit.exists():
        audit_df = pd.read_parquet(candidate_audit)

    if audit_df is not None:
        merged = (
            long_df[long_df["rule"] == top["rule"]]
            .merge(
                audit_df[["ticker", "date", "minutes_to_peak"]],
                on=["ticker", "date"],
                how="left",
            )
        )
        merged["peak_bucket"] = pd.cut(merged["minutes_to_peak"], bins=bins, labels=labels)
        agg = merged.groupby("peak_bucket", observed=False).agg(
            n=("realized_r_rule", "size"),
            base_mean=("realized_r_baseline", "mean"),
            rule_mean=("realized_r_rule", "mean"),
            base_win=("realized_r_baseline", lambda s: float((s > 0).mean())),
            rule_win=("realized_r_rule", lambda s: float((s > 0).mean())),
            base_total=("realized_r_baseline", "sum"),
            rule_total=("realized_r_rule", "sum"),
        )
        lines.append(
            "| bucket | n | base mean | rule mean | Δmean | base win | rule win | base total | rule total | Δtotal |"
        )
        lines.append("|" + "---|" * 10)
        for k, r in agg.iterrows():
            lines.append(
                f"| {k} | {int(r['n'])} | "
                f"{r['base_mean']:.3f} | {r['rule_mean']:.3f} | "
                f"{r['rule_mean'] - r['base_mean']:+.3f} | "
                f"{r['base_win']:.1%} | {r['rule_win']:.1%} | "
                f"{r['base_total']:.2f} | {r['rule_total']:.2f} | "
                f"{r['rule_total'] - r['base_total']:+.2f} |"
            )
        lines.append("")

    # Rule-axis sensitivity tables: marginal effect of one knob holding others
    for axis, col in [
        ("profit_target", "profit_target"),
        ("trailing", "trailing_activation"),
        ("partial_at", "partial_at"),
        ("time_no_new_high", "time_no_new_high"),
    ]:
        agg = df.groupby(col, dropna=False).agg(
            n=("total_r", "size"),
            best_total=("total_r", "max"),
            mean_total=("total_r", "mean"),
            median_total=("total_r", "median"),
            best_mean_r=("mean_r", "max"),
        )
        lines.append(f"### Marginal effect: {axis}")
        lines.append("")
        lines.append("| value | n_combos | best total_R | mean total_R | median total_R |")
        lines.append("|---|---|---|---|---|")
        for k, r in agg.iterrows():
            lines.append(
                f"| {k} | {int(r['n'])} | {r['best_total']:.2f} | "
                f"{r['mean_total']:.2f} | {r['median_total']:.2f} |"
            )
        lines.append("")

    return "\n".join(lines)


if __name__ == "__main__":
    main()