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=4, num_labels_l3=14, dropout=0.1, hidden_size=768,
                 id2label_l2=None, label2id_l2=None,
                 id2label_l3=None, label2id_l3=None,
                 l3_to_l2=None, loss_weight_l2=0.5, loss_weight_l3=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.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.l3_to_l2 = l3_to_l2 or []
        self.loss_weight_l2 = loss_weight_l2
        self.loss_weight_l3 = loss_weight_l3
        self.label_smoothing = label_smoothing


class MultiHeadClassifier(PreTrainedModel):
    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.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.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):
        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)

    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None,
                labels_l2=None, labels_l3=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)

        loss = None
        if labels_l2 is not None and labels_l3 is not None:
            cw2 = self.cw_l2.to(logits_l2.device).to(logits_l2.dtype)
            cw3 = self.cw_l3.to(logits_l3.device).to(logits_l3.dtype)
            ls = getattr(self.config, "label_smoothing", 0.0)
            l2 = F.cross_entropy(logits_l2, labels_l2, weight=cw2, label_smoothing=ls)
            l3 = F.cross_entropy(logits_l3, labels_l3, weight=cw3, label_smoothing=ls)
            loss = self.config.loss_weight_l2 * l2 + self.config.loss_weight_l3 * l3

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


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