#!/usr/bin/env python3
"""Production book × universe × horizon matrix (research only).

Four pre-registered arms — same live strategy knobs; only entry start date and
tradable universe change. See docs/research/prod-book-universe-horizon.md.

  A  2022-07-01  prod ~505
  B  2022-07-01  prod ∪ liquid top-1500
  C  2016-07-01  prod ~505
  D  2016-07-01  prod ∪ liquid top-1500

Example (MacBook, deep research.sqlite)
---------------------------------------
  python scripts/run_prod_book_universe_matrix.py \\
      --snapshot backtest_snapshots/research.sqlite \\
      --workers 8 --allow-spawn \\
      --candidate-cache reports/.cache/prod-book-univ-cands.pkl
"""

from __future__ import annotations

import argparse
import asyncio
import json
import os
import pickle
import sys
import time
from concurrent.futures import ProcessPoolExecutor
from datetime import date, datetime
from pathlib import Path
from typing import Any

from sqlalchemy import create_engine, text
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from app.ssl_bootstrap import bootstrap_ssl  # noqa: E402

bootstrap_ssl()

SHORT_START = date(2022, 7, 1)
LONG_START = date(2016, 7, 1)
LIQUID_TOP_N = 1500
LIQUID_MIN_PRICE = 5.0
CACHE_VERSION = "prod-book-universe-horizon-v1"

ARMS: tuple[dict[str, Any], ...] = (
    {
        "id": "A_prod_4y_505",
        "label": "Prod book · ~4y · 505 only",
        "start": SHORT_START,
        "universe": "prod_505",
    },
    {
        "id": "B_prod_4y_505_liquid",
        "label": "Prod book · ~4y · 505 + liquid top-1500",
        "start": SHORT_START,
        "universe": "prod_plus_liquid",
    },
    {
        "id": "C_prod_2016_505",
        "label": "Prod book · since 2016-07 · 505 only",
        "start": LONG_START,
        "universe": "prod_505",
    },
    {
        "id": "D_prod_2016_505_liquid",
        "label": "Prod book · since 2016-07 · 505 + liquid top-1500",
        "start": LONG_START,
        "universe": "prod_plus_liquid",
    },
)


def _parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--snapshot", default="backtest_snapshots/research.sqlite")
    p.add_argument("--workers", type=int, default=8)
    p.add_argument("--allow-spawn", action="store_true")
    p.add_argument("--quiet", action="store_true")
    p.add_argument(
        "--candidate-cache",
        default="reports/.cache/prod-book-universe-cands.pkl",
        help="Pickle cache for full GTL candidate pass (expensive).",
    )
    p.add_argument(
        "--rebuild-cache",
        action="store_true",
        help="Ignore existing candidate cache.",
    )
    p.add_argument("--out", default=None)
    p.add_argument(
        "--skip-race-guard",
        action="store_true",
        help="Allow run without completion manifest (not recommended).",
    )
    return p.parse_args()


def _sqlite_url(path: Path) -> str:
    return f"sqlite+aiosqlite:///{path.resolve().as_posix()}"


def _load_prod_and_all_symbols(snapshot: Path) -> tuple[set[str], list[str]]:
    engine = create_engine(
        f"sqlite:///{snapshot.resolve().as_posix()}",
        future=True,
    )
    try:
        with engine.connect() as conn:
            all_syms = [
                str(r[0]).upper()
                for r in conn.execute(text("SELECT symbol FROM tickers ORDER BY 1"))
            ]
            try:
                rank_only = {
                    str(r[0]).upper()
                    for r in conn.execute(text("SELECT symbol FROM research_rank_only"))
                }
            except Exception:
                rank_only = set()
    finally:
        engine.dispose()
    prod = {s for s in all_syms if s not in rank_only}
    return prod, all_syms


def _median(xs: list[float]) -> float | None:
    if len(xs) < 20:
        return None
    s = sorted(xs)
    mid = len(s) // 2
    if len(s) % 2:
        return s[mid]
    return 0.5 * (s[mid - 1] + s[mid])


