#!/usr/bin/env python3
"""Engine ablation analysis: measure each primary engine's marginal contribution
to composite performance (primary + IA + parking).

Runs N+1 backtests (baseline + leave-one-out for each enabled engine) with
presets attached and reports which engines help vs harm composite returns.

Usage:
    # Leave-one-out analysis
    python -m apps.tools.run_engine_ablation \\
        --manifest configs/experiments/return_max_long_v7.70.json \\
        --parking qqqm_low_dd \\
        --idle-alpha micro_event_alpha_plus_event_plus \\
        --output-root /tmp/ablation_v7.70

    # Greedy pruning (iteratively removes worst composite-contributor)
    python -m apps.tools.run_engine_ablation \\
        --manifest configs/experiments/return_max_long_v7.70.json \\
        --parking qqqm_low_dd --idle-alpha micro_event_alpha_plus_event_plus \\
        --greedy-prune --output-root /tmp/ablation_v7.70

    # Risk budget sweep (test 0.5x/1.0x/1.5x for each engine's risk budget)
    python -m apps.tools.run_engine_ablation \\
        --manifest configs/experiments/return_max_long_v7.70.json \\
        --parking qqqm_low_dd --idle-alpha micro_event_alpha_plus_event_plus \\
        --budget-sweep --output-root /tmp/ablation_v7.70
"""
from __future__ import annotations

import argparse
import copy
import json
import subprocess
import sys
import tempfile
from pathlib import Path


REPO_ROOT = Path(__file__).resolve().parents[2]
DEFAULT_SNAPSHOT_DIR = ""


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Engine ablation analysis with composite scoring."
    )
    parser.add_argument("--manifest", required=True, help="Base experiment manifest path.")
    parser.add_argument("--parking", default="qqqm_low_dd", help="Cash parking preset.")
    parser.add_argument("--idle-alpha", default="micro_event_alpha_plus_event_plus", help="IA sleeve preset.")
    parser.add_argument("--snapshot-dir", default=DEFAULT_SNAPSHOT_DIR)
    parser.add_argument("--output-root", required=True, help="Root directory for run artifacts.")
    parser.add_argument("--split", default="all", help="Backtest split (default: 'all' = full period merging train+valid+test).")
    parser.add_argument(
        "--greedy-prune", action="store_true",
        help="After LOO, iteratively remove worst engine until no improvement possible.",
    )
    parser.add_argument(
        "--budget-sweep", action="store_true",
        help="Instead of removing engines, test 0.5x/1.5x risk budget for each engine.",
    )
    parser.add_argument(
        "--no-presets", action="store_true",
        help="Run without IA/parking presets (primary-only ablation).",
    )
    parser.add_argument(
        "--idle-alpha-dedup", default=None, choices=["skip", "rename"],
        help="IA dedup mode: 'rename' injects IA engines even when engine_id conflicts with primary (adds __ia_sleeve suffix).",
    )
    parser.add_argument("--form4-sleeve", default=None, help="Form4 insider sleeve preset name.")
    parser.add_argument("--ownership-sleeve", default=None, help="Ownership 13D sleeve preset name.")
    return parser.parse_args()


def _load_manifest(manifest_path: str) -> dict:
    return json.loads(Path(manifest_path).read_text())


def _get_enabled_engines(manifest: dict) -> list[dict]:
    return [e for e in manifest.get("strategy_engines", []) if e.get("enabled", True) is not False]


