"""Configuration for the Dihya-5M Berber language identifier."""

from __future__ import annotations

from typing import Any

from transformers import PreTrainedConfig

CLASSES = [
    "kab_Latn",
    "shi_Latn",
    "rif_Latn",
    "taq_Latn",
    "tzm_Latn",
    "shy_Latn",
    "NOT_AMAZIGH",
]
"""Label order is the trained head's row order. Reordering it renames every prediction."""


class DihyaConfig(PreTrainedConfig):
    """The classifier's shapes, its label set, and the two inference-time corrections.

    `prior_shift` and `q_hat` are stored here rather than in a sidecar because they are
    part of what the model predicts: without the shift the argmax minimises the training
    prior's error rather than balanced error, and without the quantile there is no
    prediction set. A repository that ships the weights alone ships a different model.
    """

    model_type = "dihya"

    def __init__(
        self,
        vocab_size: int = 258,
        hidden_size: int = 256,
        intermediate_size: int = 704,
        num_attention_heads: int = 8,
        num_hidden_layers: int = 6,
        conv_kernels: list[int] | None = None,
        conv_dim: int = 128,
        max_position_embeddings: int = 256,
        rope_theta: float = 10000.0,
        dropout_prob: float = 0.1,
        logit_scale: float = 24.0,
        rms_norm_eps: float = 1e-6,
        classes: list[str] | None = None,
        prior_shift: list[float] | None = None,
        q_hat: float | None = None,
        target_coverage: float = 0.99,
        pad_token_id: int = 0,
        unk_token_id: int = 1,
        byte_offset: int = 2,
        **kwargs: Any,
    ) -> None:
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.num_attention_heads = num_attention_heads
        self.num_hidden_layers = num_hidden_layers
        self.conv_kernels = list(conv_kernels or [3, 5, 7])
        self.conv_dim = conv_dim
        self.max_position_embeddings = max_position_embeddings
        self.rope_theta = rope_theta
        self.dropout_prob = dropout_prob
        self.logit_scale = logit_scale
        self.rms_norm_eps = rms_norm_eps
        self.unk_token_id = unk_token_id
        self.byte_offset = byte_offset

        labels = list(classes or CLASSES)
        self.classes = labels
        self.prior_shift = list(prior_shift) if prior_shift is not None else [0.0] * len(labels)
        self.q_hat = q_hat
        self.target_coverage = target_coverage

        kwargs.setdefault("pad_token_id", pad_token_id)
        kwargs.setdefault("id2label", dict(enumerate(labels)))
        kwargs.setdefault("label2id", {name: i for i, name in enumerate(labels)})
        super().__init__(**kwargs)

    @property
    def head_size(self) -> int:
        return self.hidden_size // self.num_attention_heads


__all__ = ["CLASSES", "DihyaConfig"]