"""
Optimize weights command - unified optimization for any trait/task with any evaluator.

This command runs an Optuna-based optimization loop to find optimal weight
modification parameters. Works with:

- Any trait (synthetic pairs): refusal, sycophancy, verbosity, personality
- Any task (lm-eval): hellaswag, arc_easy, gsm8k, etc.
- Any evaluator: refusal detection, task accuracy, semantic similarity, LLM judge

Pipeline per trial:
1. Generate steering vector (from trait or task, or use cached)
2. Load fresh model copy
3. Apply weight modification with trial parameters
4. Evaluate using chosen evaluator
5. Return score to Optuna
6. Repeat until target reached or max trials
7. Apply best parameters and save optimized model
"""

import json
import os
import re
import tempfile
import time
from dataclasses import dataclass
from typing import Any, Callable

import torch

from wisent.core.errors import UnknownTypeError, InsufficientDataError

from wisent.core.models.wisent_model import WisentModel
from wisent.core.evaluators.steering_evaluators import (
    SteeringEvaluatorFactory,
    EvaluatorConfig,
    RefusalEvaluator,
    TaskEvaluator,
    PersonalizationEvaluator as SharedPersonalizationEvaluator,
)
from wisent.core.opti.core.atoms import HPOConfig
from wisent.core.opti.methods.opti_weights import WeightsOptimizer, WeightsOptimizerConfig
from wisent.core.utils.device import resolve_default_device, preferred_dtype


@dataclass
class OptimizationResult:
    """Final result from optimization."""
    best_params: dict
    best_score: float
    target_achieved: bool
    total_time: float
    total_trials: int


