import torch
import torch.nn as nn
from transformers import BertModel
from typing import Dict


class KoBERT(nn.Module):
    def __init__(self,
                 bert_weight: str = "kykim/bert-kor-base") -> None:
        super(KoBERT, self).__init__()
        self.bert_model = BertModel.from_pretrained(bert_weight)

    def forward(self,
                input_ids: torch.Tensor,
                attention_mask: torch.Tensor,
                token_type_ids: torch.Tensor) -> Dict[str, torch.Tensor]:
        return self.bert_model(input_ids, attention_mask, token_type_ids)


class Pooling(nn.Module):
    def __init__(self, ):
        super(Pooling, self).__init__()

    def __repr__(self):
        return "Pooling()"

    def forward(self, features: Dict[str, torch.Tensor]):
        token_embeddings = features['token_embeddings']
        attention_mask = features['attention_mask']

        # TODO 다양한 Pooling 방식 도입해보기
        """
        우선 이 부분에서는 MEAN Pooling 방식을 사용했는데,,,
        CGD 논문 보면 다양한 Pooling 방식을 사용하고, 그 출력값을 concat해서
        상당히 괜찮은 결과를 얻었으니까!

        여기에서도 좀 다양한 방식의 Pooling 방식을 적용해보고, 실험 결과를
        비교해보면 어떨까?
        """
        input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
        sum_embeddings = torch.sum(token_embeddings * input_mask_expanded, 1)

        sum_mask = input_mask_expanded.sum(1)

        sum_mask = torch.clamp(sum_mask, min=1e-9)
        output_vector = sum_embeddings / sum_mask
        features['sentence_embedding'] = output_vector
        return features
