import torch
import torch.nn as nn
from transformers import PretrainedConfig, PreTrainedModel
from transformers.modeling_outputs import SequenceClassifierOutput


class MyTextConfig(PretrainedConfig):
    model_type = "my-text-sequence-classification"

    def __init__(
        self,
        vocab_size=30522,
        max_position_embeddings=128,
        hidden_size=128,
        num_labels=2,
        hidden_dropout_prob=0.1,
        **kwargs
    ):
        super().__init__(**kwargs)

        # ## vocab_size
        # 어휘 크기(vocab size): 입력 token id의 전체 개수
        self.vocab_size = vocab_size

        # ## max_position_embeddings
        # 최대 위치 임베딩(max position embeddings): 입력 sequence의 최대 길이
        self.max_position_embeddings = max_position_embeddings

        # ## hidden_size
        # 은닉 차원(hidden size): 각 token vector의 차원
        self.hidden_size = hidden_size

        # ## num_labels
        # 분류 label 개수(num labels): 출력 class 개수
        self.num_labels = num_labels

        # ## hidden_dropout_prob
        # 드롭아웃 비율(dropout probability): 과적합(overfitting) 방지
        self.hidden_dropout_prob = hidden_dropout_prob


class MyTextSequenceClassification(PreTrainedModel):
    config_class = MyTextConfig

    def __init__(self, config):
        super().__init__(config)

        # ## Token Embedding
        # token id를 dense vector로 변환
        self.token_embedding = nn.Embedding(
            config.vocab_size,
            config.hidden_size
        )

        # ## Position Embedding
        # token의 위치(position) 정보를 vector에 추가
        self.position_embedding = nn.Embedding(
            config.max_position_embeddings,
            config.hidden_size
        )

        # ## Dropout
        # 학습 중 일부 값을 제거하여 overfitting을 줄임
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

        # ## Classifier
        # pooled output을 class별 logits로 변환
        self.classifier = nn.Linear(
            config.hidden_size,
            config.num_labels
        )

        # ## Loss Function
        # 다중 분류 문제에 사용하는 CrossEntropyLoss
        self.loss_fn = nn.CrossEntropyLoss()

        self.post_init()

    def forward(self, input_ids, attention_mask=None, labels=None):
        batch_size, seq_length = input_ids.size()

        # ## Position IDs
        # 각 token 위치에 해당하는 index 생성
        position_ids = torch.arange(
            seq_length,
            dtype=torch.long,
            device=input_ids.device
        ).unsqueeze(0).expand(batch_size, seq_length)

        # ## Embedding
        # token embedding과 position embedding을 더해서 입력 표현 생성
        x = self.token_embedding(input_ids) + self.position_embedding(position_ids)

        # ## Mean Pooling
        # sequence 전체 token vector를 하나의 문장 vector로 압축
        if attention_mask is not None:
            mask = attention_mask.unsqueeze(-1).float()
            x = x * mask
            pooled_output = x.sum(dim=1) / mask.sum(dim=1).clamp(min=1.0)
        else:
            pooled_output = x.mean(dim=1)

        pooled_output = self.dropout(pooled_output)

        # ## Logits
        # 각 label에 대한 예측 점수
        logits = self.classifier(pooled_output)

        # ## Loss
        # labels가 주어지면 classification loss 계산
        loss = None
        if labels is not None:
            loss = self.loss_fn(logits, labels)

        return SequenceClassifierOutput(
            loss=loss,
            logits=logits
        )


if __name__ == "__main__":
    config = MyTextConfig()
    model = MyTextSequenceClassification(config)

    input_ids = torch.randint(0, config.vocab_size, (2, 10))
    attention_mask = torch.ones_like(input_ids)
    labels = torch.tensor([0, 1])

    outputs = model(
        input_ids=input_ids,
        attention_mask=attention_mask,
        labels=labels
    )

    print("Loss:", outputs.loss)
    print("Logits:", outputs.logits)
    print("Logits shape:", outputs.logits.shape)