def execute_optimize_weights(args):
    """
    Execute the optimize-weights command.

    Runs Optuna optimization to find optimal weight modification parameters
    for any trait or task.
    """
    start_time = time.time()

    print(f"\n{'='*80}")
    print("WEIGHT MODIFICATION OPTIMIZATION")
    print(f"{'='*80}")
    print(f"   Model: {args.model}")
    if args.trait:
        print(f"   Mode: Trait-based (synthetic pairs)")
        print(f"   Trait: {args.trait[:60]}...")
    elif args.task:
        print(f"   Mode: Task-based (lm-eval)")
        print(f"   Task: {args.task}")
    else:
        print(f"   Mode: Pre-computed vectors")
        print(f"   Vectors: {args.steering_vectors}")
    
    # Determine evaluator type from task for display
    task_lower = (args.task or "").lower()
    if task_lower == "refusal":
        evaluator_display = "refusal"
    elif task_lower == "personalization":
        evaluator_display = f"personalization ({args.trait})"
    elif task_lower == "custom":
        evaluator_display = f"custom ({args.custom_evaluator})"
    elif "," in (args.task or ""):
        evaluator_display = "pooled (multi-benchmark)"
    else:
        evaluator_display = f"task ({args.task})"
    print(f"   Evaluator: {evaluator_display}")
    print(f"   Target: {args.target_metric} = {args.target_value}")
    print(f"   Trials: {args.trials}")
    print(f"   Output: {args.output_dir}")
    print(f"{'='*80}\n")

    # Parse search space ranges
    strength_range = tuple(float(x) for x in args.strength_range.split(","))
    max_weight_range = tuple(float(x) for x in args.max_weight_range.split(","))
    min_weight_range = tuple(float(x) for x in args.min_weight_range.split(","))
    position_range = tuple(float(x) for x in args.position_range.split(","))
    num_pairs = args.num_pairs

    # Determine optimization direction
    direction = args.direction
    if direction == "auto":
        if args.target_metric in ["refusal_rate", "kl_divergence"]:
            direction = "minimize"
        else:
            direction = "maximize"

    print(f"Search space:")
    print(f"   Strength: {strength_range}")
    print(f"   Max weight: {max_weight_range}")
    print(f"   Min weight: {min_weight_range}")
    print(f"   Position: {position_range}")
    print(f"   Num pairs: {num_pairs} (fixed)")
    print(f"   Direction: {direction}")
    print()

    # Initialize components
    device = args.device if args.device else resolve_default_device()
    print(f"Device: {device}")

    # IMPORTANT: Generate steering vectors FIRST, before loading the main model.
    # The steering vector generation pipeline loads its own model internally.
    # If we load our model first, then the pipeline loads another model, we end up
    # with two large models competing for GPU memory, causing CPU offloading.
    # By generating vectors first, that model is unloaded before we load ours.
    print(f"\nGenerating steering vectors with {num_pairs} pairs...")
    print("   (This will load the model temporarily for activation collection)\n")
    steering_vectors, positive_examples, negative_examples = _generate_steering_vectors(args, num_pairs, num_layers=None)
    print(f"\n   Steering vectors generated for {len(steering_vectors)} layers")
    print(f"   Got {len(positive_examples)} positive and {len(negative_examples)} negative examples")

    # Force garbage collection to ensure the model used for activation collection is freed
    import gc
    gc.collect()
    if device == "cuda":
        import torch
        torch.cuda.empty_cache()
    print("   Memory cleaned up\n")

    # NOW load the main model for optimization (GPU should be free now)
    print("Loading base model for optimization...")
    wisent_model = WisentModel(args.model, device=device)
    base_model = wisent_model.hf_model  # Get underlying HF model for weight modification
    tokenizer = wisent_model.tokenizer
    num_layers = wisent_model.num_layers

    print(f"   Model loaded: {num_layers} layers\n")

    # Store base model state for restoration
    base_state_dict = {k: v.clone() for k, v in base_model.state_dict().items()}

    # Initialize evaluator (pass wisent_model for personalization baseline generation, and contrastive pairs)
    evaluator = _create_evaluator(args, args.model, wisent_model=wisent_model,
                                   positive_examples=positive_examples, negative_examples=negative_examples)

    # Create optimizer config
    optimizer_config = WeightsOptimizerConfig(
        strength_range=strength_range,
        max_weight_range=max_weight_range,
        min_weight_range=min_weight_range,
        position_range=position_range,
        method=args.method,
        components=args.components,
        norm_preserve=args.norm_preserve,
        optimize_direction_index=args.optimize_direction_index,
        target_metric=args.target_metric,
        target_value=args.target_value if args.early_stop else None,
    )

    # Create the optimizer
    optimizer = WeightsOptimizer(
        model=base_model,
        base_state_dict=base_state_dict,
        steering_vectors=steering_vectors,
        evaluator=evaluator,
        tokenizer=tokenizer,
        config=optimizer_config,
        num_layers=num_layers,
    )

    # Create HPO config
    hpo_config = HPOConfig(
        n_trials=args.trials,
        direction=direction,
        sampler="tpe",
        pruner=None,
        seed=42,
    )

    # Run optimization
    print(f"\nStarting optimization ({args.trials} trials)...\n")
    result = optimizer.optimize(hpo_config)

    best_params = result.best_params
    best_value = result.best_value

    print(f"\n{'='*80}")
    print("OPTIMIZATION COMPLETE")
    print(f"{'='*80}")
    print(f"\nBest parameters:")
    for k, v in best_params.items():
        print(f"   {k}: {v:.4f}" if isinstance(v, float) else f"   {k}: {v}")
    print(f"\nBest {args.target_metric}: {best_value:.4f}")

    # Check if target achieved
    if direction == "maximize":
        target_achieved = best_value >= args.target_value
    else:
        target_achieved = best_value <= args.target_value
    print(f"\nTarget {args.target_value} achieved: {'YES' if target_achieved else 'NO'}")

    # Apply best parameters and save final model
    print(f"\n{'='*80}")
    print("SAVING OPTIMIZED MODEL")
    print(f"{'='*80}")

    optimizer.apply_best_params(best_params)

    # Save model
    os.makedirs(args.output_dir, exist_ok=True)
    print(f"\nSaving optimized model to {args.output_dir}...")
    base_model.save_pretrained(args.output_dir)
    tokenizer.save_pretrained(args.output_dir)

    # Save optimization metadata
    metadata = {
        "model": args.model,
        "task": args.task,
        "trait": getattr(args, 'trait', None),
        "evaluator_type": evaluator_display,
        "target_metric": args.target_metric,
        "target_value": args.target_value,
        "best_params": best_params,
        "best_score": best_value,
        "target_achieved": target_achieved,
        "total_trials": len(result.study.trials),
        "direction": direction,
    }

    with open(os.path.join(args.output_dir, "optimization_metadata.json"), "w") as f:
        json.dump(metadata, f, indent=2)

    print(f"   Model saved")
    print(f"   Metadata saved to optimization_metadata.json")

    # Save all trials if requested
    if args.save_trials:
        trials_data = [
            {
                "trial": t.number,
                "params": t.params,
                "score": t.value,
            }
            for t in result.study.trials
        ]
        with open(args.save_trials, "w") as f:
            json.dump(trials_data, f, indent=2)
        print(f"   Trials saved to {args.save_trials}")

    # Push to hub if requested
    if args.push_to_hub:
        if not args.repo_id:
            print("\n   ERROR: --repo-id required for --push-to-hub")
        else:
            print(f"\n   Pushing to HuggingFace Hub: {args.repo_id}...")
            base_model.push_to_hub(args.repo_id)
            tokenizer.push_to_hub(args.repo_id)
            print(f"   Pushed successfully")

    # Show/save before/after comparisons if requested
    save_comparisons_path = getattr(args, 'save_comparisons', None)
    show_comparisons_count = getattr(args, 'show_comparisons', 0)
    if show_comparisons_count > 0 or save_comparisons_path:
        _show_response_comparisons(
            base_model=base_model,
            base_state_dict=base_state_dict,
            steering_vectors=steering_vectors,
            best_params=best_params,
            num_layers=num_layers,
            model_name=args.model,
            args=args,
            optimizer_config=optimizer_config,
            num_comparisons=show_comparisons_count if show_comparisons_count > 0 else None,
            save_path=save_comparisons_path,
        )

    total_time = time.time() - start_time

    print(f"\n{'='*80}")
    print(f"Total optimization time: {total_time:.1f}s ({total_time/60:.1f} min)")
    print(f"{'='*80}\n")

    return OptimizationResult(
        best_params=best_params,
        best_score=best_value,
        target_achieved=target_achieved,
        total_time=total_time,
        total_trials=len(result.study.trials),
    )


