"""Diagnostic Experiments: Fixed-holding & entry-timing analysis.

Experiment 1 — Stop-off + fixed holding period exit (3/5/10/20d)
  Purpose: Does the raw signal carry time-based alpha?
  Method:  Use pre-computed fwd_return_Xd (no stop, no target, pure hold).
  Report:  Mean return, median, win rate, by score bucket and event type.

Experiment 2 — Reaction-day close entry vs next-open entry
  Purpose: Is the overnight gap eating the edge?
  Method:  Compare fwd_return_Xd (next-open base) with
           close_entry_return = (entry_price / event_close) * (1 + fwd_return_Xd) - 1
  Constraint: Only pre-market/intraday filings (reaction_date == event_date)
              are eligible for close entry; after-close filings shown separately.

Usage:
  python -m dev.analysis.entry_timing_experiment \
    --snapshot-dir data/datasets/snapshots/a6687401-afdd-4bb7-9bb1-4cd32c5190bb \
    --split train
"""
from __future__ import annotations

import argparse
import statistics
from pathlib import Path
from typing import Any

import pyarrow.parquet as pq

from libs.backtest.scoring import compute_entry_score


FWD_HORIZONS = ["fwd_return_3d", "fwd_return_5d", "fwd_return_10d", "fwd_return_20d"]
HORIZON_LABELS = ["3d", "5d", "10d", "20d"]


def _load_rows(snapshot_dir: Path, split: str) -> list[dict[str, Any]]:
    table = pq.read_table(snapshot_dir / f"{split}.parquet")
    return table.to_pylist()


def _stats(values: list[float]) -> dict[str, float | None]:
    if not values:
        return {"mean": None, "median": None, "win_rate": None, "n": 0}
    return {
        "mean": statistics.mean(values),
        "median": statistics.median(values),
        "win_rate": sum(1 for v in values if v > 0) / len(values),
        "n": len(values),
    }


def _fmt(val: float | None, pct: bool = True) -> str:
    if val is None:
        return "     N/A"
    if pct:
        return f"{val:+.2%}"
    return f"{val:.1%}"


def _score_bucket(score: float) -> str:
    if score < 0.3:
        return "low(<0.3)"
    if score < 0.5:
        return "mid(0.3-0.5)"
    if score < 0.7:
        return "high(0.5-0.7)"
    return "top(>=0.7)"


BUCKET_ORDER = ["low(<0.3)", "mid(0.3-0.5)", "high(0.5-0.7)", "top(>=0.7)"]


def _is_intraday_filing(row: dict) -> bool:
    """Pre-market / regular-hours filing: reaction_date == event_date."""
    return row.get("reaction_date") == row.get("event_date")


# ── Experiment 1 ──────────────────────────────────────────────────────

