"""
Gap-Down ORB Short Diagnostic v2

Improvements over v1:
  - Market regime filter: QQQ also gapped down (opens below prev_close) — only trade on
    "down market" days. This is the INVERSE of V23 (V23 needs QQQ up), creating structural
    regime-orthogonality. On QQQ-up days, gap-down stocks often recover → stops. On
    QQQ-down days, gap-down stocks tend to continue lower → wins.
  - Slightly tighter stop: ATR * 0.75 (was 1.0) to reduce loss size when wrong
  - Target: ATR * 1.5 below entry (unchanged) — still testing if stocks fall far enough
  - EOD exit as fallback (unchanged)

Hypothesis upgrade: on days when V23 is FORCED OFF (QQQ negative), gap-down shorts
should work better because the market backdrop confirms downward bias.

Gates (same):
  - WR ≥ 42%
  - avg_win / avg_loss ≥ 1.2
  - Pearson corr with V23 daily PnL ≤ 0.00 (expect negative on QQQ-down days)
  - Total PnL > 0 (added gate)
"""
from __future__ import annotations

import json
import math
import os
import sys

import numpy as np
import pandas as pd
import yaml

CACHE_DIR = "data/cache/intraday"
DAILY_CACHE_DIR = "data/cache/daily"
UNIVERSE_FILE = "configs/symbols_midlarge_snapshot_exact.yaml"
V23_BASELINE_RUN = "runs/intraday_orb/intraday_20260420_205136_5597b16d.json"
QQQ_TICKER = "QQQ"

MIN_PRICE = 10.0
MIN_GAP_DOWN = -0.02
MAX_GAP_DOWN = -0.12
MIN_AVG_DOLLAR_VOL = 25_000_000
MIN_RVOL = 1.5
ATR_STOP_MULT = 0.75     # tighter stop vs v1
ATR_TARGET_MULT = 1.5
RISK_PER_TRADE = 500.0
MAX_SIMULTANEOUS = 3
EXIT_HOUR = 15
EXIT_MIN = 55
QQQ_REGIME_THRESHOLD = 0.0   # QQQ gap must be ≤ 0 (flat or negative open)


def load_universe() -> list[str]:
    with open(UNIVERSE_FILE) as f:
        data = yaml.safe_load(f)
    return data.get("symbols", data) if isinstance(data, dict) else data


def load_daily_cache(ticker: str) -> pd.DataFrame | None:
    path = f"{DAILY_CACHE_DIR}/{ticker}.parquet"
    if not os.path.exists(path):
        return None
    try:
        return pd.read_parquet(path)
    except Exception:
        return None


def load_bars(ticker: str, date: str) -> pd.DataFrame | None:
    path = f"{CACHE_DIR}/{ticker}/{date}.parquet"
    if not os.path.exists(path):
        return None
    try:
        df = pd.read_parquet(path)
        if len(df) < 10:
            return None
        df["ts"] = pd.to_datetime(df["timestamp"]).dt.tz_convert("US/Eastern")
        df["hour"] = df["ts"].dt.hour
        df["minute"] = df["ts"].dt.minute
        df = df[(df["hour"] >= 9) & (df["hour"] < 16)].copy()
        df = df.sort_values("ts").reset_index(drop=True)
        return df
    except Exception:
        return None


def approximate_atr(df: pd.DataFrame) -> float:
    return (df["high"] - df["low"]).mean() * math.sqrt(78)


def get_qqq_gap(date: str, qqq_daily: pd.DataFrame) -> float | None:
    """QQQ gap = (today_open - prev_close) / prev_close."""
    idx_list = qqq_daily.index[qqq_daily["date"] == date].tolist()
    if not idx_list or idx_list[0] == 0:
        return None
    i = idx_list[0]
    prev_close = qqq_daily.iloc[i - 1]["close"]
    today_open = qqq_daily.iloc[i]["open"]
    if prev_close <= 0:
        return None
    return (today_open - prev_close) / prev_close