def _generate_steering_vectors(args, num_pairs: int, num_layers: int = None) -> tuple[dict[int, torch.Tensor], list[str], list[str]]:
    """Generate steering vectors from trait or task.

    Args:
        args: Command line arguments
        num_pairs: Number of contrastive pairs to generate
        num_layers: Number of model layers (optional, not used - kept for backwards compatibility)

    Returns:
        Tuple of (steering_vectors, positive_examples, negative_examples)
    """
    from argparse import Namespace

    if args.steering_vectors:
        with open(args.steering_vectors, "r") as f:
            data = json.load(f)
        vectors = {
            int(layer) - 1: torch.tensor(vec)
            for layer, vec in data["steering_vectors"].items()
        }
        # Try to load pairs from the same file if available
        positive_examples = data.get("positive_examples", [])
        negative_examples = data.get("negative_examples", [])
        return vectors, positive_examples, negative_examples

    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
        temp_output = f.name

    # Also create a temp file for pairs
    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
        temp_pairs = f.name

    try:
        # Determine task type and route accordingly
        task_lower = (args.task or "").lower().strip()
        
        # Special cases with dedicated logic
        SPECIAL_CASES = {'refusal', 'personalization', 'humanizer', 'custom'}
        
        # Check if task is a skill or risk - expand to benchmarks
        from wisent.core.task_selector import TaskSelector
        selector = TaskSelector()
        available_skills = [s.lower() for s in selector.get_available_skills()]
        available_risks = [r.lower() for r in selector.get_available_risks()]
        
        if task_lower in available_skills:
            matching_tasks = selector.find_tasks_by_tags(skills=[task_lower])
            if matching_tasks:
                args.task = ",".join(matching_tasks)
                print(f"   Expanded skill '{task_lower}' to {len(matching_tasks)} benchmarks")
        elif task_lower in available_risks:
            matching_tasks = selector.find_tasks_by_tags(risks=[task_lower])
            if matching_tasks:
                args.task = ",".join(matching_tasks)
                print(f"   Expanded risk '{task_lower}' to {len(matching_tasks)} benchmarks")
        
        # Route based on task type
        if task_lower in SPECIAL_CASES:
            # Special cases: refusal, personalization, humanizer, custom
            return _generate_vectors_special_case(args, task_lower, num_pairs, temp_output, temp_pairs)
        
        elif "," in (args.task or ""):
            # Multiple benchmarks (comma-separated or expanded): unified training
            return _generate_vectors_multi_benchmark(args, temp_output)
        
        elif _is_valid_benchmark(args.task):
            # Single benchmark: task-based generation
            return _generate_vectors_single_benchmark(args, num_pairs, temp_output, temp_pairs)
        
        else:
            # Synthetic: free-form description
            return _generate_vectors_synthetic(args, num_pairs, temp_output, temp_pairs)

    finally:
        if os.path.exists(temp_output):
            os.unlink(temp_output)
        if os.path.exists(temp_pairs):
            os.unlink(temp_pairs)


def _is_valid_benchmark(task: str) -> bool:
    """Check if task is a valid lm-eval benchmark name."""
    if not task:
        return False
    from wisent.core.task_selector import TaskSelector
    selector = TaskSelector()
    return task.lower() in [t.lower() for t in selector.tasks.keys()]


def _generate_vectors_special_case(args, task_lower: str, num_pairs: int, temp_output: str, temp_pairs: str):
    """Generate vectors for special cases: refusal, personalization, humanizer, custom.
    
    These use synthetic pair generation with predefined trait descriptions.
    """
    # Map special cases to trait descriptions
    TRAIT_DESCRIPTIONS = {
        'refusal': 'refusing harmful or dangerous requests while being helpful for legitimate ones',
        'personalization': args.trait if getattr(args, 'trait', None) else 'adapting responses to user preferences',
        'humanizer': 'generating natural, human-like responses that avoid AI-like patterns',
        'custom': args.trait if getattr(args, 'trait', None) else None,
    }
    
    trait = TRAIT_DESCRIPTIONS.get(task_lower)
    if trait is None:
        raise ValueError(f"Special case '{task_lower}' requires --trait to be specified")
    
    # Use synthetic generation with the trait description
    original_task = args.task
    args.task = trait
    try:
        return _generate_vectors_synthetic(args, num_pairs, temp_output, temp_pairs)
    finally:
        args.task = original_task


