"""
HuggingFace-compatible wrapper for XLMEduModel.

Architecture: XLM-RoBERTa-large encoder → linear classifier → CRF (torchcrf)
Task: BIO token classification for situation-entity segmentation (B-EDU / I-EDU / O)

Loading from the Hub (recommended):
    from transformers import AutoConfig, AutoModel
    config = AutoConfig.from_pretrained("xaver-krueckl/situation-entity-segmenter", trust_remote_code=True)
    model  = AutoModel.from_pretrained("xaver-krueckl/situation-entity-segmenter", trust_remote_code=True)

Inference (one pre-split sentence):
    from transformers import AutoTokenizer
    tokenizer = AutoTokenizer.from_pretrained("FacebookAI/xlm-roberta-large")
    words     = ["The", "cat", "sat", "."]
    pairs     = model.predict_text(words, tokenizer)  # [(word, tag), ...]
"""

from __future__ import annotations

from typing import Dict, List, Optional

import torch
from torchcrf import CRF
from transformers import AutoConfig, AutoModel, PretrainedConfig, PreTrainedModel


LABELS   = {"B-EDU": 0, "I-EDU": 1, "O": 2}
BIO_TAGS = [tag for tag, _ in sorted(LABELS.items(), key=lambda x: x[1])]
NUM_TAGS = len(BIO_TAGS)
O_IDX    = LABELS["O"]


class XLMEduConfig(PretrainedConfig):
    model_type = "xlm_edu"

    def __init__(
        self,
        encoder_name: str = "FacebookAI/xlm-roberta-large",
        num_tags: int = NUM_TAGS,
        label2id: Optional[Dict[str, int]] = None,
        id2label: Optional[Dict[int, str]] = None,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.encoder_name = encoder_name
        self.num_tags     = num_tags
        self.label2id     = label2id or LABELS
        raw_id2label      = id2label or {v: k for k, v in LABELS.items()}
        self.id2label     = {int(k): v for k, v in raw_id2label.items()}
        # must be an instance attribute so save_pretrained writes it into config.json
        self.auto_map = {
            "AutoConfig": "modeling_xlmedu.XLMEduConfig",
            "AutoModel":  "modeling_xlmedu.XLMEduModelHF",
        }


class XLMEduModelHF(PreTrainedModel):
    config_class = XLMEduConfig
    _tied_weights_keys = []

    @property
    def all_tied_weights_keys(self):
        return {}

    def __init__(self, config: XLMEduConfig):
        super().__init__(config)
        # from_config creates an empty encoder shell; weights are loaded by the outer from_pretrained
        encoder_config  = AutoConfig.from_pretrained(config.encoder_name)
        self.encoder    = AutoModel.from_config(encoder_config)
        self.classifier = torch.nn.Linear(self.encoder.config.hidden_size, config.num_tags)
        self.crf        = CRF(config.num_tags, batch_first=True)

    def forward(
        self,
        input_ids:      torch.Tensor,
        attention_mask: torch.Tensor,
        labels:         Optional[torch.Tensor] = None,
    ):
        """
        Returns a dict with 'loss' (if labels given) and 'logits' (emission scores).
        logits shape: (batch, seq_len, num_tags)
        """
        outputs    = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
        emissions  = self.classifier(outputs.last_hidden_state)

        loss = None
        if labels is not None:
            crf_mask   = attention_mask.bool()
            labels_crf = labels.clone()
            labels_crf[labels_crf == -100] = O_IDX
            loss = -self.crf(emissions.float(), labels_crf, mask=crf_mask, reduction="mean")

        return {"loss": loss, "logits": emissions}

    def predict(
        self,
        input_ids:      torch.Tensor,
        attention_mask: torch.Tensor,
    ) -> List[List[int]]:
        """Viterbi-decode the best tag sequence for each sample in the batch.

        Returns one integer tag per sub-token (including special tokens).
        During training, non-first sub-tokens were masked (-100 → O), so their
        predictions are meaningless. Callers should use encoding.word_ids() to
        read only the first sub-token of each word.
        """
        emissions = self.forward(input_ids, attention_mask)["logits"]
        return self.crf.decode(emissions.float(), mask=attention_mask.bool())

    def predict_text(
        self,
        words:     List[str],
        tokenizer,
    ) -> List[tuple]:
        """Tag a single pre-tokenised sentence.

        Args:
            words:     word-level tokens for one sentence (e.g. from spaCy).
            tokenizer: the HuggingFace tokenizer for this model.

        Returns:
            List of (word, tag) pairs, one per input word.
        """
        from collections import defaultdict

        encoding = tokenizer(
            words,
            return_tensors="pt",
            is_split_into_words=True,
            truncation=True,
            max_length=512,
        )

        with torch.no_grad():
            raw_tag_ids = self.predict(
                encoding["input_ids"], encoding["attention_mask"]
            )[0]

        word_ids       = encoding.word_ids(batch_index=0)
        word_token_ids: dict = defaultdict(list)
        for pos, word_idx in enumerate(word_ids):
            if word_idx is not None:
                word_token_ids[word_idx].append(
                    encoding["input_ids"][0, pos].item()
                )

        results: List[tuple] = []
        prev_word_idx = None
        for pos, word_idx in enumerate(word_ids):
            if word_idx is None:
                continue
            if word_idx != prev_word_idx:
                word = tokenizer.decode(word_token_ids[word_idx]).strip()
                tag  = self.config.id2label[raw_tag_ids[pos]]
                results.append((word, tag))
                prev_word_idx = word_idx

        return results