"""HS6 product classifier: XLM-RoBERTa encoder + one flat linear head.

HS4 and HS2 are not separate heads. They are marginals of the same HS6
distribution (logsumexp over the children of each parent), so the levels are
consistent by construction: the model cannot name one heading at 4 digits and a
code from a different heading at 6 digits.

The forward pass below must stay identical to the one used in training,
otherwise the released weights do not mean what the metrics say they mean.
"""

from dataclasses import dataclass
from typing import List, Optional

import torch
import torch.nn as nn
from transformers.modeling_outputs import ModelOutput
from transformers.modeling_utils import PreTrainedModel
from transformers.models.xlm_roberta.configuration_xlm_roberta import XLMRobertaConfig
from transformers.models.xlm_roberta.modeling_xlm_roberta import XLMRobertaModel

from .configuration_hs6 import HS6ClassifierConfig

# fields that belong to the classifier, not to the encoder
_HEAD_ONLY = ("n2", "n4", "n6", "head_dropout", "pooling", "recommended_max_length",
              "id2hs4", "id2hs2", "id2label", "label2id", "auto_map", "architectures",
              "model_type")


def _encoder_config(config: HS6ClassifierConfig) -> XLMRobertaConfig:
    """A plain XLM-R config, so the encoder does not warn about the wrapper type."""
    raw = {k: v for k, v in config.to_dict().items() if k not in _HEAD_ONLY}
    return XLMRobertaConfig(**raw)


@dataclass
class HS6ClassifierOutput(ModelOutput):
    """`logits` is the HS6 level, so the standard text-classification tooling works.

    `logits_hs4` / `logits_hs2` are the marginals over the same distribution.
    """

    loss: Optional[torch.FloatTensor] = None
    logits: Optional[torch.FloatTensor] = None
    logits_hs4: Optional[torch.FloatTensor] = None
    logits_hs2: Optional[torch.FloatTensor] = None


class HS6ClassifierModel(PreTrainedModel):
    config_class = HS6ClassifierConfig
    base_model_prefix = "encoder"
    supports_gradient_checkpointing = True

    def __init__(self, config: HS6ClassifierConfig):
        super().__init__(config)
        # no pooling layer: the head reads the CLS token of the last hidden state
        self.encoder = XLMRobertaModel(_encoder_config(config), add_pooling_layer=False)
        self.pooling = config.pooling
        self.drop = nn.Dropout(config.head_dropout)
        self.head_6 = nn.Linear(config.hidden_size, config.n6)
        self.n4, self.n2 = config.n4, config.n2
        # parent of every HS6 class: its first 4 and first 2 digits. Stored in the
        # checkpoint so the mapping cannot drift away from the trained weights.
        self.register_buffer("parent4", torch.zeros(config.n6, dtype=torch.long))
        self.register_buffer("parent2", torch.zeros(config.n6, dtype=torch.long))
        self.post_init()

    def _marginal(self, l6, parent, n_parent):
        mx = l6.max(1, keepdim=True).values
        e = (l6 - mx).exp()
        s = torch.zeros(l6.size(0), n_parent, device=l6.device, dtype=e.dtype)
        s.index_add_(1, parent, e)
        return s.clamp_min(1e-20).log() + mx

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        labels: Optional[torch.LongTensor] = None,
        return_dict: Optional[bool] = None,
        **kwargs,
    ):
        return_dict = True if return_dict is None else return_dict
        out = self.encoder(input_ids=input_ids, attention_mask=attention_mask)

        if self.pooling == "mean":
            h = out.last_hidden_state
            m = attention_mask.unsqueeze(-1).to(h.dtype)
            pooled = (h * m).sum(1) / m.sum(1).clamp(min=1)
        else:
            pooled = out.last_hidden_state[:, 0, :]

        l6 = self.head_6(self.drop(pooled))
        l4 = self._marginal(l6, self.parent4, self.n4)
        l2 = self._marginal(l6, self.parent2, self.n2)

        loss = None
        if labels is not None:
            loss = nn.functional.cross_entropy(l6, labels)

        if not return_dict:
            return (loss, l6, l4, l2) if loss is not None else (l6, l4, l2)
        return HS6ClassifierOutput(loss=loss, logits=l6, logits_hs4=l4, logits_hs2=l2)

    @torch.no_grad()
    def classify(self, texts, tokenizer, top_k: int = 5, batch_size: int = 16,
                 max_length: Optional[int] = None) -> List[List[dict]]:
        """Convenience wrapper: texts in, ranked HS6 codes with probabilities out."""
        if isinstance(texts, str):
            texts = [texts]
        max_length = max_length or self.config.recommended_max_length
        device = next(self.parameters()).device
        results = []
        for start in range(0, len(texts), batch_size):
            chunk = [t if isinstance(t, str) and t.strip() else " "
                     for t in texts[start:start + batch_size]]
            enc = tokenizer(chunk, truncation=True, max_length=max_length,
                            padding=True, return_tensors="pt").to(device)
            logits = self(**enc).logits
            probs = torch.softmax(logits.float(), dim=1)
            conf, idx = probs.topk(min(top_k, probs.size(1)), dim=1)
            for c, i in zip(conf.cpu().tolist(), idx.cpu().tolist()):
                results.append([
                    {"hs6": self.config.id2label[j], "score": round(v, 6)}
                    for j, v in zip(i, c)
                ])
        return results