def _generate_vectors_multi_benchmark(args, temp_output: str):
    """Generate unified steering vector from multiple benchmarks."""
    from argparse import Namespace
    from wisent.core.cli.train_unified_goodness import execute_train_unified_goodness
    
    temp_output_pt = temp_output.replace('.json', '.pt')
    layers_arg = args.layers if hasattr(args, 'layers') else None
    if layers_arg == 'all' or layers_arg is None:
        layers_arg = None
    
    vector_args = Namespace(
        task=args.task,
        exclude_benchmarks=None,
        max_benchmarks=getattr(args, 'max_benchmarks', None),
        cap_pairs_per_benchmark=getattr(args, 'cap_pairs_per_benchmark', None),
        train_ratio=getattr(args, 'train_ratio', 0.8),
        seed=getattr(args, 'seed', 42),
        model=args.model,
        device=args.device,
        layer=None,
        layers=layers_arg,
        token_aggregation=args.token_aggregation if hasattr(args, 'token_aggregation') else 'continuation',
        prompt_strategy='chat_template',
        method='caa',
        normalize=True,
        no_normalize=False,
        skip_evaluation=True,
        evaluate_steering_scales="0.0,1.0",
        save_pairs=None,
        save_report=None,
        output=temp_output_pt,
        verbose=False,
        timing=False,
    )

    execute_train_unified_goodness(vector_args)

    checkpoint = torch.load(temp_output_pt, map_location='cpu', weights_only=False)

    if 'all_layer_vectors' in checkpoint:
        raw_vectors = checkpoint['all_layer_vectors']
    elif 'steering_vector' in checkpoint and 'layer_index' in checkpoint:
        raw_vectors = {checkpoint['layer_index']: checkpoint['steering_vector']}
    else:
        raw_vectors = {
            k: v for k, v in checkpoint.items()
            if isinstance(k, (int, str)) and str(k).isdigit()
        }

    vectors = {}
    for layer, vec in raw_vectors.items():
        layer_idx = int(layer) if isinstance(layer, str) else layer
        vectors[layer_idx] = vec if isinstance(vec, torch.Tensor) else torch.tensor(vec)

    args._pooled_eval_pairs = checkpoint.get('eval_pairs', [])
    args._benchmarks_used = checkpoint.get('benchmarks_used', [])

    if os.path.exists(temp_output_pt):
        os.unlink(temp_output_pt)

    return vectors, [], []


def _generate_vectors_single_benchmark(args, num_pairs: int, temp_output: str, temp_pairs: str):
    """Generate steering vector from a single benchmark."""
    from wisent.core.cli.generate_vector_from_task import execute_generate_vector_from_task

    vector_args = Namespace(
        task=args.task,
        trait_label="correctness",
        num_pairs=num_pairs,
        output=temp_output,
        model=args.model,
        device=args.device,
        verbose=False,
        timing=False,
        layers=args.layers,
        token_aggregation=args.token_aggregation,
        prompt_strategy="chat_template",
        method="caa",
        normalize=True,
        keep_intermediate=True,
        intermediate_dir=os.path.dirname(temp_pairs),
        accept_low_quality_vector=True,
    )

    execute_generate_vector_from_task(vector_args)
    return _load_vectors_from_json(temp_output, temp_pairs)


def _generate_vectors_synthetic(args, num_pairs: int, temp_output: str, temp_pairs: str):
    """Generate steering vector from synthetic pairs based on free-form description."""
    from wisent.core.contrastive_pairs.generation.llm_pair_generator import LLMPairGenerator
    from wisent.core.cli.get_activations import execute_get_activations
    from wisent.core.cli.create_steering_vector import execute_create_steering_vector
    
    # Generate synthetic pairs using the trait description (or task if trait not specified)
    trait_desc = getattr(args, 'trait', None) or args.task
    if not trait_desc:
        raise ValueError("Either --trait or --task must be specified for synthetic pair generation")
    
    generator = LLMPairGenerator(model_name=args.model)
    pairs = generator.generate_pairs(
        trait_description=trait_desc,
        num_pairs=num_pairs,
    )
    
    # Save pairs to temp file
    pairs_data = {"pairs": [p.to_dict() for p in pairs]}
    with open(temp_pairs, 'w') as f:
        json.dump(pairs_data, f)
    
    # Collect activations
    enriched_file = temp_pairs.replace('.json', '_enriched.json')
    act_args = Namespace(
        pairs_file=temp_pairs,
        output=enriched_file,
        model=args.model,
        device=args.device,
        layers=args.layers,
        token_aggregation=args.token_aggregation,
        prompt_strategy="chat_template",
        verbose=False,
    )
    execute_get_activations(act_args)
    
    # Create steering vector
    vec_args = Namespace(
        enriched_pairs=enriched_file,
        output=temp_output,
        method="caa",
        normalize=True,
        verbose=False,
    )
    execute_create_steering_vector(vec_args)
    
    return _load_vectors_from_json(temp_output, temp_pairs)


