"""Steering method definitions: TETNO, GROM, PRZELOM."""

from wisent.core.control.steering_methods.registry.registry import (
    SteeringMethodDefinition,
    SteeringMethodParameter,
    SteeringMethodType,
)
from wisent.core.control.steering_methods.configs.optimal import get_optimal


TETNO_STEERING_LAYER_CONFIGS = ("single_best", "range_3", "range_5")

TETNO_DEFINITION = SteeringMethodDefinition(
    name="tetno",
    method_type=SteeringMethodType.TETNO,
    description="TETNO - Probabilistic Uncertainty-guided Layer Steering Engine. Layer-adaptive conditional steering with uncertainty-guided intensity.",
    method_class_path="wisent.core.control.steering_methods.methods.tetno.TETNOMethod",
    parameters=[
        SteeringMethodParameter(
            name="sensor_layer",
            type=int,
            default=None, required=False,
            help="Layer index where condition gating is computed (auto-computed if not set)",
            cli_flag="--tetno-sensor-layer",
        ),
        SteeringMethodParameter(
            name="steering_layers",
            type=str,
            default=None, required=False,
            help="Comma-separated layer indices where steering is applied (auto-computed if not set)",
            cli_flag="--tetno-steering-layers",
        ),
        SteeringMethodParameter(
            name="per_layer_scaling",
            type=bool,
            default=True, required=False,
            help="Learn different scaling per layer",
            action="store_true",
            cli_flag="--tetno-per-layer-scaling",
        ),
        SteeringMethodParameter(
            name="condition_threshold",
            type=float,
            help="Threshold for condition activation (0-1)",
            cli_flag="--tetno-condition-threshold",
        ),
        SteeringMethodParameter(
            name="gate_temperature",
            type=float,
            help="Temperature for sigmoid gating (lower = sharper)",
            cli_flag="--tetno-gate-temperature",
        ),
        SteeringMethodParameter(
            name="learn_threshold",
            type=bool,
            default=True, required=False,
            help="Learn optimal threshold via grid search",
            action="store_true",
            cli_flag="--tetno-learn-threshold",
        ),
        SteeringMethodParameter(
            name="use_entropy_scaling",
            type=bool,
            default=True, required=False,
            help="Enable entropy-based intensity modulation",
            action="store_true",
            cli_flag="--tetno-use-entropy-scaling",
        ),
        SteeringMethodParameter(
            name="entropy_floor",
            type=float,
            help="Minimum entropy to trigger scaling",
            cli_flag="--tetno-entropy-floor",
        ),
        SteeringMethodParameter(
            name="entropy_ceiling",
            type=float,
            help="Entropy at which max_alpha is reached",
            cli_flag="--tetno-entropy-ceiling",
        ),
        SteeringMethodParameter(
            name="max_alpha",
            type=float,
            help="Maximum steering strength",
            cli_flag="--tetno-max-alpha",
        ),
        SteeringMethodParameter(
            name="optimization_steps",
            type=int,
            help="Steps for condition vector optimization",
            cli_flag="--tetno-optimization-steps",
        ),
        SteeringMethodParameter(
            name="learning_rate",
            type=float,
            help="Learning rate for optimization",
            cli_flag="--tetno-learning-rate",
        ),
        SteeringMethodParameter(
            name="normalize",
            type=bool,
            default=True, required=False,
            help="L2-normalize vectors",
            action="store_true",
            cli_flag="--tetno-normalize",
        ),
        SteeringMethodParameter(name="threshold_search_steps", type=int,
            help="Steps for threshold grid search", cli_flag="--tetno-threshold-search-steps"),
        SteeringMethodParameter(name="condition_margin", type=float,
            help="Margin for condition activation", cli_flag="--tetno-condition-margin"),
        SteeringMethodParameter(name="min_layer_scale", type=float,
            help="Minimum per-layer scaling factor", cli_flag="--tetno-min-layer-scale"),
        SteeringMethodParameter(name="hybrid_strength_factor", type=float,
            help="Strength factor for hybrid weight export", cli_flag="--tetno-hybrid-strength-factor"),
        SteeringMethodParameter(name="gate_scale_factor", type=float,
            help="Scale factor for sigmoid gating", cli_flag="--tetno-gate-scale-factor"),
        SteeringMethodParameter(name="log_interval", type=int,
            help="Training progress log interval (every N steps)", cli_flag="--tetno-log-interval"),
    ],
)


