"""transformers-compatible config for the RadGraph-IT DyGIE++ joint NER + relation model.

Mirrors the hyperparameters of `training_v2/src/dygie/model.py`'s `DyGIEModel` plus the
label vocabulary (`training_v2/src/dygie/vocab.py`'s `Vocabulary`), so a `RadgraphModel` can
be reconstructed from `config.json` alone, matching the checkpoint trained by
`run_medbit_full_cv.sh` (single train run on the full split, not actual cross-validation --
see that script's own header comment).
"""
from transformers import PretrainedConfig


class RadgraphConfig(PretrainedConfig):
    model_type = "radgraph_it"

    def __init__(
        self,
        encoder_name: str = "IVN-RIN/medBIT-r3-plus",
        max_length: int = 512,
        max_span_width: int = 12,
        feature_size: int = 20,
        feedforward_params: dict = None,
        loss_weights: dict = None,
        relation_spans_per_word: float = 0.5,
        train_encoder: bool = True,
        span_pooling: bool = False,
        transformer_params: dict = None,
        relation_context: bool = False,
        relation_feedforward_params: dict = None,
        dataset: str = "radgraph-it",
        ner_labels: dict = None,
        relation_labels: dict = None,
        **kwargs,
    ):
        self.encoder_name = encoder_name
        self.max_length = max_length
        self.max_span_width = max_span_width
        self.feature_size = feature_size
        self.feedforward_params = feedforward_params or {"hidden_dims": [150, 150], "dropout": 0.4}
        self.loss_weights = loss_weights or {"ner": 0.2, "relation": 1.0}
        self.relation_spans_per_word = relation_spans_per_word
        self.train_encoder = train_encoder
        self.span_pooling = span_pooling
        self.transformer_params = transformer_params
        self.relation_context = relation_context
        self.relation_feedforward_params = relation_feedforward_params
        self.dataset = dataset
        # Namespace-unqualified label -> index maps, null label "" pinned to 0. Namespaced as
        # f"{dataset}__ner_labels" / f"{dataset}__relation_labels" when rebuilt into a Vocabulary
        # (see modeling_radgraph.py), matching training_v2/src/dygie/vocab.py exactly.
        self.ner_labels = ner_labels or {
            "": 0,
            "Anatomy::definitely present": 1,
            "Observation::definitely present": 2,
            "Observation::definitely absent": 3,
            "Observation::uncertain": 4,
            "Anatomy::definitely absent": 5,
            "Anatomy::uncertain": 6,
        }
        self.relation_labels = relation_labels or {
            "": 0,
            "modify": 1,
            "located_at": 2,
            "suggestive_of": 3,
        }
        super().__init__(**kwargs)