def _build_liquid_membership(
    prices: dict[str, tuple],
    *,
    top_n: int,
    min_price: float,
) -> dict[date, set[str]]:
    """For each calendar date present in any series, top-N by 63d median $vol."""
    # Collect per-symbol (date -> (close, dvol63))
    per_sym: dict[str, dict[date, tuple[float, float | None]]] = {}
    all_dates: set[date] = set()
    for sym, cols in prices.items():
        ords, _o, _h, _l, closes, vols = cols
        dates = [date.fromordinal(int(o)) for o in ords]
        n = len(dates)
        series: dict[date, tuple[float, float | None]] = {}
        for i in range(n):
            d = dates[i]
            c = float(closes[i])
            dvol = None
            if i + 1 >= 63:
                dvs = []
                for k in range(i - 62, i + 1):
                    ck = float(closes[k])
                    vk = float(vols[k] or 0)
                    if ck > 0 and vk >= 0:
                        dvs.append(ck * vk)
                dvol = _median(dvs)
            series[d] = (c, dvol)
            all_dates.add(d)
        per_sym[sym] = series

    membership: dict[date, set[str]] = {}
    for d in sorted(all_dates):
        eligible: list[tuple[float, str]] = []
        for sym, series in per_sym.items():
            row = series.get(d)
            if row is None:
                continue
            c, dvol = row
            if c < min_price or dvol is None or dvol <= 0:
                continue
            eligible.append((-dvol, sym))  # highest dvol first
        eligible.sort()
        membership[d] = {sym for _, sym in eligible[:top_n]}
    return membership


def _worker_replay(
    symbol: str,
    columns: tuple,
    config: dict,
    activation: dict,
    spy: dict,
    cadence: str,
) -> list[dict]:
    """Picklable full GTL+signals candidate replay (no signal-only)."""
    from app.services import backtest_service as bt

    cands, _series = bt._replay_and_signals(
        symbol,
        columns,
        config,
        activation,
        spy,
        bt.PRODUCTION_GTL_TARGET_MODEL,
        cadence,
        False,  # always full replay for book matrix
        None,
        None,
    )
    return cands


async def _load_or_build_candidates(
    snapshot: Path,
    *,
    cache_path: Path | None,
    rebuild: bool,
    workers: int,
    quiet: bool,
) -> tuple[list[dict], dict[str, tuple], dict, set[str], dict]:
    from app.config import settings
    from app.services import backtest_service as bt
    from app.services.admin_service import get_activation_config
    from app.services.recommendation_service import get_recommendation_config
    from app.services.paper_trade_service import get_exit_policy
    from app.services.benchmark_service import load_benchmark_closes
    from app.models.ticker import Ticker
    from sqlalchemy import select

    os.environ["BACKTEST_SNAPSHOT_OFFLINE"] = "1"
    settings.backtest_workers = max(1, workers)

    prod_set, all_syms = _load_prod_and_all_symbols(snapshot)
    print(f"Symbols: all={len(all_syms)} prod_505={len(prod_set)}")

    cache_key = {
        "version": CACHE_VERSION,
        "snapshot": str(snapshot.resolve()),
        "prod_n": len(prod_set),
        "all_n": len(all_syms),
    }
    if cache_path and cache_path.exists() and not rebuild:
        with cache_path.open("rb") as fh:
            blob = pickle.load(fh)
        if blob.get("key") == cache_key and blob.get("candidates"):
            print(f"Loaded candidate cache: {cache_path} ({len(blob['candidates'])} rows)")
            return (
                blob["candidates"],
                blob["prices"],
                blob["spy"],
                set(blob["prod_set"]),
                blob["exit_config"],
            )
        print("Cache key mismatch — rebuilding candidates")

    engine = create_async_engine(_sqlite_url(snapshot), pool_pre_ping=True)
    Session = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)

    candidates: list[dict] = []
    prices: dict[str, tuple] = {}
    try:
        async with Session() as db:
            config = await get_recommendation_config(db)
            activation = await get_activation_config(db)
            exit_config = await get_exit_policy(db)
            spy = await load_benchmark_closes(db, "SPY")
            tickers = list(
                (await db.execute(select(Ticker).order_by(Ticker.symbol))).scalars()
            )

            # Fetch all price columns first (I/O).
            for idx, t in enumerate(tickers):
                if not quiet and idx % 100 == 0:
                    print(f"  fetch prices {idx}/{len(tickers)}", end="\r", flush=True)
                cols = await bt._fetch_columns(db, t.symbol)
                if cols is not None:
                    prices[t.symbol.upper()] = cols
            if not quiet:
                print()

            # Parallel GTL replay for every symbol with prices.
            syms = sorted(prices)
            print(f"GTL replay on {len(syms)} symbols (workers={workers})…")
            t0 = time.monotonic()
            if workers <= 1:
                for i, sym in enumerate(syms):
                    if not quiet and i % 50 == 0:
                        print(f"  replay {i}/{len(syms)}", end="\r", flush=True)
                    candidates.extend(
                        _worker_replay(
                            sym, prices[sym], config, activation, spy, "weekly"
                        )
                    )
            else:
                # Process pool: pass column batches.
                import multiprocessing as mp

                ctx = mp.get_context("spawn")
                chunk = max(1, workers * 2)
                with ProcessPoolExecutor(max_workers=workers, mp_context=ctx) as pool:
                    for start in range(0, len(syms), chunk):
                        batch = syms[start : start + chunk]
                        futs = [
                            pool.submit(
                                _worker_replay,
                                sym,
                                prices[sym],
                                config,
                                activation,
                                spy,
                                "weekly",
                            )
                            for sym in batch
                        ]
                        for fut in futs:
                            try:
                                candidates.extend(fut.result())
                            except Exception as exc:
                                print(f"  worker error: {exc}")
                        if not quiet:
                            print(
                                f"  replay {min(start+chunk, len(syms))}/{len(syms)} "
                                f"cands={len(candidates)} "
                                f"elapsed={(time.monotonic()-t0)/60:.1f}m",
                                end="\r",
                                flush=True,
                            )
            if not quiet:
                print()
    finally:
        await engine.dispose()

    print(f"Total raw candidates: {len(candidates)}")
    if cache_path:
        cache_path.parent.mkdir(parents=True, exist_ok=True)
        with cache_path.open("wb") as fh:
            pickle.dump(
                {
                    "key": cache_key,
                    "candidates": candidates,
                    "prices": prices,
                    "spy": spy,
                    "prod_set": sorted(prod_set),
                    "exit_config": exit_config,
                },
                fh,
                protocol=pickle.HIGHEST_PROTOCOL,
            )
        print(f"Wrote cache {cache_path}")

    return candidates, prices, spy, prod_set, exit_config


