"""
Afternoon Momentum Diagnostic (Option 2)

Hypothesis: stocks that form a midday consolidation (10:00-12:30 ET) and then break out
in the afternoon (12:30-15:00 ET) have positive edge that is regime-orthogonal to V23.

Setup:
  - Universe: midlarge tickers with intraday cache
  - Morning/midday range: high of bars 9:30-12:30 ET
  - Afternoon breakout: first bar close above morning range high after 12:30 ET
  - Volume filter: breakout bar volume > 1.5x mean of prior 12 bars (~1 hour)
  - Entry: open of the bar following the breakout bar (next bar)
  - Stop: ATR14 * 0.75 below entry  (ATR14 approximated from 5-min ATR * sqrt(78))
  - Exit: EOD (15:55 ET bar close)
  - Risk per trade: 5% of $10,000 = $500 (same as V23)
  - Max simultaneous: 3 (same as V23)

Gates:
  - WR ≥ 50%
  - avg_win / avg_loss ≥ 1.0
  - Pearson corr with V23 daily PnL ≤ 0.25
"""
from __future__ import annotations

import glob
import json
import math
import os
import sys
from collections import defaultdict
from pathlib import Path

import numpy as np
import pandas as pd
import yaml

# ── Config ─────────────────────────────────────────────────────────────────

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

MIN_PRICE = 10.0
MIN_AVG_VOL_BARS = 5        # minimum bars in midday range
MIDDAY_END_HOUR = 12        # 12:30 ET
MIDDAY_END_MIN = 30
PM_ENTRY_HOUR = 12
PM_ENTRY_MIN = 30
PM_CUTOFF_HOUR = 15         # no new entries after 15:00 ET
EXIT_HOUR = 15
EXIT_MIN = 55
VOL_CONFIRM_MULT = 1.5      # breakout bar vol > 1.5x prior-hour avg
RISK_PER_TRADE = 500.0      # $500 = 5% of $10k
ATR_STOP_MULT = 0.75
MAX_SIMULTANEOUS = 3


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_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) < 20:
            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:
    """Estimate daily ATR from 5-min bars: sqrt(78) * avg 5-min range."""
    if len(df) < 10:
        return 0.0
    ranges = df["high"] - df["low"]
    avg_5min_range = ranges.mean()
    return avg_5min_range * math.sqrt(78)


def simulate_day(date: str, tickers: list[str]) -> dict:
    """Run afternoon momentum simulation for one day."""
    day_trades: list[dict] = []
    open_positions: list[dict] = []

    for ticker in tickers:
        bars = load_bars(ticker, date)
        if bars is None or len(bars) < 20:
            continue

        # Quality: price
        first_bar = bars.iloc[0]
        if first_bar["open"] < MIN_PRICE:
            continue

        # Split into morning/midday and afternoon
        midday_mask = (
            (bars["hour"] < MIDDAY_END_HOUR) |
            ((bars["hour"] == MIDDAY_END_HOUR) & (bars["minute"] < MIDDAY_END_MIN))
        )
        morning_bars = bars[midday_mask]
        afternoon_bars = bars[
            (bars["hour"] > PM_ENTRY_HOUR) |
            ((bars["hour"] == PM_ENTRY_HOUR) & (bars["minute"] >= PM_ENTRY_MIN))
        ]

        if len(morning_bars) < MIN_AVG_VOL_BARS or len(afternoon_bars) < 2:
            continue

        # Morning range high (consolidation ceiling)
        morning_high = morning_bars["high"].max()

        # ATR estimate for stop
        atr = approximate_atr(bars)
        if atr <= 0:
            continue

        # Scan for breakout in afternoon
        breakout_idx = None
        for i in range(len(afternoon_bars) - 1):
            bar = afternoon_bars.iloc[i]
            # Stop scanning if too late
            if bar["hour"] >= PM_CUTOFF_HOUR:
                break
            # Already have a position in this ticker?
            if any(p["ticker"] == ticker for p in open_positions):
                break
            # Check if close above morning high
            if bar["close"] > morning_high:
                # Volume confirmation: compare to prior 12 bars
                bar_idx = afternoon_bars.index[i]
                prior_start = max(0, bar_idx - 12)
                prior_vol = bars.loc[prior_start:bar_idx - 1, "volume"].mean()
                if prior_vol > 0 and bar["volume"] < VOL_CONFIRM_MULT * prior_vol:
                    continue
                # Entry on next bar open
                next_bar = afternoon_bars.iloc[i + 1]
                entry_price = next_bar["open"]
                stop_price = entry_price - ATR_STOP_MULT * atr
                if stop_price <= 0:
                    continue
                shares = RISK_PER_TRADE / (entry_price - stop_price)
                if shares <= 0 or shares * entry_price > 50000:
                    continue
                breakout_idx = i + 1
                open_positions.append({
                    "ticker": ticker,
                    "entry_price": entry_price,
                    "stop_price": stop_price,
                    "shares": shares,
                    "entry_bar_idx": afternoon_bars.index[breakout_idx],
                    "bars": afternoon_bars,
                    "atr": atr,
                })
                break

    # Limit simultaneous positions
    open_positions = open_positions[:MAX_SIMULTANEOUS]

    # Resolve positions: scan remaining afternoon bars for stop or EOD exit
    for pos in open_positions:
        bars_slice = pos["bars"]
        entry_bar_idx = pos["entry_bar_idx"]
        remaining = bars_slice[bars_slice.index > entry_bar_idx]

        exit_price = pos["entry_price"]
        exit_reason = "eod"

        for _, bar in remaining.iterrows():
            # Stop hit
            if bar["low"] <= pos["stop_price"]:
                exit_price = min(pos["stop_price"], bar["open"])
                exit_reason = "stop"
                break
            # EOD exit
            if bar["hour"] == EXIT_HOUR and bar["minute"] == EXIT_MIN:
                exit_price = bar["close"]
                exit_reason = "eod"
                break
            # Last bar fallback
            if bar["hour"] >= 15 and bar["minute"] >= 50:
                exit_price = bar["close"]
                exit_reason = "eod"
                break

        pnl = (exit_price - pos["entry_price"]) * pos["shares"]
        day_trades.append({
            "ticker": pos["ticker"],
            "date": date,
            "entry": pos["entry_price"],
            "exit": exit_price,
            "shares": pos["shares"],
            "pnl": pnl,
            "win": pnl > 0,
            "exit_reason": exit_reason,
        })

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