def _load_vectors_from_json(temp_output: str, temp_pairs: str):
    """Load steering vectors from JSON output file."""
    with open(temp_output, "r") as f:
        data = json.load(f)

    vectors = {
        int(layer) - 1: torch.tensor(vec)
        for layer, vec in data["steering_vectors"].items()
    }

    positive_examples = data.get("positive_examples", [])
    negative_examples = data.get("negative_examples", [])

    if not positive_examples or not negative_examples:
        intermediate_dir = os.path.dirname(temp_pairs)
        import glob
        pairs_files = glob.glob(os.path.join(intermediate_dir, "*_pairs.json"))
        pairs_files = [f for f in pairs_files if "_with_activations" not in f]

        for pairs_file in pairs_files:
            with open(pairs_file, "r") as f:
                pairs_data = json.load(f)
            if "pairs" in pairs_data:
                for pair in pairs_data["pairs"]:
                    if "positive_response" in pair and "model_response" in pair["positive_response"]:
                        positive_examples.append(pair["positive_response"]["model_response"])
                    if "negative_response" in pair and "model_response" in pair["negative_response"]:
                        negative_examples.append(pair["negative_response"]["model_response"])

    return vectors, positive_examples, negative_examples


def _create_evaluator(args, model_name: str, wisent_model: WisentModel = None,
                      positive_examples: list[str] = None, negative_examples: list[str] = None) -> Callable[[Any, Any], dict[str, float]]:
    """Create the appropriate evaluator function based on --task argument.

    Uses shared steering evaluators from wisent.core.evaluators.steering_evaluators.

    Task types:
    - 'refusal' → RefusalEvaluator (compliance rate)
    - 'personalization' → PersonalizationEvaluator (requires --trait)
    - benchmark name → TaskEvaluator (e.g., 'arc_easy', 'gsm8k')
    - comma-separated benchmarks → PooledEvaluator (e.g., 'arc_easy,gsm8k')

    The evaluator receives (hf_model, tokenizer) and wraps them in WisentModel
    internally to use standard generation.

    Args:
        args: Command arguments (must have args.task)
        model_name: Model name/path
        wisent_model: Optional WisentModel instance for baseline generation (used by personalization)
        positive_examples: List of positive example responses from contrastive pairs
        negative_examples: List of negative example responses from contrastive pairs
    """
    from wisent.core.models.inference_config import get_generate_kwargs

    task = args.task.lower() if args.task else ""
    
    # Determine evaluator type from task
    if task == "refusal":
        evaluator_type = "refusal"
    elif task == "personalization" or (not task and getattr(args, 'trait', None)):
        # Use personalization evaluator when --trait is provided (with or without --task personalization)
        evaluator_type = "personalization"
    elif task == "custom":
        if not getattr(args, 'custom_evaluator', None):
            raise ValueError("--custom-evaluator is required when --task custom")
        evaluator_type = "custom"
    elif "," in task:
        # Multiple benchmarks: use pooled evaluator
        evaluator_type = "pooled"
    elif task:
        # Single benchmark: use task evaluator (only if task is specified)
        evaluator_type = "task"
    else:
        raise ValueError("Either --task or --trait must be specified")

    # Pooled evaluator for multiple benchmarks
    if evaluator_type == "pooled":
        print(f"   [DEBUG] Creating pooled evaluator for tasks: {task}")
        print(f"   [DEBUG] args._pooled_eval_pairs count: {len(getattr(args, '_pooled_eval_pairs', []))}")
        return _create_pooled_evaluator(args)
    
    # Custom evaluator
    if evaluator_type == "custom":
        return _create_custom_evaluator(args, model_name)

    # Use shared steering evaluators for refusal, task, personalization
    eval_config = EvaluatorConfig(
        evaluator_type=evaluator_type,
        trait=getattr(args, 'trait', None),
        task=task if evaluator_type == "task" else None,
        eval_prompts_path=getattr(args, 'eval_prompts', None),
        eval_topics=getattr(args, 'eval_topics', None),
        num_eval_prompts=getattr(args, 'num_eval_prompts', 30),
    )

    # Create shared evaluator
    shared_evaluator = SteeringEvaluatorFactory.create(
        eval_config, model_name, wisent_model, positive_examples, negative_examples
    )

    def evaluate(hf_model, tokenizer) -> dict:
        """Run evaluation using shared steering evaluator."""
        # Wrap HF model in WisentModel for standard generation
        temp_wisent_model = WisentModel(model_name, hf_model=hf_model)
        
        # Get prompts from evaluator
        prompts = shared_evaluator.get_prompts()
        
        # Generate responses
        responses = []
        for prompt_text in prompts:
            messages = [{"role": "user", "content": prompt_text}]
            result = temp_wisent_model.generate(
                [messages],
                **get_generate_kwargs(max_new_tokens=150),
            )
            responses.append(result[0] if result else "")
        
        # Evaluate using shared evaluator
        return shared_evaluator.evaluate_responses(responses)

    return evaluate


