#!/usr/bin/env python3
"""rho-unlock: Two-axis behavioral diagnostic + contrastive decoding unlock.

Diagnose hidden model capability and unlock it via contrastive decoding.

Subcommands:
    diagnose  — Run rho-eval (Axis 1) + expression gap (Axis 2) to classify
                each behavior into four quadrants (HEALTHY/UNLOCK/RETRAIN/BOTH)
    unlock    — Apply contrastive decoding to behaviors flagged as UNLOCK,
                re-measure gap to show improvement

Usage:
    rho-unlock diagnose Qwen/Qwen2.5-7B
    rho-unlock diagnose Qwen/Qwen2.5-7B --behaviors bias,sycophancy
    rho-unlock diagnose Qwen/Qwen2.5-7B --output diagnosis.json

    rho-unlock unlock Qwen/Qwen2.5-7B
    rho-unlock unlock Qwen/Qwen2.5-7B --amateur Qwen/Qwen2.5-0.5B --alpha 0.5
"""

import argparse
import json
import sys
import time
from pathlib import Path

import numpy as np


def _json_serializer(obj):
    """Handle numpy/torch types for JSON serialization."""
    if isinstance(obj, (np.integer,)):
        return int(obj)
    if isinstance(obj, (np.floating,)):
        return float(obj)
    if isinstance(obj, np.ndarray):
        return obj.tolist()
    if isinstance(obj, bool):
        return obj
    try:
        import torch
        if isinstance(obj, torch.device):
            return str(obj)
        if isinstance(obj, torch.Tensor):
            return obj.item() if obj.numel() == 1 else obj.tolist()
    except ImportError:
        pass
    return str(obj)


def _load_model_mlx(model_id):
    """Load model + tokenizer via MLX (Apple Silicon)."""
    print(f"\033[1mLoading {model_id} (MLX)...\033[0m", flush=True)
    t0 = time.time()

    from mlx_lm import load as mlx_load

    model, tokenizer = mlx_load(model_id)

    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    print(f"  Loaded in {time.time() - t0:.1f}s", flush=True)
    return model, tokenizer


def _load_model_torch(model_id, device, dtype_str="float32"):
    """Load model + tokenizer via PyTorch."""
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    dtype = getattr(torch, dtype_str, torch.float32)
    print(f"\033[1mLoading {model_id} (PyTorch, {device})...\033[0m", flush=True)
    t0 = time.time()

    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    try:
        model = AutoModelForCausalLM.from_pretrained(
            model_id, dtype=dtype, trust_remote_code=True,
            low_cpu_mem_usage=True,
        ).to(device)
    except TypeError:
        model = AutoModelForCausalLM.from_pretrained(
            model_id, torch_dtype=dtype, trust_remote_code=True,
            low_cpu_mem_usage=True,
        ).to(device)
    model.eval()

    n_params = sum(p.numel() for p in model.parameters())
    print(f"  Loaded in {time.time() - t0:.1f}s ({n_params / 1e9:.2f}B params)", flush=True)
    return model, tokenizer


def _load_model(model_id, device=None, dtype_str="float32"):
    """Load model, auto-detecting MLX vs PyTorch.

    On macOS, tries MLX first. Falls back to PyTorch.

    Returns:
        (model, tokenizer, device_str) — device_str is "mlx" for MLX models,
        or the actual torch device string for PyTorch models.
    """
    import platform

    # Try MLX on macOS
    if platform.system() == "Darwin":
        try:
            model, tokenizer = _load_model_mlx(model_id)
            return model, tokenizer, "mlx"
        except (ImportError, Exception) as e:
            print(f"  MLX unavailable ({e}), falling back to PyTorch", flush=True)

    # Resolve device for PyTorch
    if device is None:
        try:
            from rho_eval.utils import get_device
            device = str(get_device())
        except ImportError:
            device = "cpu"

    model, tokenizer = _load_model_torch(model_id, device, dtype_str)
    return model, tokenizer, device


# ── Diagnose subcommand ──────────────────────────────────────────────

