"""ST-TENP: timestep-conditioned TENP + structural sensitivity + redundancy.

This is a separate pruning method from protocol TENP. Physical slicing is the
same; selection is not. Adaptive expert budget is not included.
"""

from __future__ import annotations

import logging

from .apply import prune_moe_layer
from .scoring import DEFAULT_MODALITY_FUSE

logger = logging.getLogger("sttenp")

METHOD = "st-tenp"
VARIANTS = ("ta", "st", "st-redundancy")
DEFAULT_VARIANT = "st"


def variant_flags(variant: str) -> tuple[str, bool]:
    """Map a named ST-TENP variant to (temporal, redundancy)."""
    if variant not in VARIANTS:
        raise ValueError(f"unknown ST-TENP variant {variant!r}; expected one of {VARIANTS}")
    if variant == "ta":
        return "ta", False
    if variant == "st":
        return "st", False
    return "st", True


def prune_sttenp_layer(
    mlp,
    collector,
    p_mid: float,
    variant: str = DEFAULT_VARIANT,
    modality_fuse: str = DEFAULT_MODALITY_FUSE,
    n_min: int = 32,
    temporal_beta: float = 0.3,
    temporal_topk: int = 2,
) -> dict:
    """Prune one SparseMoEMLP with ST-TENP. Does not change the TENP default path."""
    temporal, redundancy = variant_flags(variant)
    meta = prune_moe_layer(
        mlp,
        collector,
        p_mid=p_mid,
        modality_fuse=modality_fuse,
        score="projection",
        apply_audio_discount=False,
        compensate="none",
        temporal=temporal,
        redundancy=redundancy,
        n_min=n_min,
        temporal_beta=temporal_beta,
        temporal_topk=temporal_topk,
    )
    for spec in meta.values():
        spec["method"] = METHOD
        spec["variant"] = variant
        spec["score"] = METHOD
    logger.info("ST-TENP variant=%s temporal=%s redundancy=%s p_mid=%s", variant, temporal, redundancy, p_mid)
    return meta