from transformers import PretrainedConfig


class DualTurnConfig(PretrainedConfig):
    """Configuration for DualTurnModel.

    The backbone is a two-stream transformer (TurnTakingModel) that encodes
    dual-channel audio via Mimi and predicts per-frame turn-taking signals.

    The endpoint classifier is a lightweight sklearn model (logistic regression
    or gradient boosting) that converts the 10 backbone signals at a VAD-offset
    anchor into P(ST) — probability the user has finished their turn.
    """

    model_type = "dualturn"

    def __init__(
        self,
        # ── Backbone ──────────────────────────────────────────────────────
        backbone_input_mode: str = "continuous",    # "continuous" | "discrete"
        mimi_sample_rate: int = 24_000,             # Mimi encoder input SR
        mimi_frame_rate: float = 12.5,              # frames per second
        mimi_frame_ms: float = 80.0,                # ms per frame
        # ── Signals ───────────────────────────────────────────────────────
        signal_keys: list = None,
        # ── Endpoint classifier ───────────────────────────────────────────
        st_threshold: float = 0.30,
        vad_edge_threshold: float = 0.50,           # vad_user edge for anchor
        agent_voice_min: float = 0.15,              # min agent VAD to consider voicing
        fvad_alpha_short: float = 0.3,              # Silero smoothing (fast)
        fvad_alpha_long: float = 0.7,               # Silero smoothing (slow)
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.backbone_input_mode = backbone_input_mode
        self.mimi_sample_rate    = mimi_sample_rate
        self.mimi_frame_rate     = mimi_frame_rate
        self.mimi_frame_ms       = mimi_frame_ms
        self.signal_keys = signal_keys or [
            "vad_user", "vad_agent",
            "eot_user", "eot_agent",
            "bot_user", "bot_agent",
            "fvad_user_short", "fvad_user_long",
            "fvad_agent_short", "fvad_agent_long",
        ]
        self.st_threshold        = st_threshold
        self.vad_edge_threshold  = vad_edge_threshold
        self.agent_voice_min     = agent_voice_min
        self.fvad_alpha_short    = fvad_alpha_short
        self.fvad_alpha_long     = fvad_alpha_long