import torch
import torch.nn as nn

from typing import Optional
from dataclasses import dataclass

from transformers import PreTrainedModel
from transformers.utils import ModelOutput

from .configuration_dart2vec import Dart2VecConfig


@dataclass
class Dart2VecModelOutput(ModelOutput):
    hidden_states: torch.Tensor


@dataclass
class Dart2VecModelForFeatureExtractionOutput(ModelOutput):
    embeddings: torch.Tensor


class Dart2VecEmbeddings(nn.Module):
    def __init__(self, config: Dart2VecConfig):
        super().__init__()

        self.tag_embeddings = nn.Embedding(
            config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id
        )

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
    ):
        if inputs_embeds is not None:
            return inputs_embeds

        embeddings = self.tag_embeddings(input_ids)

        return embeddings


class Dart2VecPreTrainedModel(PreTrainedModel):
    """
    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
    models.
    """

    config_class = Dart2VecConfig
    base_model_prefix = "dart2vec"

    def _init_weights(self, module):
        """Initialize the weights"""
        if isinstance(module, nn.Embedding):
            torch.nn.init.kaiming_uniform_(module.weight)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
        elif isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=0.02)
            if module.bias is not None:
                module.bias.data.zero_()


class Dart2VecModel(Dart2VecPreTrainedModel):
    def __init__(self, config: Dart2VecConfig):
        super().__init__(config)

        self.config = config

        self.embeddings = Dart2VecEmbeddings(config)

        self.post_init()

    def get_input_embeddings(self):
        return self.embeddings.tag_embeddings

    def set_input_embeddings(self, value):
        self.embeddings.tag_embeddings = value

    def forward(
        self, input_ids: torch.Tensor
    ) -> Dart2VecModelForFeatureExtractionOutput:
        embeddings = self.embeddings(input_ids)

        return Dart2VecModelForFeatureExtractionOutput(embeddings=embeddings)