def _run_diagnose(args):
    """Run two-axis diagnostic: rho-eval + expression gap."""
    from rho_eval.audit import audit
    from rho_eval.unlock.expression_gap import (
        measure_all_gaps, MC_BEHAVIORS, BENCHMARK_ONLY, BEHAVIOR_N_CHOICES,
    )
    from rho_eval.unlock.diagnosis import diagnose, format_diagnosis_table

    # Parse behaviors
    if args.behaviors == "all":
        from rho_eval.behaviors import list_behaviors
        behavior_names = list_behaviors() + sorted(BENCHMARK_ONLY)
    else:
        behavior_names = [b.strip() for b in args.behaviors.split(",")]
    rho_behaviors = [b for b in behavior_names if b not in BENCHMARK_ONLY]
    benchmark_names = [b for b in behavior_names if b in BENCHMARK_ONLY]

    # Load model
    model, tokenizer, resolved_device = _load_model(args.model, args.device)

    # ── Axis 1: rho-eval (only for rho-eval behaviors) ────────────────
    rho_scores = {}
    if rho_behaviors:
        print(f"\n\033[1m━━━ Axis 1: Behavioral Discrimination (ρ scores) ━━━\033[0m\n", flush=True)

        report = audit(
            model=model,
            tokenizer=tokenizer,
            behaviors=rho_behaviors,
            n=args.n_probes,
            seed=args.seed,
            device=resolved_device,
        )
        report.model = args.model

        for name, result in sorted(report.behaviors.items()):
            status_colors = {"PASS": "\033[92m", "WARN": "\033[93m", "FAIL": "\033[91m"}
            reset = "\033[0m"
            c = status_colors.get(result.status, "")
            print(
                f"  {name:<14s}  "
                f"ρ={result.rho:+.4f}  "
                f"{result.positive_count:>3d}/{result.total:<4d}  "
                f"[{c}{result.status}{reset}]  "
                f"{result.elapsed:.1f}s",
                flush=True,
            )

        rho_scores = {name: result.rho for name, result in report.behaviors.items()}

    # ── Axis 2: expression gap ────────────────────────────────────────
    mc_behaviors = [b for b in behavior_names if b in MC_BEHAVIORS]

    print(f"\n\033[1m━━━ Axis 2: Expression Gap (logit vs generation) ━━━\033[0m\n", flush=True)

    if mc_behaviors:
        gap_results = measure_all_gaps(
            model, tokenizer,
            behaviors=mc_behaviors,
            n_probes=args.n_probes,
            seed=args.seed,
            device=resolved_device,
        )
    else:
        gap_results = {}
        print("  No MC-based behaviors selected — expression gap N/A", flush=True)

    # ── Combine into diagnosis ────────────────────────────────────────
    print(f"\n\033[1m━━━ Diagnosis ━━━\033[0m\n", flush=True)
    expression_gaps = {}
    parse_rates = {}
    logit_accs = {}
    gen_accs = {}

    # Separate benchmark scores from behavioral ρ scores
    benchmark_scores = {}
    benchmark_n_choices = {}

    for name in behavior_names:
        if name in gap_results and gap_results[name].supports_gap:
            expression_gaps[name] = gap_results[name].gap
            parse_rates[name] = gap_results[name].parse_rate
            logit_accs[name] = gap_results[name].logit_accuracy
            gen_accs[name] = gap_results[name].gen_accuracy

            # Route benchmark logit accuracy to benchmark_scores
            if name in BENCHMARK_ONLY:
                benchmark_scores[name] = gap_results[name].logit_accuracy
                benchmark_n_choices[name] = BEHAVIOR_N_CHOICES.get(name, 4)
        else:
            expression_gaps[name] = None

    diagnoses = diagnose(
        rho_scores=rho_scores,
        expression_gaps=expression_gaps,
        parse_rates=parse_rates,
        logit_accuracies=logit_accs,
        gen_accuracies=gen_accs,
        benchmark_scores=benchmark_scores,
        benchmark_n_choices=benchmark_n_choices,
        rho_threshold=args.rho_threshold,
        above_chance_margin=args.above_chance_margin,
        gap_threshold=args.gap_threshold,
    )

    print(format_diagnosis_table(diagnoses))

    # Show unlock candidates
    unlock_candidates = [
        name for name, d in diagnoses.items()
        if d.quadrant.value == "UNLOCK"
    ]
    if unlock_candidates:
        print(f"\n  \033[93m→ Unlock candidates: {', '.join(unlock_candidates)}\033[0m")
        print(f"  Run: rho-unlock unlock {args.model}", flush=True)

    # ── Save output ───────────────────────────────────────────────────
    if args.output:
        output = {
            "model": args.model,
            "rho_scores": rho_scores,
            "benchmark_scores": benchmark_scores,
            "expression_gaps": {k: v for k, v in expression_gaps.items()},
            "parse_rates": {k: v for k, v in parse_rates.items()},
            "logit_accuracies": {k: v for k, v in logit_accs.items()},
            "gen_accuracies": {k: v for k, v in gen_accs.items()},
            "diagnoses": {k: v.to_dict() for k, v in diagnoses.items()},
            "thresholds": {
                "rho": args.rho_threshold,
                "above_chance_margin": args.above_chance_margin,
                "gap": args.gap_threshold,
            },
        }
        out_path = Path(args.output)
        out_path.parent.mkdir(parents=True, exist_ok=True)
        with open(out_path, "w") as f:
            json.dump(output, f, indent=2, default=_json_serializer)
        print(f"\n  Saved: {out_path}", flush=True)