def _run_backtest(
    manifest_dict: dict,
    parking: str | None,
    idle_alpha: str | None,
    split: str | None,
    output_root: Path,
    run_label: str,
    snapshot_dir: str,
    idle_alpha_dedup: str | None = None,
    form4_sleeve: str | None = None,
    ownership_sleeve: str | None = None,
) -> dict[str, object]:
    """Write a temp manifest JSON and run the backtest, returning key metrics."""
    run_dir = output_root / run_label
    run_dir.mkdir(parents=True, exist_ok=True)

    with tempfile.NamedTemporaryFile(
        mode="w", suffix=".json", dir=run_dir, delete=False
    ) as f:
        json.dump(manifest_dict, f, indent=2)
        temp_manifest = f.name

    cmd = [
        sys.executable, "-m", "apps.backtester.run",
        "--manifest", temp_manifest,
        "--output-root", str(run_dir),
    ]
    if snapshot_dir:
        cmd += ["--snapshot-dir", snapshot_dir]
    if split:
        cmd += ["--split", split]
    if parking:
        cmd += ["--parking", parking]
    if idle_alpha:
        cmd += ["--idle-alpha", idle_alpha]
    if idle_alpha_dedup:
        cmd += ["--idle-alpha-dedup", idle_alpha_dedup]
    if form4_sleeve:
        cmd += ["--form4-sleeve", form4_sleeve]
    if ownership_sleeve:
        cmd += ["--ownership-sleeve", ownership_sleeve]

    print(f"  [{run_label}]...", end=" ", flush=True)
    result = subprocess.run(
        cmd,
        cwd=str(REPO_ROOT),
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        check=False,
    )
    if result.returncode != 0:
        print("FAILED")
        lines = (result.stdout or "").splitlines()
        for line in lines[-4:]:
            print(f"    {line}")
        return {"error": True, "run_label": run_label, "total_return_pct": None, "composite_amplification": None}

    # Find latest run dir
    bt_dirs = sorted(run_dir.glob("bt_*"), key=lambda p: p.stat().st_mtime)
    if not bt_dirs:
        print("NO_RUN_DIR")
        return {"error": True, "run_label": run_label, "total_return_pct": None}
    latest = bt_dirs[-1]

    metrics_path = latest / "metrics" / "metrics_summary.json"
    sleeve_path = latest / "metrics" / "sleeve_decomposition.json"
    summary = json.loads(metrics_path.read_text()) if metrics_path.exists() else {}
    sleeve = json.loads(sleeve_path.read_text()) if sleeve_path.exists() else {}

    ret = summary.get("total_return_pct")
    sharpe = summary.get("sharpe_ratio")
    dd = summary.get("max_drawdown_pct")
    amp = sleeve.get("composite_amplification")
    idle = sleeve.get("avg_idle_fraction_pct")
    print(f"CW={ret:.2f}% Sharpe={sharpe:.2f} DD={dd:.2f}% Idle={idle:.1f}%" if ret is not None else "no metrics")

    return {
        "run_label": run_label,
        "total_return_pct": ret,
        "sharpe_ratio": sharpe,
        "max_drawdown_pct": dd,
        "composite_amplification": amp,
        "avg_idle_fraction_pct": idle,
        "trade_count": summary.get("trade_count", 0),
        "error": False,
    }


def run_loo_ablation(
    base_manifest: dict,
    parking: str | None,
    idle_alpha: str | None,
    split: str | None,
    output_root: Path,
    snapshot_dir: str,
    idle_alpha_dedup: str | None = None,
    form4_sleeve: str | None = None,
    ownership_sleeve: str | None = None,
) -> tuple[dict, list[dict]]:
    """Run baseline + leave-one-out for each enabled engine. Returns (baseline, loo_results)."""
    engines = _get_enabled_engines(base_manifest)
    print(f"\n=== Leave-One-Out Ablation: {len(engines)} enabled engines ===")
    print(f"Total runs: {len(engines) + 1}\n")

    print("Running baseline...")
    baseline = _run_backtest(
        base_manifest, parking, idle_alpha, split,
        output_root / "loo", "baseline", snapshot_dir, idle_alpha_dedup,
        form4_sleeve, ownership_sleeve
    )

    loo_results: list[dict] = []
    for i, engine in enumerate(engines):
        engine_id = engine.get("engine_id", f"engine_{i}")
        # Build manifest with this engine removed
        mod = copy.deepcopy(base_manifest)
        mod["strategy_engines"] = [
            e for e in mod["strategy_engines"]
            if e.get("engine_id") != engine_id
        ]
        mod["experiment_name"] = f"{base_manifest.get('experiment_name', 'ablation')}_loo_{engine_id}"

        result = _run_backtest(
            mod, parking, idle_alpha, split,
            output_root / "loo", f"loo_{engine_id}", snapshot_dir, idle_alpha_dedup,
            form4_sleeve, ownership_sleeve
        )
        base_ret = baseline.get("total_return_pct")
        loo_ret = result.get("total_return_pct")
        delta = (loo_ret - base_ret) if (loo_ret is not None and base_ret is not None) else None
        result["engine_id"] = engine_id
        result["delta_return_pct"] = delta
        result["engine_risk_budget_pct"] = engine.get("engine_risk_budget_pct")
        loo_results.append(result)

    return baseline, loo_results