def _candidate_eligible(
    cand: dict,
    *,
    prod_set: set[str],
    universe: str,
    liquid_by_date: dict[date, set[str]],
) -> bool:
    if cand.get("direction") != "long":
        return False
    sym = str(cand.get("symbol") or "").upper()
    if not sym:
        return False
    if universe == "prod_505":
        return sym in prod_set
    # prod_plus_liquid
    if sym in prod_set:
        return True
    try:
        d = date.fromisoformat(str(cand["date"])[:10])
    except Exception:
        return False
    return sym in (liquid_by_date.get(d) or set())


def _run_arm(
    arm: dict[str, Any],
    *,
    all_candidates: list[dict],
    prices: dict[str, tuple],
    spy: dict,
    prod_set: set[str],
    liquid_by_date: dict[date, set[str]],
    exit_config: dict,
) -> dict[str, Any]:
    from app.services import backtest_service as bt

    start: date = arm["start"]
    universe: str = arm["universe"]

    filtered: list[dict] = []
    for c in all_candidates:
        try:
            d = date.fromisoformat(str(c["date"])[:10])
        except Exception:
            continue
        if d < start:
            continue
        if not _candidate_eligible(
            c, prod_set=prod_set, universe=universe, liquid_by_date=liquid_by_date
        ):
            continue
        filtered.append(dict(c))

    # Re-rank inside this arm's universe (production percentile logic).
    bt._assign_momentum_percentiles(filtered)
    bt._assign_residual_momentum_percentiles(filtered)
    bt._assign_low_volatility_percentiles(filtered)
    bt._assign_activation_momentum_percentiles(filtered)
    bt._assign_residual_high_vol_blend(filtered)
    for c in filtered:
        c["qualified"] = bt._momentum_qualifies(c, 80.0)

    longs = [
        c for c in filtered if c.get("qualified") and c.get("direction") == "long"
    ]

    strategy = next(s for s in bt.PORTFOLIO_MONITOR_STRATEGIES if s.get("is_production"))
    entry_cfg = bt._entry_variant_config(str(strategy["entry_variant"]))
    assert entry_cfg is not None
    ranking_key = str(
        entry_cfg.get("ranking_key") or entry_cfg["percentile_key"]
    )
    exit_policy = bt.LIVE_EXIT_MODE_TO_SIM.get(
        str(exit_config.get("mode", "atr_trailing")), "atr_trail3"
    )
    hold_days = int(exit_config.get("hold_days", 30))
    trail = float(exit_config.get("atr_multiplier", bt.ATR_TRAIL_MULTIPLIER))
    risk = float(entry_cfg["risk_per_trade"])
    max_pos = int(entry_cfg["max_positions"])

    reentry = bt._make_gate_reset_reentry_fn(
        longs, prices, cadence="weekly", ranking_key=ranking_key
    )
    sim = bt._simulate_portfolio(
        longs,
        prices,
        spy,
        exit_policy,
        hold_days,
        ranking_key=ranking_key,
        max_positions=max_pos,
        risk_per_trade=risk,
        atr_trail_multiplier=trail,
        post_stop_reentry_fn=reentry,
        start_date=start,
        end_date=None,
        fill_mode=bt.FILL_MODE_CLOSE,
        include_trades=False,
    )
    if sim is None:
        return {
            "id": arm["id"],
            "label": arm["label"],
            "start": start.isoformat(),
            "universe": universe,
            "n_candidates": len(filtered),
            "n_qualified_longs": 0,
            "error": "no_trades",
        }

    keep = {
        k: sim.get(k)
        for k in (
            "sharpe",
            "sharpe_se",
            "cagr_pct",
            "max_drawdown_pct",
            "total_return_pct",
            "calmar",
            "trades",
            "win_rate",
            "n_returns",
            "psr",
            "start_date",
            "end_date",
            "spy_return_pct",
            "final_equity",
        )
    }
    return {
        "id": arm["id"],
        "label": arm["label"],
        "start": start.isoformat(),
        "universe": universe,
        "n_candidates": len(filtered),
        "n_qualified_longs": len(longs),
        "fill_mode": "close",
        "ranking_key": ranking_key,
        "exit_policy": exit_policy,
        "hold_days": hold_days,
        **keep,
    }


