"""
Parser for optimize-weights command.

This command runs an optimization loop to find optimal weight modification
parameters for any trait or task, using any evaluator.

Unified pipeline:
1. Generate steering vector (from trait or task)
2. Apply weight modification with trial parameters
3. Evaluate using chosen evaluator
4. Optuna adjusts parameters
5. Repeat until target reached or max trials
6. Save optimized model
"""

import argparse


def setup_optimize_weights_parser(parser: argparse.ArgumentParser) -> None:
    """
    Set up argument parser for optimize-weights command.

    This command optimizes weight modification parameters to achieve a target
    metric value. Works with any trait (synthetic pairs) or task (lm-eval),
    and any evaluator (refusal, task accuracy, personality, custom).
    """

    # ==========================================================================
    # INPUT SOURCE: What to optimize for (trait, task, tags, or pre-computed)
    # ==========================================================================
    input_group = parser.add_mutually_exclusive_group(required=True)
    input_group.add_argument(
        "--trait",
        type=str,
        help="Trait description for synthetic pair generation (e.g., 'refusal behavior', 'sycophantic responses')"
    )
    input_group.add_argument(
        "--task",
        type=str,
        help="LM-eval task name for task-based optimization (e.g., 'hellaswag', 'arc_easy')"
    )
    input_group.add_argument(
        "--tags",
        type=str,
        nargs="+",
        help="Train on ALL benchmarks with these tags (e.g., --tags coding mathematics). "
             "Pools data from all matching benchmarks to train unified steering vectors."
    )
    input_group.add_argument(
        "--steering-vectors",
        type=str,
        help="Path to pre-computed steering vectors JSON file (skip vector generation)"
    )
    
    # Additional options for --tags mode
    parser.add_argument(
        "--cap-pairs-per-benchmark",
        type=int,
        default=None,
        help="Cap pairs per benchmark (for --tags mode). Benchmarks with more pairs get randomly sampled."
    )
    parser.add_argument(
        "--max-benchmarks",
        type=int,
        default=None,
        help="Maximum number of benchmarks to use (for --tags mode)"
    )
    parser.add_argument(
        "--train-ratio",
        type=float,
        default=0.8,
        help="Fraction of pairs for training vs evaluation (for --tags mode, default: 0.8)"
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=42,
        help="Random seed for reproducibility (default: 42)"
    )

    # ==========================================================================
    # MODEL CONFIGURATION
    # ==========================================================================
    parser.add_argument(
        "--model",
        type=str,
        required=True,
        help="HuggingFace model name or path"
    )
    parser.add_argument(
        "--device",
        type=str,
        default=None,
        help="Device to use (e.g., 'cuda', 'cuda:0', 'cpu'). Default: auto"
    )

    # ==========================================================================
    # OUTPUT
    # ==========================================================================
    parser.add_argument(
        "--output-dir",
        type=str,
        required=True,
        help="Directory to save the optimized model"
    )
    parser.add_argument(
        "--save-trials",
        type=str,
        default=None,
        help="Path to save all trial results as JSON (optional)"
    )

    # ==========================================================================
    # EVALUATOR CONFIGURATION
    # ==========================================================================
    eval_group = parser.add_argument_group("evaluation")
    eval_group.add_argument(
        "--evaluator",
        type=str,
        default="auto",
        choices=["auto", "refusal", "task", "semantic", "keyword", "llm_judge", "personalization", "custom"],
        help="Evaluator to use. 'auto' selects based on trait/task. 'custom' requires --custom-evaluator. Default: auto"
    )
    eval_group.add_argument(
        "--custom-evaluator",
        type=str,
        default=None,
        help=(
            "Custom evaluator specification. Can be: "
            "(1) Python module path e.g. 'my_evaluators.gptzero', "
            "(2) File path with function e.g. './my_eval.py:score_fn', "
            "(3) Built-in example e.g. 'wisent.core.evaluators.custom.examples.gptzero'. "
            "The module must define 'create_evaluator', 'evaluator', or 'evaluate' function."
        )
    )
    eval_group.add_argument(
        "--custom-evaluator-kwargs",
        type=str,
        default=None,
        help="JSON string of kwargs for custom evaluator, e.g. '{\"api_key\": \"xxx\", \"optimize_for\": \"human_prob\"}'"
    )
    eval_group.add_argument(
        "--eval-prompts",
        type=str,
        default=None,
        help="Path to custom evaluation prompts JSON file (for refusal/semantic evaluation)"
    )
    eval_group.add_argument(
        "--eval-topics",
        type=str,
        default=None,
        help="Comma-separated UncensorBench topics (for refusal evaluation)"
    )
    eval_group.add_argument(
        "--num-eval-prompts",
        type=int,
        default=30,
        help="Number of evaluation prompts per trial. Default: 30"
    )

    # ==========================================================================
    # OPTIMIZATION TARGET
    # ==========================================================================
    target_group = parser.add_argument_group("optimization target")
    target_group.add_argument(
        "--target-metric",
        type=str,
        default="compliance_rate",
        help="Metric to optimize: compliance_rate, refusal_rate, accuracy. Default: compliance_rate"
    )
    target_group.add_argument(
        "--target-value",
        type=float,
        default=0.95,
        help="Target value for the metric. Default: 0.95"
    )
    target_group.add_argument(
        "--direction",
        type=str,
        default="auto",
        choices=["auto", "maximize", "minimize"],
        help="Optimization direction. 'auto' infers from metric. Default: auto"
    )

    # ==========================================================================
    # OPTIMIZATION PARAMETERS
    # ==========================================================================
    optim_group = parser.add_argument_group("optimization")
    optim_group.add_argument(
        "--trials",
        type=int,
        default=300,
        help="Number of optimization trials. Default: 300"
    )
    optim_group.add_argument(
        "--startup-trials",
        type=int,
        default=10,
        help="Number of random startup trials before TPE. Default: 10"
    )
    optim_group.add_argument(
        "--early-stop",
        action="store_true",
        help="Stop early if target value is reached"
    )
    optim_group.add_argument(
        "--early-stop-patience",
        type=int,
        default=10,
        help="Stop if no improvement for N trials. Default: 10"
    )

    # ==========================================================================
    # SEARCH SPACE: Parameters to optimize
    # ==========================================================================
    search_group = parser.add_argument_group("search space")
    search_group.add_argument(
        "--strength-range",
        type=str,
        default="0.3,2.0",
        help="Min,max range for directional projection strength. Default: 0.3,2.0"
    )
    search_group.add_argument(
        "--max-weight-range",
        type=str,
        default="0.5,3.0",
        help="Min,max range for kernel max weight. Default: 0.5,3.0"
    )
    search_group.add_argument(
        "--min-weight-range",
        type=str,
        default="0.0,0.5",
        help="Min,max range for kernel min weight. Default: 0.0,0.5"
    )
    search_group.add_argument(
        "--position-range",
        type=str,
        default="0.3,0.7",
        help="Min,max range for kernel peak position (as ratio 0-1). Default: 0.3,0.7"
    )
    search_group.add_argument(
        "--num-pairs",
        type=int,
        default=100,
        help="Number of contrastive pairs to generate (fixed, not optimized). Default: 100"
    )
    search_group.add_argument(
        "--optimize-direction-index",
        action="store_true",
        help="Optimize float direction index (interpolate between layer directions)"
    )

    # ==========================================================================
    # WEIGHT MODIFICATION METHOD
    # ==========================================================================
    method_group = parser.add_argument_group("weight modification")
    method_group.add_argument(
        "--method",
        type=str,
        default="directional",
        choices=["directional", "additive"],
        help="Weight modification method. Default: directional"
    )
    method_group.add_argument(
        "--components",
        type=str,
        nargs="+",
        default=["self_attn.o_proj", "mlp.down_proj"],
        help="Components to modify. Default: self_attn.o_proj mlp.down_proj"
    )
    method_group.add_argument(
        "--norm-preserve",
        action="store_true",
        default=True,
        help="Use norm-preserving directional projection. Default: True"
    )
    method_group.add_argument(
        "--no-norm-preserve",
        action="store_false",
        dest="norm_preserve",
        help="Disable norm-preserving directional projection"
    )

    # ==========================================================================
    # STEERING VECTOR GENERATION
    # ==========================================================================
    vector_group = parser.add_argument_group("steering vector generation")
    vector_group.add_argument(
        "--layers",
        type=str,
        default="all",
        help="Layers for activation collection: 'all' or comma-separated indices. Default: all"
    )
    vector_group.add_argument(
        "--token-aggregation",
        type=str,
        default="average",
        choices=["average", "last", "first", "max"],
        help="Token aggregation strategy. Default: average"
    )
    vector_group.add_argument(
        "--similarity-threshold",
        type=float,
        default=0.8,
        help="Similarity threshold for synthetic pair filtering. Default: 0.8"
    )
    vector_group.add_argument(
        "--pairs-cache-dir",
        type=str,
        default=None,
        help="Directory to cache/load generated pairs"
    )

    # ==========================================================================
    # HUB EXPORT
    # ==========================================================================
    hub_group = parser.add_argument_group("huggingface hub")
    hub_group.add_argument(
        "--push-to-hub",
        action="store_true",
        help="Push optimized model to HuggingFace Hub"
    )
    hub_group.add_argument(
        "--repo-id",
        type=str,
        default=None,
        help="HuggingFace repo ID (required if --push-to-hub)"
    )

    # ==========================================================================
    # DISPLAY
    # ==========================================================================
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="Enable verbose output"
    )
    parser.add_argument(
        "--timing",
        action="store_true",
        help="Show timing information"
    )
    parser.add_argument(
        "--show-comparisons",
        type=int,
        default=0,
        metavar="N",
        help="Show N before/after response comparisons with biggest score changes in console. Default: 0 (disabled)"
    )
    parser.add_argument(
        "--save-comparisons",
        type=str,
        default=None,
        metavar="PATH",
        help="Save all comparisons to JSON file (use with --show-comparisons to also display in console)"
    )
