"""
Gap-Down Catalyst Short Diagnostic — Three-Arm Comparison

Hypothesis: gap-down stocks with bad-news catalysts (earnings miss, impairment,
officer departure, etc.) continue lower reliably, while gap-downs with no catalyst
tend to recover. Separating them lifts WR above break-even.

Arms:
  A — no catalyst filter (baseline from diag_gapdown_short.py v1)
  B — require bad-news event on T or T-1 (CONTINUATION hypothesis)
  C — exclude bad-news event / no catalyst (RECOVERY control group)

Gates (Arm B):
  G1: WR ≥ 50%
  G2: avg_win / avg_loss ≥ 1.2
  G3: Pearson corr with V23 daily PnL ≤ 0.00
  G4: trades ≥ 30 (sufficient sample)
  G5: Arm B WR − Arm C WR ≥ 5pp (filter truly separates populations)
"""
from __future__ import annotations

import json
import math
import os
import sys
from datetime import date as _date, timedelta

import numpy as np
import pandas as pd
import yaml

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

MIN_PRICE = 10.0
MIN_GAP_DOWN = -0.02
MAX_GAP_DOWN = -0.10
MIN_AVG_DOLLAR_VOL = 25_000_000
MIN_RVOL = 1.5
ATR_STOP_MULT = 0.75
ATR_TARGET_MULT = 1.5
RISK_PER_TRADE = 500.0
MAX_SIMULTANEOUS = 3
EXIT_HOUR = 15
EXIT_MIN = 55

# Bad-news 8-K items that signal continuation (negative catalyst)
BAD_NEWS_ITEMS = {"2.02", "2.06", "4.02", "1.03", "1.02"}
# event_types that are definitively bad news regardless of item
BAD_NEWS_EVENT_TYPES = {"earnings_result", "financial_restatement", "contract_termination"}
# For item 8.01 (other_material_event) / 5.02 (management_change): keyword filter on title
BAD_NEWS_ITEM_8_01_KEYWORDS = {"downgrade", "lawsuit", "investigation", "regulatory", "probe",
                                "violation", "fraud", "warning", "recall", "suspend"}
BAD_NEWS_MANAGEMENT_KEYWORDS = {"resign", "terminat", "depart", "step", "remov", "dismiss"}


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 load_catalyst_events(ticker: str, date: str) -> list[dict] | None:
    """Load filing events for ticker from disk cache. Returns None if uncached."""
    import gzip
    path = f"{CATALYST_CACHE_DIR}/{ticker.upper()}.json.gz"
    if not os.path.exists(path):
        return None
    try:
        with gzip.open(path, "rt", encoding="utf-8") as fh:
            payload = json.load(fh)
    except Exception:
        return None

    coverage_start = str(payload.get("coverage_start") or "")
    coverage_end = str(payload.get("coverage_end") or "")
    if not coverage_start or not coverage_end:
        return None

    # Window: T-1 through T
    prev_date = str(_date.fromisoformat(date) - timedelta(days=1))
    if prev_date < coverage_start or date > coverage_end:
        return None

    events = payload.get("events", [])
    return [
        e for e in events
        if prev_date <= str(e.get("filing_date", ""))[:10] <= date
    ]


def classify_bad_news(events: list[dict]) -> bool:
    """Return True if any event in the list is a bad-news catalyst."""
    if not events:
        return False
    for e in events:
        item = str(e.get("item_number", "")).strip()
        etype = str(e.get("event_type", "")).strip().lower()
        title = str(e.get("title", "")).strip().lower()

        # Definitive bad-news items
        if item in BAD_NEWS_ITEMS:
            return True
        # Definitive bad-news event types
        if etype in BAD_NEWS_EVENT_TYPES:
            return True
        # item 8.01 with bad keywords
        if item == "8.01" and any(kw in title for kw in BAD_NEWS_ITEM_8_01_KEYWORDS):
            return True
        # item 5.02 management change with departure keywords
        if item == "5.02" and any(kw in title for kw in BAD_NEWS_MANAGEMENT_KEYWORDS):
            return True

    return False


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


def simulate_trade(ticker: str, date: str, gap: float, prev_close: float,
                   bars: pd.DataFrame) -> dict | None:
    """Simulate one gap-down short trade. Returns trade dict or None if no entry."""
    if bars.iloc[0]["open"] < MIN_PRICE:
        return None

    orb_bars = bars[(bars["hour"] == 9) & (bars["minute"] >= 30) & (bars["minute"] < 35)]
    if len(orb_bars) == 0:
        return None
    orb = orb_bars.iloc[0]
    if orb["close"] >= orb["open"]:  # must be bearish ORB
        return None

    atr = approximate_atr(bars)
    if atr <= 0:
        return None

    stop_dist = ATR_STOP_MULT * atr
    entry_trigger = orb["low"]
    stop_price = entry_trigger + stop_dist
    shares = RISK_PER_TRADE / stop_dist
    if shares <= 0 or shares * entry_trigger > 50000:
        return None

    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:
        return None

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

    pnl = (entry_price - exit_price) * shares

    return {
        "ticker": ticker,
        "date": date,
        "gap": gap,
        "entry": entry_price,
        "exit": exit_price,
        "pnl": pnl,
        "win": pnl > 0,
        "exit_reason": exit_reason,
    }