def run_budget_sweep(
    base_manifest: dict,
    parking: str | None,
    idle_alpha: str | None,
    split: str | None,
    output_root: Path,
    snapshot_dir: str,
    idle_alpha_dedup: str | None = None,
    form4_sleeve: str | None = None,
    ownership_sleeve: str | None = None,
) -> tuple[dict, list[dict]]:
    """Test 0.5x and 1.5x risk budget for each engine vs baseline."""
    engines = _get_enabled_engines(base_manifest)
    print(f"\n=== Budget Sweep: {len(engines)} engines × 2 levels ===")
    print(f"Total additional runs: {len(engines) * 2}\n")

    print("Running baseline...")
    baseline = _run_backtest(
        base_manifest, parking, idle_alpha, split,
        output_root / "budget_sweep", "baseline", snapshot_dir, idle_alpha_dedup,
        form4_sleeve, ownership_sleeve
    )

    sweep_results: list[dict] = []
    for i, engine in enumerate(engines):
        engine_id = engine.get("engine_id", f"engine_{i}")
        orig_budget = engine.get("engine_risk_budget_pct", 1.0)
        for scale, label in [(0.5, "half"), (1.5, "boost")]:
            mod = copy.deepcopy(base_manifest)
            for e in mod["strategy_engines"]:
                if e.get("engine_id") == engine_id:
                    e["engine_risk_budget_pct"] = orig_budget * scale
            mod["experiment_name"] = f"sweep_{engine_id}_{label}"

            result = _run_backtest(
                mod, parking, idle_alpha, split,
                output_root / "budget_sweep", f"{engine_id}_{label}", snapshot_dir, idle_alpha_dedup,
                form4_sleeve, ownership_sleeve
            )
            base_ret = baseline.get("total_return_pct")
            sweep_ret = result.get("total_return_pct")
            delta = (sweep_ret - base_ret) if (sweep_ret is not None and base_ret is not None) else None
            result["engine_id"] = engine_id
            result["scale"] = scale
            result["scale_label"] = label
            result["delta_return_pct"] = delta
            result["orig_budget"] = orig_budget
            result["new_budget"] = orig_budget * scale
            sweep_results.append(result)

    return baseline, sweep_results


def run_greedy_prune(
    base_manifest: dict,
    parking: str | None,
    idle_alpha: str | None,
    split: str | None,
    output_root: Path,
    snapshot_dir: str,
    idle_alpha_dedup: str | None = None,
    form4_sleeve: str | None = None,
    ownership_sleeve: str | None = None,
) -> list[str]:
    """Iteratively remove the engine with highest positive delta (least hurts or improves composite).
    Returns list of removed engine_ids in removal order."""
    current_manifest = copy.deepcopy(base_manifest)
    removed: list[str] = []
    round_num = 0

    print("\n=== Greedy Pruning ===")
    while True:
        round_num += 1
        print(f"\n--- Round {round_num} ---")
        baseline, loo_results = run_loo_ablation(
            current_manifest, parking, idle_alpha, split,
            output_root / f"greedy_round_{round_num}", snapshot_dir, idle_alpha_dedup,
            form4_sleeve, ownership_sleeve
        )
        valid = [r for r in loo_results if not r.get("error") and r.get("delta_return_pct") is not None]
        if not valid:
            print("No valid LOO results, stopping.")
            break
        # Find engine whose removal gives highest delta (most beneficial or least harmful to remove)
        best = max(valid, key=lambda r: r["delta_return_pct"])
        if best["delta_return_pct"] <= 0:
            print(f"Best removal delta={best['delta_return_pct']:.2f}pp (non-positive). Stopping.")
            break
        engine_id = best["engine_id"]
        delta = best["delta_return_pct"]
        print(f"Removing '{engine_id}' (delta={delta:+.2f}pp)")
        removed.append(engine_id)
        current_manifest["strategy_engines"] = [
            e for e in current_manifest["strategy_engines"]
            if e.get("engine_id") != engine_id
        ]

    print(f"\nGreedy pruning complete. Removed {len(removed)} engines: {removed}")
    # Save the pruned manifest
    pruned_name = base_manifest.get("experiment_name", "ablation") + "_pruned"
    current_manifest["experiment_name"] = pruned_name
    pruned_path = output_root / f"{pruned_name}.json"
    pruned_path.write_text(json.dumps(current_manifest, indent=2))
    print(f"Pruned manifest saved to: {pruned_path}")
    return removed