def simulate_day(date: str, tickers: list[str],
                 daily_data: dict[str, pd.DataFrame],
                 qqq_daily: pd.DataFrame) -> dict:
    # Regime gate: QQQ must have opened flat or negative
    qqq_gap = get_qqq_gap(date, qqq_daily)
    if qqq_gap is None or qqq_gap > QQQ_REGIME_THRESHOLD:
        return {"date": date, "trades": [], "day_pnl": 0.0, "skip_reason": "qqq_up"}

    day_trades: list[dict] = []
    open_positions: int = 0

    candidates = []
    for ticker in tickers:
        if ticker not in daily_data:
            continue
        df = daily_data[ticker]
        idx = df.index[df["date"] == date].tolist()
        if not idx or idx[0] == 0:
            continue
        row_idx = idx[0]
        today = df.iloc[row_idx]
        prev = df.iloc[row_idx - 1]
        if prev["close"] <= 0:
            continue

        gap = (today["open"] - prev["close"]) / prev["close"]
        if gap > MIN_GAP_DOWN or gap < MAX_GAP_DOWN:
            continue

        avg_dvol = (
            df["close"].iloc[max(0, row_idx - 20):row_idx].mean() *
            df["volume"].iloc[max(0, row_idx - 20):row_idx].mean()
        )
        if avg_dvol < MIN_AVG_DOLLAR_VOL:
            continue

        candidates.append((ticker, gap, prev["close"], today["open"]))

    candidates.sort(key=lambda x: x[1])  # largest gap-down first

    for ticker, gap, prev_close, today_open in candidates:
        if open_positions >= MAX_SIMULTANEOUS:
            break

        bars = load_bars(ticker, date)
        if bars is None or len(bars) < 10:
            continue

        if bars.iloc[0]["open"] < MIN_PRICE:
            continue

        orb_bars = bars[(bars["hour"] == 9) & (bars["minute"] >= 30) & (bars["minute"] < 35)]
        if len(orb_bars) == 0:
            continue
        orb = orb_bars.iloc[0]

        if orb["close"] >= orb["open"]:  # must be bearish ORB
            continue

        avg_daily = daily_data[ticker]["volume"].mean() if ticker in daily_data else 0
        first_vol = bars.iloc[0]["volume"]
        rvol = first_vol / (avg_daily / 78) if avg_daily > 0 else 0
        if rvol < MIN_RVOL:
            continue

        atr = approximate_atr(bars)
        if atr <= 0:
            continue

        orb_low = orb["low"]
        orb_high = orb["high"]
        entry_trigger = orb_low

        stop_dist = ATR_STOP_MULT * atr
        stop_price = entry_trigger + stop_dist

        if stop_price <= entry_trigger:
            continue
        shares = RISK_PER_TRADE / stop_dist
        if shares <= 0 or shares * entry_trigger > 50000:
            continue

        # Scan post-ORB bars
        post_orb = bars[bars.index > orb_bars.index[0]]
        entry_price = None
        exit_price = None
        exit_reason = "no_entry"
        profit_target: float = 0.0

        for _, bar in post_orb.iterrows():
            if entry_price is None:
                if bar["low"] <= entry_trigger:
                    entry_price = min(entry_trigger, bar["open"])
                    profit_target = entry_price - ATR_TARGET_MULT * atr
                    if bar["high"] >= stop_price:
                        exit_price = stop_price
                        exit_reason = "stop"
                        break
                    continue
            else:
                if bar["hour"] >= EXIT_HOUR and bar["minute"] >= EXIT_MIN:
                    exit_price = bar["close"]
                    exit_reason = "eod"
                    break
                if bar["high"] >= stop_price or bar["close"] >= prev_close:
                    exit_price = max(stop_price, bar["open"])
                    exit_reason = "stop"
                    break
                if bar["low"] <= profit_target:
                    exit_price = max(profit_target, bar["open"])
                    exit_reason = "target"
                    break

        if entry_price is None:
            continue

        if exit_price is None:
            exit_price = bars.iloc[-1]["close"]
            exit_reason = "eod"

        pnl = (entry_price - exit_price) * shares  # short pnl

        day_trades.append({
            "ticker": ticker,
            "date": date,
            "gap": gap,
            "qqq_gap": qqq_gap,
            "entry": entry_price,
            "exit": exit_price,
            "pnl": pnl,
            "win": pnl > 0,
            "exit_reason": exit_reason,
        })
        open_positions += 1

    day_pnl = sum(t["pnl"] for t in day_trades)
    return {"date": date, "trades": day_trades, "day_pnl": day_pnl, "skip_reason": None}


