import torch
import torch.nn as nn


class DualEmbeddings(nn.Module):
    def __init__(self, config):
        super().__init__()
        d = config.word_char_emb_dim

        self.char_embeddings = nn.Embedding(
            config.vocab_char_size, d, padding_idx=config.pad_token_id
        )
        self.word_embeddings = nn.Embedding(
            config.vocab_word_size, d, padding_idx=0
        )
        self.projection = nn.Linear(2 * d, config.hidden_size, bias=False)
        self.position_embeddings = nn.Embedding(
            config.max_position_embeddings, config.hidden_size
        )
        self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

        self.register_buffer(
            "position_ids", torch.arange(config.max_position_embeddings).unsqueeze(0), persistent=True
        )

    def forward(self, input_ids, word_ids):
        bsz, seq_len = input_ids.shape
        pos_ids = self.position_ids[:, :seq_len]

        c = self.char_embeddings(input_ids)
        w = self.word_embeddings(word_ids)

        x = torch.cat([c, w], dim=-1)
        x = self.projection(x)
        x = x + self.position_embeddings(pos_ids)
        x = self.layer_norm(x)
        x = self.dropout(x)
        return x