OBLITERATUS/obliteratus/cli.py

"""CLI entry point for Obliteratus — Master Ablation Suite."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

from rich.console import Console

console = Console()

_BANNER = r"""
[bold red]
  ░▒█▀▀▀█ ░▒█▀▀▄ ░▒█░░░ ▀█▀ ▀▀█▀▀ ░▒█▀▀▀ ░▒█▀▀█ ▒█▀▀█ ▀▀█▀▀ ░▒█░░▒█ ░▒█▀▀▀█
  ░▒█░░▒█ ░▒█▀▀▄ ░▒█░░░ ░█░ ░░█░░ ░▒█▀▀▀ ░▒█▄▄▀ ▒█▄▄█ ░░█░░ ░▒█░░▒█ ░░▀▀▀▄▄
  ░▒█▄▄▄█ ░▒█▄▄▀ ░▒█▄▄█ ▄█▄ ░░▀░░ ░▒█▄▄▄ ░▒█░▒█ ▒█░▒█ ░░▀░░ ░░▒█▄▄█ ░▒█▄▄▄█
[/bold red]
[dim]  ════════════════════════════════════════════════════════════════════[/dim]
[bold white]   MASTER ABLATION SUITE[/bold white] [dim]//[/dim] [bold red]Break the chains. Free the mind.[/bold red]
[dim]  ════════════════════════════════════════════════════════════════════[/dim]
"""


def _add_gpu_args(parser):
    """Add --gpus flag for multi-GPU control."""
    gpu_group = parser.add_argument_group("GPU selection")
    gpu_group.add_argument(
        "--gpus", type=str, default=None, metavar="IDS",
        help=(
            "Comma-separated GPU IDs to use (e.g. '0,1,2,3' or 'all'). "
            "Sets CUDA_VISIBLE_DEVICES. By default uses all available GPUs. "
            "Models are automatically split across selected GPUs via accelerate."
        ),
    )


def _add_remote_args(parser):
    """Add --remote execution flags to a subcommand parser."""
    remote_group = parser.add_argument_group("remote execution")
    remote_group.add_argument(
        "--remote", type=str, default=None, metavar="[USER@]HOST",
        help="Run on a remote GPU node via SSH (e.g. root@gpu-node or just gpu-node)",
    )
    remote_group.add_argument(
        "--ssh-key", type=str, default=None,
        help="Path to SSH private key (default: use SSH agent or ~/.ssh/id_rsa)",
    )
    remote_group.add_argument(
        "--ssh-port", type=int, default=22,
        help="SSH port on remote host (default: 22)",
    )
    remote_group.add_argument(
        "--remote-dir", type=str, default="/tmp/obliteratus_run",
        help="Working directory on the remote machine (default: /tmp/obliteratus_run)",
    )
    remote_group.add_argument(
        "--remote-python", type=str, default="python3",
        help="Python binary on the remote machine (default: python3)",
    )
    remote_group.add_argument(
        "--no-sync", action="store_true", default=False,
        help="Don't copy results back to local machine after remote run",
    )


def _apply_gpu_selection(args):
    """Set CUDA_VISIBLE_DEVICES based on --gpus flag (for local runs only)."""
    import os

    gpus = getattr(args, "gpus", None)
    if gpus is None or getattr(args, "remote", None):
        return  # skip for remote runs (handled by remote runner)

    if gpus.lower() == "all":
        return  # use all GPUs (default behavior)

    # Validate: should be comma-separated integers
    try:
        gpu_ids = [int(g.strip()) for g in gpus.split(",")]
    except ValueError:
        console.print(f"[red]Invalid --gpus value: {gpus!r}. Expected comma-separated integers or 'all'.[/]")
        raise SystemExit(1)

    os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(str(g) for g in gpu_ids)
    console.print(f"[dim]Using GPUs: {gpu_ids} (CUDA_VISIBLE_DEVICES={os.environ['CUDA_VISIBLE_DEVICES']})[/dim]")


def main(argv: list[str] | None = None):
    console.print(_BANNER)
    parser = argparse.ArgumentParser(
        prog="obliteratus",
        description="Master Ablation Suite for HuggingFace transformers",
    )
    subparsers = parser.add_subparsers(dest="command", required=True)

    # --- run ---
    run_parser = subparsers.add_parser("run", help="Run an ablation from a YAML config")
    run_parser.add_argument("config", type=str, help="Path to YAML config file")
    run_parser.add_argument("--output-dir", type=str, default=None, help="Override output dir")
    run_parser.add_argument(
        "--preset",
        type=str,
        default=None,
        help="Apply a preset (e.g. quick, full, attention, jailbreak, guardrail)",
    )
    _add_gpu_args(run_parser)
    _add_remote_args(run_parser)

    # --- info ---
    info_parser = subparsers.add_parser("info", help="Print model architecture info")
    info_parser.add_argument("model", type=str, help="HuggingFace model name/path")
    info_parser.add_argument("--task", type=str, default="causal_lm", choices=["causal_lm", "classification"])
    info_parser.add_argument("--device", type=str, default="cpu")
    info_parser.add_argument("--dtype", type=str, default="float32")

    # --- interactive ---
    subparsers.add_parser(
        "interactive",
        help="Guided setup — pick hardware, model, and preset interactively",
    )

    # --- models ---
    models_parser = subparsers.add_parser("models", help="Browse curated models by compute tier")
    models_parser.add_argument(
        "--tier",
        type=str,
        default=None,
        choices=["tiny", "small", "medium", "large", "frontier"],
        help="Filter by compute tier",
    )

    # --- presets ---
    subparsers.add_parser("presets", help="Browse ablation presets (quick, full, jailbreak, etc.)")

    # --- strategies ---
    subparsers.add_parser("strategies", help="List available ablation strategies")

    # --- ui ---
    ui_parser = subparsers.add_parser(
        "ui",
        help="Launch the Gradio web UI locally (same UI as the HuggingFace Space)",
    )
    ui_parser.add_argument(
        "--port", type=int, default=7860, help="Server port (default: 7860)",
    )
    ui_parser.add_argument(
        "--host", type=str, default="0.0.0.0", help="Server host (default: 0.0.0.0)",
    )
    ui_parser.add_argument(
        "--share", action="store_true", help="Create a public Gradio share link",
    )
    ui_parser.add_argument(
        "--no-browser", action="store_true", help="Don't auto-open browser on launch",
    )
    ui_parser.add_argument(
        "--auth", type=str, default=None,
        help="Basic auth as user:pass",
    )
    ui_parser.add_argument(
        "--quiet", action="store_true", help="Suppress the startup banner",
    )

    # --- obliterate (primary) + abliterate (backward-compat alias) ---
    def _add_obliterate_args(p):
        p.add_argument("model", type=str, help="HuggingFace model name/path")
        p.add_argument("--output-dir", type=str, default=None, help="Where to save the obliterated model")
        p.add_argument("--device", type=str, default="auto")
        p.add_argument("--dtype", type=str, default="float16")
        p.add_argument(
            "--method", type=str, default="advanced",
            choices=[
                "basic", "advanced", "aggressive", "spectral_cascade",
                "informed", "surgical", "optimized", "inverted", "nuclear",
            ],
            help="Liberation method (default: advanced)",
        )
        p.add_argument("--n-directions", type=int, default=None, help="Override: number of refusal directions to extract")
        p.add_argument(
            "--direction-method", type=str, default=None,
            choices=["diff_means", "svd", "leace"],
            help="Direction extraction method: diff_means (simple, robust), svd (multi-direction), leace (optimal erasure)",
        )
        p.add_argument("--regularization", type=float, default=None, help="Override: fraction to preserve (0.0-1.0)")
        p.add_argument("--refinement-passes", type=int, default=None, help="Override: number of iterative passes")
        p.add_argument(
            "--quantization", type=str, default=None, choices=["4bit", "8bit"],
            help="Load model with quantization (4bit or 8bit). Requires bitsandbytes.",
        )
        p.add_argument(
            "--large-model", action="store_true", default=False,
            help="Enable conservative defaults for 120B+ models (fewer directions, 1 pass, lower SAE expansion).",
        )
        p.add_argument(
            "--verify-sample-size", type=int, default=None,
            help="Number of harmful prompts to test for refusal rate (default: 30). "
                 "Increase for tighter confidence intervals (e.g. 100 for ~1%% resolution).",
        )
        p.add_argument(
            "--contribute", action="store_true", default=False,
            help="Save a community contribution record after the run completes.",
        )
        p.add_argument(
            "--contribute-notes", type=str, default="",
            help="Optional notes to include with the community contribution.",
        )

    abl_parser = subparsers.add_parser(
        "obliterate",
        help="One-click: remove refusal directions from a model (SOTA multi-technique)",
    )
    _add_obliterate_args(abl_parser)
    _add_gpu_args(abl_parser)
    _add_remote_args(abl_parser)
    # Backward-compat alias (hidden from help)
    abl_alias = subparsers.add_parser("abliterate", help=argparse.SUPPRESS)
    _add_obliterate_args(abl_alias)
    _add_gpu_args(abl_alias)
    _add_remote_args(abl_alias)

    # --- report ---
    report_parser = subparsers.add_parser("report", help="Regenerate report from saved results")
    report_parser.add_argument("results_json", type=str, help="Path to results.json")
    report_parser.add_argument("--output-dir", type=str, default=None)

    # --- aggregate ---
    aggregate_parser = subparsers.add_parser("aggregate", help="Aggregate community contribution results")
    aggregate_parser.add_argument(
        "--dir", type=str, default="community_results",
        help="Directory containing contribution JSON files",
    )

    # --- tourney ---
    tourney_parser = subparsers.add_parser(
        "tourney",
        help="March Madness tournament — pit all methods against each other, push winner to Hub",
    )
    tourney_parser.add_argument("model", type=str, help="HuggingFace model name/path")
    tourney_parser.add_argument("--hub-org", type=str, default=None, help="HF org to push winner (e.g. my-org)")
    tourney_parser.add_argument("--hub-repo", type=str, default=None, help="Full HF repo ID (overrides --hub-org)")
    tourney_parser.add_argument("--device", type=str, default="auto")
    tourney_parser.add_argument("--dtype", type=str, default="float16")
    tourney_parser.add_argument("--dataset", type=str, default="builtin", help="Dataset source (default: builtin)")
    tourney_parser.add_argument(
        "--quantization", type=str, default=None, choices=["4bit", "8bit"],
        help="Load model with quantization",
    )
    tourney_parser.add_argument("--output-dir", type=str, default="/tmp/obliteratus_tourney")
    tourney_parser.add_argument(
        "--methods", type=str, nargs="+", default=None,
        help="Override: only run these methods (space-separated)",
    )
    _add_gpu_args(tourney_parser)
    _add_remote_args(tourney_parser)

    # --- recommend ---
    recommend_parser = subparsers.add_parser(
        "recommend",
        help="Show telemetry-driven best method + hyperparams for a model",
    )
    recommend_parser.add_argument("model", type=str, help="HuggingFace model name/path")
    recommend_parser.add_argument("--device", type=str, default="cpu")
    recommend_parser.add_argument("--dtype", type=str, default="float32")
    recommend_parser.add_argument(
        "--insights", action="store_true", default=False,
        help="Also show global cross-architecture insights",
    )

    # --- gpu-calc ---
    calc_parser = subparsers.add_parser(
        "gpu-calc",
        help="Estimate minimum GPUs needed for a model",
    )
    calc_parser.add_argument(
        "model", type=str, nargs="?", default=None,
        help="HuggingFace model name/path (auto-fetches param counts)",
    )
    calc_parser.add_argument(
        "--params", type=float, default=None, metavar="B",
        help="Total parameters in billions (overrides auto-detection)",
    )
    calc_parser.add_argument(
        "--active-params", type=float, default=None, metavar="B",
        help="Active parameters in billions (for MoE models; defaults to --params)",
    )
    calc_parser.add_argument(
        "--dtype", type=str, default="bfloat16",
        choices=["float32", "float16", "bfloat16", "int8", "int4"],
        help="Data type for model weights (default: bfloat16)",
    )
    calc_parser.add_argument(
        "--gpu-mem", type=float, default=80.0, metavar="GB",
        help="VRAM per GPU in GB (default: 80 for A100-80GB)",
    )

    args = parser.parse_args(argv)

    # Apply GPU selection early (before any CUDA init)
    _apply_gpu_selection(args)

    if args.command == "gpu-calc":
        _cmd_gpu_calc(args)
        return
    elif args.command == "run":
        if getattr(args, "remote", None):
            _cmd_remote_run(args)
        else:
            _cmd_run(args)
    elif args.command == "interactive":
        _cmd_interactive()
    elif args.command == "models":
        _cmd_models(args)
    elif args.command == "presets":
        _cmd_presets()
    elif args.command == "info":
        _cmd_info(args)
    elif args.command == "strategies":
        _cmd_strategies()
    elif args.command == "report":
        _cmd_report(args)
    elif args.command == "aggregate":
        _cmd_aggregate(args)
    elif args.command == "ui":
        _cmd_ui(args)
    elif args.command == "recommend":
        _cmd_recommend(args)
    elif args.command == "tourney":
        if getattr(args, "remote", None):
            _cmd_remote_tourney(args)
        else:
            _cmd_tourney(args)
    elif args.command in ("obliterate", "abliterate"):
        if getattr(args, "remote", None):
            _cmd_remote_abliterate(args)
        else:
            _cmd_abliterate(args)


def _cmd_ui(args):
    from obliteratus.local_ui import launch_local_ui

    auth = tuple(args.auth.split(":", 1)) if args.auth else None
    launch_local_ui(
        host=args.host,
        port=args.port,
        share=args.share,
        open_browser=not args.no_browser,
        auth=auth,
        quiet=args.quiet,
    )


def _cmd_interactive():
    from obliteratus.interactive import run_interactive
    run_interactive()


def _cmd_models(args):
    from rich.table import Table

    from obliteratus.presets import get_presets_by_tier, list_all_presets

    presets = get_presets_by_tier(args.tier) if args.tier else list_all_presets()

    table = Table(title="Model Library — Curated Targets")
    table.add_column("Model", style="green")
    table.add_column("HuggingFace ID", style="cyan")
    table.add_column("Params", justify="right")
    table.add_column("Tier", style="yellow")
    table.add_column("Dtype")
    table.add_column("Quant")
    table.add_column("Description")

    for p in presets:
        table.add_row(
            p.name,
            p.hf_id,
            p.params,
            p.tier.upper(),
            p.recommended_dtype,
            p.recommended_quantization or "—",
            p.description,
        )

    console.print(table)
    console.print(
        "\n[dim]Tiers: TINY = CPU/laptop | SMALL = 4-8GB | "
        "MEDIUM = 8-16GB | LARGE = 24GB+ | FRONTIER = multi-GPU/cloud[/dim]"
    )


def _cmd_presets():
    from rich.table import Table

    from obliteratus.study_presets import list_study_presets

    presets = list_study_presets()

    table = Table(title="Ablation Presets")
    table.add_column("Key", style="cyan", min_width=12)
    table.add_column("Name", style="green")
    table.add_column("Strategies", style="yellow")
    table.add_column("Samples", justify="right")
    table.add_column("Description", max_width=55)

    for p in presets:
        strats = ", ".join(s["name"] for s in p.strategies)
        table.add_row(p.key, p.name, strats, str(p.max_samples), p.description)

    console.print(table)
    console.print(
        "\n[dim]Usage:  obliteratus run config.yaml --preset quick\n"
        "   or:  set preset: quick in your YAML file[/dim]"
    )


def _cmd_run(args):
    from obliteratus.config import StudyConfig
    from obliteratus.runner import run_study

    config = StudyConfig.from_yaml(args.config)
    # If --preset flag given, inject it so from_dict picks it up
    if args.preset:
        import yaml

        raw = yaml.safe_load(Path(args.config).read_text())
        raw["preset"] = args.preset
        config = StudyConfig.from_dict(raw)
    if args.output_dir:
        config.output_dir = args.output_dir

    # If YAML has a remote: section, dispatch to remote runner
    if config.remote is not None:
        from obliteratus.remote import RemoteConfig as _RC, RemoteRunner

        rc = _RC(
            host=config.remote.host,
            user=config.remote.user,
            port=config.remote.port,
            ssh_key=config.remote.ssh_key,
            remote_dir=config.remote.remote_dir,
            python=config.remote.python,
            sync_results=config.remote.sync_results,
            gpus=config.remote.gpus,
        )
        runner = RemoteRunner(rc)
        result_path = runner.run_config(
            local_config_path=args.config,
            local_output_dir=config.output_dir,
            preset=args.preset,
        )
        if result_path:
            console.print(f"\n[bold green]Remote run complete.[/] Results at: [cyan]{result_path}[/]")
        else:
            console.print("[red]Remote run failed. Check logs above.[/]")
            raise SystemExit(1)
        return

    run_study(config)


def _cmd_info(args):
    from obliteratus.models.loader import load_model

    console.print(f"[bold cyan]Loading model:[/bold cyan] {args.model}")
    handle = load_model(
        model_name=args.model,
        task=args.task,
        device=args.device,
        dtype=args.dtype,
    )
    summary = handle.summary()
    for key, val in summary.items():
        if isinstance(val, int) and val > 1000:
            console.print(f"  {key}: {val:,}")
        else:
            console.print(f"  {key}: {val}")


def _cmd_strategies():
    from obliteratus.strategies import STRATEGY_REGISTRY

    console.print("[bold]Available ablation strategies:[/bold]\n")
    for name, cls in sorted(STRATEGY_REGISTRY.items()):
        doc = (cls.__doc__ or "").strip().split("\n")[0]
        console.print(f"  [cyan]{name}[/cyan]  — {doc}")


def _cmd_report(args):
    from obliteratus.reporting.report import AblationReport, AblationResult

    path = Path(args.results_json)
    data = json.loads(path.read_text())

    report = AblationReport(model_name=data["model_name"])
    report.add_baseline(data["baseline_metrics"])
    for r in data["results"]:
        report.add_result(
            AblationResult(
                strategy=r["strategy"],
                component=r["component"],
                description=r["description"],
                metrics=r["metrics"],
                metadata=r.get("metadata"),
            )
        )

    report.print_summary()

    output_dir = Path(args.output_dir) if args.output_dir else path.parent
    metric_name = list(data["baseline_metrics"].keys())[0]
    try:
        report.plot_impact(metric=metric_name, output_path=output_dir / "impact.png")
        report.plot_heatmap(output_path=output_dir / "heatmap.png")
        console.print(f"\nPlots saved to {output_dir}/")
    except Exception as e:
        console.print(f"[yellow]Could not generate plots: {e}[/yellow]")


def _cmd_aggregate(args):
    from obliteratus.community import aggregate_results, load_contributions

    contrib_dir = args.dir
    records = load_contributions(contrib_dir)
    if not records:
        console.print(f"[yellow]No contributions found in {contrib_dir}[/yellow]")
        return

    aggregated = aggregate_results(records)

    from rich.table import Table

    table = Table(title="Aggregated Community Results")
    table.add_column("Model", style="green")
    table.add_column("Method", style="cyan")
    table.add_column("Runs", justify="right")
    table.add_column("Mean Refusal", justify="right")
    table.add_column("Mean Perplexity", justify="right")

    for model_name, methods in sorted(aggregated.items()):
        for method_name, stats in sorted(methods.items()):
            refusal = stats.get("refusal_rate", {}).get("mean", "N/A")
            ppl = stats.get("perplexity", {}).get("mean", "N/A")
            if isinstance(refusal, float):
                refusal = f"{refusal:.4f}"
            if isinstance(ppl, float):
                ppl = f"{ppl:.2f}"
            table.add_row(
                model_name.split("/")[-1] if "/" in model_name else model_name,
                method_name,
                str(stats["n_runs"]),
                str(refusal),
                str(ppl),
            )

    console.print(table)


def _cmd_recommend(args):
    from rich.markdown import Markdown
    from rich.panel import Panel

    from obliteratus.architecture_profiles import detect_architecture, enhance_profile_with_telemetry
    from obliteratus.adaptive_defaults import format_recommendation, get_global_insights

    model_name = args.model
    console.print(f"\nAnalyzing [bold]{model_name}[/]...")

    # Detect architecture
    try:
        from transformers import AutoConfig
        config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
        num_layers = getattr(config, "num_hidden_layers", 0)
        hidden_size = getattr(config, "hidden_size", 0)
    except Exception:
        config = None
        num_layers = 0
        hidden_size = 0

    profile = detect_architecture(model_name, config, num_layers, hidden_size)
    profile, rec = enhance_profile_with_telemetry(profile)

    console.print(Panel(
        f"[bold]{profile.profile_label}[/]\n"
        f"Architecture: {profile.arch_class.value} | Reasoning: {profile.reasoning_class.value}\n"
        f"Params: ~{profile.total_params_b:.1f}B | Layers: {profile.num_layers} | "
        f"Hidden: {profile.hidden_size}",
        title="Architecture Profile",
        border_style="cyan",
    ))

    if rec:
        md = format_recommendation(rec)
        console.print(Markdown(md))
    else:
        console.print("\n[yellow]Could not fetch telemetry — using research-grounded defaults.[/]")

    console.print(f"\n[bold green]Research default method:[/] {profile.recommended_method}")
    if profile.method_overrides:
        console.print("[bold green]Overrides:[/]")
        for k, v in sorted(profile.method_overrides.items()):
            console.print(f"  {k}: {v}")

    if args.insights:
        console.print("\n")
        console.rule("[bold magenta]Global Telemetry Insights")
        insights = get_global_insights()
        console.print(f"Total records analyzed: {insights['total_records']}")
        if insights["overall_best_methods"]:
            console.print("\n[bold]Overall method ranking (all architectures):[/]")
            for entry in insights["overall_best_methods"][:10]:
                console.print(
                    f"  {entry['method']}: {entry['mean_score']:.4f} "
                    f"({entry['n_runs']} runs)"
                )
        if insights["architecture_breakdown"]:
            console.print("\n[bold]Per-architecture breakdown:[/]")
            for label, info in insights["architecture_breakdown"].items():
                console.print(
                    f"  {label}: best={info['best_method']} "
                    f"({info['best_score']:.4f}), "
                    f"{info['n_methods_tested']} methods tested, "
                    f"{info['total_runs']} runs"
                )


def _cmd_tourney(args):
    from obliteratus.tourney import TourneyRunner, render_bracket

    def on_log(msg):
        console.print(msg)

    def on_round(rnd):
        console.print()
        console.rule(f"[bold green]Round {rnd.round_num} complete — "
                     f"{len(rnd.advanced_to)} advance, {len(rnd.eliminated)} eliminated")

    runner = TourneyRunner(
        model_name=args.model,
        hub_org=args.hub_org,
        hub_repo=args.hub_repo,
        device=args.device,
        dtype=args.dtype,
        dataset_key=args.dataset,
        quantization=args.quantization,
        methods=args.methods,
        output_dir=args.output_dir,
        on_log=on_log,
        on_round=on_round,
    )

    result = runner.run()

    if result.winner:
        console.print()
        console.rule("[bold magenta]TOURNAMENT CHAMPION", style="magenta")
        console.print(f"  [bold]{result.winner.method}[/] — score {result.winner.score:.4f}")
        console.print(f"  Refusal rate: {result.winner.metrics.get('refusal_rate', '?')}")
        console.print(f"  Coherence:    {result.winner.metrics.get('coherence', '?')}")
        if result.hub_repo:
            console.print(f"  Pushed to:    [link=https://huggingface.co/{result.hub_repo}]{result.hub_repo}[/link]")
        console.print(f"\n  Full bracket: {args.output_dir}/tourney_bracket.md")


def _cmd_abliterate(args):
    from rich.live import Live
    from rich.panel import Panel
    from rich.table import Table
    from rich.text import Text

    from obliteratus.abliterate import METHODS, STAGES, AbliterationPipeline

    model_name = args.model
    output_dir = args.output_dir or f"abliterated/{model_name.replace('/', '_')}"
    method = args.method
    method_label = METHODS.get(method, {}).get("label", method)

    # Stage state tracking
    stage_status = {s.key: "waiting" for s in STAGES}
    stage_msgs = {s.key: "" for s in STAGES}
    log_lines: list[str] = []

    def make_display():
        table = Table(show_header=False, expand=True, border_style="green")
        table.add_column("", width=6)
        table.add_column("Stage", min_width=10)
        table.add_column("Status", min_width=50)
        for i, s in enumerate(STAGES):
            st = stage_status[s.key]
            if st == "done":
                icon = "[bold green]✓[/]"
                bar = "[green]" + "█" * 20 + "[/]"
            elif st == "running":
                icon = "[bold yellow]⚡[/]"
                bar = "[yellow]" + "▓" * 10 + "░" * 10 + "[/]"
            else:
                icon = "[dim]○[/]"
                bar = "[dim]" + "░" * 20 + "[/]"
            msg = stage_msgs.get(s.key, "")
            table.add_row(
                f"[cyan][{i + 1}/6][/]",
                f"{icon} [bold]{s.name}[/]",
                f"{bar}  {msg}",
            )

        header = Text.from_markup(
            f"[bold green]OBLITERATUS — ABLITERATION PIPELINE[/]\n"
            f"[dim]Target:[/] [cyan]{model_name}[/]  →  [cyan]{output_dir}[/]\n"
            f"[dim]Method:[/] [magenta]{method_label}[/]"
        )

        # Last 12 log lines
        recent = log_lines[-12:] if log_lines else ["Initializing..."]
        log_text = "\n".join(f"[dim]>[/] {line}" for line in recent)

        return Panel(
            f"{header}\n\n{table}\n\n[dim]─── LOG ───[/]\n{log_text}",
            border_style="green",
            title="[bold green]⚗ ABLITERATE ⚗[/]",
        )

    def on_stage(result):
        stage_status[result.stage] = result.status
        stage_msgs[result.stage] = result.message
        if live:
            live.update(make_display())

    def on_log(msg):
        log_lines.append(msg)
        if live:
            live.update(make_display())

    live = None
    pipeline = AbliterationPipeline(
        model_name=model_name,
        output_dir=output_dir,
        device=args.device,
        dtype=args.dtype,
        method=method,
        n_directions=args.n_directions,
        direction_method=getattr(args, "direction_method", None),
        regularization=args.regularization,
        refinement_passes=args.refinement_passes,
        quantization=args.quantization,
        large_model_mode=getattr(args, "large_model", False),
        verify_sample_size=getattr(args, "verify_sample_size", None),
        on_stage=on_stage,
        on_log=on_log,
    )

    with Live(make_display(), console=console, refresh_per_second=4) as live_ctx:
        live = live_ctx
        try:
            result_path = pipeline.run()
            live.update(make_display())
        except Exception as e:
            log_lines.append(f"[red]ERROR: {e}[/]")
            live.update(make_display())
            raise

    # ── Telemetry: send pipeline report to community leaderboard ──
    try:
        from obliteratus.telemetry import maybe_send_pipeline_report
        maybe_send_pipeline_report(pipeline)
    except Exception:
        pass  # Telemetry is best-effort

    # ── Community contribution (--contribute flag) ──
    contrib_path = None
    if getattr(args, "contribute", False):
        try:
            from obliteratus.community import save_contribution
            contrib_path = save_contribution(
                pipeline,
                model_name=model_name,
                notes=getattr(args, "contribute_notes", ""),
            )
        except Exception as e:
            console.print(f"[yellow]Could not save contribution: {e}[/yellow]")

    console.print()
    contrib_line = ""
    if contrib_path:
        contrib_line = f"\n  Contribution: [cyan]{contrib_path}[/]"
    console.print(
        Panel(
            f"[bold green]Abliteration complete![/]\n\n"
            f"  Model saved to: [cyan]{result_path}[/]\n"
            f"  Metadata: [cyan]{result_path}/abliteration_metadata.json[/]"
            f"{contrib_line}\n\n"
            f"  [dim]Load with:[/] AutoModelForCausalLM.from_pretrained('{result_path}')",
            border_style="green",
            title="[bold green]✓ REBIRTH COMPLETE[/]",
        )
    )


def _cmd_gpu_calc(args):
    import math

    from rich.panel import Panel
    from rich.table import Table

    BYTES_PER_PARAM = {
        "float32": 4,
        "float16": 2,
        "bfloat16": 2,
        "int8": 1,
        "int4": 0.5,
    }

    # Resolve param counts
    total_params_b = args.params
    active_params_b = args.active_params

    if total_params_b is None:
        if args.model is None:
            console.print("[red]Provide either a model name or --params.[/]")
            raise SystemExit(1)
        console.print(f"Fetching config for [cyan]{args.model}[/]...")
        try:
            from transformers import AutoConfig
            config = AutoConfig.from_pretrained(args.model, trust_remote_code=True)
        except Exception as e:
            console.print(f"[red]Could not load config: {e}[/]")
            raise SystemExit(1)

        # Total params: prefer explicit num_parameters, else estimate from config
        total_params_b = _estimate_total_params_b(config)

        # Active params for MoE
        if active_params_b is None:
            active_params_b = _estimate_active_params_b(config, total_params_b)

    if active_params_b is None:
        active_params_b = total_params_b

    bpp = BYTES_PER_PARAM[args.dtype]
    gpu_mem_gb = args.gpu_mem

    # Model weight memory (use base-10 GB to match HF/nvidia conventions)
    weight_gb = total_params_b * bpp

    # Activation overhead during forward passes (PROBE/VERIFY).
    # Scales with active params, not total. Empirical from benchmarks:
    # - DeepSeek-70B (149GB): failed at 160GB (2 GPUs), OK at 240GB (3 GPUs)
    # - GPT-OSS-120B (234GB): failed at 240GB (3 GPUs), OK at 320GB (4 GPUs)
    # This implies ~15-35% overhead. We use 20% as a reasonable middle ground.
    active_weight_gb = active_params_b * bpp
    activation_overhead_gb = active_weight_gb * 0.20

    # CUDA context + fragmentation overhead: ~1.5 GB per GPU (fixed cost)
    cuda_overhead_per_gpu = 1.5

    # Total memory needed (before splitting across GPUs)
    total_needed_gb = weight_gb + activation_overhead_gb

    # Find minimum GPUs: we need total_needed / (gpu_mem - cuda_overhead) GPUs
    usable_per_gpu = gpu_mem_gb - cuda_overhead_per_gpu
    if usable_per_gpu <= 0:
        console.print("[red]GPU memory too small after CUDA overhead.[/]")
        raise SystemExit(1)

    min_gpus = math.ceil(total_needed_gb / usable_per_gpu)
    min_gpus = max(min_gpus, 1)

    # Show results for a range of GPU counts
    is_moe = active_params_b < total_params_b * 0.99

    table = Table(title="GPU Configurations", show_edge=True)
    table.add_column("GPUs", justify="right", style="cyan")
    table.add_column("VRAM/GPU", justify="right")
    table.add_column("Total VRAM", justify="right")
    table.add_column("Headroom", justify="right")
    table.add_column("Verdict", min_width=20)

    # Show from min_gpus-1 (to show why it fails) up to 8
    low = max(1, min_gpus - 1)
    high = max(min_gpus + 3, 8)
    for n in range(low, high + 1):
        total_vram = n * gpu_mem_gb
        usable_vram = n * usable_per_gpu
        headroom = usable_vram - total_needed_gb
        headroom_pct = headroom / total_needed_gb * 100
        vram_per = total_needed_gb / n

        if headroom < 0:
            verdict = "[red]INSUFFICIENT[/]"
        elif headroom_pct < 15:
            verdict = "[yellow]TIGHT — may fail[/]"
        elif n == min_gpus:
            verdict = "[bold green]MINIMUM (recommended)[/]"
        else:
            verdict = "[green]OK[/] [dim](more GPUs = slower)[/]"

        table.add_row(
            str(n),
            f"{vram_per:.1f} GB",
            f"{total_vram:.0f} GB",
            f"{headroom:+.1f} GB ({headroom_pct:+.0f}%)",
            verdict,
        )

    model_label = args.model or f"{total_params_b:.1f}B params"
    moe_line = ""
    if is_moe:
        moe_line = f"\n  Active params:  [cyan]{active_params_b:.1f}B[/] ({active_params_b/total_params_b*100:.0f}% of total — MoE)"

    console.print(Panel(
        f"  Model:          [cyan]{model_label}[/]\n"
        f"  Total params:   [cyan]{total_params_b:.1f}B[/]"
        f"{moe_line}\n"
        f"  Dtype:          [cyan]{args.dtype}[/] ({bpp} bytes/param)\n"
        f"  Weight memory:  [cyan]{weight_gb:.1f} GB[/]\n"
        f"  Activation est: [cyan]{activation_overhead_gb:.1f} GB[/]\n"
        f"  Total needed:   [bold]{total_needed_gb:.1f} GB[/]\n"
        f"  GPU VRAM:       [cyan]{gpu_mem_gb:.0f} GB[/] per device",
        title="[bold]GPU Calculator[/]",
        border_style="cyan",
    ))
    console.print(table)
    console.print(
        f"\n  [bold green]Minimum GPUs: {min_gpus}[/]"
        f"  ({min_gpus} x {gpu_mem_gb:.0f} GB = {min_gpus * gpu_mem_gb:.0f} GB)\n"
    )
    console.print(
        "[dim]Note: fewer GPUs = faster (pipeline parallel has cross-device overhead).\n"
        "Estimates are conservative. Actual memory may vary with sequence length\n"
        "and model architecture. See 'obliteratus obliterate --help' for runtime options.[/]\n"
    )


def _estimate_total_params_b(config) -> float:
    """Estimate total parameter count in billions from a HuggingFace config."""
    # Some configs have explicit param counts
    for attr in ("num_parameters", "n_params"):
        val = getattr(config, attr, None)
        if val and val > 1000:
            return val / 1e9

    # Estimate from architecture dimensions
    h = getattr(config, "hidden_size", 0)
    L = getattr(config, "num_hidden_layers", 0)
    V = getattr(config, "vocab_size", 0)
    i = getattr(config, "intermediate_size", h * 4)

    if h == 0 or L == 0:
        console.print("[red]Cannot determine model size from config. Use --params.[/]")
        raise SystemExit(1)

    n_heads = getattr(config, "num_attention_heads", None) or (h // 128)
    head_dim = getattr(config, "head_dim", None) or (h // n_heads if n_heads else 128)
    kv_heads = getattr(config, "num_key_value_heads", None) or n_heads

    # Attention: Q + K + V projections + output projection
    attn_params = h * (n_heads * head_dim) + h * (kv_heads * head_dim) * 2 + (n_heads * head_dim) * h

    # FFN (MoE or dense)
    n_experts = getattr(config, "num_local_experts", getattr(config, "num_experts", 1)) or 1
    # MoE models often have a separate intermediate size for expert FFNs
    moe_i = getattr(config, "moe_intermediate_size", i)
    # gate + up + down projections per expert
    ffn_per_expert = h * moe_i * 3
    ffn_params = ffn_per_expert * n_experts
    # Some architectures (Qwen, DeepSeek) also have a shared/dense FFN per layer
    if n_experts > 1 and hasattr(config, "moe_intermediate_size"):
        # The dense FFN uses the main intermediate_size
        ffn_params += h * i * 3
    # Router
    if n_experts > 1:
        ffn_params += h * n_experts

    # Per-layer: attention + FFN + layernorms
    layer_params = attn_params + ffn_params + h * 4  # 2 layernorms, 2 params each

    # Embedding + LM head
    embed_params = V * h * 2  # input + output embeddings (may be tied but counts for memory)

    total = L * layer_params + embed_params
    return total / 1e9


def _estimate_active_params_b(config, total_params_b: float) -> float:
    """For MoE models, estimate active parameters per forward pass."""
    n_experts = getattr(config, "num_local_experts", getattr(config, "num_experts", 1)) or 1
    if n_experts <= 1:
        return total_params_b

    top_k = getattr(config, "num_experts_per_tok", getattr(config, "top_k", 2)) or 2

    h = getattr(config, "hidden_size", 0)
    i = getattr(config, "intermediate_size", h * 4)
    moe_i = getattr(config, "moe_intermediate_size", i)
    L = getattr(config, "num_hidden_layers", 0)

    # FFN per expert (uses moe_intermediate_size if available)
    ffn_per_expert = h * moe_i * 3
    # Active FFN = top_k experts instead of all n_experts
    ffn_all = ffn_per_expert * n_experts * L
    ffn_active = ffn_per_expert * top_k * L
    # Non-FFN params (includes any shared/dense FFN)
    non_ffn = total_params_b * 1e9 - ffn_all
    active = non_ffn + ffn_active
    return max(active / 1e9, 0.1)


def _make_remote_runner(args):
    """Create a RemoteRunner from CLI --remote flags."""
    from obliteratus.remote import RemoteConfig, RemoteRunner

    rc = RemoteConfig.from_cli_args(
        args.remote,
        port=args.ssh_port,
        ssh_key=args.ssh_key,
        remote_dir=args.remote_dir,
        python=args.remote_python,
        sync_results=not args.no_sync,
        gpus=getattr(args, "gpus", None),
    )
    return RemoteRunner(rc)


def _cmd_remote_abliterate(args):
    from rich.panel import Panel

    runner = _make_remote_runner(args)

    kwargs = {}
    if args.method:
        kwargs["method"] = args.method
    if args.device:
        kwargs["device"] = args.device
    if args.dtype:
        kwargs["dtype"] = args.dtype
    if args.quantization:
        kwargs["quantization"] = args.quantization
    if args.n_directions is not None:
        kwargs["n_directions"] = args.n_directions
    if getattr(args, "direction_method", None):
        kwargs["direction_method"] = args.direction_method
    if args.regularization is not None:
        kwargs["regularization"] = args.regularization
    if args.refinement_passes is not None:
        kwargs["refinement_passes"] = args.refinement_passes
    if getattr(args, "large_model", False):
        kwargs["large_model"] = True
    if getattr(args, "verify_sample_size", None) is not None:
        kwargs["verify_sample_size"] = args.verify_sample_size

    result_path = runner.run_obliterate(
        model=args.model,
        local_output_dir=args.output_dir,
        **kwargs,
    )

    if result_path:
        console.print(
            Panel(
                f"[bold green]Remote abliteration complete![/]\n\n"
                f"  Results at: [cyan]{result_path}[/]\n\n"
                f"  [dim]Load with:[/] AutoModelForCausalLM.from_pretrained('{result_path}')",
                border_style="green",
                title="[bold green]REBIRTH COMPLETE (remote)[/]",
            )
        )
    else:
        console.print("[red]Remote abliteration failed. Check logs above.[/]")
        raise SystemExit(1)


def _cmd_remote_run(args):
    runner = _make_remote_runner(args)
    result_path = runner.run_config(
        local_config_path=args.config,
        local_output_dir=args.output_dir,
        preset=args.preset,
    )
    if result_path:
        console.print(f"\n[bold green]Remote run complete.[/] Results at: [cyan]{result_path}[/]")
    else:
        console.print("[red]Remote run failed. Check logs above.[/]")
        raise SystemExit(1)


def _cmd_remote_tourney(args):
    from rich.panel import Panel

    runner = _make_remote_runner(args)
    result_path = runner.run_tourney(
        model=args.model,
        local_output_dir=args.output_dir,
        device=args.device,
        dtype=args.dtype,
        quantization=args.quantization,
        methods=args.methods,
        hub_org=args.hub_org,
        hub_repo=args.hub_repo,
        dataset=args.dataset,
    )
    if result_path:
        console.print(
            Panel(
                f"[bold green]Remote tournament complete![/]\n\n"
                f"  Results at: [cyan]{result_path}[/]",
                border_style="green",
                title="[bold green]TOURNAMENT COMPLETE (remote)[/]",
            )
        )
    else:
        console.print("[red]Remote tournament failed. Check logs above.[/]")
        raise SystemExit(1)


if __name__ == "__main__":
    main()