def main() -> None:
    print("=== Gap-Down ORB Short Diagnostic v2 (QQQ-down regime) ===\n")

    with open(V23_BASELINE_RUN) as f:
        v23_data = json.load(f)
    v23_dates = [d["date"] for d in v23_data["daily_summary"]]
    v23_pnl = {d["date"]: d["daily_pnl"] for d in v23_data["daily_summary"]}

    universe = load_universe()
    print(f"Universe: {len(universe)} tickers")
    print("Loading daily caches...", flush=True)

    daily_data: dict[str, pd.DataFrame] = {}
    for ticker in universe:
        df = load_daily_cache(ticker)
        if df is not None and len(df) >= 2:
            daily_data[ticker] = df

    qqq_daily = load_daily_cache(QQQ_TICKER)
    if qqq_daily is None:
        print("ERROR: QQQ daily cache not found")
        return

    print(f"Daily cache loaded: {len(daily_data)} tickers")
    print(f"V23 window: {v23_dates[0]} → {v23_dates[-1]} ({len(v23_dates)} days)\n")

    all_results = []
    for i, date in enumerate(v23_dates):
        result = simulate_day(date, universe, daily_data, qqq_daily)
        all_results.append(result)
        if (i + 1) % 20 == 0:
            print(f"  {i+1}/{len(v23_dates)} days processed...", flush=True)

    all_trades = [t for r in all_results for t in r["trades"]]
    total_trades = len(all_trades)
    active_days = sum(1 for r in all_results if r.get("skip_reason") != "qqq_up")
    trade_days = sum(1 for r in all_results if r["trades"])
    skipped_days = sum(1 for r in all_results if r.get("skip_reason") == "qqq_up")

    wins = sum(1 for t in all_trades if t["win"])
    win_rate = wins / total_trades if total_trades else 0.0
    avg_win = float(np.mean([t["pnl"] for t in all_trades if t["pnl"] > 0])) if wins > 0 else 0.0
    avg_loss = float(np.mean([abs(t["pnl"]) for t in all_trades if t["pnl"] <= 0])) if total_trades - wins > 0 else 0.0
    total_pnl = sum(t["pnl"] for t in all_trades)
    exit_reasons = {}
    for t in all_trades:
        exit_reasons[t["exit_reason"]] = exit_reasons.get(t["exit_reason"], 0) + 1

    gd_daily = {r["date"]: r["day_pnl"] for r in all_results}
    gd_series = [gd_daily.get(d, 0.0) for d in v23_dates]
    v23_series = [v23_pnl.get(d, 0.0) for d in v23_dates]
    corr = float(np.corrcoef(gd_series, v23_series)[0, 1]) if len(v23_dates) > 1 else 0.0

    print("=" * 60)
    print(f"  QQQ-down days (eligible):  {active_days} / {len(v23_dates)}")
    print(f"  QQQ-up days (skipped):     {skipped_days}")
    print(f"  Total trades:              {total_trades}")
    print(f"  Trade days:                {trade_days}")
    print(f"  Win rate:                  {win_rate*100:.1f}%  (gate: ≥42%)")
    print(f"  Avg win:                   ${avg_win:.2f}")
    print(f"  Avg loss:                  ${avg_loss:.2f}")
    if avg_loss > 0:
        print(f"  Win/Loss ratio:            {avg_win/avg_loss:.2f}  (gate: ≥1.2)")
    print(f"  Total PnL:                 ${total_pnl:+.2f}  (gate: >0)")
    print(f"  Corr vs V23:               {corr:.3f}  (gate: ≤0.00)")
    print(f"  Exit breakdown:            {exit_reasons}")
    print("=" * 60)

    g1 = win_rate >= 0.42
    g2 = (avg_win / avg_loss >= 1.2) if avg_loss > 0 else False
    g3 = corr <= 0.00
    g4 = total_trades >= 30
    g5 = total_pnl > 0
    print(f"\nGate G1 (WR ≥ 42%):         {'PASS' if g1 else 'FAIL'}")
    print(f"Gate G2 (W/L ≥ 1.2):        {'PASS' if g2 else 'FAIL'}")
    print(f"Gate G3 (corr ≤ 0.00):      {'PASS' if g3 else 'FAIL'}")
    print(f"Gate G4 (trades ≥ 30):      {'PASS' if g4 else 'FAIL'}")
    print(f"Gate G5 (PnL > 0):          {'PASS' if g5 else 'FAIL'}")
    passed = sum([g1, g2, g3, g4, g5])
    verdict = "PROCEED TO ENGINE BUILD" if passed == 5 else f"ABORT ({passed}/5 gates passed)"
    print(f"\nVERDICT: {verdict}")

    out = {
        "total_trades": total_trades,
        "active_days": active_days,
        "trade_days": trade_days,
        "skipped_days": skipped_days,
        "win_rate": win_rate,
        "avg_win": avg_win,
        "avg_loss": avg_loss,
        "total_pnl": total_pnl,
        "corr_v23": corr,
        "exit_reasons": exit_reasons,
        "trades": all_trades,
        "daily_pnl": [{"date": r["date"], "pnl": r["day_pnl"]} for r in all_results],
    }
    out_path = "runs/intraday_orb/diag_gapdown_short_v2.json"
    with open(out_path, "w") as f:
        json.dump(out, f, indent=2, default=str)
    print(f"\nResults saved to {out_path}")


if __name__ == "__main__":
    sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../..")))
    main()