def main() -> None:
    print("=== Afternoon Momentum Diagnostic ===\n")

    # Load V23 baseline
    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(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)
        all_results.append(result)
        if (i + 1) % 20 == 0:
            print(f"  {i+1}/{len(v23_dates)} days processed...", flush=True)

    # Aggregate
    all_trades = [t for r in all_results for t in r["trades"]]
    total_trades = len(all_trades)
    wins = sum(1 for t in all_trades if t["win"])
    win_rate = wins / total_trades if total_trades else 0.0
    avg_win = np.mean([t["pnl"] for t in all_trades if t["pnl"] > 0]) if wins > 0 else 0.0
    avg_loss = 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)

    # Correlation with V23
    pm_daily = {r["date"]: r["day_pnl"] for r in all_results}
    pm_series = [pm_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(pm_series, v23_series)[0, 1]) if len(v23_dates) > 1 else 0.0

    trade_days = sum(1 for r in all_results if r["trades"])

    print("=" * 55)
    print(f"  Total trades:     {total_trades}")
    print(f"  Trade days:       {trade_days} / {len(v23_dates)}")
    print(f"  Win rate:         {win_rate*100:.1f}%  (gate: ≥50%)")
    print(f"  Avg win:          ${avg_win:.2f}")
    print(f"  Avg loss:         ${avg_loss:.2f}")
    print(f"  Win/Loss ratio:   {avg_win/avg_loss:.2f}  (gate: ≥1.0)" if avg_loss > 0 else "  Win/Loss ratio:   N/A")
    print(f"  Total PnL:        ${total_pnl:+.2f}")
    print(f"  Corr vs V23:      {corr:.3f}  (gate: ≤0.25)")
    print("=" * 55)

    # Gate verdict
    g1 = win_rate >= 0.50
    g2 = (avg_win / avg_loss >= 1.0) if avg_loss > 0 else False
    g3 = corr <= 0.25
    g4 = total_trades >= 30
    print(f"\nGate G1 (WR ≥ 50%):       {'PASS' if g1 else 'FAIL'}")
    print(f"Gate G2 (W/L ≥ 1.0):      {'PASS' if g2 else 'FAIL'}")
    print(f"Gate G3 (corr ≤ 0.25):    {'PASS' if g3 else 'FAIL'}")
    print(f"Gate G4 (trades ≥ 30):    {'PASS' if g4 else 'FAIL'}")
    verdict = "PROCEED TO ENGINE BUILD" if (g1 and g2 and g3 and g4) else "ABORT — edge not confirmed"
    print(f"\nVERDICT: {verdict}")

    # Save results
    out = {
        "total_trades": total_trades,
        "trade_days": trade_days,
        "win_rate": win_rate,
        "avg_win": avg_win,
        "avg_loss": avg_loss,
        "total_pnl": total_pnl,
        "corr_v23": corr,
        "trades": all_trades,
        "daily_pnl": [{"date": r["date"], "pnl": r["day_pnl"]} for r in all_results],
    }
    out_path = "runs/intraday_orb/diag_afternoon_momentum.json"
    import json as _j
    with open(out_path, "w") as f:
        _j.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()