def experiment_1(rows: list[dict[str, Any]]) -> None:
    """Stop-off + fixed holding period: raw forward returns."""
    print(f"\n{'='*100}")
    print("EXPERIMENT 1: Stop-off + Fixed Holding Period Exit")
    print(f"{'='*100}")

    # Enrich with score
    for r in rows:
        r["_score"] = compute_entry_score(r)
        r["_bucket"] = _score_bucket(r["_score"])

    # ── 1a. ALL events ──
    print(f"\n--- 1a. ALL events (N={len(rows)}) ---")
    _print_horizon_table("ALL", rows)

    # ── 1b. By score bucket ──
    print(f"\n--- 1b. By score bucket ---")
    by_bucket: dict[str, list] = {b: [] for b in BUCKET_ORDER}
    for r in rows:
        by_bucket[r["_bucket"]].append(r)

    header = f"{'Bucket':<16} {'N':>5}"
    for lbl in HORIZON_LABELS:
        header += f"  {lbl+'_mean':>8} {lbl+'_med':>8} {lbl+'_wr':>7}"
    print(header)
    print("-" * len(header))

    for bucket in BUCKET_ORDER:
        items = by_bucket[bucket]
        line = f"{bucket:<16} {len(items):>5}"
        for h in FWD_HORIZONS:
            vals = [float(r[h]) for r in items if r.get(h) is not None]
            s = _stats(vals)
            line += f"  {_fmt(s['mean']):>8} {_fmt(s['median']):>8} {_fmt(s['win_rate'], pct=False):>7}"
        print(line)

    # ── 1c. By event_type ──
    print(f"\n--- 1c. By event_type ---")
    by_type: dict[str, list] = {}
    for r in rows:
        et = r.get("event_type", "unknown")
        by_type.setdefault(et, []).append(r)

    header = f"{'EventType':<25} {'N':>5} {'AvgScr':>7}"
    for lbl in HORIZON_LABELS:
        header += f"  {lbl+'_mean':>8} {lbl+'_wr':>7}"
    print(header)
    print("-" * len(header))

    for et in sorted(by_type, key=lambda k: -len(by_type[k])):
        items = by_type[et]
        avg_score = statistics.mean(r["_score"] for r in items)
        line = f"{et:<25} {len(items):>5} {avg_score:>7.3f}"
        for h in FWD_HORIZONS:
            vals = [float(r[h]) for r in items if r.get(h) is not None]
            s = _stats(vals)
            line += f"  {_fmt(s['mean']):>8} {_fmt(s['win_rate'], pct=False):>7}"
        print(line)

    # ── 1d. MFE/MAE analysis (reward:risk) ──
    print(f"\n--- 1d. MFE/MAE analysis (max favorable / max adverse excursion) ---")
    mfe_mae_horizons = [("3d", "mfe_3d", "mae_3d"),
                        ("5d", "mfe_5d", "mae_5d"),
                        ("10d", "mfe_10d", "mae_10d"),
                        ("20d", "mfe_20d", "mae_20d")]

    header = f"{'Horizon':<10} {'MFE_mean':>10} {'MAE_mean':>10} {'R:R':>8} {'MFE_med':>10} {'MAE_med':>10}"
    print(header)
    print("-" * len(header))

    for lbl, mfe_col, mae_col in mfe_mae_horizons:
        mfe_vals = [float(r[mfe_col]) for r in rows if r.get(mfe_col) is not None]
        mae_vals = [float(r[mae_col]) for r in rows if r.get(mae_col) is not None]
        if mfe_vals and mae_vals:
            mfe_mean = statistics.mean(mfe_vals)
            mae_mean = statistics.mean(mae_vals)
            rr = mfe_mean / abs(mae_mean) if mae_mean != 0 else 0
            mfe_med = statistics.median(mfe_vals)
            mae_med = statistics.median(mae_vals)
            print(f"{lbl:<10} {mfe_mean:>+10.2%} {mae_mean:>+10.2%} {rr:>8.2f} {mfe_med:>+10.2%} {mae_med:>+10.2%}")

    # ── 1e. Score monotonicity per horizon ──
    print(f"\n--- 1e. Score monotonicity (does higher score → better return?) ---")
    for h, lbl in zip(FWD_HORIZONS, HORIZON_LABELS):
        bucket_means = []
        for bucket in BUCKET_ORDER:
            vals = [float(r[h]) for r in by_bucket[bucket] if r.get(h) is not None]
            bucket_means.append(statistics.mean(vals) if vals else None)

        valid_means = [m for m in bucket_means if m is not None]
        monotonic = all(a <= b for a, b in zip(valid_means, valid_means[1:])) if len(valid_means) >= 2 else False
        direction = "MONOTONIC" if monotonic else "NOT monotonic"
        means_str = " → ".join(f"{m:+.2%}" if m is not None else "N/A" for m in bucket_means)
        print(f"  {lbl}: {means_str}  [{direction}]")


# ── Experiment 2 ──────────────────────────────────────────────────────