# ── Unlock subcommand ────────────────────────────────────────────────

def _run_unlock(args):
    """Apply contrastive decoding to unlock hidden capability."""
    from rho_eval.unlock.contrastive import (
        detect_amateur, contrastive_logit_classify, contrastive_generate,
        get_answer_token_ids,
    )
    from rho_eval.unlock.expression_gap import (
        measure_all_gaps, MC_BEHAVIORS, BEHAVIOR_N_CHOICES, BENCHMARK_ONLY,
    )
    from rho_eval.unlock.diagnosis import diagnose, format_diagnosis_table

    # Parse behaviors
    if args.behaviors == "all":
        behavior_names = sorted(MC_BEHAVIORS)
    else:
        behavior_names = [b.strip() for b in args.behaviors.split(",")]
        behavior_names = [b for b in behavior_names if b in MC_BEHAVIORS]

    if not behavior_names:
        print("Error: No MC-based behaviors selected. "
              "Supported: bias, sycophancy, mmlu.", file=sys.stderr)
        sys.exit(1)

    # Detect or validate amateur model
    amateur_id = args.amateur or detect_amateur(args.model)
    if amateur_id is None:
        print(
            f"Error: No amateur model found for '{args.model}'.\n"
            f"  Specify one with: --amateur <model_id>",
            file=sys.stderr,
        )
        sys.exit(1)

    alpha = args.alpha

    # Load both models
    print(f"\n\033[1m━━━ Loading Models ━━━\033[0m\n", flush=True)
    print(f"  Expert:  {args.model}")
    print(f"  Amateur: {amateur_id} {'(auto-detected)' if not args.amateur else ''}")
    print(f"  α = {alpha}\n")

    expert_model, expert_tokenizer, resolved_device = _load_model(args.model, args.device)
    amateur_model, amateur_tokenizer, _ = _load_model(amateur_id, args.device)

    # ── Baseline: measure expression gap without CD ───────────────────
    print(f"\n\033[1m━━━ Baseline (no contrastive decoding) ━━━\033[0m\n", flush=True)

    baseline_gaps = measure_all_gaps(
        expert_model, expert_tokenizer,
        behaviors=behavior_names,
        n_probes=args.n_probes,
        seed=args.seed,
        device=resolved_device,
    )

    # ── Check for meaningful gaps before applying CD ─────────────────
    max_gap = 0.0
    small_gap_behaviors = []
    for beh_name in behavior_names:
        if beh_name in baseline_gaps and baseline_gaps[beh_name].supports_gap:
            gap = baseline_gaps[beh_name].gap
            max_gap = max(max_gap, gap)
            if gap < 0.02:
                small_gap_behaviors.append((beh_name, gap))

    if max_gap < 0.02:
        print(f"\n\033[33m⚠  WARNING: No significant expression gaps detected.\033[0m")
        print(f"  All gaps are below 2%. Contrastive decoding is unlikely to help")
        print(f"  and may hurt accuracy (observed -4.6% on Qwen 7B with no gap).")
        for beh_name, gap in small_gap_behaviors:
            print(f"    {beh_name}: gap = {gap:+.1%}")
        if not getattr(args, 'force', False):
            print(f"\n  To proceed anyway, use --force")
            print(f"  Tip: Run 'rho-unlock diagnose {args.model}' to check model status.\n")
            sys.exit(0)
        print(f"  --force specified, proceeding anyway...\n")
    elif small_gap_behaviors:
        print(f"\n\033[33m⚠  Note: Some behaviors have small gaps (<2%):\033[0m")
        for beh_name, gap in small_gap_behaviors:
            print(f"    {beh_name}: gap = {gap:+.1%} — CD may not help here")
        print()

    # ── Apply contrastive decoding ────────────────────────────────────
    print(f"\n\033[1m━━━ Contrastive Decoding (α={alpha}) ━━━\033[0m\n", flush=True)

    cd_results = {}

    for beh_name in behavior_names:
        if beh_name not in baseline_gaps or not baseline_gaps[beh_name].supports_gap:
            continue

        baseline = baseline_gaps[beh_name]
        n_choices = BEHAVIOR_N_CHOICES.get(beh_name, 3)
        answer_ids = get_answer_token_ids(expert_tokenizer, n_choices=n_choices)

        print(f"  [cd] Applying CD to {beh_name} "
              f"({baseline.n_probes} probes, {n_choices} choices)...", flush=True)

        cd_correct = 0
        t0 = time.time()

        if beh_name in BENCHMARK_ONLY:
            # Benchmark: use expression_gap's loader + formatter
            from rho_eval.unlock.expression_gap import (
                _load_mmlu, _load_truthfulqa, _load_arc, _load_hellaswag,
                _format_mmlu_prompt, _format_truthfulqa_prompt,
                _format_arc_prompt, _format_hellaswag_prompt,
            )
            loader = {
                "mmlu": _load_mmlu, "truthfulqa": _load_truthfulqa,
                "arc": _load_arc, "hellaswag": _load_hellaswag,
            }[beh_name]
            formatter = {
                "mmlu": _format_mmlu_prompt,
                "truthfulqa": _format_truthfulqa_prompt,
                "arc": _format_arc_prompt,
                "hellaswag": _format_hellaswag_prompt,
            }[beh_name]

            full_probes = loader(
                n=args.n_probes or baseline.n_probes,
                seed=args.seed,
            )
            letters = "ABCD"[:n_choices]

            for i, probe in enumerate(full_probes):
                prompt = formatter(expert_tokenizer, probe)
                correct_letter = letters[probe["answer_idx"]]

                best, cd_scores = contrastive_logit_classify(
                    expert_model, expert_tokenizer,
                    amateur_model, amateur_tokenizer,
                    prompt, answer_ids, alpha=alpha,
                    device=resolved_device,
                )

                if best == correct_letter:
                    cd_correct += 1

                if (i + 1) % 50 == 0:
                    print(f"    [{i+1}/{len(full_probes)}] "
                          f"cd_acc={cd_correct/(i+1):.1%}", flush=True)

            n = len(full_probes)
        else:
            # rho-eval behavior: reload via behavior plugin
            from rho_eval.behaviors import get_behavior
            behavior = get_behavior(beh_name)
            full_probes = behavior.load_probes(n=args.n_probes, seed=args.seed)

            for probe in full_probes:
                prompt = probe["text"]

                best, cd_scores = contrastive_logit_classify(
                    expert_model, expert_tokenizer,
                    amateur_model, amateur_tokenizer,
                    prompt, answer_ids, alpha=alpha,
                    device=resolved_device,
                )

                correct = probe.get("correct_answer",
                                    probe.get("truthful_answer", ""))
                if best == correct:
                    cd_correct += 1

            n = len(full_probes)

        cd_acc = cd_correct / n if n > 0 else 0.0
        elapsed = time.time() - t0

        cd_results[beh_name] = {
            "cd_accuracy": cd_acc,
            "baseline_gen_accuracy": baseline.gen_accuracy,
            "baseline_logit_accuracy": baseline.logit_accuracy,
            "improvement": cd_acc - (baseline.gen_accuracy or 0.0),
            "n_probes": n,
            "elapsed": elapsed,
        }

        print(
            f"  [cd] {beh_name}: "
            f"baseline={baseline.gen_accuracy:.1%} → CD={cd_acc:.1%} "
            f"(Δ={cd_acc - (baseline.gen_accuracy or 0.0):+.1%})",
            flush=True,
        )

    # ── Summary table ─────────────────────────────────────────────────
    print(f"\n\033[1m━━━ Results ━━━\033[0m\n", flush=True)

    header = f"  {'Behavior':<14s} {'Before':>8s} {'After':>8s} {'Δ':>8s}"
    sep = "  " + "─" * (len(header) - 2)
    print(header)
    print(sep)

    for name, res in sorted(cd_results.items()):
        before = res["baseline_gen_accuracy"] or 0.0
        after = res["cd_accuracy"]
        delta = res["improvement"]
        print(f"  {name:<14s} {before:>7.1%} {after:>7.1%} {delta:>+7.1%}")

    # ── Save output ───────────────────────────────────────────────────
    if args.output:
        output = {
            "model": args.model,
            "amateur": amateur_id,
            "alpha": alpha,
            "baseline": {
                k: v.to_dict() for k, v in baseline_gaps.items()
                if v.supports_gap
            },
            "contrastive_decoding": cd_results,
        }
        out_path = Path(args.output)
        out_path.parent.mkdir(parents=True, exist_ok=True)
        with open(out_path, "w") as f:
            json.dump(output, f, indent=2, default=_json_serializer)
        print(f"\n  Saved: {out_path}", flush=True)