def simulate_day(date: str, tickers: list[str],
                 daily_data: dict[str, pd.DataFrame]) -> dict:
    """Simulate one day for all three arms."""
    trades_a = []
    trades_b = []
    trades_c = []
    open_a = open_b = open_c = 0
    missing_catalyst_cache = 0
    total_candidates = 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"]))
        total_candidates += 1

    candidates.sort(key=lambda x: x[1])

    for ticker, gap, prev_close in candidates:
        bars = load_bars(ticker, date)
        if bars is None or len(bars) < 10:
            continue

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

        # Load catalyst events for Arm B / C classification
        events = load_catalyst_events(ticker, date)
        if events is None:
            missing_catalyst_cache += 1
            has_bad_news = None  # unknown — include in Arm A but exclude from B/C
        else:
            has_bad_news = classify_bad_news(events)

        # Arm A: no filter
        if open_a < MAX_SIMULTANEOUS:
            trade = simulate_trade(ticker, date, gap, prev_close, bars)
            if trade is not None:
                trades_a.append(trade)
                open_a += 1

        # Arm B: require bad-news
        if has_bad_news is True and open_b < MAX_SIMULTANEOUS:
            trade = simulate_trade(ticker, date, gap, prev_close, bars)
            if trade is not None:
                trades_b.append({**trade, "arm": "B"})
                open_b += 1

        # Arm C: exclude bad-news (no catalyst OR non-bad catalyst)
        if has_bad_news is False and open_c < MAX_SIMULTANEOUS:
            trade = simulate_trade(ticker, date, gap, prev_close, bars)
            if trade is not None:
                trades_c.append({**trade, "arm": "C"})
                open_c += 1

    return {
        "date": date,
        "trades_a": trades_a, "pnl_a": sum(t["pnl"] for t in trades_a),
        "trades_b": trades_b, "pnl_b": sum(t["pnl"] for t in trades_b),
        "trades_c": trades_c, "pnl_c": sum(t["pnl"] for t in trades_c),
        "missing_catalyst_cache": missing_catalyst_cache,
        "total_candidates": total_candidates,
    }