GROM_DEFINITION = SteeringMethodDefinition(
    name="grom",
    method_type=SteeringMethodType.GROM,
    description="GROM - Joint optimization of manifold, gating, and intensity.",
    method_class_path="wisent.core.control.steering_methods.methods.grom.GROMMethod",
    parameters=[
        SteeringMethodParameter(name="num_directions", type=int,
            help="Number of directions per layer", cli_flag="--grom-num-directions"),
        SteeringMethodParameter(name="steering_layers", type=str, default=None, required=False,
            help="Layer indices for steering (auto if not set)", cli_flag="--grom-steering-layers"),
        SteeringMethodParameter(name="sensor_layer", type=int, default=None, required=False,
            help="Primary layer for gating (auto if not set)", cli_flag="--grom-sensor-layer"),
        SteeringMethodParameter(name="gate_hidden_dim", type=int, default=None, required=False,
            help="Gate network hidden dim (auto if not set)", cli_flag="--grom-gate-hidden-dim"),
        SteeringMethodParameter(name="intensity_hidden_dim", type=int, default=None, required=False,
            help="Intensity network hidden dim (auto if not set)", cli_flag="--grom-intensity-hidden-dim"),
        SteeringMethodParameter(name="optimization_steps", type=int,
            help="Total optimization steps", cli_flag="--grom-optimization-steps"),
        SteeringMethodParameter(name="learning_rate", type=float,
            help="Learning rate for all components", cli_flag="--grom-learning-rate"),
        SteeringMethodParameter(name="warmup_steps", type=int,
            help="Warmup steps for LR scheduler", cli_flag="--grom-warmup-steps"),
        SteeringMethodParameter(name="behavior_weight", type=float,
            help="Weight for behavior loss", cli_flag="--grom-behavior-weight"),
        SteeringMethodParameter(name="retain_weight", type=float,
            help="Weight for retain loss", cli_flag="--grom-retain-weight"),
        SteeringMethodParameter(name="sparse_weight", type=float,
            help="Weight for sparsity loss", cli_flag="--grom-sparse-weight"),
        SteeringMethodParameter(name="smooth_weight", type=float,
            help="Weight for smoothness loss", cli_flag="--grom-smooth-weight"),
        SteeringMethodParameter(name="independence_weight", type=float,
            help="Weight for direction independence loss", cli_flag="--grom-independence-weight"),
        SteeringMethodParameter(name="max_alpha", type=float,
            help="Maximum steering intensity", cli_flag="--grom-max-alpha"),
        SteeringMethodParameter(name="gate_temperature", type=float,
            help="Temperature for sigmoid gating", cli_flag="--grom-gate-temperature"),
        SteeringMethodParameter(name="max_grad_norm", type=float,
            help="Max gradient norm for clipping", cli_flag="--grom-max-grad-norm"),
        SteeringMethodParameter(name="eta_min_factor", type=float,
            help="Min LR factor for cosine annealing", cli_flag="--grom-eta-min-factor"),
        SteeringMethodParameter(name="linear_threshold", type=float,
            help="Threshold to classify geometry as linear", cli_flag="--grom-linear-threshold"),
        SteeringMethodParameter(name="adapt_cone_threshold", type=float,
            help="Geometry score for cone adaptation", cli_flag="--grom-adapt-cone-threshold"),
        SteeringMethodParameter(name="adapt_manifold_threshold", type=float,
            help="Geometry score for manifold adaptation", cli_flag="--grom-adapt-manifold-threshold"),
        SteeringMethodParameter(name="adapt_linear_directions", type=int,
            help="Num directions for linear geometry", cli_flag="--grom-adapt-linear-directions"),
        SteeringMethodParameter(name="adapt_complex_directions", type=int,
            help="Num directions for complex geometry", cli_flag="--grom-adapt-complex-directions"),
        SteeringMethodParameter(name="adapt_max_directions", type=int,
            help="Max directions after adaptation", cli_flag="--grom-adapt-max-directions"),
        SteeringMethodParameter(name="significant_directions_default", type=int,
            help="Default significant directions", cli_flag="--grom-significant-directions-default"),
        SteeringMethodParameter(name="min_adapted_directions", type=int,
            help="Minimum directions after adaptation", cli_flag="--grom-min-adapted-directions"),
        SteeringMethodParameter(name="caa_similarity_skip", type=float,
            help="Skip GROM if CAA similarity exceeds this", cli_flag="--grom-caa-similarity-skip"),
        SteeringMethodParameter(name="contrastive_margin", type=float,
            help="Margin for contrastive loss", cli_flag="--grom-contrastive-margin"),
        SteeringMethodParameter(name="contrastive_weight", type=float,
            help="Weight for contrastive loss", cli_flag="--grom-contrastive-weight"),
        SteeringMethodParameter(name="utility_weight", type=float,
            help="Weight for utility loss", cli_flag="--grom-utility-weight"),
        SteeringMethodParameter(name="concentration_weight", type=float,
            help="Weight for concentration loss", cli_flag="--grom-concentration-weight"),
        SteeringMethodParameter(name="gate_warmup_weight", type=float,
            help="Weight for gate warmup loss", cli_flag="--grom-gate-warmup-weight"),
        SteeringMethodParameter(name="caa_alignment_weight", type=float,
            help="Weight for CAA alignment loss", cli_flag="--grom-caa-alignment-weight"),
        SteeringMethodParameter(name="gate_dim_min", type=int,
            help="Min gate network hidden dim", cli_flag="--grom-gate-dim-min"),
        SteeringMethodParameter(name="gate_dim_max", type=int,
            help="Max gate network hidden dim", cli_flag="--grom-gate-dim-max"),
        SteeringMethodParameter(name="gate_dim_divisor", type=int,
            help="Divisor for gate dim snapping", cli_flag="--grom-gate-dim-divisor"),
        SteeringMethodParameter(name="intensity_dim_min", type=int,
            help="Min intensity network hidden dim", cli_flag="--grom-intensity-dim-min"),
        SteeringMethodParameter(name="intensity_dim_max", type=int,
            help="Max intensity network hidden dim", cli_flag="--grom-intensity-dim-max"),
        SteeringMethodParameter(name="intensity_dim_divisor", type=int,
            help="Divisor for intensity dim snapping", cli_flag="--grom-intensity-dim-divisor"),
        SteeringMethodParameter(name="gate_shrink_factor", type=int,
            help="Internal shrink factor for gate network bottleneck layer", cli_flag="--grom-gate-shrink-factor"),
        SteeringMethodParameter(name="log_interval", type=int,
            help="Training log interval (every N steps)", cli_flag="--grom-log-interval"),
        SteeringMethodParameter(name="normalize", type=bool, default=get_optimal("normalize"), required=False,
            help="L2-normalize directions", action="store_true", cli_flag="--grom-normalize"),
        SteeringMethodParameter(name="weight_decay", type=float,
            help="Weight decay for optimizer", cli_flag="--grom-weight-decay"),
        SteeringMethodParameter(name="min_cosine_sim", type=float,
            help="Minimum cosine similarity for cone analysis", cli_flag="--grom-min-cosine-sim"),
        SteeringMethodParameter(name="max_cosine_sim", type=float,
            help="Maximum cosine similarity threshold", cli_flag="--grom-max-cosine-sim"),
        SteeringMethodParameter(name="create_noise_scale", type=float,
            help="Noise scale for steering object creation", cli_flag="--grom-create-noise-scale"),
        SteeringMethodParameter(name="create_gate_threshold", type=float,
            help="Gate threshold for steering object creation", cli_flag="--grom-create-gate-threshold"),
    ],
)


PRZELOM_DEFINITION = SteeringMethodDefinition(
    name="przelom",
    method_type=SteeringMethodType.PRZELOM,
    description="Attention-Transport steering - computes desired transport plan via EOT and inverts through Q-projection pseudoinverse.",
    method_class_path="wisent.core.control.steering_methods.methods.przelom.PrzelomMethod",
    parameters=[
        SteeringMethodParameter(
            name="epsilon",
            type=float,
            help="Entropic regularization (softmax temperature) for transport plan",
            cli_flag="--przelom-epsilon",
        ),
        SteeringMethodParameter(
            name="target_mode",
            type=str,
            required=True,
            help="Target transport plan: uniform or nearest (must be specified explicitly)",
            cli_flag="--przelom-target-mode",
            choices=["uniform", "nearest"],
        ),
        SteeringMethodParameter(
            name="regularization",
            type=float,
            help="Tikhonov regularization for pseudoinverse",
            cli_flag="--przelom-regularization",
        ),
        SteeringMethodParameter(
            name="inference_k",
            type=int,
            help="Nearest source points for inference interpolation",
            cli_flag="--przelom-inference-k",
        ),
    ],
)
