#!/usr/bin/env python3
"""Equal-weight vs ATR-weight ablation for the Clenow weekly recipe.

Question: the ranked book draws down DEEPER than random-among-eligible
monkeys (gated MaxDD p2.5 QQQ / p9 SPX). Is that concentration from the ATR
inverse-vol SIZING, or from the momentum SELECTION (correlated trenders falling
together)? A 2×2 decomposes it — ranking {clenow, random} × weighting {atr,
equal}, regime gate ON throughout:

  real_atr   = (clenow,  atr)    the system
  real_ew    = (clenow,  equal)  ranked book, vol-sizing removed
  monkey_atr = (random,  atr)    K=200
  monkey_ew  = (random,  equal)  K=200

Then:
  - Sizing effect (hold selection fixed): real_atr vs real_ew; and the
    monkey_atr-median vs monkey_ew-median MaxDD.
  - Selection effect (hold weighting fixed): real's percentile within the
    matching monkey distribution — atr-real in atr-monkey, ew-real in ew-monkey.
    If the ranked book is still a low-MaxDD-percentile outlier under EQUAL
    weight, the deep drawdown is SELECTION, not sizing.

Usage:
    python ablation_weight.py --index QQQ_Nasdaq-100 --regime-etf QQQ
    python ablation_weight.py --index SPY_SandP_500  --regime-etf SPY
"""
import argparse
import random
from datetime import date

import numpy as np

from prototype import prepare_run, simulate_portfolio
from monkey_baseline import metrics_of, pct_rank, pctiles


def real_run(prep, weight_mode):
    return metrics_of(simulate_portfolio(
        prep, rank_mode="clenow", gate=True, weight_mode=weight_mode, verbose=False))


def monkey_dist(prep, weight_mode, k):
    out = []
    for seed in range(k):
        res = simulate_portfolio(
            prep, rank_mode="random", rng=random.Random(seed),
            gate=True, weight_mode=weight_mode, verbose=False)
        out.append(metrics_of(res))
        if (seed + 1) % 25 == 0:
            print(f"    {weight_mode} {seed+1}/{k}", flush=True)
    return out


def cell(name, real, monk):
    """Print one (real vs its matching monkey) cell."""
    print(f"\n--- {name} ---")
    print(f"  real: CAGR {real['cagr']:+.2%}  Sharpe {real['sharpe']:.2f}  "
          f"MaxDD {real['max_dd']:.2%}  final ${real['final']:,.0f}")
    for metric, fmt in [("cagr", "{:+.2%}"), ("sharpe", "{:.2f}"), ("max_dd", "{:.2%}")]:
        dist = [m[metric] for m in monk]
        pc = pctiles(dist)
        rp = pct_rank(real[metric], dist)
        line = "  ".join(f"p{p}={fmt.format(v)}" for p, v in pc.items())
        print(f"  {metric:>7}: [{line}]  real_pctile={rp:.1f}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--index", default="QQQ_Nasdaq-100")
    ap.add_argument("--regime-etf", default="QQQ")
    ap.add_argument("--start", default="2005-01-01")
    ap.add_argument("--end", default="2026-05-15")
    ap.add_argument("--k", type=int, default=200)
    args = ap.parse_args()

    start = date.fromisoformat(args.start)
    end = date.fromisoformat(args.end)

    print(f"Preparing {args.index} (gate ETF {args.regime_etf})...")
    prep = prepare_run(args.index, start, end, args.regime_etf, verbose=True)

    real_atr = real_run(prep, "atr")
    real_ew = real_run(prep, "equal")
    print(f"\nReal (clenow) ATR : CAGR {real_atr['cagr']:+.2%}  Sharpe {real_atr['sharpe']:.2f}  MaxDD {real_atr['max_dd']:.2%}")
    print(f"Real (clenow) EQUAL: CAGR {real_ew['cagr']:+.2%}  Sharpe {real_ew['sharpe']:.2f}  MaxDD {real_ew['max_dd']:.2%}")

    print(f"\nRunning {args.k} ATR-weight monkeys...")
    monk_atr = monkey_dist(prep, "atr", args.k)
    print(f"\nRunning {args.k} EQUAL-weight monkeys...")
    monk_ew = monkey_dist(prep, "equal", args.k)

    print(f"\n========== {args.index} 2x2 ABLATION (gated, K={args.k}) ==========")
    cell("clenow ATR  vs ATR-monkey", real_atr, monk_atr)
    cell("clenow EQUAL vs EQUAL-monkey", real_ew, monk_ew)

    # Sizing effect on RANDOM books (selection held = random)
    atr_dd_med = float(np.median([m["max_dd"] for m in monk_atr]))
    ew_dd_med = float(np.median([m["max_dd"] for m in monk_ew]))
    atr_cagr_med = float(np.median([m["cagr"] for m in monk_atr]))
    ew_cagr_med = float(np.median([m["cagr"] for m in monk_ew]))
    print("\n--- Sizing effect on RANDOM books (median monkey) ---")
    print(f"  MaxDD: ATR {atr_dd_med:.2%}  ->  EQUAL {ew_dd_med:.2%}  (Δ {(ew_dd_med-atr_dd_med)*100:+.1f} pp)")
    print(f"  CAGR : ATR {atr_cagr_med:+.2%}  ->  EQUAL {ew_cagr_med:+.2%}")
    print("\n--- Sizing effect on the RANKED book ---")
    print(f"  MaxDD: ATR {real_atr['max_dd']:.2%}  ->  EQUAL {real_ew['max_dd']:.2%}  (Δ {(real_ew['max_dd']-real_atr['max_dd'])*100:+.1f} pp)")
    print(f"  CAGR : ATR {real_atr['cagr']:+.2%}  ->  EQUAL {real_ew['cagr']:+.2%}")


if __name__ == "__main__":
    main()