import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel, PretrainedConfig, AutoModel, AutoConfig


class MultiHeadConfig(PretrainedConfig):
    model_type = "multihead_text_classifier"

    def __init__(self, encoder_name="cointegrated/LaBSE-en-ru",
                 num_labels_l2=5, num_labels_l3=11, num_labels_l4=41,
                 dropout=0.1, hidden_size=768,
                 id2label_l2=None, label2id_l2=None,
                 id2label_l3=None, label2id_l3=None,
                 id2label_l4=None, label2id_l4=None,
                 l3_to_l2=None, l4_to_l3=None, merged_labels=None,
                 loss_weight_l2=0.25, loss_weight_l3=0.5, loss_weight_l4=1.0,
                 label_smoothing=0.0, **kwargs):
        super().__init__(**kwargs)
        self.encoder_name = encoder_name
        self.num_labels_l2 = num_labels_l2
        self.num_labels_l3 = num_labels_l3
        self.num_labels_l4 = num_labels_l4
        self.dropout = dropout
        self.hidden_size = hidden_size
        self.id2label_l2 = id2label_l2 or {}
        self.label2id_l2 = label2id_l2 or {}
        self.id2label_l3 = id2label_l3 or {}
        self.label2id_l3 = label2id_l3 or {}
        self.id2label_l4 = id2label_l4 or {}
        self.label2id_l4 = label2id_l4 or {}
        self.l3_to_l2 = l3_to_l2 or []          # index = l3_id -> l2_id
        self.l4_to_l3 = l4_to_l3 or []          # index = l4_id -> l3_id; -1 у мёрж-классов TC1/TC2
        self.merged_labels = merged_labels or []  # метки, требующие слоя правил (TC1/TC2)
        self.loss_weight_l2 = loss_weight_l2
        self.loss_weight_l3 = loss_weight_l3
        self.loss_weight_l4 = loss_weight_l4
        self.label_smoothing = label_smoothing


class MultiHeadClassifier(PreTrainedModel):
    """LaBSE -> [CLS] -> dropout -> три линейные головы (L2 / L3 / L4)."""
    config_class = MultiHeadConfig

    def __init__(self, config):
        super().__init__(config)
        enc_config = AutoConfig.from_pretrained(config.encoder_name)
        self.encoder = AutoModel.from_config(enc_config)   # только архитектура (веса грузятся отдельно)
        h = enc_config.hidden_size
        self.dropout = nn.Dropout(config.dropout)
        self.head_l2 = nn.Linear(h, config.num_labels_l2)
        self.head_l3 = nn.Linear(h, config.num_labels_l3)
        self.head_l4 = nn.Linear(h, config.num_labels_l4)
        # веса классов (не сохраняются в чекпойнт): по умолчанию единицы
        self.register_buffer("cw_l2", torch.ones(config.num_labels_l2), persistent=False)
        self.register_buffer("cw_l3", torch.ones(config.num_labels_l3), persistent=False)
        self.register_buffer("cw_l4", torch.ones(config.num_labels_l4), persistent=False)
        self.post_init()

    @classmethod
    def from_encoder(cls, config):
        """Инициализация для обучения: подгружаем претренированные веса энкодера."""
        model = cls(config)
        model.encoder = AutoModel.from_pretrained(config.encoder_name)
        return model

    def set_class_weights(self, w_l2, w_l3, w_l4):
        self.register_buffer("cw_l2", torch.as_tensor(w_l2, dtype=torch.float), persistent=False)
        self.register_buffer("cw_l3", torch.as_tensor(w_l3, dtype=torch.float), persistent=False)
        self.register_buffer("cw_l4", torch.as_tensor(w_l4, dtype=torch.float), persistent=False)

    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None,
                labels_l2=None, labels_l3=None, labels_l4=None, **kwargs):
        out = self.encoder(input_ids=input_ids, attention_mask=attention_mask,
                           token_type_ids=token_type_ids)
        cls = out.last_hidden_state[:, 0]          # [CLS]
        cls = self.dropout(cls)
        logits_l2 = self.head_l2(cls)
        logits_l3 = self.head_l3(cls)
        logits_l4 = self.head_l4(cls)

        loss = None
        if labels_l2 is not None and labels_l3 is not None and labels_l4 is not None:
            ls = getattr(self.config, "label_smoothing", 0.0)
            def _ce(logits, labels, cw):
                # ignore_index=-100: у мёрж-классов (TC1/TC2) нет однозначного L2/L3, и лосс этих
                # голов на таких строках выключается маркером -100 в labels (см. Шаг 3).
                w = cw.to(logits.device).to(logits.dtype)
                # если ВЕСЬ батч замаскирован, cross_entropy вернул бы NaN (деление на нулевую
                # сумму весов) и отравил бы обучение -> отдаём ровно нулевой вклад, сохраняя граф.
                if int((labels != -100).sum()) == 0:
                    return logits.sum() * 0.0
                return F.cross_entropy(logits, labels, weight=w, label_smoothing=ls,
                                       ignore_index=-100)
            l2 = _ce(logits_l2, labels_l2, self.cw_l2)
            l3 = _ce(logits_l3, labels_l3, self.cw_l3)
            l4 = _ce(logits_l4, labels_l4, self.cw_l4)
            loss = (self.config.loss_weight_l2 * l2
                    + self.config.loss_weight_l3 * l3
                    + self.config.loss_weight_l4 * l4)

        return {"loss": loss, "logits_l2": logits_l2, "logits_l3": logits_l3, "logits_l4": logits_l4}


MultiHeadConfig.register_for_auto_class()
MultiHeadClassifier.register_for_auto_class("AutoModel")