"""Parser for the 'zwiad' command - geometry analysis with concept decomposition."""


def setup_zwiad_parser(parser):
    """
    Set up the zwiad command parser.

    Usage:
        wisent zwiad --from-database --model meta-llama/Llama-3.2-1B-Instruct --task truthfulqa_custom
        wisent zwiad --from-cache cache/optimize/meta-llama_Llama-3.2-1B-Instruct/activations_xxx.pkl
        wisent zwiad --from-database --task truthfulqa_custom --layers 8-16 --visualizations
    """
    # Data source - mutually exclusive
    data_source = parser.add_mutually_exclusive_group(required=True)
    data_source.add_argument(
        "--from-database",
        action="store_true",
        help="Load activations from Supabase database"
    )
    data_source.add_argument(
        "--from-cache",
        type=str,
        metavar="PATH",
        help="Load activations from local cache pickle file"
    )

    # Model and task (required only for --from-database)
    parser.add_argument(
        "--model",
        type=str,
        required=True,
        help="Model name in database"
    )
    parser.add_argument(
        "--task",
        type=str,
        default=None,
        help="Task/benchmark name in database (required for --from-database, optional for --from-cache to enable LLM concept naming)"
    )

    # Layer selection
    parser.add_argument(
        "--layers",
        type=str,
        default=None,
        help="Layers to analyze: range (8-16) or comma-separated (8,10,12). Default: all available layers in database"
    )

    # Database options
    parser.add_argument(
        "--prompt-format",
        type=str,
        required=True,
        choices=["chat", "completion"],
        help="Prompt format used for activations"
    )
    parser.add_argument(
        "--extraction-strategy",
        type=str,
        required=True,
        choices=["last_token", "first_token", "chat_last"],
        help="Token extraction strategy"
    )
    parser.add_argument(
        "--extraction-component",
        type=str,
        required=True,
        choices=["residual_stream", "attn_output", "mlp_output", "per_head",
                 "mlp_intermediate", "post_attn_residual", "pre_attn_layernorm",
                 "embedding_output", "final_layernorm", "q_proj", "k_proj",
                 "v_proj", "mlp_gate_activation", "attention_scores", "logits"],
        help="Transformer component to extract from"
    )
    parser.add_argument(
        "--database-url",
        type=str,
        default=None,
        help="Database URL (default: DATABASE_URL env var)"
    )

    # Protocol steps
    parser.add_argument(
        "--steps",
        type=str,
        required=True,
        help="Protocol steps to run: 'all', 'signal', 'geometry', 'decomposition', 'intervention', 'editability', or comma-separated (e.g., 'signal,geometry,editability')"
    )

    # Analysis options
    parser.add_argument(
        "--no-visualizations",
        action="store_true",
        default=False,
        help="Disable visualization generation (visualizations are on by default)"
    )
    parser.add_argument(
        "--llm-model",
        type=str,
        required=True,
        help="LLM model for concept naming"
    )
    parser.add_argument(
        "--concept-naming-n-samples",
        type=int,
        required=True,
        help="Number of samples to use for LLM concept naming prompts"
    )

    # Output
    parser.add_argument(
        "--output",
        type=str,
        default=None,
        help="Output file path for JSON results"
    )
    parser.add_argument("--visualizations-dir", type=str, required=True, help="Directory to save visualization PNGs")

    # Geometry metric design-choice parameters (non-derivable from data)
    parser.add_argument("--probe-small-hidden", type=int, required=True, dest="probe_small_hidden", help="Hidden dim for small probe")
    parser.add_argument("--probe-mlp-hidden", type=int, required=True, dest="probe_mlp_hidden", help="Hidden dim for MLP probe")
    parser.add_argument("--probe-mlp-alpha", type=float, required=True, dest="probe_mlp_alpha", help="L2 regularization alpha for MLP probe")
    parser.add_argument("--spectral-n-neighbors", type=int, required=True, dest="spectral_n_neighbors", help="Neighbor count for spectral/manifold metrics")
    parser.add_argument("--direction-n-bootstrap", type=int, required=True, dest="direction_n_bootstrap", help="Bootstrap iterations for direction stability")
    parser.add_argument("--direction-subset-fraction", type=float, required=True, dest="direction_subset_fraction", help="Fraction per bootstrap sample")
    parser.add_argument("--direction-std-penalty", type=float, required=True, dest="direction_std_penalty", help="Penalty weight for direction std")
    parser.add_argument("--consistency-w-cosine", type=float, required=True, dest="consistency_w_cosine", help="Weight for cosine in consistency")
    parser.add_argument("--consistency-w-positive", type=float, required=True, dest="consistency_w_positive", help="Weight for positive-fraction in consistency")
    parser.add_argument("--consistency-w-high-sim", type=float, required=True, dest="consistency_w_high_sim", help="Weight for high-similarity in consistency")
    parser.add_argument("--sparsity-threshold-fraction", type=float, required=True, dest="sparsity_threshold_fraction", help="Fraction threshold for sparsity")
    parser.add_argument("--detection-threshold", type=float, required=True, dest="detection_threshold", help="Accuracy threshold for multi-direction detection")
    parser.add_argument("--direction-moderate-similarity", type=float, required=True, dest="direction_moderate_similarity", help="Threshold for moderate similarity")
