"""Parser setup for the unified 'optimize' command.

Consolidates optimize-weights, optimize-steering, and optimize-classification
into a single command with --target flag.

Usage:
    wisent optimize --target weights --task coding --model Qwen/Qwen3-8B
    wisent optimize --target steering --task coding --model Qwen/Qwen3-8B
    wisent optimize --target classification --task coding --model Qwen/Qwen3-8B
"""

import argparse


def setup_optimize_parser(subparsers: argparse._SubParsersAction) -> None:
    """Set up the unified optimize command parser."""
    
    optimize_parser = subparsers.add_parser(
        "optimize",
        help="Optimize model parameters (weights, steering, or classification)",
        description=(
            "Unified optimization command.\n\n"
            "Examples:\n"
            "  wisent optimize --target weights --task coding --model Qwen/Qwen3-8B\n"
            "  wisent optimize --target steering --task coding --model Qwen/Qwen3-8B\n"
            "  wisent optimize --target classification --model Qwen/Qwen3-8B"
        ),
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    
    # Target selection (required)
    optimize_parser.add_argument(
        "--target",
        type=str,
        required=True,
        choices=["weights", "steering", "classification"],
        help="What to optimize: weights, steering, or classification"
    )
    
    # Common arguments shared across all targets
    optimize_parser.add_argument(
        "--model",
        type=str,
        required=True,
        help="Model name or path"
    )
    
    optimize_parser.add_argument(
        "--task",
        type=str,
        default=None,
        help=(
            "Task to optimize for. Can be: "
            "skill name (coding, math), "
            "risk name (hallucination, toxicity), "
            "special case (refusal, personalization, humanizer, custom), "
            "benchmark name (arc_easy, gsm8k), "
            "comma-separated benchmarks (arc_easy,gsm8k), "
            "or free-form description ('be more concise')"
        )
    )
    
    optimize_parser.add_argument(
        "--device",
        type=str,
        default=None,
        help="Device to run on (default: auto)"
    )
    
    optimize_parser.add_argument(
        "--verbose",
        action="store_true",
        help="Enable verbose output"
    )
    
    # Weights-specific arguments
    weights_group = optimize_parser.add_argument_group("weights optimization")
    
    weights_group.add_argument(
        "--output-dir",
        type=str,
        default="./optimized_model",
        help="Output directory for optimized model (weights target)"
    )
    
    weights_group.add_argument(
        "--n-trials",
        type=int,
        default=300,
        help="Number of Optuna trials (default: 300)"
    )
    
    weights_group.add_argument(
        "--target-metric",
        type=float,
        default=0.95,
        help="Target metric value to achieve (default: 0.95)"
    )
    
    weights_group.add_argument(
        "--num-pairs",
        type=int,
        default=100,
        help="Number of contrastive pairs for vector generation (default: 100)"
    )
    
    weights_group.add_argument(
        "--cap-pairs-per-benchmark",
        type=int,
        default=None,
        help="Cap pairs per benchmark for multi-benchmark training"
    )
    
    weights_group.add_argument(
        "--trait",
        type=str,
        default=None,
        help="Trait description (required for --task personalization or custom)"
    )
    
    weights_group.add_argument(
        "--steering-vectors",
        type=str,
        default=None,
        help="Path to pre-computed steering vectors JSON (skip generation)"
    )
    
    # Steering-specific arguments
    steering_group = optimize_parser.add_argument_group("steering optimization")
    
    steering_group.add_argument(
        "--steering-action",
        type=str,
        choices=["comprehensive", "compare-methods", "optimize-layer", "optimize-strength", "auto"],
        default="comprehensive",
        help="Steering optimization action (default: comprehensive)"
    )
    
    steering_group.add_argument(
        "--limit",
        type=int,
        default=100,
        help="Sample limit per task (default: 100)"
    )
    
    steering_group.add_argument(
        "--no-save",
        action="store_true",
        help="Don't save results to model config"
    )
    
    # Classification-specific arguments  
    classification_group = optimize_parser.add_argument_group("classification optimization")
    
    classification_group.add_argument(
        "--optimization-metric",
        type=str,
        default="f1",
        choices=["f1", "accuracy", "precision", "recall"],
        help="Metric to optimize for classification (default: f1)"
    )
    
    classification_group.add_argument(
        "--layer-range",
        type=str,
        default=None,
        help="Layer range to test (e.g., '10-20')"
    )
    
    classification_group.add_argument(
        "--threshold-range",
        type=float,
        nargs="+",
        default=[0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9],
        help="Detection thresholds to test"
    )