# ── CLI entry point ──────────────────────────────────────────────────

def main(argv=None):
    parser = argparse.ArgumentParser(
        prog="rho-unlock",
        description="Two-axis behavioral diagnostic + contrastive decoding unlock for LLMs",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Diagnose: measure ρ + expression gap, classify into quadrants
  rho-unlock diagnose Qwen/Qwen2.5-7B

  # Unlock: apply contrastive decoding to rescue hidden capability
  rho-unlock unlock Qwen/Qwen2.5-7B --alpha 0.5

  # Specify amateur model explicitly
  rho-unlock unlock Qwen/Qwen2.5-7B --amateur Qwen/Qwen2.5-0.5B

  # Save results to JSON
  rho-unlock diagnose Qwen/Qwen2.5-7B --output diagnosis.json
""",
    )

    subparsers = parser.add_subparsers(dest="command", help="Subcommand")

    # ── diagnose ──────────────────────────────────────────────────────
    diag_parser = subparsers.add_parser(
        "diagnose",
        help="Run two-axis diagnostic (ρ + expression gap)",
        description="Measure behavioral discrimination (ρ) and expression gap per behavior.",
    )
    diag_parser.add_argument("model", help="HuggingFace model ID or local path")
    diag_parser.add_argument(
        "--behaviors", default="bias,sycophancy,factual,toxicity",
        help="Comma-separated behavior names, or 'all' (default: bias,sycophancy,factual,toxicity)",
    )
    diag_parser.add_argument("--n-probes", type=int, default=None, help="Number of probes per behavior")
    diag_parser.add_argument("--seed", type=int, default=42, help="Random seed (default: 42)")
    diag_parser.add_argument("--device", default=None, help="Torch device (auto-detected if omitted)")
    diag_parser.add_argument("--output", "-o", help="Save JSON report to this path")
    diag_parser.add_argument(
        "--rho-threshold", type=float, default=0.3,
        help="ρ above this = 'knows' for behavioral dimensions (default: 0.3)",
    )
    diag_parser.add_argument(
        "--above-chance-margin", type=float, default=0.15,
        help="Margin above chance level for benchmark 'knows' (default: 0.15, "
             "e.g. 4-choice → 25%% + 15%% = 40%%)",
    )
    diag_parser.add_argument(
        "--gap-threshold", type=float, default=0.05,
        help="Gap above this = 'can't express' (default: 0.05 = 5%%)",
    )
    diag_parser.set_defaults(func=_run_diagnose)

    # ── unlock ────────────────────────────────────────────────────────
    unlock_parser = subparsers.add_parser(
        "unlock",
        help="Apply contrastive decoding to unlock hidden capability",
        description="Subtract amateur model logits to rescue expression bottleneck.",
    )
    unlock_parser.add_argument("model", help="Expert model (HuggingFace ID or local path)")
    unlock_parser.add_argument(
        "--amateur", default=None,
        help="Amateur model ID (auto-detected if omitted)",
    )
    unlock_parser.add_argument(
        "--alpha", type=float, default=0.5,
        help="Contrastive strength (default: 0.5)",
    )
    unlock_parser.add_argument(
        "--behaviors", default="bias,sycophancy",
        help="Comma-separated MC behaviors to unlock (default: bias,sycophancy)",
    )
    unlock_parser.add_argument("--n-probes", type=int, default=None, help="Number of probes per behavior")
    unlock_parser.add_argument("--seed", type=int, default=42, help="Random seed (default: 42)")
    unlock_parser.add_argument("--device", default=None, help="Torch device (auto-detected if omitted)")
    unlock_parser.add_argument("--output", "-o", help="Save JSON results to this path")
    unlock_parser.add_argument("--force", action="store_true",
                               help="Run CD even if no significant expression gap detected")
    unlock_parser.set_defaults(func=_run_unlock)

    # Parse
    args = parser.parse_args(argv)
    if not args.command:
        parser.print_help()
        sys.exit(1)

    args.func(args)


if __name__ == "__main__":
    main()