def experiment_2(rows: list[dict[str, Any]]) -> None:
    """Reaction-close entry vs next-open entry."""
    print(f"\n{'='*100}")
    print("EXPERIMENT 2: Reaction-day Close Entry vs Next-Open Entry")
    print(f"{'='*100}")

    # Split by filing timing
    intraday = [r for r in rows if _is_intraday_filing(r)]
    after_close = [r for r in rows if not _is_intraday_filing(r)]

    print(f"\nFiling timing split:")
    print(f"  Pre-market / intraday (reaction_date == event_date): {len(intraday)}")
    print(f"  After-close (reaction_date > event_date):            {len(after_close)}")

    # ── 2a. Overnight gap analysis (intraday filings only) ──
    print(f"\n--- 2a. Overnight gap cost (intraday filings, N={len(intraday)}) ---")
    gaps = []
    for r in intraday:
        ep = r.get("entry_price")
        ec = r.get("event_close")
        if ep and ec and ec > 0:
            gap = ep / ec - 1.0
            gaps.append(gap)

    if gaps:
        s = _stats(gaps)
        print(f"  Overnight gap (next_open / reaction_close - 1):")
        print(f"    Mean:   {_fmt(s['mean'])}")
        print(f"    Median: {_fmt(s['median'])}")
        print(f"    Gap up rate (open > close): {_fmt(s['win_rate'], pct=False)}")
        print(f"    Std:    {statistics.stdev(gaps):+.2%}" if len(gaps) > 1 else "")

        # Distribution
        gap_buckets = {"< -2%": 0, "-2% to 0%": 0, "0% to +1%": 0,
                       "+1% to +3%": 0, "+3% to +5%": 0, "> +5%": 0}
        for g in gaps:
            if g < -0.02:
                gap_buckets["< -2%"] += 1
            elif g < 0:
                gap_buckets["-2% to 0%"] += 1
            elif g < 0.01:
                gap_buckets["0% to +1%"] += 1
            elif g < 0.03:
                gap_buckets["+1% to +3%"] += 1
            elif g < 0.05:
                gap_buckets["+3% to +5%"] += 1
            else:
                gap_buckets["> +5%"] += 1

        print(f"\n  Gap distribution:")
        for label, count in gap_buckets.items():
            pct = count / len(gaps) * 100
            bar = "#" * int(pct / 2)
            print(f"    {label:>12}: {count:>4} ({pct:5.1f}%) {bar}")

    # ── 2b. Close entry vs next-open returns (intraday only) ──
    print(f"\n--- 2b. Close entry vs Next-open entry returns (intraday, N={len(intraday)}) ---")

    header = f"{'Horizon':<10} {'NextOpen_mean':>13} {'NextOpen_wr':>11} {'CloseEntry_mean':>15} {'CloseEntry_wr':>13} {'Diff':>8}"
    print(header)
    print("-" * len(header))

    for h, lbl in zip(FWD_HORIZONS, HORIZON_LABELS):
        next_open_vals = []
        close_entry_vals = []

        for r in intraday:
            fwd = r.get(h)
            ep = r.get("entry_price")
            ec = r.get("event_close")
            if fwd is not None and ep and ec and ec > 0:
                fwd_f = float(fwd)
                next_open_vals.append(fwd_f)
                # close_entry_return = (entry_price / event_close) * (1 + fwd) - 1
                close_ret = (ep / ec) * (1.0 + fwd_f) - 1.0
                close_entry_vals.append(close_ret)

        no_s = _stats(next_open_vals)
        ce_s = _stats(close_entry_vals)
        diff = (ce_s["mean"] - no_s["mean"]) if ce_s["mean"] is not None and no_s["mean"] is not None else None
        diff_str = _fmt(diff) if diff is not None else "     N/A"

        print(f"{lbl:<10} {_fmt(no_s['mean']):>13} {_fmt(no_s['win_rate'], pct=False):>11}"
              f" {_fmt(ce_s['mean']):>15} {_fmt(ce_s['win_rate'], pct=False):>13} {diff_str:>8}")

    # ── 2c. After-close filings (separate analysis) ──
    if after_close:
        print(f"\n--- 2c. After-close filings (N={len(after_close)}) ---")
        print(f"  (These events filed after market close; reaction_date = next trading day)")

        header = f"{'Horizon':<10} {'NextOpen_mean':>13} {'NextOpen_wr':>11}"
        print(f"  {header}")
        print(f"  {'-' * len(header)}")

        for h, lbl in zip(FWD_HORIZONS, HORIZON_LABELS):
            vals = [float(r[h]) for r in after_close if r.get(h) is not None]
            s = _stats(vals)
            print(f"  {lbl:<10} {_fmt(s['mean']):>13} {_fmt(s['win_rate'], pct=False):>11}")

    # ── 2d. Close entry by event type (intraday only) ──
    print(f"\n--- 2d. Close-entry returns by event type (intraday only) ---")
    by_type: dict[str, list] = {}
    for r in intraday:
        et = r.get("event_type", "unknown")
        by_type.setdefault(et, []).append(r)

    header = f"{'EventType':<25} {'N':>5} {'Gap_mean':>9}"
    for lbl in HORIZON_LABELS:
        header += f"  {lbl+'_CE':>8} {lbl+'_NO':>8}"
    print(header)
    print("-" * len(header))

    for et in sorted(by_type, key=lambda k: -len(by_type[k])):
        items = by_type[et]
        # Gap
        type_gaps = []
        for r in items:
            ep = r.get("entry_price")
            ec = r.get("event_close")
            if ep and ec and ec > 0:
                type_gaps.append(ep / ec - 1.0)

        gap_mean = statistics.mean(type_gaps) if type_gaps else None
        line = f"{et:<25} {len(items):>5} {_fmt(gap_mean):>9}"

        for h in FWD_HORIZONS:
            ce_vals = []
            no_vals = []
            for r in items:
                fwd = r.get(h)
                ep = r.get("entry_price")
                ec = r.get("event_close")
                if fwd is not None and ep and ec and ec > 0:
                    fwd_f = float(fwd)
                    no_vals.append(fwd_f)
                    ce_vals.append((ep / ec) * (1.0 + fwd_f) - 1.0)
            ce_s = _stats(ce_vals)
            no_s = _stats(no_vals)
            line += f"  {_fmt(ce_s['mean']):>8} {_fmt(no_s['mean']):>8}"
        print(line)

    # ── 2e. Score-filtered close entry (score >= 0.5, intraday) ──
    filtered = [r for r in intraday if r.get("_score", 0) >= 0.5]
    if filtered:
        print(f"\n--- 2e. Score-filtered (>= 0.5) close entry (intraday, N={len(filtered)}) ---")

        header = f"{'Horizon':<10} {'CloseEntry_mean':>15} {'CE_wr':>8} {'NextOpen_mean':>14} {'NO_wr':>8}"
        print(header)
        print("-" * len(header))

        for h, lbl in zip(FWD_HORIZONS, HORIZON_LABELS):
            ce_vals = []
            no_vals = []
            for r in filtered:
                fwd = r.get(h)
                ep = r.get("entry_price")
                ec = r.get("event_close")
                if fwd is not None and ep and ec and ec > 0:
                    fwd_f = float(fwd)
                    no_vals.append(fwd_f)
                    ce_vals.append((ep / ec) * (1.0 + fwd_f) - 1.0)
            ce_s = _stats(ce_vals)
            no_s = _stats(no_vals)
            print(f"{lbl:<10} {_fmt(ce_s['mean']):>15} {_fmt(ce_s['win_rate'], pct=False):>8}"
                  f" {_fmt(no_s['mean']):>14} {_fmt(no_s['win_rate'], pct=False):>8}")


