"""Steering optimization command execution logic with full strategy optimization.

Results are persisted to ~/.wisent/configs/ via WisentConfigManager
so they can be automatically loaded on subsequent runs.
"""

import json
import sys
import time

import numpy as np

from wisent.core.evaluators.rotator import EvaluatorRotator
from wisent.core.models.inference_config import get_generate_kwargs
from wisent.core.config_manager import (
    get_config_manager,
    save_steering_config,
    get_steering_config,
    get_cached_optimization,
    store_optimization,
    SteeringConfig,
)


def execute_optimize_steering(args):
    """
    Execute the optimize-steering command.

    Supports multiple subcommands:
    - comprehensive: Run comprehensive steering optimization
    - compare-methods: Compare different steering methods
    - optimize-layer: Find optimal steering layer
    - optimize-strength: Find optimal steering strength
    - auto: Automatically optimize based on classification config
    """
    from wisent.core.data_loaders.loaders.lm_loader import LMEvalDataLoader
    from wisent.core.models.wisent_model import WisentModel

    # Check which subcommand was called
    if not hasattr(args, "steering_action") or args.steering_action is None:
        print("\n✗ No steering optimization action specified")
        print("Available actions: comprehensive, compare-methods, optimize-layer, optimize-strength, auto")
        sys.exit(1)

    print(f"\n{'=' * 80}")
    print(f"🎯 STEERING PARAMETER OPTIMIZATION: {args.steering_action.upper()}")
    print(f"{'=' * 80}")
    print(f"   Model: {args.model}")
    print(f"   Device: {args.device or 'auto'}")
    print(f"{'=' * 80}\n")

    # Load model
    print("📦 Loading model...")
    model = WisentModel(args.model, device=args.device)
    print(f"   ✓ Model loaded with {model.num_layers} layers\n")

    # Initialize data loader
    loader = LMEvalDataLoader()

    # Execute based on subcommand and return results
    if args.steering_action == "comprehensive":
        return execute_comprehensive(args, model, loader)
    if args.steering_action == "compare-methods":
        return execute_compare_methods(args, model, loader)
    if args.steering_action == "optimize-layer":
        return execute_optimize_layer(args, model, loader)
    if args.steering_action == "optimize-strength":
        return execute_optimize_strength(args, model, loader)
    if args.steering_action == "auto":
        return execute_auto(args, model, loader)
    if args.steering_action == "personalization":
        return execute_personalization(args, model)
    if args.steering_action == "multi-personalization":
        return execute_multi_personalization(args, model)
    print(f"\n✗ Unknown steering action: {args.steering_action}")
    sys.exit(1)