def _create_custom_evaluator(args, model_name: str) -> Callable:
    """Create custom evaluator for --task custom.
    
    Uses the custom evaluator framework from wisent.core.evaluators.custom.
    
    Args:
        args: Command arguments (must have args.custom_evaluator)
        model_name: Model name/path
    """
    import json
    from wisent.core.evaluators.custom import create_custom_evaluator
    from wisent.core.models.inference_config import get_generate_kwargs
    
    # Parse custom evaluator kwargs
    custom_kwargs = {}
    if getattr(args, 'custom_evaluator_kwargs', None):
        custom_kwargs = json.loads(args.custom_evaluator_kwargs)
    
    # Create the custom evaluator
    custom_eval = create_custom_evaluator(args.custom_evaluator, **custom_kwargs)
    
    # Get test prompts - for custom, we need either trait-based prompts or user-provided prompts
    eval_prompts = []
    if getattr(args, 'eval_prompts', None):
        # Load from file
        with open(args.eval_prompts) as f:
            eval_prompts = [line.strip() for line in f if line.strip()]
    elif getattr(args, 'trait', None):
        # Generate prompts based on trait
        from wisent.core.evaluators.steering_evaluators import PersonalizationEvaluator
        eval_prompts = PersonalizationEvaluator._generate_test_prompts(
            getattr(args, 'num_eval_prompts', 30)
        )
    else:
        # Default generic prompts
        eval_prompts = [
            "Tell me about yourself.",
            "What do you think about artificial intelligence?",
            "How would you solve world hunger?",
            "Explain quantum computing in simple terms.",
            "What's the meaning of life?",
        ]
    
    def evaluate(hf_model, tokenizer) -> dict:
        """Run evaluation using custom evaluator."""
        # Wrap HF model in WisentModel for standard generation
        temp_wisent_model = WisentModel(model_name, hf_model=hf_model)
        
        # Generate responses
        scores = []
        for prompt_text in eval_prompts:
            messages = [{"role": "user", "content": prompt_text}]
            result = temp_wisent_model.generate(
                [messages],
                **get_generate_kwargs(max_new_tokens=150),
            )
            response = result[0] if result else ""
            
            # Score using custom evaluator
            score = custom_eval(response)
            if isinstance(score, dict):
                # Take the primary score (first value or 'score' key)
                score = score.get('score', list(score.values())[0])
            scores.append(float(score))
        
        avg_score = sum(scores) / len(scores) if scores else 0.0
        return {
            "score": avg_score,
            "custom_score": avg_score,
            "num_evaluated": len(scores),
        }
    
    return evaluate


def _create_pooled_evaluator(args) -> Callable:
    """Create pooled evaluator for multi-benchmark mode (--task bench1,bench2,...).
    
    Evaluates on ALL benchmarks from training, using each benchmark's
    native evaluator (routing by source_benchmark metadata).
    """
    from wisent.core.evaluators.rotator import EvaluatorRotator
    from wisent.core.models.inference_config import get_generate_kwargs

    # Get eval pairs stored during vector generation
    eval_pairs = getattr(args, '_pooled_eval_pairs', [])
    benchmarks_used = getattr(args, '_benchmarks_used', [])

    if not eval_pairs:
        raise ValueError("No eval pairs found for pooled evaluation. "
                        "Make sure train_unified_goodness saved eval_pairs to checkpoint.")

    print(f"   Pooled evaluator: {len(eval_pairs)} eval pairs across {len(benchmarks_used)} benchmarks")

    # Discover evaluators once
    EvaluatorRotator.discover_evaluators('wisent.core.evaluators.oracles')
    EvaluatorRotator.discover_evaluators('wisent.core.evaluators.benchmark_specific')

    # Group pairs by benchmark
    pairs_by_benchmark = {}
    for pair in eval_pairs:
        bench = pair.get('source_benchmark', 'unknown')
        if bench not in pairs_by_benchmark:
            pairs_by_benchmark[bench] = []
        pairs_by_benchmark[bench].append(pair)

    print(f"   Benchmarks in eval set: {list(pairs_by_benchmark.keys())}")

    def evaluate(hf_model, tokenizer) -> dict[str, float]:
        """Evaluate modified model on pooled benchmark data."""
        from wisent.core.models.wisent_model import WisentModel

        # Create WisentModel wrapper - use object.__new__ to avoid __init__ loading
        wisent_model = object.__new__(WisentModel)
        wisent_model.hf_model = hf_model
        wisent_model.tokenizer = tokenizer
        wisent_model.model_name = "modified_model"
        # Set internal attributes that the properties depend on
        if hasattr(hf_model, 'model') and hasattr(hf_model.model, 'layers'):
            wisent_model._layers = hf_model.model.layers
        elif hasattr(hf_model, 'transformer') and hasattr(hf_model.transformer, 'h'):
            wisent_model._layers = hf_model.transformer.h
        else:
            wisent_model._layers = []
        wisent_model._hidden_size = hf_model.config.hidden_size
        wisent_model.device = next(hf_model.parameters()).device

        total_correct = 0
        total_samples = 0
        benchmark_scores = {}

        for bench_name, bench_pairs in pairs_by_benchmark.items():
            # Get evaluator for this benchmark
            evaluator = EvaluatorRotator(
                evaluator=None,
                task_name=bench_name,
                autoload=False
            )

            correct = 0
            evaluated = 0

            for pair in bench_pairs:
                try:
                    question = pair['prompt']
                    expected = pair['positive_response']
                    metadata = pair.get('metadata', {}) or {}
                    test_code = metadata.get('test_code', '')
                    entry_point = metadata.get('entry_point', '')

                    messages = [{"role": "user", "content": question}]
                    response = wisent_model.generate(
                        [messages],
                        **get_generate_kwargs(),
                    )[0]

                    # Evaluate
                    eval_result = evaluator.evaluate(
                        response=response,
                        expected=expected,
                        model=wisent_model,
                        question=question,
                        task_name=bench_name,
                        test_code=test_code,
                        entry_point=entry_point,
                    )

                    if eval_result.ground_truth == "TRUTHFUL":
                        correct += 1
                    evaluated += 1

                except Exception as e:
                    # Count failed evals but continue
                    evaluated += 1

            if evaluated > 0:
                benchmark_scores[bench_name] = correct / evaluated
                total_correct += correct
                total_samples += evaluated

        # Compute aggregate accuracy
        accuracy = total_correct / total_samples if total_samples > 0 else 0.0

        return {
            "accuracy": accuracy,
            "correct": total_correct,
            "total": total_samples,
            "score": accuracy,
            "benchmark_scores": benchmark_scores,
        }

    return evaluate