def _write_outputs(payload: dict, out_json: Path, doc_path: Path) -> None:
    out_json.parent.mkdir(parents=True, exist_ok=True)
    out_json.write_text(
        json.dumps(payload, indent=2, default=str) + "\n", encoding="utf-8"
    )

    lines = [
        "# Production book × universe × horizon — results",
        "",
        f"Generated: `{payload.get('generated_at')}`",
        "",
        "> Survivorship: today's constituents backfilled. Compare arms relatively; "
        "do not treat deep CAGR/Sharpe levels as deployable forecasts.",
        "",
        "## Arms",
        "",
        "| arm | universe | entry start | Sharpe | SE | CAGR % | max DD % | ret % | trades | qual longs | span |",
        "|---|---|---|---:|---:|---:|---:|---:|---:|---:|---|",
    ]
    for row in payload.get("arms") or []:
        if row.get("error"):
            lines.append(
                f"| {row.get('id')} | {row.get('universe')} | {row.get('start')} | "
                f"ERR | | | | | | {row.get('n_qualified_longs')} | {row.get('error')} |"
            )
            continue
        lines.append(
            f"| {row.get('id')} | {row.get('universe')} | {row.get('start')} | "
            f"{row.get('sharpe')} | {row.get('sharpe_se')} | {row.get('cagr_pct')} | "
            f"{row.get('max_drawdown_pct')} | {row.get('total_return_pct')} | "
            f"{row.get('trades')} | {row.get('n_qualified_longs')} | "
            f"{row.get('start_date')}→{row.get('end_date')} |"
        )
    lines.extend([
        "",
        "## Config (production, unchanged)",
        "",
        f"```json\n{json.dumps(payload.get('strategy') or {}, indent=2)}\n```",
        "",
        "## Snapshot",
        "",
        f"```json\n{json.dumps(payload.get('snapshot_meta') or {}, indent=2, default=str)}\n```",
        "",
        "PENDING_HUMAN — descriptive matrix only; no auto promotion.",
        "",
        f"JSON: `{out_json.as_posix()}`",
        "",
    ])
    out_json.with_suffix(".md").write_text("\n".join(lines) + "\n", encoding="utf-8")

    # Fill results section of the research doc.
    if doc_path.exists():
        text = doc_path.read_text(encoding="utf-8")
        marker = "## Results"
        idx = text.find(marker)
        header = text[:idx] if idx >= 0 else text
        # Drop old results/verdict tail
        for m in ("## Results", "## Verdict"):
            pass
        body = [
            header.rstrip(),
            "",
            "## Results",
            "",
            f"Generated: `{payload.get('generated_at')}`",
            "",
            "| arm | universe | entry start | Sharpe | SE | CAGR % | max DD % | trades |",
            "|---|---|---|---:|---:|---:|---:|---:|",
        ]
        for row in payload.get("arms") or []:
            body.append(
                f"| {row.get('id')} | {row.get('universe')} | {row.get('start')} | "
                f"{row.get('sharpe', '')} | {row.get('sharpe_se', '')} | "
                f"{row.get('cagr_pct', '')} | {row.get('max_drawdown_pct', '')} | "
                f"{row.get('trades', '')} |"
            )
        body.extend([
            "",
            f"Full report: `{out_json.as_posix()}`",
            "",
            "## Verdict",
            "",
            "**PENDING_HUMAN** — descriptive only; production knobs unchanged.",
            "",
        ])
        doc_path.write_text("\n".join(body) + "\n", encoding="utf-8")