def execute_comprehensive(args, model, loader):
    """Execute comprehensive steering optimization with generation-based evaluation."""
    import torch

    from wisent.core.activations.activations_collector import ActivationCollector
    from wisent.core.activations.core.atoms import ActivationAggregationStrategy
    from wisent.core.activations.prompt_construction_strategy import PromptConstructionStrategy
    from wisent.core.models.core.atoms import SteeringPlan
    from wisent.core.cli.steering_method_trainer import create_steering_method
    from wisent.core.cli.steering_search_space import (
        get_search_space_from_args,
        print_search_space_summary,
        CAASearchSpace,
        PRISMSearchSpace,
        PULSESearchSpace,
        TITANSearchSpace,
    )

    print("🔍 Running comprehensive steering optimization...")
    print("   Optimizing: Method-specific search space (layer, strength, strategy, + method params)")

    # Determine tasks to optimize
    if args.tasks:
        task_list = args.tasks
    else:
        task_list = ["arc_easy", "hellaswag", "winogrande", "gsm8k"]

    # Check for cached results if --use-cached is specified
    use_cached = getattr(args, "use_cached", False)
    save_as_default = getattr(args, "save_as_default", False)

    if use_cached:
        print("\n📦 Checking optimization cache...")
        cached_results = {}
        tasks_to_run = []

        for task_name in task_list:
            for method in args.methods:
                cached = get_cached_optimization(args.model, task_name, method)
                if cached:
                    print(
                        f"   ✓ Found cached result for {task_name}/{method}: layer={cached.layer}, strength={cached.strength}, score={cached.score:.3f}"
                    )
                    cached_results[f"{task_name}::{method}"] = cached
                else:
                    if task_name not in tasks_to_run:
                        tasks_to_run.append(task_name)

        if cached_results and not tasks_to_run:
            print("\n✅ All tasks have cached results. Returning cached configurations.")
            # Convert cached results to the expected return format
            all_results = {}
            for key, cached in cached_results.items():
                task_name, method = key.split("::")
                if task_name not in all_results:
                    all_results[task_name] = {}
                all_results[task_name][method] = {
                    "best_layer": cached.layer,
                    "best_strength": cached.strength,
                    "best_score": cached.score,
                    "token_aggregation": cached.token_aggregation,
                    "prompt_strategy": cached.prompt_strategy,
                    "from_cache": True,
                }
            return all_results

        if tasks_to_run:
            print(f"   Tasks needing optimization: {', '.join(tasks_to_run)}")
            task_list = tasks_to_run

    print(f"   Tasks: {', '.join(task_list)}")
    print(f"   Methods: {', '.join(args.methods)}")
    print(f"   Limit: {args.limit} samples per task")
    quick_search = getattr(args, 'quick_search', False)
    print(f"   Quick search: {quick_search}")
    print("   Time limit: DISABLED (no time limit)\n")

    all_results = {}

    # Get search spaces for each method and print summary
    method_search_spaces = {}
    total_all_methods = 0
    for method_name in args.methods:
        search_space = get_search_space_from_args(method_name, args, model.num_layers)
        method_search_spaces[method_name] = search_space
        print_search_space_summary(search_space, method_name)
        total_all_methods += search_space.get_total_configs()
    
    print(f"\n   Total configurations across all methods: {total_all_methods:,}\n")
    
    # For backward compatibility, also set up the legacy variables
    # These are used by some code paths that haven't been fully migrated
    first_method = args.methods[0] if args.methods else "CAA"
    first_space = method_search_spaces.get(first_method)
    if isinstance(first_space, (CAASearchSpace, PRISMSearchSpace)):
        layers_to_test = first_space.layers
    else:
        # PULSE/TITAN don't use direct layers, compute defaults
        layers_to_test = list(range(model.num_layers // 2, model.num_layers - 2, 2))
    
    strengths_to_test = first_space.strengths if first_space else [0.5, 1.0, 1.5, 2.0]
    strategies_to_test = first_space.strategies if first_space else ["constant", "initial_only", "diminishing"]
    
    # Convert string token aggregations to enum
    token_agg_map = {
        "last_token": ActivationAggregationStrategy.LAST_TOKEN,
        "mean_pooling": ActivationAggregationStrategy.MEAN_POOLING,
        "first_token": ActivationAggregationStrategy.FIRST_TOKEN,
        "max_pooling": ActivationAggregationStrategy.MAX_POOLING,
        "choice_token": ActivationAggregationStrategy.CHOICE_TOKEN,
        "continuation_token": ActivationAggregationStrategy.CONTINUATION_TOKEN,
    }
    token_aggregations_to_test = [
        token_agg_map.get(t, ActivationAggregationStrategy.LAST_TOKEN) 
        for t in (first_space.token_aggregations if first_space else ["last_token", "mean_pooling"])
    ]
    
    # Convert string prompt constructions to enum
    prompt_const_map = {
        "chat_template": PromptConstructionStrategy.CHAT_TEMPLATE,
        "direct_completion": PromptConstructionStrategy.DIRECT_COMPLETION,
        "multiple_choice": PromptConstructionStrategy.MULTIPLE_CHOICE,
        "role_playing": PromptConstructionStrategy.ROLE_PLAYING,
        "instruction_following": PromptConstructionStrategy.INSTRUCTION_FOLLOWING,
    }
    prompt_constructions_to_test = [
        prompt_const_map.get(p, PromptConstructionStrategy.CHAT_TEMPLATE)
        for p in (first_space.prompt_constructions if first_space else ["chat_template", "direct_completion"])
    ]
    
    # For legacy code paths
    total_configs = first_space.get_total_configs() if first_space else 100

    for task_idx, task_name in enumerate(task_list, 1):
        print(f"\n{'=' * 80}")
        print(f"Task {task_idx}/{len(task_list)}: {task_name}")
        print(f"{'=' * 80}")

        task_start_time = time.time()

        try:
            # Load task data
            print("  📊 Loading task data...")
            result = loader._load_one_task(
                task_name=task_name, split_ratio=0.8, seed=42, limit=args.limit, training_limit=None, testing_limit=None
            )

            train_pairs = result["train_qa_pairs"]
            test_pairs = result["test_qa_pairs"]

            print(f"      ✓ Loaded {len(train_pairs.pairs)} train, {len(test_pairs.pairs)} test pairs")

            # Initialize evaluator for this task (auto-select based on task_name)
            EvaluatorRotator.discover_evaluators("wisent.core.evaluators.benchmark_specific")
            evaluator = EvaluatorRotator(evaluator=None, task_name=task_name)  # None = auto-select
            print(f"      ✓ Using evaluator: {evaluator._plugin.name} (auto-selected for {task_name})")

            # Compute baseline (unsteered) results if requested
            baseline_results = {}
            if hasattr(args, "compute_baseline") and args.compute_baseline:
                print("\n  📊 Computing BASELINE (unsteered) accuracy...")
                baseline_scores = []
                baseline_per_problem = []

                for pair_idx, pair in enumerate(test_pairs.pairs):
                    try:
                        # Prepare choices for multiple choice evaluation
                        choices = [pair.negative_response.model_response, pair.positive_response.model_response]
                        expected = pair.positive_response.model_response

                        # Evaluate WITHOUT steering
                        test_code = pair.metadata.get("test_code") if pair.metadata else None
                        eval_result = evaluator.evaluate(
                            response="",
                            expected=expected,
                            model=model,
                            question=pair.prompt,
                            choices=choices,
                            steering_plan=None,  # No steering for baseline
                            test_code=test_code,
                            task_name=task_name,
                        )

                        is_correct = eval_result.ground_truth == "TRUTHFUL"
                        baseline_scores.append(1.0 if is_correct else 0.0)

                        # Store per-problem baseline result with details
                        baseline_per_problem.append(
                            {
                                "pair_index": pair_idx,
                                "prompt": pair.prompt,
                                "expected": expected,
                                "baseline_correct": is_correct,
                                "ground_truth": eval_result.ground_truth,
                                "method_used": eval_result.method_used,
                                "confidence": eval_result.confidence,
                            }
                        )

                        if (pair_idx + 1) % 10 == 0:
                            print(
                                f"      Evaluated {pair_idx + 1}/{len(test_pairs.pairs)} baseline samples...", end="\r"
                            )

                    except Exception as e:
                        print(f"\n❌ Baseline evaluation failed for pair {pair_idx}:")
                        print(f"   Error: {e}")
                        raise

                baseline_accuracy = np.mean(baseline_scores) if baseline_scores else 0.0
                print(
                    f"\n      ✓ Baseline accuracy: {baseline_accuracy:.3f} ({sum(baseline_scores):.0f}/{len(baseline_scores)} correct)"
                )

                baseline_results = {
                    "accuracy": baseline_accuracy,
                    "per_problem": baseline_per_problem,
                    "num_correct": int(sum(baseline_scores)),
                    "num_total": len(baseline_scores),
                }

            print(
                "\n  🔍 Testing CAA method across layers, strengths, strategies, token aggregations, prompt constructions..."
            )
            print(f"      Total configurations: {total_configs}")

            best_score = 0
            best_config = None
            method_results = {}
            configs_tested = 0
            all_generation_examples = []  # Store generation examples for all configs

            # Prepare test prompts if generating examples for all configs
            if args.save_all_generation_examples or args.save_generation_examples:
                num_examples = min(args.num_generation_examples, len(test_pairs.pairs))
                example_pairs = test_pairs.pairs[:num_examples]
                print(f"  📝 Will generate {num_examples} example responses per configuration")

            for layer in layers_to_test:
                for strength in strengths_to_test:
                    for strategy in strategies_to_test:
                        for token_agg in token_aggregations_to_test:
                            for prompt_const in prompt_constructions_to_test:
                                # Time limit disabled - run all configurations

                                try:
                                    configs_tested += 1
                                    layer_str = str(layer)

                                    # Step 1: Generate steering vector using CAA with current token aggregation
                                    collector = ActivationCollector(model=model, store_device="cpu")

                                    pos_acts = []
                                    neg_acts = []

                                    for pair in train_pairs.pairs:
                                        updated_pair = collector.collect_for_pair(
                                            pair,
                                            layers=[layer_str],
                                            aggregation=token_agg,  # Use current token aggregation strategy
                                            return_full_sequence=False,
                                            normalize_layers=False,
                                        )

                                        if (
                                            updated_pair.positive_response.layers_activations
                                            and layer_str in updated_pair.positive_response.layers_activations
                                        ):
                                            act = updated_pair.positive_response.layers_activations[layer_str]
                                            if act is not None:
                                                pos_acts.append(act)

                                        if (
                                            updated_pair.negative_response.layers_activations
                                            and layer_str in updated_pair.negative_response.layers_activations
                                        ):
                                            act = updated_pair.negative_response.layers_activations[layer_str]
                                            if act is not None:
                                                neg_acts.append(act)

                                    if len(pos_acts) == 0 or len(neg_acts) == 0:
                                        continue

                                    # Create CAA steering vector
                                    # Use the selected method (first from args.methods or default to CAA)
                                    method_name = args.methods[0] if args.methods else "CAA"
                                    steering_method = create_steering_method(method_name, args)
                                    steering_vector = steering_method.train_for_layer(pos_acts, neg_acts)

                                    # Step 2: Evaluate with ACTUAL GENERATION and task evaluator
                                    # Create steering plan
                                    from wisent.core.models.core.atoms import SteeringPlan, SteeringVector

                                    steering_vec = SteeringVector(vector=steering_vector, scale=strength)
                                    steering_plan = SteeringPlan(
                                        layers={layer_str: steering_vec},
                                        layers_description=[
                                            f"CAA L{layer} S{strength} {strategy} T:{token_agg.value} P:{prompt_const.value}"
                                        ],
                                    )

                                    # Apply steering to model
                                    model.apply_steering(steering_plan)

                                    test_scores = []
                                    detailed_results = []  # Store full evaluation details
                                    delta_tracking = []  # Track improved/regressed/unchanged per problem

                                    for pair_idx, pair in enumerate(test_pairs.pairs):
                                        try:
                                            # Prepare choices for multiple choice evaluation
                                            # ContrastivePair uses: prompt, positive_response.model_response, negative_response.model_response
                                            choices = [
                                                pair.negative_response.model_response,
                                                pair.positive_response.model_response,
                                            ]
                                            expected = pair.positive_response.model_response

                                            # Use the Wisent evaluator to check correctness
                                            # The evaluator will use log likelihood if possible,
                                            # otherwise fall back to generation
                                            # Pass test_code from metadata for coding tasks
                                            test_code = pair.metadata.get("test_code") if pair.metadata else None
                                            eval_result = evaluator.evaluate(
                                                response="",  # Not used for log likelihood eval
                                                expected=expected,
                                                model=model,
                                                question=pair.prompt,
                                                choices=choices,
                                                steering_plan=steering_plan,
                                                test_code=test_code,
                                                task_name=task_name,
                                            )

                                            # Convert TRUTHFUL/UNTRUTHFUL to 1.0/0.0
                                            is_correct = eval_result.ground_truth == "TRUTHFUL"
                                            test_scores.append(1.0 if is_correct else 0.0)

                                            # Save full evaluation details
                                            detailed_results.append(
                                                {
                                                    "prompt": pair.prompt,
                                                    "choices": choices,
                                                    "expected": expected,
                                                    "ground_truth": eval_result.ground_truth,
                                                    "method_used": eval_result.method_used,
                                                    "confidence": eval_result.confidence,
                                                    "details": eval_result.details,
                                                    "meta": dict(eval_result.meta) if eval_result.meta else {},
                                                    "is_correct": is_correct,
                                                }
                                            )

                                            # Track delta if baseline was computed
                                            if baseline_results and "per_problem" in baseline_results:
                                                baseline_correct = baseline_results["per_problem"][pair_idx][
                                                    "baseline_correct"
                                                ]
                                                if not baseline_correct and is_correct:
                                                    delta_status = "improved"
                                                elif baseline_correct and not is_correct:
                                                    delta_status = "regressed"
                                                else:
                                                    delta_status = "unchanged"

                                                delta_tracking.append(
                                                    {
                                                        "pair_index": pair_idx,
                                                        "prompt": pair.prompt,
                                                        "expected": expected,
                                                        "baseline_correct": baseline_correct,
                                                        "steered_correct": is_correct,
                                                        "delta_status": delta_status,
                                                    }
                                                )

                                        except Exception as e:
                                            # NO FALLBACK - raise the error immediately
                                            print("\n❌ Evaluation failed for test pair:")
                                            print(f"   Prompt: {pair.prompt[:100]}")
                                            print(f"   Error: {e}")
                                            raise

                                    # Clear steering
                                    model.clear_steering()

                                    if len(test_scores) > 0:
                                        avg_score = np.mean(test_scores)

                                        # Generate examples for this configuration if requested
                                        if args.save_all_generation_examples:
                                            config_examples = []
                                            # Get inference config settings
                                            for idx, pair in enumerate(example_pairs):
                                                prompt = pair.prompt
                                                try:
                                                    # Generate without steering (only once per prompt, reuse if already generated)
                                                    unsteered_response = model.generate(
                                                        [[{"role": "user", "content": prompt}]],
                                                        **get_generate_kwargs(max_new_tokens=100),
                                                        use_steering=False,
                                                    )[0]

                                                    # Create steering plan for this config
                                                    from wisent.core.models.core.atoms import (
                                                        SteeringPlan,
                                                        SteeringVector,
                                                    )

                                                    steering_vec = SteeringVector(
                                                        vector=steering_vector, scale=strength
                                                    )
                                                    steering_plan = SteeringPlan(
                                                        layers={layer_str: steering_vec},
                                                        layers_description=[
                                                            f"CAA steering layer={layer}, strength={strength}, strategy={strategy}"
                                                        ],
                                                    )

                                                    # Generate with steering
                                                    model.apply_steering(steering_plan)
                                                    steered_response = model.generate(
                                                        [[{"role": "user", "content": prompt}]],
                                                        **get_generate_kwargs(max_new_tokens=100),
                                                        use_steering=True,
                                                        steering_plan=steering_plan,
                                                    )[0]
                                                    model.clear_steering()

                                                    config_examples.append(
                                                        {
                                                            "prompt": prompt,
                                                            "correct_answer": pair.positive_response.model_response,
                                                            "incorrect_answer": pair.negative_response.model_response,
                                                            "unsteered_generation": unsteered_response,
                                                            "steered_generation": steered_response,
                                                        }
                                                    )
                                                except Exception as e:
                                                    if args.verbose:
                                                        print(
                                                            f"      ⚠️ Failed to generate example for config layer={layer}, strength={strength}, strategy={strategy}: {e}"
                                                        )

                                            # Store this config's examples
                                            all_generation_examples.append(
                                                {
                                                    "layer": layer,
                                                    "strength": strength,
                                                    "strategy": strategy,
                                                    "accuracy": avg_score,
                                                    "examples": config_examples,
                                                }
                                            )

                                        # Compute delta summary if baseline was computed
                                        delta_summary = {}
                                        if delta_tracking:
                                            improved = sum(1 for d in delta_tracking if d["delta_status"] == "improved")
                                            regressed = sum(
                                                1 for d in delta_tracking if d["delta_status"] == "regressed"
                                            )
                                            unchanged = sum(
                                                1 for d in delta_tracking if d["delta_status"] == "unchanged"
                                            )
                                            delta_summary = {
                                                "improved": improved,
                                                "regressed": regressed,
                                                "unchanged": unchanged,
                                                "net_change": improved - regressed,
                                            }

                                        # Store detailed results for this configuration
                                        config_key = (
                                            f"L{layer}_S{strength}_{strategy}_{token_agg.value}_{prompt_const.value}"
                                        )
                                        method_results[config_key] = {
                                            "layer": layer,
                                            "strength": strength,
                                            "strategy": strategy,
                                            "token_aggregation": token_agg.value,
                                            "prompt_construction": prompt_const.value,
                                            "accuracy": avg_score,
                                            "num_test_samples": len(test_scores),
                                            "detailed_results": detailed_results,  # Save all eval details
                                            "delta_tracking": delta_tracking if delta_tracking else None,
                                            "delta_summary": delta_summary if delta_summary else None,
                                        }

                                        if avg_score > best_score:
                                            best_score = avg_score
                                            best_config = {
                                                "layer": layer,
                                                "strength": strength,
                                                "strategy": strategy,
                                                "token_aggregation": token_agg.value,
                                                "prompt_construction": prompt_const.value,
                                                "accuracy": avg_score,
                                            }

                                    if configs_tested % 10 == 0 and args.verbose:
                                        print(f"      Tested {configs_tested} configurations...", end="\r")

                                except Exception as e:
                                    # NO FALLBACK - raise the error immediately
                                    print("\n❌ Configuration test failed:")
                                    print(f"   Layer: {layer}")
                                    print(f"   Strength: {strength}")
                                    print(f"   Strategy: {strategy}")
                                    print(f"   Error: {e}")
                                    raise

            if best_config:
                print("\n  ✅ Best configuration found:")
                print("      Method: CAA")
                print(f"      Layer: {best_config['layer']}")
                print(f"      Strength: {best_config['strength']}")
                print(f"      Strategy: {best_config['strategy']} ⭐")
                print(f"      Token Aggregation: {best_config['token_aggregation']}")
                print(f"      Prompt Construction: {best_config['prompt_construction']}")
                print(f"      Accuracy: {best_config['accuracy']:.3f}")

                method_results["CAA"] = {
                    "optimal_layer": best_config["layer"],
                    "optimal_strength": best_config["strength"],
                    "optimal_strategy": best_config["strategy"],
                    "optimal_token_aggregation": best_config["token_aggregation"],
                    "optimal_prompt_construction": best_config["prompt_construction"],
                    "accuracy": best_config["accuracy"],
                    "f1": best_config["accuracy"],
                }

                # Save baseline comparison results if computed
                if hasattr(args, "compute_baseline") and args.compute_baseline and baseline_results:
                    import os

                    baseline_dir = (
                        args.baseline_output_dir if hasattr(args, "baseline_output_dir") else "./baseline_comparison"
                    )
                    os.makedirs(baseline_dir, exist_ok=True)

                    # Get delta tracking for best config
                    best_config_key = f"L{best_config['layer']}_S{best_config['strength']}_{best_config['strategy']}_{best_config['token_aggregation']}_{best_config['prompt_construction']}"
                    best_config_results = method_results.get(best_config_key, {})
                    best_delta_tracking = best_config_results.get("delta_tracking", [])
                    best_delta_summary = best_config_results.get("delta_summary", {})

                    # Separate improved, regressed, unchanged for inspection
                    improved_examples = [d for d in best_delta_tracking if d.get("delta_status") == "improved"]
                    regressed_examples = [d for d in best_delta_tracking if d.get("delta_status") == "regressed"]
                    unchanged_examples = [d for d in best_delta_tracking if d.get("delta_status") == "unchanged"]

                    baseline_comparison_data = {
                        "task": task_name,
                        "model": args.model,
                        "baseline_accuracy": baseline_results["accuracy"],
                        "best_steered_accuracy": best_config["accuracy"],
                        "delta": best_config["accuracy"] - baseline_results["accuracy"],
                        "best_config": best_config,
                        "summary": best_delta_summary,
                        "improved_examples": improved_examples,
                        "regressed_examples": regressed_examples,
                        "unchanged_examples": unchanged_examples,
                        "baseline_per_problem": baseline_results["per_problem"],
                    }

                    comparison_path = os.path.join(baseline_dir, f"{task_name}_baseline_comparison.json")
                    with open(comparison_path, "w") as f:
                        json.dump(baseline_comparison_data, f, indent=2)

                    print("\n  📊 Baseline Comparison Summary:")
                    print(f"      Baseline (unsteered) accuracy: {baseline_results['accuracy']:.3f}")
                    print(f"      Best steered accuracy: {best_config['accuracy']:.3f}")
                    print(f"      Delta: {(best_config['accuracy'] - baseline_results['accuracy']) * 100:+.1f}%")
                    if best_delta_summary:
                        print(f"      Improved: {best_delta_summary.get('improved', 0)} problems")
                        print(f"      Regressed: {best_delta_summary.get('regressed', 0)} problems")
                        print(f"      Unchanged: {best_delta_summary.get('unchanged', 0)} problems")
                        print(f"      Net change: {best_delta_summary.get('net_change', 0)} problems")
                    print(f"      💾 Saved comparison to: {comparison_path}")

                # Save best steering vector if requested
                if args.save_best_vector:
                    import os

                    vector_dir = args.save_best_vector
                    os.makedirs(vector_dir, exist_ok=True)

                    # Recreate the best steering vector with optimal token aggregation
                    best_layer_str = str(best_config["layer"])
                    best_token_agg = ActivationAggregationStrategy(best_config["token_aggregation"])
                    pos_acts_best = []
                    neg_acts_best = []

                    for pair in train_pairs.pairs:
                        updated_pair = collector.collect_for_pair(
                            pair,
                            layers=[best_layer_str],
                            aggregation=best_token_agg,  # Use optimal token aggregation
                            return_full_sequence=False,
                            normalize_layers=False,
                        )

                        if (
                            updated_pair.positive_response.layers_activations
                            and best_layer_str in updated_pair.positive_response.layers_activations
                        ):
                            act = updated_pair.positive_response.layers_activations[best_layer_str]
                            if act is not None:
                                pos_acts_best.append(act)

                        if (
                            updated_pair.negative_response.layers_activations
                            and best_layer_str in updated_pair.negative_response.layers_activations
                        ):
                            act = updated_pair.negative_response.layers_activations[best_layer_str]
                            if act is not None:
                                neg_acts_best.append(act)

                    # Create and save steering vector
                    method_name = args.methods[0] if args.methods else "CAA"
                    steering_method = create_steering_method(method_name, args)
                    best_steering_vector = steering_method.train_for_layer(pos_acts_best, neg_acts_best)

                    vector_path = os.path.join(vector_dir, f"{task_name}_layer{best_config['layer']}.pt")
                    torch.save(
                        {
                            "steering_vector": best_steering_vector,
                            "vector": best_steering_vector,  # Legacy key
                            "layer": best_config["layer"],
                            "layer_index": best_config["layer"],  # Legacy key
                            "strength": best_config["strength"],
                            "strategy": best_config["strategy"],
                            "token_aggregation": best_config["token_aggregation"],
                            "prompt_construction": best_config["prompt_construction"],
                            "method": "CAA",
                            "task": task_name,
                            "model": args.model,
                            "accuracy": best_config["accuracy"],
                        },
                        vector_path,
                    )
                    print(f"      💾 Saved steering vector to: {vector_path}")

                # Save generation examples
                if args.save_all_generation_examples:
                    # Save examples for ALL configurations
                    examples_path = os.path.join(
                        args.save_best_vector if args.save_best_vector else "./optimization_results",
                        f"{task_name}_all_generation_examples.json",
                    )
                    os.makedirs(os.path.dirname(examples_path), exist_ok=True)

                    with open(examples_path, "w") as f:
                        json.dump(
                            {
                                "task": task_name,
                                "model": args.model,
                                "best_config": best_config,
                                "configurations": all_generation_examples,
                            },
                            f,
                            indent=2,
                        )

                    print(
                        f"\n  💾 Saved generation examples for {len(all_generation_examples)} configurations to: {examples_path}"
                    )

                # Generate examples for --save-generation-examples, --show-comparisons, or --save-comparisons
                show_comparisons = getattr(args, 'show_comparisons', 0)
                save_comparisons = getattr(args, 'save_comparisons', None)
                need_generation = args.save_generation_examples or show_comparisons > 0 or save_comparisons

                if need_generation:
                    # Save examples only for the best configuration
                    print("\n  📝 Generating example responses for best configuration...")

                    # Get a few test examples to generate from
                    num_examples = min(args.num_generation_examples, len(test_pairs.pairs))
                    example_pairs = test_pairs.pairs[:num_examples]

                    generation_examples = []

                    # Get inference config settings
                    gen_kwargs = get_generate_kwargs()

                    for idx, pair in enumerate(example_pairs):
                        # Create prompt from the question
                        prompt = pair.prompt

                        try:
                            # Generate without steering
                            unsteered_response = model.generate(
                                [[{"role": "user", "content": prompt}]],
                                **get_generate_kwargs(max_new_tokens=100),
                                use_steering=False,
                            )[0]

                            # Recreate best steering vector for generation
                            best_layer_str = str(best_config["layer"])
                            pos_acts_gen = []
                            neg_acts_gen = []

                            # Collect activations again for steering
                            for train_pair in train_pairs.pairs[:20]:  # Use subset for speed
                                updated_pair = collector.collect_for_pair(
                                    train_pair,
                                    layers=[best_layer_str],
                                    aggregation=ActivationAggregationStrategy.MEAN_POOLING,
                                    return_full_sequence=False,
                                    normalize_layers=False,
                                )

                                if (
                                    updated_pair.positive_response.layers_activations
                                    and best_layer_str in updated_pair.positive_response.layers_activations
                                ):
                                    act = updated_pair.positive_response.layers_activations[best_layer_str]
                                    if act is not None:
                                        pos_acts_gen.append(act)

                                if (
                                    updated_pair.negative_response.layers_activations
                                    and best_layer_str in updated_pair.negative_response.layers_activations
                                ):
                                    act = updated_pair.negative_response.layers_activations[best_layer_str]
                                    if act is not None:
                                        neg_acts_gen.append(act)

                            # Create steering vector
                            method_name_gen = args.methods[0] if args.methods else "CAA"
                            steering_method_gen = create_steering_method(method_name_gen, args)
                            steering_vector_gen = steering_method_gen.train_for_layer(pos_acts_gen, neg_acts_gen)

                            # Create SteeringPlan
                            from wisent.core.models.core.atoms import SteeringPlan, SteeringVector

                            steering_vec = SteeringVector(vector=steering_vector_gen, scale=best_config["strength"])
                            steering_plan = SteeringPlan(
                                layers={best_layer_str: steering_vec},
                                layers_description=[f"CAA steering for {task_name}"],
                            )

                            # Generate with steering
                            model.apply_steering(steering_plan)
                            steered_response = model.generate(
                                [[{"role": "user", "content": prompt}]],
                                **get_generate_kwargs(max_new_tokens=100),
                                use_steering=True,
                                steering_plan=steering_plan,
                            )[0]
                            model.detach()

                            generation_examples.append(
                                {
                                    "question": prompt,
                                    "correct_answer": pair.positive_response.model_response,
                                    "incorrect_answer": pair.negative_response.model_response,
                                    "unsteered_generation": unsteered_response,
                                    "steered_generation": steered_response,
                                }
                            )

                            print(f"      Generated example {idx + 1}/{num_examples}")

                        except Exception as e:
                            print(f"      ⚠️ Failed to generate example {idx + 1}: {e}")
                            if args.verbose:
                                import traceback

                                traceback.print_exc()

                    # Save examples to JSON only if --save-generation-examples is set
                    if args.save_generation_examples:
                        examples_path = os.path.join(
                            args.save_best_vector if args.save_best_vector else "./optimization_results",
                            f"{task_name}_generation_examples.json",
                        )
                        os.makedirs(os.path.dirname(examples_path), exist_ok=True)

                        with open(examples_path, "w") as f:
                            json.dump(
                                {
                                    "task": task_name,
                                    "model": args.model,
                                    "best_config": best_config,
                                    "examples": generation_examples,
                                },
                                f,
                                indent=2,
                            )

                        print(f"      💾 Saved {len(generation_examples)} generation examples to: {examples_path}")

                    # Handle --show-comparisons and --save-comparisons flags
                    if (show_comparisons > 0 or save_comparisons) and generation_examples:
                        # Build comparisons list from generation_examples
                        comparisons = []
                        for ex in generation_examples:
                            comparisons.append({
                                "prompt": ex["question"],
                                "baseline_response": ex["unsteered_generation"],
                                "optimized_response": ex["steered_generation"],
                                "correct_answer": ex.get("correct_answer", ""),
                                "incorrect_answer": ex.get("incorrect_answer", ""),
                            })

                        # Save to JSON if requested
                        if save_comparisons:
                            os.makedirs(os.path.dirname(save_comparisons) if os.path.dirname(save_comparisons) else ".", exist_ok=True)
                            with open(save_comparisons, "w") as f:
                                json.dump({
                                    "model": args.model,
                                    "task": task_name,
                                    "best_config": best_config,
                                    "comparisons": comparisons,
                                }, f, indent=2)
                            print(f"      💾 Saved comparisons to: {save_comparisons}")

                        # Display in console if requested
                        if show_comparisons > 0:
                            print(f"\n  📊 Top {min(show_comparisons, len(comparisons))} Baseline vs Optimized Comparisons:\n")
                            for i, comp in enumerate(comparisons[:show_comparisons]):
                                print(f"{'─'*80}")
                                print(f"Comparison {i+1}/{min(show_comparisons, len(comparisons))}")
                                print(f"{'─'*80}")
                                print(f"PROMPT: {comp['prompt'][:200]}{'...' if len(comp['prompt']) > 200 else ''}")
                                print()
                                print(f"BASELINE (unsteered):")
                                print(f"  {comp['baseline_response'][:300]}{'...' if len(comp['baseline_response']) > 300 else ''}")
                                print()
                                print(f"OPTIMIZED (steered):")
                                print(f"  {comp['optimized_response'][:300]}{'...' if len(comp['optimized_response']) > 300 else ''}")
                                print()

            else:
                print("\n  ⚠️  No valid configuration found")
                method_results["CAA"] = {
                    "optimal_layer": 8,
                    "optimal_strength": 1.0,
                    "optimal_strategy": "constant",
                    "optimal_token_aggregation": "last_token",
                    "optimal_prompt_construction": "chat_template",
                    "accuracy": 0.5,
                    "f1": 0.5,
                }

            all_results[task_name] = {
                "methods": method_results,
                "best_method": "CAA",
                "best_layer": method_results["CAA"]["optimal_layer"],
                "best_strength": method_results["CAA"]["optimal_strength"],
                "best_strategy": method_results["CAA"]["optimal_strategy"],
                "best_token_aggregation": method_results["CAA"]["optimal_token_aggregation"],
                "best_prompt_construction": method_results["CAA"]["optimal_prompt_construction"],
            }

            task_time = time.time() - task_start_time
            print(f"\n  ⏱️  Task completed in {task_time:.1f}s (tested {configs_tested} configurations)")

        except Exception as e:
            # NO FALLBACK - raise the error immediately
            print(f"\n❌ Task '{task_name}' optimization failed:")
            print(f"   Error: {e}")
            import traceback

            traceback.print_exc()
            raise

    # Save results
    print(f"\n{'=' * 80}")
    print("📊 COMPREHENSIVE OPTIMIZATION COMPLETE")
    print(f"{'=' * 80}\n")

    results_file = f"./optimization_results/steering_comprehensive_{args.model.replace('/', '_')}.json"
    import os

    os.makedirs(os.path.dirname(results_file), exist_ok=True)

    output_data = {
        "model": args.model,
        "tasks": all_results,
        "methods_tested": args.methods,
        "limit": args.limit,
        "optimization_dimensions": ["layer", "strength", "strategy", "token_aggregation", "prompt_construction"],
    }

    with open(results_file, "w") as f:
        json.dump(output_data, f, indent=2)

    print(f"✅ Results saved to: {results_file}\n")

    # Print summary
    print("📋 SUMMARY BY TASK:")
    print("-" * 140)
    for task_name, config in all_results.items():
        print(
            f"  {task_name:20s} | L{config['best_layer']:2d} S{config['best_strength']:.1f} | {config['best_strategy']:12s} | T:{config['best_token_aggregation']:12s} | P:{config['best_prompt_construction']:18s}"
        )
    print("-" * 140 + "\n")

    # Store results in optimization cache
    save_as_default = getattr(args, "save_as_default", False)
    print("💾 Storing results in optimization cache...")
    for task_name, config in all_results.items():
        # Skip results that came from cache
        if config.get("from_cache"):
            continue

        # Get best score from methods if available
        best_score = 0.0
        if "methods" in config and "CAA" in config["methods"]:
            best_score = config["methods"]["CAA"].get("accuracy", 0.0)

        # Get best method name
        best_method_name = config.get("best_method", "CAA")
        
        # Get method-specific parameters from the best config
        method_config = config.get("methods", {}).get(best_method_name, {})
        
        cache_key = store_optimization(
            model=args.model,
            task=task_name,
            layer=config["best_layer"],
            strength=config["best_strength"],
            method=best_method_name,
            token_aggregation=config.get("best_token_aggregation", "last_token"),
            prompt_strategy=config.get("best_prompt_construction", "chat_template"),
            strategy=config.get("best_strategy", "constant"),
            score=best_score,
            metric="accuracy",
            metadata={"limit": args.limit},
            set_as_default=save_as_default,
            # PRISM parameters
            num_directions=method_config.get("num_directions", 1),
            direction_weighting=method_config.get("direction_weighting", "primary_only"),
            retain_weight=method_config.get("retain_weight", 0.0),
            independence_weight=method_config.get("independence_weight", 0.05),
            prism_optimization_steps=method_config.get("optimization_steps", 100),
            # PULSE parameters
            sensor_layer=method_config.get("sensor_layer", -1),
            steering_layers=method_config.get("steering_layers", ""),
            condition_threshold=method_config.get("condition_threshold", 0.5),
            gate_temperature=method_config.get("gate_temperature", 0.5),
            per_layer_scaling=method_config.get("per_layer_scaling", True),
            use_entropy_scaling=method_config.get("use_entropy_scaling", False),
            max_alpha=method_config.get("max_alpha", 2.0),
            # TITAN parameters
            gate_hidden_dim=method_config.get("gate_hidden_dim", 64),
            intensity_hidden_dim=method_config.get("intensity_hidden_dim", 32),
            behavior_weight=method_config.get("behavior_weight", 1.0),
            sparse_weight=method_config.get("sparse_weight", 0.05),
            titan_optimization_steps=method_config.get("titan_optimization_steps", 200),
            titan_learning_rate=method_config.get("titan_learning_rate", 0.005),
            # Store all method params as generic dict
            method_params=method_config,
        )
        print(f"   ✓ Cached {task_name}: {cache_key}")

    if save_as_default:
        print("   ✓ Results set as default configurations")

    # Return results for programmatic access
    return {
        "model": args.model,
        "action": "comprehensive",
        "methods_tested": args.methods,
        "tasks_optimized": list(all_results.keys()),
        "results": all_results,
        "results_file": results_file,
        "optimization_dimensions": ["layer", "strength", "strategy", "token_aggregation", "prompt_construction"],
    }


def get_strategy_weight(strategy: str, position: float) -> float:
    """
    Calculate steering weight based on strategy and token position.

    Args:
        strategy: Steering strategy name
        position: Token position as fraction (0.0 = start, 1.0 = end)

    Returns:
        Weight multiplier for steering vector
    """
    if strategy == "last_only":
        return 1.0 if position >= 0.9 else 0.0
    if strategy == "first_only":
        return 1.0 if position <= 0.1 else 0.0
    if strategy == "all_equal":
        return 1.0
    if strategy == "exponential_decay":
        return np.exp(-3.0 * position)  # Decay rate of 3
    if strategy == "exponential_growth":
        return np.exp(3.0 * position)
    if strategy == "linear_decay":
        return 1.0 - position
    if strategy == "linear_growth":
        return position
    return 1.0  # Default to all_equal


def execute_compare_methods(args, model, loader):
    """Execute method comparison - currently only CAA is implemented."""
    import matplotlib.pyplot as plt
    from wisent_plots import LineChart

    from wisent.core.activations.activations_collector import ActivationCollector
    from wisent.core.activations.core.atoms import ActivationAggregationStrategy
    from wisent.core.models.core.atoms import SteeringPlan, SteeringVector
    from wisent.core.cli.steering_method_trainer import create_steering_method

    # Check for cached results if --use-cached is specified
    use_cached = getattr(args, "use_cached", False)
    save_as_default = getattr(args, "save_as_default", False)

    if use_cached:
        print(f"\n📦 Checking optimization cache for {args.task}...")
        for method in args.methods:
            cached = get_cached_optimization(args.model, args.task, method)
            if cached:
                print(
                    f"   ✓ Found cached result for {method}: layer={cached.layer}, strength={cached.strength}, score={cached.score:.3f}"
                )
                return {
                    "model": args.model,
                    "action": "compare-methods",
                    "task": args.task,
                    "best_method": method,
                    "best_layer": cached.layer,
                    "best_strength": cached.strength,
                    "best_score": cached.score,
                    "from_cache": True,
                }
        print("   No cached results found. Running optimization...")

    print(f"🔍 Comparing steering methods for task: {args.task}\n")
    print(f"   Methods: {', '.join(args.methods)}")
    print(f"   Limit: {args.limit} samples")
    print(f"   Layer: {args.layer}")
    print(f"   Strength: {args.strength}\n")

    # Load task data
    print("📊 Loading task data...")
    result = loader._load_one_task(
        task_name=args.task, split_ratio=0.8, seed=42, limit=args.limit, training_limit=None, testing_limit=None
    )

    train_pairs = result["train_qa_pairs"]
    test_pairs = result["test_qa_pairs"]
    print(f"   ✓ Loaded {len(train_pairs.pairs)} train, {len(test_pairs.pairs)} test pairs\n")

    # Initialize evaluator
    EvaluatorRotator.discover_evaluators("wisent.core.evaluators.benchmark_specific")
    evaluator = EvaluatorRotator(evaluator=None, task_name=args.task)
    print(f"   ✓ Using evaluator: {evaluator._plugin.name}\n")

    # Collect activations once for all methods
    layer_str = str(args.layer)
    collector = ActivationCollector(model=model, store_device="cpu")

    print("🎯 Collecting training activations (ONCE)...")
    pos_acts = []
    neg_acts = []

    for i, pair in enumerate(train_pairs.pairs):
        if i % 10 == 0:
            print(f"   Processing train pair {i + 1}/{len(train_pairs.pairs)}...", end="\r")

        updated_pair = collector.collect_for_pair(
            pair,
            layers=[layer_str],
            aggregation=ActivationAggregationStrategy.MEAN_POOLING,
            return_full_sequence=False,
            normalize_layers=False,
        )

        if (
            updated_pair.positive_response.layers_activations
            and layer_str in updated_pair.positive_response.layers_activations
        ):
            act = updated_pair.positive_response.layers_activations[layer_str]
            if act is not None:
                pos_acts.append(act)

        if (
            updated_pair.negative_response.layers_activations
            and layer_str in updated_pair.negative_response.layers_activations
        ):
            act = updated_pair.negative_response.layers_activations[layer_str]
            if act is not None:
                neg_acts.append(act)

    print(f"   Processing train pair {len(train_pairs.pairs)}/{len(train_pairs.pairs)}... Done!")
    print(f"   ✓ Collected {len(pos_acts)} positive, {len(neg_acts)} negative activations\n")

    # Test each method
    print("🧪 Testing methods...")
    method_results = {}

    # Only CAA is implemented for now
    if "CAA" in args.methods:
        print("\n   Testing CAA method...")

        # Train steering vector using selected method
        method_name = args.methods[0] if args.methods else "CAA"
        steering_method = create_steering_method(method_name, args)
        steering_vector = steering_method.train_for_layer(pos_acts, neg_acts)

        # Create steering plan
        steering_vec = SteeringVector(vector=steering_vector, scale=args.strength)
        steering_plan = SteeringPlan(
            layers={layer_str: steering_vec},
            layers_description=[f"CAA steering layer={args.layer}, strength={args.strength}"],
        )

        # Apply steering and evaluate
        model.apply_steering(steering_plan)

        test_scores = []
        detailed_results = []
        for pair in test_pairs.pairs:
            choices = [pair.negative_response.model_response, pair.positive_response.model_response]
            expected = pair.positive_response.model_response
            test_code = pair.metadata.get("test_code") if pair.metadata else None

            eval_result = evaluator.evaluate(
                response="",
                expected=expected,
                model=model,
                question=pair.prompt,
                choices=choices,
                steering_plan=steering_plan,
                test_code=test_code,
                task_name=args.task,
            )

            is_correct = eval_result.ground_truth == "TRUTHFUL"
            test_scores.append(1.0 if is_correct else 0.0)

            # Save full evaluation details
            detailed_results.append(
                {
                    "question": pair.prompt,
                    "choices": choices,
                    "expected": expected,
                    "ground_truth": eval_result.ground_truth,
                    "method_used": eval_result.method_used,
                    "confidence": eval_result.confidence,
                    "details": eval_result.details,
                    "meta": dict(eval_result.meta) if eval_result.meta else {},
                    "is_correct": is_correct,
                }
            )

        model.clear_steering()

        caa_accuracy = np.mean(test_scores) if len(test_scores) > 0 else 0.0
        method_results["CAA"] = {
            "accuracy": caa_accuracy,
            "num_test_samples": len(test_scores),
            "detailed_results": detailed_results,
        }

        print(f"      ✓ CAA: accuracy={caa_accuracy:.3f}")

    # Other methods are not yet implemented
    for method in args.methods:
        if method not in ["CAA"]:
            print(f"      ⚠️  {method}: not yet implemented")
            method_results[method] = {"accuracy": 0.0, "status": "not_implemented"}

    # Save results
    print(f"\n{'=' * 80}")
    print("📊 METHOD COMPARISON COMPLETE")
    print(f"{'=' * 80}\n")

    results_file = f"./optimization_results/steering_compare_methods_{args.task}_{args.model.replace('/', '_')}.json"
    import os

    os.makedirs(os.path.dirname(results_file), exist_ok=True)

    output_data = {
        "model": args.model,
        "task": args.task,
        "layer": args.layer,
        "strength": args.strength,
        "methods": method_results,
        "limit": args.limit,
    }

    with open(results_file, "w") as f:
        json.dump(output_data, f, indent=2)

    print(f"✅ Results saved to: {results_file}\n")

    # Create comparison plot if we have results
    implemented_methods = [m for m in method_results if method_results[m].get("accuracy", 0) > 0]
    if len(implemented_methods) > 1 and args.save_plot:
        plot_path_svg = f"steering_compare_methods_{args.task}_{args.model.replace('/', '_')}.svg"
        plot_path_png = f"steering_compare_methods_{args.task}_{args.model.replace('/', '_')}.png"

        method_names = list(implemented_methods)
        accuracies = [method_results[m]["accuracy"] for m in method_names]

        chart = LineChart(style=1, figsize=(10, 6), show_markers=True)
        fig, ax = plt.subplots(1, 1, figsize=(10, 6))

        ax.bar(method_names, accuracies, color="#3498db", alpha=0.8)
        ax.set_xlabel("Steering Method")
        ax.set_ylabel("Accuracy")
        ax.set_title(f"Steering Method Comparison\n{args.model} on {args.task}")
        ax.set_ylim(0, 1)

        fig.savefig(plot_path_svg, format="svg", bbox_inches="tight")
        fig.savefig(plot_path_png, dpi=150, bbox_inches="tight")
        plt.close(fig)

        print("💾 Comparison plot saved to:")
        print(f"   SVG: {plot_path_svg}")
        print(f"   PNG: {plot_path_png}\n")

    # Store best result in cache
    save_as_default = getattr(args, "save_as_default", False)
    best_method = max(method_results.keys(), key=lambda m: method_results[m].get("accuracy", 0))
    best_accuracy = method_results[best_method].get("accuracy", 0)

    if best_accuracy > 0:
        print("💾 Storing best result in optimization cache...")
        cache_key = store_optimization(
            model=args.model,
            task=args.task,
            layer=args.layer,
            strength=args.strength,
            method=best_method,
            strategy="constant",
            score=best_accuracy,
            metric="accuracy",
            metadata={"limit": args.limit},
            set_as_default=save_as_default,
        )
        print(f"   ✓ Cached: {cache_key}")
        if save_as_default:
            print("   ✓ Set as default configuration")

    return {"action": "compare-methods", "task": args.task, "methods": method_results, "results_file": results_file}


def execute_optimize_layer(args, model, loader):
    """Execute layer optimization - find the best layer for steering."""
    import matplotlib.pyplot as plt
    from wisent_plots import LineChart

    from wisent.core.activations.activations_collector import ActivationCollector
    from wisent.core.activations.core.atoms import ActivationAggregationStrategy
    from wisent.core.models.core.atoms import SteeringPlan, SteeringVector
    from wisent.core.cli.steering_method_trainer import create_steering_method

    # Check for cached results if --use-cached is specified
    use_cached = getattr(args, "use_cached", False)
    save_as_default = getattr(args, "save_as_default", False)

    if use_cached:
        print(f"\n📦 Checking optimization cache for {args.task}/{args.method}...")
        cached = get_cached_optimization(args.model, args.task, args.method)
        if cached:
            print(
                f"   ✓ Found cached result: layer={cached.layer}, strength={cached.strength}, score={cached.score:.3f}"
            )
            return {
                "model": args.model,
                "action": "optimize-layer",
                "task": args.task,
                "method": args.method,
                "best_layer": cached.layer,
                "best_strength": cached.strength,
                "best_accuracy": cached.score,
                "from_cache": True,
            }
        print("   No cached results found. Running optimization...")

    print(f"🎯 Optimizing steering layer for task: {args.task}\n")
    print(f"   Method: {args.method}")
    print(f"   Strength: {args.strength}")
    print(f"   Limit: {args.limit} samples\n")

    # Load task data
    print("📊 Loading task data...")
    result = loader._load_one_task(
        task_name=args.task, split_ratio=0.8, seed=42, limit=args.limit, training_limit=None, testing_limit=None
    )

    train_pairs = result["train_qa_pairs"]
    test_pairs = result["test_qa_pairs"]
    print(f"   ✓ Loaded {len(train_pairs.pairs)} train, {len(test_pairs.pairs)} test pairs\n")

    # Initialize evaluator
    EvaluatorRotator.discover_evaluators("wisent.core.evaluators.benchmark_specific")
    evaluator = EvaluatorRotator(evaluator=None, task_name=args.task)
    print(f"   ✓ Using evaluator: {evaluator._plugin.name}\n")

    # Determine layers to test
    if args.layers:
        layers_to_test = args.layers
    else:
        # Test all layers from 0 to num_layers-1
        layers_to_test = list(range(model.num_layers))

    print(f"🔍 Testing {len(layers_to_test)} layers: {layers_to_test[:5]}{'...' if len(layers_to_test) > 5 else ''}\n")

    collector = ActivationCollector(model=model, store_device="cpu")
    layer_results = {}
    best_layer = None
    best_accuracy = 0.0

    for layer_idx, layer in enumerate(layers_to_test, 1):
        layer_str = str(layer)
        print(f"   [{layer_idx}/{len(layers_to_test)}] Testing layer {layer}...", end=" ")

        try:
            # Collect activations for this layer
            pos_acts = []
            neg_acts = []

            for pair in train_pairs.pairs:
                updated_pair = collector.collect_for_pair(
                    pair,
                    layers=[layer_str],
                    aggregation=ActivationAggregationStrategy.MEAN_POOLING,
                    return_full_sequence=False,
                    normalize_layers=False,
                )

                if (
                    updated_pair.positive_response.layers_activations
                    and layer_str in updated_pair.positive_response.layers_activations
                ):
                    act = updated_pair.positive_response.layers_activations[layer_str]
                    if act is not None:
                        pos_acts.append(act)

                if (
                    updated_pair.negative_response.layers_activations
                    and layer_str in updated_pair.negative_response.layers_activations
                ):
                    act = updated_pair.negative_response.layers_activations[layer_str]
                    if act is not None:
                        neg_acts.append(act)

            if len(pos_acts) == 0 or len(neg_acts) == 0:
                print("⚠️  No activations collected")
                continue

            # Train steering vector using selected method
            steering_method = create_steering_method(args.method, args)
            steering_vector = steering_method.train_for_layer(pos_acts, neg_acts)
            if False:  # Compatibility placeholder
                print(f"⚠️  Method {args.method} not supported")
                continue

            # Create steering plan
            steering_vec = SteeringVector(vector=steering_vector, scale=args.strength)
            steering_plan = SteeringPlan(
                layers={layer_str: steering_vec}, layers_description=[f"{args.method} steering layer={layer}"]
            )

            # Evaluate
            model.apply_steering(steering_plan)

            test_scores = []
            detailed_results = []
            for pair in test_pairs.pairs:
                choices = [pair.negative_response.model_response, pair.positive_response.model_response]
                expected = pair.positive_response.model_response
                test_code = pair.metadata.get("test_code") if pair.metadata else None

                eval_result = evaluator.evaluate(
                    response="",
                    expected=expected,
                    model=model,
                    question=pair.prompt,
                    choices=choices,
                    steering_plan=steering_plan,
                    test_code=test_code,
                    task_name=task_name,
                )

                is_correct = eval_result.ground_truth == "TRUTHFUL"
                test_scores.append(1.0 if is_correct else 0.0)

                # Save full evaluation details
                detailed_results.append(
                    {
                        "question": pair.prompt,
                        "choices": choices,
                        "expected": expected,
                        "ground_truth": eval_result.ground_truth,
                        "method_used": eval_result.method_used,
                        "confidence": eval_result.confidence,
                        "details": eval_result.details,
                        "meta": dict(eval_result.meta) if eval_result.meta else {},
                        "is_correct": is_correct,
                    }
                )

            model.clear_steering()

            accuracy = np.mean(test_scores) if len(test_scores) > 0 else 0.0
            layer_results[layer] = {
                "accuracy": accuracy,
                "num_test_samples": len(test_scores),
                "detailed_results": detailed_results,
            }

            print(f"accuracy={accuracy:.3f}")

            if accuracy > best_accuracy:
                best_accuracy = accuracy
                best_layer = layer

        except Exception as e:
            print(f"❌ Error: {e}")
            if args.verbose:
                import traceback

                traceback.print_exc()

    # Results
    print(f"\n{'=' * 80}")
    print("📊 LAYER OPTIMIZATION COMPLETE")
    print(f"{'=' * 80}")
    print(f"   Best layer: {best_layer}")
    print(f"   Best accuracy: {best_accuracy:.4f}")
    print(f"{'=' * 80}\n")

    # Save results
    results_file = f"./optimization_results/steering_optimize_layer_{args.task}_{args.model.replace('/', '_')}.json"
    import os

    os.makedirs(os.path.dirname(results_file), exist_ok=True)

    output_data = {
        "model": args.model,
        "task": args.task,
        "method": args.method,
        "strength": args.strength,
        "best_layer": best_layer,
        "best_accuracy": best_accuracy,
        "layer_results": {str(k): v for k, v in layer_results.items()},
        "limit": args.limit,
    }

    with open(results_file, "w") as f:
        json.dump(output_data, f, indent=2)

    print(f"✅ Results saved to: {results_file}\n")

    # Create plot
    if args.save_plot and len(layer_results) > 0:
        plot_path_svg = f"steering_optimize_layer_{args.task}_{args.model.replace('/', '_')}.svg"
        plot_path_png = f"steering_optimize_layer_{args.task}_{args.model.replace('/', '_')}.png"

        layers = sorted(layer_results.keys())
        accuracies = [layer_results[l]["accuracy"] for l in layers]

        chart = LineChart(style=1, figsize=(10, 6), show_markers=True)
        fig, ax = plt.subplots(1, 1, figsize=(10, 6))

        chart.plot_multiple(
            x=layers,
            y_series=[accuracies],
            labels=["Accuracy"],
            title=f"Layer Optimization\n{args.model} on {args.task}",
            xlabel="Layer",
            ylabel="Accuracy",
            fig=fig,
            ax=ax,
            output_format="png",
        )

        # Add vertical line for optimal layer
        ax.axvline(
            x=best_layer, color="#2ecc71", linestyle="--", linewidth=2, label=f"Best: Layer {best_layer}", alpha=0.7
        )
        ax.legend()

        fig.savefig(plot_path_svg, format="svg", bbox_inches="tight")
        fig.savefig(plot_path_png, dpi=150, bbox_inches="tight")
        plt.close(fig)

        print("💾 Layer optimization plot saved to:")
        print(f"   SVG: {plot_path_svg}")
        print(f"   PNG: {plot_path_png}\n")

    # Store result in cache
    save_as_default = getattr(args, "save_as_default", False)
    if best_layer is not None and best_accuracy > 0:
        print("💾 Storing result in optimization cache...")
        cache_key = store_optimization(
            model=args.model,
            task=args.task,
            layer=best_layer,
            strength=args.strength,
            method=args.method,
            strategy="constant",
            score=best_accuracy,
            metric="accuracy",
            metadata={"limit": args.limit},
            set_as_default=save_as_default,
        )
        print(f"   ✓ Cached: {cache_key}")
        if save_as_default:
            print("   ✓ Set as default configuration")

    return {
        "action": "optimize-layer",
        "task": args.task,
        "method": args.method,
        "best_layer": best_layer,
        "best_accuracy": best_accuracy,
        "results_file": results_file,
    }


def execute_optimize_strength(args, model, loader):
    """Execute strength optimization - find the best steering strength."""
    import matplotlib.pyplot as plt
    from wisent_plots import LineChart

    from wisent.core.activations.activations_collector import ActivationCollector
    from wisent.core.activations.core.atoms import ActivationAggregationStrategy
    from wisent.core.models.core.atoms import SteeringPlan, SteeringVector
    from wisent.core.cli.steering_method_trainer import create_steering_method

    # Check for cached results if --use-cached is specified
    use_cached = getattr(args, "use_cached", False)
    save_as_default = getattr(args, "save_as_default", False)

    if use_cached:
        print(f"\n📦 Checking optimization cache for {args.task}/{args.method}...")
        cached = get_cached_optimization(args.model, args.task, args.method)
        if cached:
            print(
                f"   ✓ Found cached result: layer={cached.layer}, strength={cached.strength}, score={cached.score:.3f}"
            )
            return {
                "model": args.model,
                "action": "optimize-strength",
                "task": args.task,
                "method": args.method,
                "best_layer": cached.layer,
                "best_strength": cached.strength,
                "best_accuracy": cached.score,
                "from_cache": True,
            }
        print("   No cached results found. Running optimization...")

    print(f"💪 Optimizing steering strength for task: {args.task}\n")
    print(f"   Method: {args.method}")
    print(f"   Layer: {args.layer}")
    print(f"   Strength range: {args.strength_range[0]} to {args.strength_range[1]}")
    print(f"   Num steps: {args.num_strength_steps}")
    print(f"   Limit: {args.limit} samples\n")

    # Load task data
    print("📊 Loading task data...")
    result = loader._load_one_task(
        task_name=args.task, split_ratio=0.8, seed=42, limit=args.limit, training_limit=None, testing_limit=None
    )

    train_pairs = result["train_qa_pairs"]
    test_pairs = result["test_qa_pairs"]
    print(f"   ✓ Loaded {len(train_pairs.pairs)} train, {len(test_pairs.pairs)} test pairs\n")

    # Initialize evaluator
    EvaluatorRotator.discover_evaluators("wisent.core.evaluators.benchmark_specific")
    evaluator = EvaluatorRotator(evaluator=None, task_name=args.task)
    print(f"   ✓ Using evaluator: {evaluator._plugin.name}\n")

    # Collect activations ONCE
    layer_str = str(args.layer)
    collector = ActivationCollector(model=model, store_device="cpu")

    print("🎯 Collecting training activations (ONCE)...")
    pos_acts = []
    neg_acts = []

    for i, pair in enumerate(train_pairs.pairs):
        if i % 10 == 0:
            print(f"   Processing train pair {i + 1}/{len(train_pairs.pairs)}...", end="\r")

        updated_pair = collector.collect_for_pair(
            pair,
            layers=[layer_str],
            aggregation=ActivationAggregationStrategy.MEAN_POOLING,
            return_full_sequence=False,
            normalize_layers=False,
        )

        if (
            updated_pair.positive_response.layers_activations
            and layer_str in updated_pair.positive_response.layers_activations
        ):
            act = updated_pair.positive_response.layers_activations[layer_str]
            if act is not None:
                pos_acts.append(act)

        if (
            updated_pair.negative_response.layers_activations
            and layer_str in updated_pair.negative_response.layers_activations
        ):
            act = updated_pair.negative_response.layers_activations[layer_str]
            if act is not None:
                neg_acts.append(act)

    print(f"   Processing train pair {len(train_pairs.pairs)}/{len(train_pairs.pairs)}... Done!")
    print(f"   ✓ Collected {len(pos_acts)} positive, {len(neg_acts)} negative activations\n")

    # Train steering vector using selected method
    steering_method = create_steering_method(args.method, args)
    steering_vector = steering_method.train_for_layer(pos_acts, neg_acts)
    if False:  # Compatibility placeholder
        print(f"❌ Method {args.method} not supported")
        return {
            "action": "optimize-strength",
            "task": args.task,
            "method": args.method,
            "status": "method_not_supported",
        }

    # Generate strength values to test
    min_strength, max_strength = args.strength_range
    strengths_to_test = np.linspace(min_strength, max_strength, args.num_strength_steps)

    print(
        f"🔍 Testing {len(strengths_to_test)} strength values: {strengths_to_test[0]:.2f} to {strengths_to_test[-1]:.2f}\n"
    )

    strength_results = {}
    best_strength = None
    best_accuracy = 0.0

    for strength_idx, strength in enumerate(strengths_to_test, 1):
        print(f"   [{strength_idx}/{len(strengths_to_test)}] Testing strength {strength:.2f}...", end=" ")

        try:
            # Create steering plan with this strength
            steering_vec = SteeringVector(vector=steering_vector, scale=float(strength))
            steering_plan = SteeringPlan(
                layers={layer_str: steering_vec}, layers_description=[f"{args.method} steering strength={strength:.2f}"]
            )

            # Evaluate
            model.apply_steering(steering_plan)

            test_scores = []
            detailed_results = []
            for pair in test_pairs.pairs:
                choices = [pair.negative_response.model_response, pair.positive_response.model_response]
                expected = pair.positive_response.model_response
                test_code = pair.metadata.get("test_code") if pair.metadata else None

                eval_result = evaluator.evaluate(
                    response="",
                    expected=expected,
                    model=model,
                    question=pair.prompt,
                    choices=choices,
                    steering_plan=steering_plan,
                    test_code=test_code,
                    task_name=task_name,
                )

                is_correct = eval_result.ground_truth == "TRUTHFUL"
                test_scores.append(1.0 if is_correct else 0.0)

                # Save full evaluation details
                detailed_results.append(
                    {
                        "question": pair.prompt,
                        "choices": choices,
                        "expected": expected,
                        "ground_truth": eval_result.ground_truth,
                        "method_used": eval_result.method_used,
                        "confidence": eval_result.confidence,
                        "details": eval_result.details,
                        "meta": dict(eval_result.meta) if eval_result.meta else {},
                        "is_correct": is_correct,
                    }
                )

            model.clear_steering()

            accuracy = np.mean(test_scores) if len(test_scores) > 0 else 0.0
            strength_results[float(strength)] = {
                "accuracy": accuracy,
                "num_test_samples": len(test_scores),
                "detailed_results": detailed_results,
            }

            print(f"accuracy={accuracy:.3f}")

            if accuracy > best_accuracy:
                best_accuracy = accuracy
                best_strength = float(strength)

        except Exception as e:
            print(f"❌ Error: {e}")
            if args.verbose:
                import traceback

                traceback.print_exc()

    # Results
    print(f"\n{'=' * 80}")
    print("📊 STRENGTH OPTIMIZATION COMPLETE")
    print(f"{'=' * 80}")
    print(f"   Best strength: {best_strength:.2f}")
    print(f"   Best accuracy: {best_accuracy:.4f}")
    print(f"{'=' * 80}\n")

    # Save results
    results_file = f"./optimization_results/steering_optimize_strength_{args.task}_{args.model.replace('/', '_')}.json"
    import os

    os.makedirs(os.path.dirname(results_file), exist_ok=True)

    output_data = {
        "model": args.model,
        "task": args.task,
        "method": args.method,
        "layer": args.layer,
        "best_strength": best_strength,
        "best_accuracy": best_accuracy,
        "strength_results": {str(k): v for k, v in strength_results.items()},
        "limit": args.limit,
    }

    with open(results_file, "w") as f:
        json.dump(output_data, f, indent=2)

    print(f"✅ Results saved to: {results_file}\n")

    # Create plot
    if args.save_plot and len(strength_results) > 0:
        plot_path_svg = f"steering_optimize_strength_{args.task}_{args.model.replace('/', '_')}.svg"
        plot_path_png = f"steering_optimize_strength_{args.task}_{args.model.replace('/', '_')}.png"

        strengths = sorted(strength_results.keys())
        accuracies = [strength_results[s]["accuracy"] for s in strengths]

        chart = LineChart(style=1, figsize=(10, 6), show_markers=True)
        fig, ax = plt.subplots(1, 1, figsize=(10, 6))

        chart.plot_multiple(
            x=strengths,
            y_series=[accuracies],
            labels=["Accuracy"],
            title=f"Strength Optimization\n{args.model} on {args.task}",
            xlabel="Steering Strength",
            ylabel="Accuracy",
            fig=fig,
            ax=ax,
            output_format="png",
        )

        # Add vertical line for optimal strength
        ax.axvline(
            x=best_strength, color="#2ecc71", linestyle="--", linewidth=2, label=f"Best: {best_strength:.2f}", alpha=0.7
        )
        ax.legend()

        fig.savefig(plot_path_svg, format="svg", bbox_inches="tight")
        fig.savefig(plot_path_png, dpi=150, bbox_inches="tight")
        plt.close(fig)

        print("💾 Strength optimization plot saved to:")
        print(f"   SVG: {plot_path_svg}")
        print(f"   PNG: {plot_path_png}\n")

    # Store result in cache
    save_as_default = getattr(args, "save_as_default", False)
    if best_strength is not None and best_accuracy > 0:
        print("💾 Storing result in optimization cache...")
        cache_key = store_optimization(
            model=args.model,
            task=args.task,
            layer=args.layer,
            strength=best_strength,
            method=args.method,
            strategy="constant",
            score=best_accuracy,
            metric="accuracy",
            metadata={"limit": args.limit, "strength_range": args.strength_range},
            set_as_default=save_as_default,
        )
        print(f"   ✓ Cached: {cache_key}")
        if save_as_default:
            print("   ✓ Set as default configuration")

    return {
        "action": "optimize-strength",
        "task": args.task,
        "method": args.method,
        "best_strength": best_strength,
        "best_accuracy": best_accuracy,
        "results_file": results_file,
    }


def execute_auto(args, model, loader):
    """Execute automatic optimization - optimizes layer AND strength together."""
    import matplotlib.pyplot as plt

    from wisent.core.activations.activations_collector import ActivationCollector
    from wisent.core.activations.core.atoms import ActivationAggregationStrategy
    from wisent.core.models.core.atoms import SteeringPlan, SteeringVector
    from wisent.core.cli.steering_method_trainer import create_steering_method

    # Check for cached results if --use-cached is specified
    use_cached = getattr(args, "use_cached", False)
    save_as_default = getattr(args, "save_as_default", False)
    task_name = args.task or "default"

    if use_cached:
        print(f"\n📦 Checking optimization cache for {task_name}...")
        for method in args.methods:
            cached = get_cached_optimization(args.model, task_name, method)
            if cached:
                print(
                    f"   ✓ Found cached result for {method}: layer={cached.layer}, strength={cached.strength}, score={cached.score:.3f}"
                )
                return {
                    "model": args.model,
                    "action": "auto",
                    "task": task_name,
                    "best_method": method,
                    "best_layer": cached.layer,
                    "best_strength": cached.strength,
                    "best_accuracy": cached.score,
                    "from_cache": True,
                }
        print("   No cached results found. Running optimization...")

    print("🤖 Running automatic steering optimization...\n")
    print(f"   Task: {args.task}")
    print(f"   Methods: {', '.join(args.methods)}")
    print(f"   Strength range: {args.strength_range}")
    print(f"   Limit: {args.limit} samples\n")

    # Load task data
    print("📊 Loading task data...")
    result = loader._load_one_task(
        task_name=args.task, split_ratio=0.8, seed=42, limit=args.limit, training_limit=None, testing_limit=None
    )

    train_pairs = result["train_qa_pairs"]
    test_pairs = result["test_qa_pairs"]
    print(f"   ✓ Loaded {len(train_pairs.pairs)} train, {len(test_pairs.pairs)} test pairs\n")

    # Initialize evaluator
    EvaluatorRotator.discover_evaluators("wisent.core.evaluators.benchmark_specific")
    evaluator = EvaluatorRotator(evaluator=None, task_name=args.task)
    print(f"   ✓ Using evaluator: {evaluator._plugin.name}\n")

    # Define search space
    layers_to_test = list(
        range(max(0, model.num_layers // 2 - 2), min(model.num_layers, model.num_layers // 2 + 3))
    )  # Test 5 layers around middle
    min_strength, max_strength = args.strength_range
    strengths_to_test = np.linspace(min_strength, max_strength, 5)  # 5 strength values

    print("🔍 Auto-optimizing layer and strength...")
    print(f"   Testing {len(layers_to_test)} layers: {layers_to_test}")
    print(f"   Testing {len(strengths_to_test)} strengths: {strengths_to_test[0]:.2f} to {strengths_to_test[-1]:.2f}")
    print(f"   Total configurations: {len(layers_to_test) * len(strengths_to_test)}\n")

    collector = ActivationCollector(model=model, store_device="cpu")
    all_results = {}
    best_config = None
    best_accuracy = 0.0

    config_count = 0
    total_configs = len(layers_to_test) * len(strengths_to_test)

    for layer in layers_to_test:
        layer_str = str(layer)

        # Collect activations for this layer
        print(f"   Collecting activations for layer {layer}...")
        pos_acts = []
        neg_acts = []

        for pair in train_pairs.pairs:
            updated_pair = collector.collect_for_pair(
                pair,
                layers=[layer_str],
                aggregation=ActivationAggregationStrategy.MEAN_POOLING,
                return_full_sequence=False,
                normalize_layers=False,
            )

            if (
                updated_pair.positive_response.layers_activations
                and layer_str in updated_pair.positive_response.layers_activations
            ):
                act = updated_pair.positive_response.layers_activations[layer_str]
                if act is not None:
                    pos_acts.append(act)

            if (
                updated_pair.negative_response.layers_activations
                and layer_str in updated_pair.negative_response.layers_activations
            ):
                act = updated_pair.negative_response.layers_activations[layer_str]
                if act is not None:
                    neg_acts.append(act)

        if len(pos_acts) == 0 or len(neg_acts) == 0:
            print(f"      ⚠️  No activations collected for layer {layer}")
            continue

        # Train steering vector for this layer using selected method
        method_name = args.methods[0] if args.methods else "CAA"
        steering_method = create_steering_method(method_name, args)
        steering_vector = steering_method.train_for_layer(pos_acts, neg_acts)
        if False:  # Compatibility placeholder
            print("      ⚠️  Only CAA method is supported")
            continue

        # Test different strengths for this layer
        for strength in strengths_to_test:
            config_count += 1
            print(f"      [{config_count}/{total_configs}] Layer {layer}, Strength {strength:.2f}...", end=" ")

            try:
                # Create steering plan
                steering_vec = SteeringVector(vector=steering_vector, scale=float(strength))
                steering_plan = SteeringPlan(
                    layers={layer_str: steering_vec}, layers_description=[f"CAA layer={layer}, strength={strength:.2f}"]
                )

                # Evaluate
                model.apply_steering(steering_plan)

                test_scores = []
                detailed_results = []
                for pair in test_pairs.pairs:
                    choices = [pair.negative_response.model_response, pair.positive_response.model_response]
                    expected = pair.positive_response.model_response
                    test_code = pair.metadata.get("test_code") if pair.metadata else None

                    eval_result = evaluator.evaluate(
                        response="",
                        expected=expected,
                        model=model,
                        question=pair.prompt,
                        choices=choices,
                        steering_plan=steering_plan,
                        test_code=test_code,
                        task_name=task_name,
                    )

                    is_correct = eval_result.ground_truth == "TRUTHFUL"
                    test_scores.append(1.0 if is_correct else 0.0)

                    # Save full evaluation details
                    detailed_results.append(
                        {
                            "question": pair.prompt,
                            "choices": choices,
                            "expected": expected,
                            "ground_truth": eval_result.ground_truth,
                            "method_used": eval_result.method_used,
                            "confidence": eval_result.confidence,
                            "details": eval_result.details,
                            "meta": dict(eval_result.meta) if eval_result.meta else {},
                            "is_correct": is_correct,
                        }
                    )

                model.clear_steering()

                accuracy = np.mean(test_scores) if len(test_scores) > 0 else 0.0
                all_results[(layer, float(strength))] = {
                    "accuracy": accuracy,
                    "num_test_samples": len(test_scores),
                    "detailed_results": detailed_results,
                }

                print(f"accuracy={accuracy:.3f}")

                if accuracy > best_accuracy:
                    best_accuracy = accuracy
                    best_config = {"layer": layer, "strength": float(strength), "accuracy": accuracy}

            except Exception as e:
                print(f"❌ Error: {e}")
                if args.verbose:
                    import traceback

                    traceback.print_exc()

    # Results
    print(f"\n{'=' * 80}")
    print("📊 AUTO OPTIMIZATION COMPLETE")
    print(f"{'=' * 80}")
    if best_config:
        print(f"   Best layer: {best_config['layer']}")
        print(f"   Best strength: {best_config['strength']:.2f}")
        print(f"   Best accuracy: {best_config['accuracy']:.4f}")
    else:
        print("   ⚠️  No valid configuration found")
    print(f"{'=' * 80}\n")

    # Save results
    results_file = f"./optimization_results/steering_auto_{args.task}_{args.model.replace('/', '_')}.json"
    import os

    os.makedirs(os.path.dirname(results_file), exist_ok=True)

    output_data = {
        "model": args.model,
        "task": args.task,
        "methods": args.methods,
        "best_config": best_config,
        "all_results": {f"layer{k[0]}_strength{k[1]:.2f}": v for k, v in all_results.items()},
        "limit": args.limit,
    }

    with open(results_file, "w") as f:
        json.dump(output_data, f, indent=2)

    print(f"✅ Results saved to: {results_file}\n")

    # Create heatmap plot
    if args.save_plot and len(all_results) > 0 and best_config:
        plot_path_svg = f"steering_auto_{args.task}_{args.model.replace('/', '_')}.svg"
        plot_path_png = f"steering_auto_{args.task}_{args.model.replace('/', '_')}.png"

        # Prepare data for heatmap
        layers = sorted(set(k[0] for k in all_results))
        strengths = sorted(set(k[1] for k in all_results))

        # Create accuracy matrix
        accuracy_matrix = np.zeros((len(strengths), len(layers)))
        for i, strength in enumerate(strengths):
            for j, layer in enumerate(layers):
                if (layer, strength) in all_results:
                    accuracy_matrix[i, j] = all_results[(layer, strength)]["accuracy"]

        fig, ax = plt.subplots(1, 1, figsize=(10, 8))

        im = ax.imshow(accuracy_matrix, cmap="viridis", aspect="auto")

        # Set ticks and labels
        ax.set_xticks(np.arange(len(layers)))
        ax.set_yticks(np.arange(len(strengths)))
        ax.set_xticklabels(layers)
        ax.set_yticklabels([f"{s:.2f}" for s in strengths])

        # Labels
        ax.set_xlabel("Layer")
        ax.set_ylabel("Strength")
        ax.set_title(f"Auto Optimization Heatmap\n{args.model} on {args.task}")

        # Colorbar
        cbar = plt.colorbar(im, ax=ax)
        cbar.set_label("Accuracy", rotation=270, labelpad=15)

        # Mark best configuration
        best_layer_idx = layers.index(best_config["layer"])
        best_strength_idx = strengths.index(best_config["strength"])
        ax.plot(
            best_layer_idx,
            best_strength_idx,
            "r*",
            markersize=20,
            label=f"Best: L{best_config['layer']}, S{best_config['strength']:.2f}",
        )
        ax.legend()

        fig.savefig(plot_path_svg, format="svg", bbox_inches="tight")
        fig.savefig(plot_path_png, dpi=150, bbox_inches="tight")
        plt.close(fig)

        print("💾 Auto optimization heatmap saved to:")
        print(f"   SVG: {plot_path_svg}")
        print(f"   PNG: {plot_path_png}\n")

    # Store result in cache
    save_as_default = getattr(args, "save_as_default", False)
    if best_config and best_config.get("accuracy", 0) > 0:
        print("💾 Storing result in optimization cache...")
        cache_key = store_optimization(
            model=args.model,
            task=args.task or "auto",
            layer=best_config["layer"],
            strength=best_config["strength"],
            method=best_config.get("method", "CAA"),
            strategy=best_config.get("strategy", "constant"),
            score=best_config["accuracy"],
            metric="accuracy",
            metadata={"limit": args.limit, "strength_range": list(args.strength_range)},
            set_as_default=save_as_default,
        )
        print(f"   ✓ Cached: {cache_key}")
        if save_as_default:
            print("   ✓ Set as default configuration")

    return {
        "action": "auto",
        "task": args.task,
        "methods": args.methods,
        "best_config": best_config,
        "results_file": results_file,
    }


def execute_personalization(args, model):
    """
    Execute personalization optimization - find optimal parameters for trait steering.

    This optimizes ALL steering parameters for personality/trait vectors by:
    1. Generating synthetic contrastive pairs for the trait
    2. Testing all combinations of:
       - Layers (where to apply steering)
       - Strengths (how strong the steering signal is)
       - Token aggregation strategies (LAST_TOKEN, MEAN_POOLING, FIRST_TOKEN)
       - Prompt construction strategies (CHAT_TEMPLATE, DIRECT_COMPLETION)
    3. Evaluating each configuration using personalization metrics:
       - Difference: Is the steered response different from baseline?
       - Quality: Is the response coherent (not lobotomized)?
       - Alignment: Does the response match the intended trait?
    4. Selecting the configuration with the highest overall score
    """
    import os

    import torch

    from wisent.core.activations.activations_collector import ActivationCollector
    from wisent.core.activations.core.atoms import ActivationAggregationStrategy
    from wisent.core.activations.prompt_construction_strategy import PromptConstructionStrategy
    from wisent.core.evaluators.steering_evaluators import PersonalizationEvaluator
    from wisent.core.models.core.atoms import SteeringPlan, SteeringVector
    from wisent.core.cli.steering_method_trainer import create_steering_method
    from wisent.core.synthetic.cleaners.pairs_cleaner import PairsCleaner
    from wisent.core.synthetic.db_instructions.mini_dp import Default_DB_Instructions
    from wisent.core.synthetic.generators.diversities.methods.fast_diversity import FastDiversity
    from wisent.core.synthetic.generators.pairs_generator import SyntheticContrastivePairsGenerator

    trait = args.trait
    trait_name = args.trait_name or trait.split()[0].lower()

    print(f"\n{'=' * 80}", flush=True)
    print("🎭 PERSONALIZATION OPTIMIZATION (COMPREHENSIVE)", flush=True)
    print(f"{'=' * 80}", flush=True)
    print(f"   Trait: {trait}", flush=True)
    print(f"   Trait Name: {trait_name}", flush=True)
    print(f"   Model: {args.model}", flush=True)
    print(f"   Num Pairs: {args.num_pairs}", flush=True)
    print(f"   Num Test Prompts: {args.num_test_prompts}", flush=True)
    print(f"   Output Directory: {args.output_dir}", flush=True)
    print(f"{'=' * 80}\n", flush=True)

    # Create output directory
    os.makedirs(args.output_dir, exist_ok=True)
    os.makedirs(os.path.join(args.output_dir, "vectors"), exist_ok=True)

    # Determine layers to test - ALL layers by default
    if args.layers:
        layers_to_test = args.layers
    else:
        # Test ALL layers (1-indexed, since activation collector uses 1-based indexing)
        num_layers = model.num_layers
        layers_to_test = list(range(1, num_layers + 1))

    # Determine strengths to test
    min_strength, max_strength = args.strength_range
    strengths_to_test = np.linspace(min_strength, max_strength, args.num_strength_steps)

    # Token aggregation strategies to test - ALL strategies
    token_aggregations_to_test = [
        ActivationAggregationStrategy.LAST_TOKEN,
        ActivationAggregationStrategy.MEAN_POOLING,
        ActivationAggregationStrategy.FIRST_TOKEN,
        ActivationAggregationStrategy.MAX_POOLING,
    ]

    # Prompt construction strategies to test - ALL strategies
    prompt_constructions_to_test = [
        PromptConstructionStrategy.CHAT_TEMPLATE,
        PromptConstructionStrategy.DIRECT_COMPLETION,
        PromptConstructionStrategy.INSTRUCTION_FOLLOWING,
        PromptConstructionStrategy.ROLE_PLAYING,
        PromptConstructionStrategy.MULTIPLE_CHOICE,
    ]

    # Steering application strategies to test - ALL strategies
    steering_strategies_to_test = ["constant", "initial_only", "diminishing", "all_equal"]

    total_configs = (
        len(layers_to_test)
        * len(strengths_to_test)
        * len(steering_strategies_to_test)
        * len(token_aggregations_to_test)
        * len(prompt_constructions_to_test)
    )

    print("📊 Search Space:", flush=True)
    print(f"   Layers: {layers_to_test} ({len(layers_to_test)} total)", flush=True)
    print(f"   Strengths: {[f'{s:.2f}' for s in strengths_to_test]}", flush=True)
    print(f"   Steering Strategies: {steering_strategies_to_test}", flush=True)
    print(f"   Token Aggregations: {[t.value for t in token_aggregations_to_test]}", flush=True)
    print(f"   Prompt Constructions: {[p.value for p in prompt_constructions_to_test]}", flush=True)
    print(f"   Total configurations: {total_configs}\n", flush=True)

    # Step 1: Generate synthetic contrastive pairs
    print(f"🔧 Step 1: Generating {args.num_pairs} synthetic contrastive pairs...", flush=True)

    # Create the pair generator with required components
    cleaner = PairsCleaner(steps=[])  # Minimal cleaning
    diversity = FastDiversity()
    db_instructions = Default_DB_Instructions()

    pair_generator = SyntheticContrastivePairsGenerator(
        model=model,
        generation_config=get_generate_kwargs(max_new_tokens=150),
        contrastive_set_name=f"{trait_name}_pairs",
        trait_description=trait,
        trait_label=trait_name,
        db_instructions=db_instructions,
        cleaner=cleaner,
        diversity=diversity,
    )

    pair_set, generation_report = pair_generator.generate(num_pairs=args.num_pairs)
    pairs = pair_set.pairs

    print(f"   ✓ Generated {len(pairs)} contrastive pairs\n", flush=True)

    # Generate test prompts for evaluation
    test_prompts = [
        "What's your favorite food?",
        "How do you spend your weekends?",
        "What motivates you in life?",
        "How do you handle setbacks?",
        "What's your opinion on teamwork?",
    ][: args.num_test_prompts]

    print("📝 Test prompts for evaluation:", flush=True)
    for i, prompt in enumerate(test_prompts, 1):
        print(f"   {i}. {prompt}", flush=True)
    print(flush=True)

    # Initialize activation collector
    collector = ActivationCollector(model=model, store_device="cpu")

    # Track results for all configurations
    all_results = {}
    best_config = None
    best_score = -1.0
    best_steering_vector = None

    # Cache for steering vectors per (layer, token_agg, prompt_const) combination
    # to avoid recomputing activations unnecessarily
    steering_vector_cache = {}

    # Step 2: Test all configurations
    print(f"🎯 Step 2: Testing {total_configs} configurations...", flush=True)

    config_count = 0

    # Initialize file for saving generation examples if requested
    examples_file_path = None
    if args.save_all_generation_examples:
        os.makedirs(args.output_dir, exist_ok=True)
        examples_file_path = os.path.join(args.output_dir, f"{trait_name}_all_generation_examples.jsonl")
        # Write header line with metadata
        with open(examples_file_path, "w") as f:
            f.write(json.dumps({"_header": True, "trait": trait, "trait_name": trait_name, "model": args.model}) + "\n")
        print(f"   📝 Will save generation examples to: {examples_file_path}", flush=True)

    # Pre-generate baseline responses ONCE (they don't depend on any loop variables)
    print("   📊 Pre-generating baseline responses for test prompts...", flush=True)
    baseline_responses_cache = {}
    for prompt in test_prompts:
        baseline = model.generate(
            [[{"role": "user", "content": prompt}]],
            **get_generate_kwargs(max_new_tokens=args.max_new_tokens),
            use_steering=False,
        )[0]
        baseline_responses_cache[prompt] = baseline
    print(f"   ✓ Generated {len(baseline_responses_cache)} baseline responses", flush=True)

    for token_agg in token_aggregations_to_test:
        for prompt_const in prompt_constructions_to_test:
            print(
                f"\n   📊 Token Aggregation: {token_agg.value}, Prompt Construction: {prompt_const.value}", flush=True
            )

            for layer in layers_to_test:
                layer_str = str(layer)

                # Check if we already have activations for this (layer, token_agg) combo
                cache_key = (layer, token_agg.value, prompt_const.value)

                if cache_key not in steering_vector_cache:
                    print(f"\n      📍 Layer {layer}: Collecting activations...", flush=True)

                    # Collect activations for this layer with current token_agg and prompt_const
                    pos_acts = []
                    neg_acts = []

                    for pair in pairs:
                        updated_pair = collector.collect_for_pair(
                            pair,
                            layers=[layer_str],
                            aggregation=token_agg,
                            prompt_strategy=prompt_const,
                            return_full_sequence=False,
                            normalize_layers=False,
                        )

                        if (
                            updated_pair.positive_response.layers_activations
                            and layer_str in updated_pair.positive_response.layers_activations
                        ):
                            act = updated_pair.positive_response.layers_activations[layer_str]
                            if act is not None:
                                pos_acts.append(act)

                        if (
                            updated_pair.negative_response.layers_activations
                            and layer_str in updated_pair.negative_response.layers_activations
                        ):
                            act = updated_pair.negative_response.layers_activations[layer_str]
                            if act is not None:
                                neg_acts.append(act)

                    if len(pos_acts) == 0 or len(neg_acts) == 0:
                        print(f"         ⚠️ No activations collected for layer {layer}", flush=True)
                        steering_vector_cache[cache_key] = None
                        continue

                    print(
                        f"         ✓ Collected {len(pos_acts)} positive, {len(neg_acts)} negative activations",
                        flush=True,
                    )

                    # Create steering vector using selected method
                    steering_method = create_steering_method("CAA", args)
                    steering_vector = steering_method.train_for_layer(pos_acts, neg_acts)
                    steering_vector_cache[cache_key] = steering_vector

                    print(
                        f"         ✓ Created steering vector (norm: {torch.norm(steering_vector).item():.4f})",
                        flush=True,
                    )
                else:
                    steering_vector = steering_vector_cache[cache_key]
                    if steering_vector is None:
                        continue

                # Test different strengths and steering strategies
                for strength in strengths_to_test:
                    for steering_strategy in steering_strategies_to_test:
                        config_count += 1
                        config_desc = f"L{layer} S{strength:.2f} St:{steering_strategy} T:{token_agg.value} P:{prompt_const.value}"
                        print(f"      [{config_count}/{total_configs}] Testing {config_desc}...", end=" ")

                        # Create steering plan
                        steering_vec = SteeringVector(vector=steering_vector, scale=float(strength))
                        steering_plan = SteeringPlan(
                            layers={layer_str: steering_vec}, layers_description=[f"Personalization {config_desc}"]
                        )

                        # Get baseline from cache and generate steered responses
                        baseline_responses = [baseline_responses_cache[prompt] for prompt in test_prompts]
                        steered_responses = []

                        for prompt in test_prompts:
                            # Generate steered response
                            model.apply_steering(steering_plan)
                            steered = model.generate(
                                [[{"role": "user", "content": prompt}]],
                                **get_generate_kwargs(max_new_tokens=args.max_new_tokens),
                                use_steering=True,
                                steering_plan=steering_plan,
                            )[0]
                            model.clear_steering()
                            steered_responses.append(steered)

                        # Evaluate using personalization metrics
                        evaluator = PersonalizationEvaluator()

                        # Calculate difference score
                        difference_score = evaluator._evaluate_difference(baseline_responses, steered_responses)

                        # Calculate quality score
                        quality_score = evaluator._evaluate_quality(steered_responses)

                        # Calculate alignment score using simple keyword matching
                        # (Full alignment needs model-based judge which is expensive)
                        alignment_score = PersonalizationEvaluator.estimate_alignment(steered_responses, trait)

                        # Calculate overall score (weighted average)
                        # Only count if difference > 0.3 (steering is actually doing something)
                        if difference_score < 0.3:
                            overall_score = 0.0
                        else:
                            overall_score = 0.2 * difference_score + 0.3 * quality_score + 0.5 * alignment_score

                        print(
                            f"diff={difference_score:.2f} qual={quality_score:.2f} align={alignment_score:.2f} overall={overall_score:.2f}"
                        )

                        # Store results with full config key
                        config_key = f"L{layer}_S{strength:.2f}_St:{steering_strategy}_T:{token_agg.value}_P:{prompt_const.value}"
                        all_results[config_key] = {
                            "layer": layer,
                            "strength": float(strength),
                            "steering_strategy": steering_strategy,
                            "token_aggregation": token_agg.value,
                            "prompt_construction": prompt_const.value,
                            "difference_score": float(difference_score),
                            "quality_score": float(quality_score),
                            "alignment_score": float(alignment_score),
                            "overall_score": float(overall_score),
                            "sample_baseline": baseline_responses[0][:200] if baseline_responses else "",
                            "sample_steered": steered_responses[0][:200] if steered_responses else "",
                        }

                        # Save generation examples if requested
                        if args.save_all_generation_examples and examples_file_path:
                            example_record = {
                                "layer": layer,
                                "strength": float(strength),
                                "steering_strategy": steering_strategy,
                                "token_aggregation": token_agg.value,
                                "prompt_construction": prompt_const.value,
                                "overall_score": float(overall_score),
                                "difference_score": float(difference_score),
                                "quality_score": float(quality_score),
                                "alignment_score": float(alignment_score),
                                "examples": [
                                    {
                                        "prompt": test_prompts[i],
                                        "baseline_response": baseline_responses[i],
                                        "steered_response": steered_responses[i],
                                    }
                                    for i in range(len(test_prompts))
                                ],
                            }
                            with open(examples_file_path, "a") as f:
                                f.write(json.dumps(example_record) + "\n")

                        # Track best configuration
                        if overall_score > best_score:
                            best_score = overall_score
                            best_config = {
                                "layer": layer,
                                "strength": float(strength),
                                "steering_strategy": steering_strategy,
                                "token_aggregation": token_agg.value,
                                "prompt_construction": prompt_const.value,
                                "difference_score": float(difference_score),
                                "quality_score": float(quality_score),
                                "alignment_score": float(alignment_score),
                                "overall_score": float(overall_score),
                            }
                            best_steering_vector = steering_vector

    # Step 3: Save results
    print(f"\n{'=' * 80}")
    print("📊 OPTIMIZATION COMPLETE")
    print(f"{'=' * 80}")

    vector_path = None
    if best_config:
        print("\n✅ Best Configuration:")
        print(f"   Layer: {best_config['layer']}")
        print(f"   Strength: {best_config['strength']:.2f}")
        print(f"   Steering Strategy: {best_config['steering_strategy']}")
        print(f"   Token Aggregation: {best_config['token_aggregation']}")
        print(f"   Prompt Construction: {best_config['prompt_construction']}")
        print(f"   Difference Score: {best_config['difference_score']:.3f}")
        print(f"   Quality Score: {best_config['quality_score']:.3f}")
        print(f"   Alignment Score: {best_config['alignment_score']:.3f}")
        print(f"   Overall Score: {best_config['overall_score']:.3f}")

        # Save best steering vector
        vector_path = os.path.join(args.output_dir, "vectors", f"{trait_name}_optimal.pt")
        torch.save(
            {
                "steering_vector": best_steering_vector,
                "layer": best_config["layer"],
                "layer_index": best_config["layer"],
                "strength": best_config["strength"],
                "steering_strategy": best_config["steering_strategy"],
                "token_aggregation": best_config["token_aggregation"],
                "prompt_construction": best_config["prompt_construction"],
                "trait": trait,
                "trait_name": trait_name,
                "model": args.model,
                "method": "CAA",
                "optimization_scores": {
                    "difference": best_config["difference_score"],
                    "quality": best_config["quality_score"],
                    "alignment": best_config["alignment_score"],
                    "overall": best_config["overall_score"],
                },
            },
            vector_path,
        )
        print(f"\n💾 Saved optimal steering vector to: {vector_path}")
    else:
        print("\n⚠️ No valid configuration found")

    # Save full results to JSON
    results_file = os.path.join(args.output_dir, f"{trait_name}_optimization_results.json")

    # best_config doesn't have steering_vector anymore (it's in best_steering_vector)
    best_config_json = best_config

    output_data = {
        "model": args.model,
        "trait": trait,
        "trait_name": trait_name,
        "num_pairs": args.num_pairs,
        "num_test_prompts": args.num_test_prompts,
        "layers_tested": layers_to_test,
        "strengths_tested": [float(s) for s in strengths_to_test],
        "steering_strategies_tested": steering_strategies_to_test,
        "token_aggregations_tested": [t.value for t in token_aggregations_to_test],
        "prompt_constructions_tested": [p.value for p in prompt_constructions_to_test],
        "best_config": best_config_json,
        "all_results": all_results,
    }

    with open(results_file, "w") as f:
        json.dump(output_data, f, indent=2)

    print(f"💾 Saved full results to: {results_file}")

    if args.save_all_generation_examples and examples_file_path:
        print(f"💾 Generation examples saved iteratively to: {examples_file_path}")

    # Print usage example
    print("\n📝 Usage Example:")
    if best_config:
        print("   python -m wisent.core.main multi-steer \\")
        print(f"       --vector {vector_path}:{best_config['strength']:.1f} \\")
        print(f"       --model {args.model} \\")
        print(f"       --layer {best_config['layer']} \\")
        print('       --prompt "Your prompt here"')

    print(f"\n{'=' * 80}\n")

    return {
        "action": "personalization",
        "trait": trait,
        "trait_name": trait_name,
        "best_config": best_config_json,
        "results_file": results_file,
        "vector_path": vector_path if best_config else None,
    }


def execute_multi_personalization(args, model):
    """
    Execute multi-trait joint personalization optimization.

    This finds a SINGLE optimal configuration (layer, token_aggregation, prompt_construction)
    that works well for ALL traits, then optimizes strength per-trait individually.

    The approach:
    1. Generate synthetic contrastive pairs for each trait
    2. For each (layer, token_agg, prompt_const) configuration:
       - Compute steering vectors for ALL traits
       - Find optimal strength for each trait individually
       - Compute combined score = mean(trait_scores)
    3. Select the configuration with highest combined score
    4. Return: shared (layer, token_agg, prompt_const) + per-trait strength
    """
    import os

    import torch

    from wisent.core.activations.activations_collector import ActivationCollector
    from wisent.core.activations.core.atoms import ActivationAggregationStrategy
    from wisent.core.activations.prompt_construction_strategy import PromptConstructionStrategy
    from wisent.core.evaluators.steering_evaluators import PersonalizationEvaluator
    from wisent.core.models.core.atoms import SteeringPlan, SteeringVector
    from wisent.core.cli.steering_method_trainer import create_steering_method
    from wisent.core.synthetic.cleaners.pairs_cleaner import PairsCleaner
    from wisent.core.synthetic.db_instructions.mini_dp import Default_DB_Instructions
    from wisent.core.synthetic.generators.diversities.methods.fast_diversity import FastDiversity
    from wisent.core.synthetic.generators.pairs_generator import SyntheticContrastivePairsGenerator

    traits = args.traits
    trait_names = args.trait_names or [t.split()[0].lower() for t in traits]

    if len(trait_names) != len(traits):
        print(f"Error: Number of --trait-name args ({len(trait_names)}) must match --trait args ({len(traits)})")
        return None

    print(f"\n{'=' * 80}", flush=True)
    print("🎭 MULTI-TRAIT JOINT PERSONALIZATION OPTIMIZATION", flush=True)
    print(f"{'=' * 80}", flush=True)
    print(f"   Model: {args.model}", flush=True)
    print(f"   Traits: {len(traits)}", flush=True)
    for i, (trait, name) in enumerate(zip(traits, trait_names)):
        print(f"      {i + 1}. {name}: {trait[:50]}...", flush=True)
    print(f"   Num Pairs per trait: {args.num_pairs}", flush=True)
    print(f"   Num Test Prompts: {args.num_test_prompts}", flush=True)
    print(f"   Output Directory: {args.output_dir}", flush=True)
    print(f"{'=' * 80}\n", flush=True)

    # Create output directory
    os.makedirs(args.output_dir, exist_ok=True)
    os.makedirs(os.path.join(args.output_dir, "vectors"), exist_ok=True)

    # Determine layers to test - default to middle 50% of layers where steering works best
    if args.layers:
        layers_to_test = args.layers
    else:
        num_layers = model.num_layers
        # Test middle 50% of layers (e.g., layers 8-20 for a 28-layer model)
        start_layer = max(1, num_layers // 4)
        end_layer = min(num_layers, 3 * num_layers // 4)
        layers_to_test = list(range(start_layer, end_layer + 1))

    # Determine strengths to test
    min_strength, max_strength = args.strength_range
    strengths_to_test = np.linspace(min_strength, max_strength, args.num_strength_steps)

    # Token aggregation strategies to test
    token_aggregations_to_test = [
        ActivationAggregationStrategy.LAST_TOKEN,
        ActivationAggregationStrategy.MEAN_POOLING,
        ActivationAggregationStrategy.FIRST_TOKEN,
        ActivationAggregationStrategy.MAX_POOLING,
    ]

    # Prompt construction strategies to test
    prompt_constructions_to_test = [
        PromptConstructionStrategy.CHAT_TEMPLATE,
        PromptConstructionStrategy.DIRECT_COMPLETION,
        PromptConstructionStrategy.INSTRUCTION_FOLLOWING,
        PromptConstructionStrategy.ROLE_PLAYING,
        PromptConstructionStrategy.MULTIPLE_CHOICE,
    ]

    # Use a fixed steering strategy (initial_only works well for multi-trait)
    steering_strategy = "initial_only"

    total_shared_configs = len(layers_to_test) * len(token_aggregations_to_test) * len(prompt_constructions_to_test)

    print("📊 Search Space:", flush=True)
    print(f"   Shared configs (layer × token_agg × prompt_const): {total_shared_configs}", flush=True)
    print(f"   Strengths per trait: {len(strengths_to_test)}", flush=True)
    print(f"   Steering strategy: {steering_strategy} (fixed)", flush=True)
    print("\n", flush=True)

    # Step 1: Generate synthetic contrastive pairs for each trait
    print(f"🔧 Step 1: Generating synthetic pairs for {len(traits)} traits...", flush=True)

    trait_pairs = {}
    for trait, name in zip(traits, trait_names):
        print(f"\n   Generating pairs for '{name}'...", flush=True)

        cleaner = PairsCleaner(steps=[])
        diversity = FastDiversity()
        db_instructions = Default_DB_Instructions()

        pair_generator = SyntheticContrastivePairsGenerator(
            model=model,
            generation_config=get_generate_kwargs(max_new_tokens=150),
            contrastive_set_name=f"{name}_pairs",
            trait_description=trait,
            trait_label=name,
            db_instructions=db_instructions,
            cleaner=cleaner,
            diversity=diversity,
        )

        pair_set, _ = pair_generator.generate(num_pairs=args.num_pairs)
        trait_pairs[name] = {"trait": trait, "pairs": pair_set.pairs}
        print(f"      ✓ Generated {len(pair_set.pairs)} pairs for '{name}'", flush=True)

    # Test prompts for evaluation
    test_prompts = [
        "What's your favorite food?",
        "How do you spend your weekends?",
        "What motivates you in life?",
        "How do you handle setbacks?",
        "What's your opinion on teamwork?",
    ][: args.num_test_prompts]

    print(f"\n📝 Test prompts: {test_prompts}", flush=True)

    # Initialize collector
    collector = ActivationCollector(model=model, store_device="cpu")

    # Track results
    all_results = {}
    best_shared_config = None
    best_combined_score = -1.0
    best_per_trait_strengths = {}
    best_steering_vectors = {}
    best_overall_sample_responses = []

    # Step 2: Test each shared configuration
    print(f"\n🎯 Step 2: Testing {total_shared_configs} shared configurations...", flush=True)

    config_count = 0

    for token_agg in token_aggregations_to_test:
        for prompt_const in prompt_constructions_to_test:
            for layer in layers_to_test:
                config_count += 1
                layer_str = str(layer)
                shared_config_key = f"L{layer}_T:{token_agg.value}_P:{prompt_const.value}"

                print(f"\n[{config_count}/{total_shared_configs}] {shared_config_key}", flush=True)

                # Compute steering vectors for each trait with this config
                trait_vectors = {}
                for name, data in trait_pairs.items():
                    pairs = data["pairs"]

                    pos_acts = []
                    neg_acts = []

                    for pair in pairs:
                        updated_pair = collector.collect_for_pair(
                            pair,
                            layers=[layer_str],
                            aggregation=token_agg,
                            prompt_strategy=prompt_const,
                            return_full_sequence=False,
                            normalize_layers=False,
                        )

                        if (
                            updated_pair.positive_response.layers_activations
                            and layer_str in updated_pair.positive_response.layers_activations
                        ):
                            act = updated_pair.positive_response.layers_activations[layer_str]
                            if act is not None:
                                pos_acts.append(act)

                        if (
                            updated_pair.negative_response.layers_activations
                            and layer_str in updated_pair.negative_response.layers_activations
                        ):
                            act = updated_pair.negative_response.layers_activations[layer_str]
                            if act is not None:
                                neg_acts.append(act)

                    if len(pos_acts) == 0 or len(neg_acts) == 0:
                        print(f"      ⚠️ No activations for '{name}' - skipping config", flush=True)
                        trait_vectors = None
                        break

                    steering_method = create_steering_method("CAA", args)
                    steering_vector = steering_method.train_for_layer(pos_acts, neg_acts)
                    trait_vectors[name] = steering_vector

                if trait_vectors is None:
                    continue

                # Use Latin Hypercube Sampling to efficiently explore strength space
                # Instead of testing all N^T combinations, sample ~20 representative points
                import random
                from itertools import product

                best_combined_score_for_config = -1.0
                best_strengths_for_config = dict.fromkeys(trait_names, strengths_to_test[0])
                best_sample_responses = []

                # Generate strength combinations - use sampling to reduce search space
                all_strength_combos = list(product(strengths_to_test, repeat=len(trait_names)))

                # If too many combinations, sample a subset that includes edges and random middle points
                max_samples = 25  # Test at most 25 combinations per config
                if len(all_strength_combos) > max_samples:
                    # Always include corner cases (min, max for each trait)
                    corners = [
                        tuple([strengths_to_test[0]] * len(trait_names)),  # All min
                        tuple([strengths_to_test[-1]] * len(trait_names)),  # All max
                        tuple([strengths_to_test[len(strengths_to_test) // 2]] * len(trait_names)),  # All mid
                    ]
                    # Add some diagonal samples
                    for i in range(len(strengths_to_test)):
                        corners.append(tuple([strengths_to_test[i]] * len(trait_names)))

                    # Randomly sample the rest
                    remaining = max_samples - len(set(corners))
                    random.seed(42)  # For reproducibility
                    other_combos = [c for c in all_strength_combos if c not in corners]
                    sampled = random.sample(other_combos, min(remaining, len(other_combos)))

                    strength_combos = list(set(corners)) + sampled
                else:
                    strength_combos = all_strength_combos

                # Pre-generate baseline responses ONCE per config (they don't depend on steering)
                baseline_responses = []
                for prompt in test_prompts:
                    baseline = model.generate(
                        [[{"role": "user", "content": prompt}]],
                        **get_generate_kwargs(max_new_tokens=args.max_new_tokens),
                        use_steering=False,
                    )[0]
                    baseline_responses.append(baseline)

                num_strength_combos = len(strength_combos)
                for combo_idx, strength_combo in enumerate(strength_combos):
                    if args.verbose and combo_idx % 5 == 0:
                        import sys

                        sys.stdout.write(f"\r      Testing strength {combo_idx + 1}/{num_strength_combos}...")
                        sys.stdout.flush()
                    current_strengths = dict(zip(trait_names, strength_combo))

                    # Create COMBINED steering plan with ALL vectors at once
                    combined_vector = None
                    for name, strength in current_strengths.items():
                        scaled_vector = trait_vectors[name] * float(strength)
                        if combined_vector is None:
                            combined_vector = scaled_vector.clone()
                        else:
                            combined_vector = combined_vector + scaled_vector

                    steering_vec = SteeringVector(vector=combined_vector, scale=1.0)
                    steering_plan = SteeringPlan(
                        layers={layer_str: steering_vec},
                        layers_description=[f"Multi-trait combined: {'+'.join(trait_names)}"],
                    )

                    # Generate only steered responses (baselines were pre-generated)
                    steered_responses = []
                    for prompt in test_prompts:
                        model.apply_steering(steering_plan)
                        steered = model.generate(
                            [[{"role": "user", "content": prompt}]],
                            **get_generate_kwargs(max_new_tokens=args.max_new_tokens),
                            use_steering=True,
                            steering_plan=steering_plan,
                        )[0]
                        model.clear_steering()
                        steered_responses.append(steered)

                    # Evaluate combined output against ALL traits together
                    evaluator = PersonalizationEvaluator()
                    difference_score = evaluator._evaluate_difference(baseline_responses, steered_responses)
                    quality_score = evaluator._evaluate_quality(steered_responses)

                    # Compute alignment score against COMBINED trait description
                    combined_trait_description = " AND ".join([trait_pairs[name]["trait"] for name in trait_names])
                    alignment_score = PersonalizationEvaluator.estimate_alignment(steered_responses, combined_trait_description)

                    if difference_score < 0.3:
                        overall_score = 0.0
                    else:
                        overall_score = 0.2 * difference_score + 0.3 * quality_score + 0.5 * alignment_score

                    if overall_score > best_combined_score_for_config:
                        best_combined_score_for_config = overall_score
                        best_strengths_for_config = current_strengths.copy()
                        best_sample_responses = list(zip(test_prompts, baseline_responses, steered_responses))

                # Store per-trait strengths from best combo
                trait_best_strengths = best_strengths_for_config
                combined_score = best_combined_score_for_config

                print(
                    f"      Strengths: {', '.join([f'{n}={s:.2f}' for n, s in trait_best_strengths.items()])}",
                    flush=True,
                )
                print(f"      → Combined score (all traits at once): {combined_score:.3f}", flush=True)

                # Show sample responses for this config
                if best_sample_responses and args.verbose:
                    print("\n      📝 Sample responses:", flush=True)
                    for prompt, baseline, steered in best_sample_responses[:2]:
                        print(f"         Prompt: {prompt}", flush=True)
                        print(f"         Baseline: {baseline[:100]}...", flush=True)
                        print(f"         Steered:  {steered[:100]}...", flush=True)

                # Store result
                all_results[shared_config_key] = {
                    "layer": layer,
                    "token_aggregation": token_agg.value,
                    "prompt_construction": prompt_const.value,
                    "steering_strategy": steering_strategy,
                    "per_trait_strengths": trait_best_strengths,
                    "combined_score": float(combined_score),
                    "sample_responses": [
                        {"prompt": p, "baseline": b, "steered": s} for p, b, s in best_sample_responses
                    ]
                    if best_sample_responses
                    else [],
                }

                if combined_score > best_combined_score:
                    best_combined_score = combined_score
                    best_shared_config = {
                        "layer": layer,
                        "token_aggregation": token_agg.value,
                        "prompt_construction": prompt_const.value,
                        "steering_strategy": steering_strategy,
                    }
                    best_per_trait_strengths = trait_best_strengths.copy()
                    best_steering_vectors = {name: v.clone() for name, v in trait_vectors.items()}
                    best_overall_sample_responses = best_sample_responses.copy() if best_sample_responses else []

    # Step 3: Save results
    print(f"\n{'=' * 80}", flush=True)
    print("📊 MULTI-TRAIT OPTIMIZATION COMPLETE", flush=True)
    print(f"{'=' * 80}", flush=True)

    vector_paths = {}
    if best_shared_config:
        print("\n✅ Best Shared Configuration:", flush=True)
        print(f"   Layer: {best_shared_config['layer']}", flush=True)
        print(f"   Token Aggregation: {best_shared_config['token_aggregation']}", flush=True)
        print(f"   Prompt Construction: {best_shared_config['prompt_construction']}", flush=True)
        print(f"   Steering Strategy: {best_shared_config['steering_strategy']}", flush=True)
        print("\n✅ Per-Trait Optimal Strengths:", flush=True)
        for name, strength in best_per_trait_strengths.items():
            print(f"      {name}: {strength:.2f}", flush=True)
        print(f"\n   Combined Score: {best_combined_score:.3f}", flush=True)

        # Print sample responses from the best configuration
        if best_overall_sample_responses:
            print(f"\n{'=' * 80}", flush=True)
            print("📝 SAMPLE RESPONSES (Best Configuration)", flush=True)
            print(f"{'=' * 80}", flush=True)
            for prompt, baseline, steered in best_overall_sample_responses:
                print(f"\n   🗣️ Prompt: {prompt}", flush=True)
                print("\n   📄 Baseline Response:", flush=True)
                print(f"      {baseline}", flush=True)
                print("\n   🎯 Steered Response (evil + italian):", flush=True)
                print(f"      {steered}", flush=True)
                print(f"\n   {'-' * 70}", flush=True)

        # Save steering vectors for each trait
        for name in trait_names:
            vector_path = os.path.join(args.output_dir, "vectors", f"{name}_optimal.pt")
            torch.save(
                {
                    "steering_vector": best_steering_vectors[name],
                    "layer": best_shared_config["layer"],
                    "layer_index": best_shared_config["layer"],
                    "strength": best_per_trait_strengths[name],
                    "steering_strategy": best_shared_config["steering_strategy"],
                    "token_aggregation": best_shared_config["token_aggregation"],
                    "prompt_construction": best_shared_config["prompt_construction"],
                    "trait": trait_pairs[name]["trait"],
                    "trait_name": name,
                    "model": args.model,
                    "method": "CAA",
                    "multi_trait_optimization": True,
                },
                vector_path,
            )
            vector_paths[name] = vector_path
            print(f"\n💾 Saved {name} vector to: {vector_path}", flush=True)
    else:
        print("\n⚠️ No valid configuration found", flush=True)

    # Save full results
    results_file = os.path.join(args.output_dir, "multi_trait_optimization_results.json")

    output_data = {
        "model": args.model,
        "traits": {name: trait_pairs[name]["trait"] for name in trait_names},
        "num_pairs_per_trait": args.num_pairs,
        "num_test_prompts": args.num_test_prompts,
        "layers_tested": layers_to_test,
        "strengths_tested": [float(s) for s in strengths_to_test],
        "token_aggregations_tested": [t.value for t in token_aggregations_to_test],
        "prompt_constructions_tested": [p.value for p in prompt_constructions_to_test],
        "best_shared_config": best_shared_config,
        "best_per_trait_strengths": best_per_trait_strengths,
        "best_combined_score": best_combined_score,
        "best_sample_responses": [
            {"prompt": p, "baseline": b, "steered": s} for p, b, s in best_overall_sample_responses
        ]
        if best_overall_sample_responses
        else [],
        "all_results": all_results,
    }

    with open(results_file, "w") as f:
        json.dump(output_data, f, indent=2)

    print(f"\n💾 Saved full results to: {results_file}", flush=True)

    # Print usage example
    if best_shared_config and vector_paths:
        print("\n📝 Usage Example:", flush=True)
        print("   python -m wisent.core.main multi-steer \\", flush=True)
        for name in trait_names:
            print(f"       --vector {vector_paths[name]}:{best_per_trait_strengths[name]:.1f} \\", flush=True)
        print(f"       --model {args.model} \\", flush=True)
        print(f"       --layer {best_shared_config['layer']} \\", flush=True)
        print('       --prompt "Your prompt here"', flush=True)

    print(f"\n{'=' * 80}\n", flush=True)

    return {
        "action": "multi-personalization",
        "traits": trait_names,
        "best_shared_config": best_shared_config,
        "best_per_trait_strengths": best_per_trait_strengths,
        "best_combined_score": best_combined_score,
        "results_file": results_file,
        "vector_paths": vector_paths,
    }