def arm_stats(all_trades: list[dict], v23_dates: list[str],
              v23_pnl_map: dict, day_pnl_map: dict) -> dict:
    if not all_trades:
        return {
            "trades": 0, "win_rate": 0.0, "avg_win": 0.0, "avg_loss": 0.0,
            "wl_ratio": 0.0, "total_pnl": 0.0, "corr_v23": 0.0,
            "exit_reasons": {},
        }
    wins = [t for t in all_trades if t["win"]]
    losses = [t for t in all_trades if not t["win"]]
    avg_win = float(np.mean([t["pnl"] for t in wins])) if wins else 0.0
    avg_loss = float(np.mean([abs(t["pnl"]) for t in losses])) if losses else 0.0
    wl = avg_win / avg_loss if avg_loss > 0 else 0.0
    total_pnl = sum(t["pnl"] for t in all_trades)
    exit_reasons: dict[str, int] = {}
    for t in all_trades:
        exit_reasons[t["exit_reason"]] = exit_reasons.get(t["exit_reason"], 0) + 1

    gd_series = [day_pnl_map.get(d, 0.0) for d in v23_dates]
    v23_series = [v23_pnl_map.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

    return {
        "trades": len(all_trades),
        "wins": len(wins),
        "win_rate": len(wins) / len(all_trades),
        "avg_win": avg_win,
        "avg_loss": avg_loss,
        "wl_ratio": wl,
        "total_pnl": total_pnl,
        "corr_v23": corr,
        "exit_reasons": exit_reasons,
    }


def main() -> None:
    print("=== Gap-Down Catalyst Short Diagnostic (Three-Arm) ===\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_map = {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

    print(f"Daily cache loaded: {len(daily_data)} tickers")

    catalyst_covered = sum(
        1 for t in universe
        if os.path.exists(f"{CATALYST_CACHE_DIR}/{t.upper()}.json.gz")
    )
    print(f"Catalyst cache coverage: {catalyst_covered}/{len(universe)} tickers")
    print(f"V23 window: {v23_dates[0]} → {v23_dates[-1]} ({len(v23_dates)} days)\n")

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

    # Aggregate trades per arm
    all_a = [t for r in all_results for t in r["trades_a"]]
    all_b = [t for r in all_results for t in r["trades_b"]]
    all_c = [t for r in all_results for t in r["trades_c"]]

    # Daily PnL series per arm (for correlation)
    pnl_a = {r["date"]: r["pnl_a"] for r in all_results}
    pnl_b = {r["date"]: r["pnl_b"] for r in all_results}
    pnl_c = {r["date"]: r["pnl_c"] for r in all_results}

    sa = arm_stats(all_a, v23_dates, v23_pnl_map, pnl_a)
    sb = arm_stats(all_b, v23_dates, v23_pnl_map, pnl_b)
    sc = arm_stats(all_c, v23_dates, v23_pnl_map, pnl_c)

    missing_pct = total_missing / total_candidates * 100 if total_candidates else 0

    print("=" * 70)
    print(f"  Catalyst skip rate: {total_missing}/{total_candidates} candidates ({missing_pct:.1f}% uncached)")
    print()
    print(f"  {'Metric':<30}  {'Arm A (baseline)':<20}  {'Arm B (bad-news)':<20}  {'Arm C (no-catalyst)':<20}")
    print(f"  {'-'*30}  {'-'*20}  {'-'*20}  {'-'*20}")
    print(f"  {'Trades':<30}  {sa['trades']:<20}  {sb['trades']:<20}  {sc['trades']:<20}")

    def wr(s):
        return f"{s['win_rate']*100:.1f}%"
    def wl(s):
        return f"{s['wl_ratio']:.2f}" if s['wl_ratio'] else "N/A"
    def pnl(s):
        return f"${s['total_pnl']:+,.0f}"
    def cr(s):
        return f"{s['corr_v23']:.3f}"

    print(f"  {'Win rate':<30}  {wr(sa):<20}  {wr(sb):<20}  {wr(sc):<20}")
    print(f"  {'Avg win':<30}  ${sa['avg_win']:<19.2f}  ${sb['avg_win']:<19.2f}  ${sc['avg_win']:<19.2f}")
    print(f"  {'Avg loss':<30}  ${sa['avg_loss']:<19.2f}  ${sb['avg_loss']:<19.2f}  ${sc['avg_loss']:<19.2f}")
    print(f"  {'Win/Loss ratio':<30}  {wl(sa):<20}  {wl(sb):<20}  {wl(sc):<20}")
    print(f"  {'Total PnL':<30}  {pnl(sa):<20}  {pnl(sb):<20}  {pnl(sc):<20}")
    print(f"  {'Corr vs V23':<30}  {cr(sa):<20}  {cr(sb):<20}  {cr(sc):<20}")
    print(f"  {'Exit reasons':<30}  {str(sa['exit_reasons']):<20}  {str(sb['exit_reasons']):<20}  {str(sc['exit_reasons']):<20}")
    print("=" * 70)

    # Gate evaluation (Arm B)
    wr_b = sb["win_rate"]
    wr_c = sc["win_rate"]
    wl_b = sb["wl_ratio"]
    corr_b = sb["corr_v23"]
    trades_b = sb["trades"]
    pnl_b_val = sb["total_pnl"]

    g1 = wr_b >= 0.50
    g2 = wl_b >= 1.2
    g3 = corr_b <= 0.00
    g4 = trades_b >= 30
    g5 = (wr_b - wr_c) >= 0.05 if sc["trades"] > 0 else False

    print(f"\n  Gate G1 (Arm B WR ≥ 50%):            {'PASS' if g1 else 'FAIL'}  ({wr_b*100:.1f}%)")
    print(f"  Gate G2 (Arm B W/L ≥ 1.2):           {'PASS' if g2 else 'FAIL'}  ({wl_b:.2f})")
    print(f"  Gate G3 (Arm B corr ≤ 0.00):         {'PASS' if g3 else 'FAIL'}  ({corr_b:.3f})")
    print(f"  Gate G4 (Arm B trades ≥ 30):         {'PASS' if g4 else 'FAIL'}  ({trades_b})")
    print(f"  Gate G5 (B WR − C WR ≥ 5pp):        {'PASS' if g5 else 'FAIL'}  ({(wr_b-wr_c)*100:.1f}pp)")

    passed = sum([g1, g2, g3, g4, g5])
    if passed == 5:
        verdict = "PROCEED TO ENGINE BUILD (Phase 2)"
    else:
        verdict = f"ABORT ({passed}/5 gates passed)"
    print(f"\n  VERDICT: {verdict}")

    out = {
        "arm_a": {**sa, "trades": all_a},
        "arm_b": {**sb, "trades": all_b},
        "arm_c": {**sc, "trades": all_c},
        "gates": {"g1": g1, "g2": g2, "g3": g3, "g4": g4, "g5": g5, "passed": passed},
        "missing_catalyst_pct": missing_pct,
        "daily_pnl_a": [{"date": r["date"], "pnl": r["pnl_a"]} for r in all_results],
        "daily_pnl_b": [{"date": r["date"], "pnl": r["pnl_b"]} for r in all_results],
        "daily_pnl_c": [{"date": r["date"], "pnl": r["pnl_c"]} for r in all_results],
    }
    out_path = "runs/intraday_orb/diag_gapdown_catalyst_short.json"
    with open(out_path, "w") as f:
        json.dump(out, f, indent=2, default=str)
    print(f"\n  Results saved to {out_path}")


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