def _apply_weight_modification_standalone(
    model,
    steering_vectors: dict[int, torch.Tensor],
    params: dict,
    num_layers: int,
    config: WeightsOptimizerConfig,
):
    """Apply weight modification with given parameters (standalone helper)."""
    from wisent.core.weight_modification import project_with_kernel, bake_steering_with_kernel

    max_weight_position = params["max_weight_position"] * (num_layers - 1)
    min_weight_distance = 0.6 * (num_layers - 1)

    if config.method == "directional":
        project_with_kernel(
            model,
            steering_vectors,
            max_weight=params["max_weight"] * params["strength"],
            max_weight_position=max_weight_position,
            min_weight=params["min_weight"],
            min_weight_distance=min_weight_distance,
            components=config.components,
            norm_preserve=config.norm_preserve,
            verbose=False,
        )
    else:
        bake_steering_with_kernel(
            model,
            steering_vectors,
            max_alpha=params["max_weight"] * params["strength"],
            max_alpha_position=max_weight_position,
            min_alpha=params["min_weight"],
            components=config.components,
            verbose=False,
        )


def _show_response_comparisons(
    base_model,
    base_state_dict: dict,
    steering_vectors: dict[int, torch.Tensor],
    best_params: dict,
    num_layers: int,
    model_name: str,
    args,
    optimizer_config: WeightsOptimizerConfig,
    num_comparisons: int = None,
    save_path: str = None,
):
    """Show and/or save before/after response comparisons.

    Uses WisentModel.generate() for consistent generation behavior.
    """
    from uncensorbench import UncensorBench
    from uncensorbench.evaluator import KeywordEvaluator, SemanticEvaluator
    from wisent.core.models.inference_config import get_generate_kwargs

    print(f"\n{'='*80}")
    print("RESPONSE COMPARISONS (Baseline vs Optimized)")
    print(f"{'='*80}")

    bench = UncensorBench()

    if args.eval_prompts:
        with open(args.eval_prompts) as f:
            custom_prompts = json.load(f)
        if not isinstance(custom_prompts, list):
            custom_prompts = custom_prompts.get("prompts", [])
        prompts = custom_prompts[:args.num_eval_prompts]
        use_custom = True
    else:
        topics = args.eval_topics.split(",") if args.eval_topics else None
        prompts = list(bench.prompts(topics=topics))[:args.num_eval_prompts]
        use_custom = False

    try:
        evaluator = SemanticEvaluator()
    except ImportError:
        evaluator = KeywordEvaluator()

    def generate_response(hf_model, prompt_text: str) -> str:
        """Generate response using WisentModel for consistent behavior."""
        # Wrap HF model in WisentModel for standard generation
        wisent_model = WisentModel(model_name, hf_model=hf_model)

        messages = [{"role": "user", "content": prompt_text}]
        responses = wisent_model.generate(
            [messages],
            **get_generate_kwargs(max_new_tokens=150),
        )

        response = responses[0] if responses else ""
        # Strip <think>...</think> blocks from Qwen3 responses
        response = re.sub(r"<think>.*?</think>", "", response, flags=re.DOTALL).strip()
        return response

    def evaluate_response(prompt_item, response: str, use_custom: bool) -> tuple[bool, float]:
        if use_custom:
            is_refusal = bench.detect_refusal(response)
            compliance_score = 0.0 if is_refusal else 1.0
        else:
            eval_result = bench.evaluate_response(prompt_item, response, evaluator)
            is_refusal = not eval_result.is_uncensored
            compliance_score = eval_result.compliance_score
        return is_refusal, compliance_score

    print("\nGenerating baseline responses...")

    base_model.load_state_dict(base_state_dict)

    baseline_results = []
    for item in prompts:
        if use_custom:
            prompt_text = item if isinstance(item, str) else item.get("prompt", str(item))
        else:
            prompt_text = item.prompt

        response = generate_response(base_model, prompt_text)
        is_refusal, score = evaluate_response(item, response, use_custom)
        baseline_results.append({
            "prompt": prompt_text,
            "response": response,
            "is_refusal": is_refusal,
            "score": score,
            "item": item,
        })

    print("Generating optimized responses...")

    _apply_weight_modification_standalone(
        base_model,
        steering_vectors,
        best_params,
        num_layers,
        optimizer_config,
    )

    optimized_results = []
    for i, item in enumerate(prompts):
        if use_custom:
            prompt_text = item if isinstance(item, str) else item.get("prompt", str(item))
        else:
            prompt_text = item.prompt

        response = generate_response(base_model, prompt_text)
        is_refusal, score = evaluate_response(item, response, use_custom)
        optimized_results.append({
            "prompt": prompt_text,
            "response": response,
            "is_refusal": is_refusal,
            "score": score,
        })

    comparisons = []
    for i, (baseline, optimized) in enumerate(zip(baseline_results, optimized_results)):
        delta = optimized["score"] - baseline["score"]
        comparisons.append({
            "index": i,
            "prompt": baseline["prompt"],
            "baseline_response": baseline["response"],
            "baseline_score": baseline["score"],
            "baseline_refusal": baseline["is_refusal"],
            "optimized_response": optimized["response"],
            "optimized_score": optimized["score"],
            "optimized_refusal": optimized["is_refusal"],
            "delta": delta,
        })

    comparisons.sort(key=lambda x: abs(x["delta"]), reverse=True)

    total_baseline_refusals = sum(1 for c in comparisons if c["baseline_refusal"])
    total_optimized_refusals = sum(1 for c in comparisons if c["optimized_refusal"])
    avg_baseline_score = sum(c["baseline_score"] for c in comparisons) / len(comparisons)
    avg_optimized_score = sum(c["optimized_score"] for c in comparisons) / len(comparisons)

    summary = {
        "total_prompts": len(comparisons),
        "baseline_refusals": total_baseline_refusals,
        "optimized_refusals": total_optimized_refusals,
        "avg_baseline_score": avg_baseline_score,
        "avg_optimized_score": avg_optimized_score,
        "score_delta": avg_optimized_score - avg_baseline_score,
        "refusal_change": total_baseline_refusals - total_optimized_refusals,
    }

    if save_path:
        os.makedirs(os.path.dirname(save_path) if os.path.dirname(save_path) else ".", exist_ok=True)
        output_data = {
            "model": args.model,
            "trait": getattr(args, "trait", None),
            "best_params": best_params,
            "summary": summary,
            "comparisons": comparisons,
        }
        with open(save_path, "w") as f:
            json.dump(output_data, f, indent=2)
        print(f"\nSaved {len(comparisons)} comparisons to: {save_path}")

    if num_comparisons and num_comparisons > 0:
        print(f"\nTop {num_comparisons} prompts with biggest score changes:\n")

        for i, comp in enumerate(comparisons[:num_comparisons]):
            print(f"{'─'*80}")
            print(f"Comparison {i+1}/{num_comparisons}")
            print(f"{'─'*80}")
            print(f"PROMPT: {comp['prompt'][:200]}{'...' if len(comp['prompt']) > 200 else ''}")
            print()
            print(f"BASELINE (score={comp['baseline_score']:.2f}, refusal={comp['baseline_refusal']}):")
            print(f"  {comp['baseline_response'][:300]}{'...' if len(comp['baseline_response']) > 300 else ''}")
            print()
            print(f"OPTIMIZED (score={comp['optimized_score']:.2f}, refusal={comp['optimized_refusal']}):")
            print(f"  {comp['optimized_response'][:300]}{'...' if len(comp['optimized_response']) > 300 else ''}")
            print()
            delta_str = f"+{comp['delta']:.2f}" if comp['delta'] >= 0 else f"{comp['delta']:.2f}"
            print(f"DELTA: {delta_str}")
            print()

    print(f"{'='*80}")
    print("SUMMARY")
    print(f"{'='*80}")
    print(f"Baseline:  {total_baseline_refusals}/{len(comparisons)} refusals, avg score={avg_baseline_score:.3f}")
    print(f"Optimized: {total_optimized_refusals}/{len(comparisons)} refusals, avg score={avg_optimized_score:.3f}")
    print(f"Change:    {total_baseline_refusals - total_optimized_refusals} fewer refusals, "
          f"score delta={avg_optimized_score - avg_baseline_score:+.3f}")