async def _main() -> None:
    args = _parse_args()
    snapshot = Path(args.snapshot)
    if not snapshot.exists():
        raise SystemExit(f"Missing snapshot: {snapshot}")
    if args.allow_spawn:
        os.environ["BACKTEST_ALLOW_SPAWN"] = "1"

    if not args.skip_race_guard:
        try:
            from scripts.research_snapshot_manifest import (
                assert_research_snapshot_complete,
            )

            manifest = assert_research_snapshot_complete(snapshot)
            print(
                f"Race guard OK: tickers={manifest.get('ticker_count')} "
                f"ohlcv={manifest.get('ohlcv_row_count')}"
            )
        except SystemExit as exc:
            # Prod-only snapshot without manifest: allow with warning if ~505.
            engine = create_engine(
                f"sqlite:///{snapshot.resolve().as_posix()}",
                future=True,
            )
            try:
                with engine.connect() as conn:
                    n = int(conn.execute(text("SELECT COUNT(*) FROM tickers")).scalar_one())
            finally:
                engine.dispose()
            if n < 400:
                raise
            print(f"WARNING: no research manifest ({exc}); proceeding n_tickers={n}")

    cache = Path(args.candidate_cache) if args.candidate_cache else None
    candidates, prices, spy, prod_set, exit_config = await _load_or_build_candidates(
        snapshot,
        cache_path=cache,
        rebuild=args.rebuild_cache,
        workers=args.workers,
        quiet=args.quiet,
    )

    print("Building PIT liquid membership (top-1500, price≥5)…")
    t0 = time.monotonic()
    liquid_by_date = _build_liquid_membership(
        prices, top_n=LIQUID_TOP_N, min_price=LIQUID_MIN_PRICE
    )
    print(
        f"  liquid dates={len(liquid_by_date)} "
        f"elapsed={(time.monotonic()-t0)/60:.1f}m"
    )

    arms_out = []
    for arm in ARMS:
        print(f"Running arm {arm['id']}…")
        row = _run_arm(
            arm,
            all_candidates=candidates,
            prices=prices,
            spy=spy,
            prod_set=prod_set,
            liquid_by_date=liquid_by_date,
            exit_config=exit_config,
        )
        arms_out.append(row)
        print(
            f"  Sharpe={row.get('sharpe')} CAGR={row.get('cagr_pct')} "
            f"DD={row.get('max_drawdown_pct')} trades={row.get('trades')} "
            f"qual={row.get('n_qualified_longs')}"
        )

    stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
    out = (
        Path(args.out)
        if args.out
        else Path("reports") / f"prod-book-universe-horizon-{stamp}.json"
    )
    payload = {
        "generated_at": datetime.now().isoformat(),
        "snapshot": str(snapshot.resolve()),
        "snapshot_meta": {
            "prod_universe_n": len(prod_set),
            "price_symbols_n": len(prices),
            "raw_candidates": len(candidates),
            "liquid_top_n": LIQUID_TOP_N,
            "liquid_min_price": LIQUID_MIN_PRICE,
            "short_start": SHORT_START.isoformat(),
            "long_start": LONG_START.isoformat(),
        },
        "strategy": {
            "note": "Live production knobs — no modifications",
            "momentum": "residual_12_1 gate 80",
            "rank": "residual_high_vol_blend_80_20",
            "fill_mode": "close",
            "cost_per_side": 0.001,
            "exit": exit_config,
            "max_positions": 10,
            "risk_per_trade": 0.01,
            "reentry": "gate_reset",
        },
        "arms": arms_out,
        "survivorship_banner": (
            "Today's constituents backfilled. Relative arm comparison only."
        ),
        "pending_human": True,
    }
    _write_outputs(
        payload,
        out,
        Path("docs/research/prod-book-universe-horizon.md"),
    )
    print(f"Wrote {out}")
    print(f"Wrote {out.with_suffix('.md')}")


if __name__ == "__main__":
    asyncio.run(_main())