def _print_horizon_table(label: str, items: list[dict]) -> None:
    header = f"{'Horizon':<10} {'Mean':>10} {'Median':>10} {'WinRate':>8} {'N':>6} {'StdDev':>10}"
    print(header)
    print("-" * len(header))
    for h, lbl in zip(FWD_HORIZONS, HORIZON_LABELS):
        vals = [float(r[h]) for r in items if r.get(h) is not None]
        s = _stats(vals)
        std = statistics.stdev(vals) if len(vals) > 1 else 0
        print(f"{lbl:<10} {_fmt(s['mean']):>10} {_fmt(s['median']):>10}"
              f" {_fmt(s['win_rate'], pct=False):>8} {s['n']:>6} {std:>10.2%}")


def main() -> None:
    parser = argparse.ArgumentParser(description="Entry Timing Experiments")
    parser.add_argument("--snapshot-dir", required=True)
    parser.add_argument("--split", default="train")
    args = parser.parse_args()

    rows = _load_rows(Path(args.snapshot_dir), args.split)
    if not rows:
        print("No data found.")
        return

    print(f"Loaded {len(rows)} events from {args.split} split")

    experiment_1(rows)
    experiment_2(rows)

    # ── Summary ──
    print(f"\n{'='*100}")
    print("KEY DIAGNOSTIC QUESTIONS")
    print(f"{'='*100}")
    print("""
  Exp 1 — Time-based alpha:
    • Are any fixed-horizon returns consistently positive?
    • Does score monotonicity hold at any horizon?
    • Which event types carry actual alpha (mean > 0, win rate > 50%)?

  Exp 2 — Entry timing:
    • Is the overnight gap systematically positive (eating the edge)?
    • Does close-entry improve returns vs next-open?
    • Is the gap effect uniform across event types, or concentrated?
""")


if __name__ == "__main__":
    main()