from typing import Any

from transformers import PreTrainedConfig


class SaniGecConfig(PreTrainedConfig):
    """
    Configuration for SaniGec, the non-autoregressive byte-encoder GEC tagger.

    A pretrained ByT5 encoder over raw UTF-8 bytes is pooled back to codepoint resolution and feeds
    two heads: a binary error-detection head and an edit-correction head over the tag vocabulary
    (KEEP/DELETE/APPEND_c/REPLACE_c). `backbone_name` is the HF checkpoint the encoder is first
    transferred from (via `SaniGecModel.from_backbone`); `encoder_config` is that encoder's
    serialized architecture, persisted so a reload rebuilds the encoder to the right shape offline
    and the saved state-dict fills it — the model never re-fetches the backbone on reload.
    `hidden_size` mirrors the encoder's `d_model` so the heads and the pooling anchor are sized to
    it. The loss hyperparameters live here so a checkpoint records the objective it trained under.
    """

    model_type = "sani_gec"

    def __init__(
        self,
        backbone_name: str = "google/byt5-base",
        encoder_config: dict[str, Any] | None = None,
        hidden_size: int = 1536,
        num_correction_tags: int = 512,
        correction_class_weights: list[float] | None = None,
        dropout: float = 0.1,
        detection_loss_weight: float = 1.0,
        correction_loss_weight: float = 1.0,
        focal_gamma: float = 2.0,
        label_smoothing: float = 0.1,
        keep_confidence_margin: float = 0.0,
        max_correction_passes: int = 12,
        pad_token_id: int = 0,
        **kwargs: Any,
    ) -> None:
        self.backbone_name = backbone_name
        self.encoder_config = encoder_config
        # The encoder's d_model is authoritative when its architecture is known; the default is only
        # a placeholder for a bare config built before a backbone is attached.
        self.hidden_size = encoder_config["d_model"] if encoder_config else hidden_size
        self.num_correction_tags = num_correction_tags
        self.correction_class_weights = correction_class_weights
        self.dropout = dropout
        self.detection_loss_weight = detection_loss_weight
        self.correction_loss_weight = correction_loss_weight
        self.focal_gamma = focal_gamma
        self.label_smoothing = label_smoothing
        # Inference-time decode rule (the iterative corrector reads these), persisted so a reloaded
        # checkpoint corrects with no external arguments.
        self.keep_confidence_margin = keep_confidence_margin
        self.max_correction_passes = max_correction_passes
        # Forward via a dict so the special tokens ride in **kwargs: the transformers v5 type stub
        # does not name pad_token_id on PreTrainedConfig.__init__, though the runtime accepts it.
        init_kwargs: dict[str, Any] = {"pad_token_id": pad_token_id, **kwargs}
        super().__init__(**init_kwargs)