def _fmt(v: object, decimals: int = 2) -> str:
    if v is None:
        return "N/A"
    try:
        return f"{float(v):.{decimals}f}"  # type: ignore[arg-type]
    except (TypeError, ValueError):
        return str(v)


def print_loo_table(baseline: dict, loo_results: list[dict]) -> None:
    base_ret = baseline.get("total_return_pct")
    print(f"\n{'Engine':<55} {'Without%':>9} {'Baseline%':>10} {'Δpp':>7} {'Action':>8}")
    print("-" * 95)
    sorted_loo = sorted(
        loo_results,
        key=lambda r: (r.get("delta_return_pct") is None, r.get("delta_return_pct") or 0.0),
        reverse=True,
    )
    for r in sorted_loo:
        engine_id = str(r.get("engine_id", "?"))[:54]
        without = _fmt(r.get("total_return_pct"))
        delta = r.get("delta_return_pct")
        delta_str = f"{delta:+.2f}" if delta is not None else "N/A"
        action = ""
        if delta is not None:
            action = "REMOVE!" if delta > 0 else "KEEP"
        print(f"{engine_id:<55} {without:>9} {_fmt(base_ret):>10} {delta_str:>7} {action:>8}")
    print("-" * 95)
    print(f"{'BASELINE':<55} {_fmt(base_ret):>9}")


def print_budget_sweep_table(baseline: dict, sweep_results: list[dict]) -> None:
    base_ret = baseline.get("total_return_pct")
    print(f"\n{'Engine':<45} {'Scale':>6} {'Budget':>7} {'CW%':>8} {'Δpp':>7}")
    print("-" * 80)
    for r in sweep_results:
        engine_id = str(r.get("engine_id", "?"))[:44]
        scale = r.get("scale_label", "?")
        budget = _fmt(r.get("new_budget"), decimals=4)
        cw = _fmt(r.get("total_return_pct"))
        delta = r.get("delta_return_pct")
        delta_str = f"{delta:+.2f}" if delta is not None else "N/A"
        print(f"{engine_id:<45} {scale:>6} {budget:>7} {cw:>8} {delta_str:>7}")
    print("-" * 80)
    print(f"{'BASELINE':>58} {_fmt(base_ret):>8}")


def main() -> None:
    args = parse_args()
    output_root = Path(args.output_root)
    output_root.mkdir(parents=True, exist_ok=True)

    base_manifest = _load_manifest(args.manifest)
    exp_name = base_manifest.get("experiment_name", Path(args.manifest).stem)
    parking = None if args.no_presets else args.parking
    idle_alpha = None if args.no_presets else getattr(args, "idle_alpha")

    print(f"Engine Ablation: {exp_name}")
    print(f"Presets: parking={parking}, idle_alpha={idle_alpha}")
    engines = _get_enabled_engines(base_manifest)
    print(f"Enabled engines: {len(engines)}")

    idle_alpha_dedup = None if args.no_presets else getattr(args, "idle_alpha_dedup", None)
    form4_sleeve = None if args.no_presets else getattr(args, "form4_sleeve", None)
    ownership_sleeve = None if args.no_presets else getattr(args, "ownership_sleeve", None)

    if args.greedy_prune:
        removed = run_greedy_prune(
            base_manifest, parking, idle_alpha, args.split,
            output_root, args.snapshot_dir, idle_alpha_dedup,
            form4_sleeve, ownership_sleeve
        )
        print(f"\nFinal removed engines: {removed}")
    elif args.budget_sweep:
        baseline, sweep_results = run_budget_sweep(
            base_manifest, parking, idle_alpha, args.split,
            output_root, args.snapshot_dir, idle_alpha_dedup,
            form4_sleeve, ownership_sleeve
        )
        print_budget_sweep_table(baseline, sweep_results)
        # Save results
        out = output_root / "budget_sweep_results.json"
        out.write_text(json.dumps({"baseline": baseline, "sweep": sweep_results}, indent=2))
        print(f"\nResults saved to: {out}")
    else:
        # Default: LOO ablation
        baseline, loo_results = run_loo_ablation(
            base_manifest, parking, idle_alpha, args.split,
            output_root, args.snapshot_dir, idle_alpha_dedup,
            form4_sleeve, ownership_sleeve
        )
        print_loo_table(baseline, loo_results)
        # Save results
        out = output_root / "loo_results.json"
        out.write_text(json.dumps({"baseline": baseline, "loo": loo_results}, indent=2))
        print(f"\nResults saved to: {out}")


if __name__ == "__main__":
    main()