import math

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel

try:
    from .configuration_custom import CustomConfig
except ImportError:
    from configuration_custom import CustomConfig


class MultiHeadAttention(nn.Module):
    """멀티 헤드 어텐션(Multi-Head Attention)을 구현합니다."""

    def __init__(self, hidden_size, num_heads, dropout=0.1):
        super().__init__()
        if hidden_size % num_heads != 0:
            raise ValueError("hidden_size must be divisible by num_heads")

        self.num_heads = num_heads
        self.head_size = hidden_size // num_heads

        # Query, Key, Value를 만들기 위한 선형 계층(linear layers)입니다.
        self.linear_q = nn.Linear(hidden_size, hidden_size)
        self.linear_k = nn.Linear(hidden_size, hidden_size)
        self.linear_v = nn.Linear(hidden_size, hidden_size)
        self.linear_out = nn.Linear(hidden_size, hidden_size)
        self.dropout = nn.Dropout(dropout)

    def _split_heads(self, x, batch_size, seq_len):
        # hidden_size를 num_heads와 head_size로 나누어 head별 표현(representation)을 만듭니다.
        return x.view(batch_size, seq_len, self.num_heads, self.head_size).transpose(1, 2)

    def forward(self, x, pad_mask=None):
        batch_size, seq_len, hidden_size = x.size()

        # 셀프 어텐션(self-attention)이므로 query, key, value가 모두 같은 입력에서 나옵니다.
        query = self._split_heads(self.linear_q(x), batch_size, seq_len)
        key = self._split_heads(self.linear_k(x), batch_size, seq_len)
        value = self._split_heads(self.linear_v(x), batch_size, seq_len)

        # 스케일드 닷 프로덕트 어텐션(scaled dot-product attention) 점수를 계산합니다.
        scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(self.head_size)

        if pad_mask is not None:
            # 패딩 마스크(padding mask)가 0인 위치는 attention에서 제외합니다.
            pad_mask = pad_mask.unsqueeze(1).unsqueeze(2)
            scores = scores.masked_fill(pad_mask == 0, torch.finfo(scores.dtype).min)

        attn_weights = F.softmax(scores, dim=-1)
        attn_weights = self.dropout(attn_weights)

        context_by_heads = torch.matmul(attn_weights, value)
        context = context_by_heads.transpose(1, 2).contiguous().view(batch_size, seq_len, hidden_size)
        context = self.linear_out(context)
        return context, attn_weights, scores


class PositionWiseFeedForward(nn.Module):
    """위치별 피드포워드 네트워크(Position-Wise Feed-Forward Network, FFN)를 구현합니다."""

    def __init__(self, hidden_size, intermediate_size, dropout=0.1):
        super().__init__()

        # 첫 번째 선형 변환(linear transformation): hidden_size를 intermediate_size로 확장(expand)합니다.
        self.linear_expand = nn.Linear(hidden_size, intermediate_size)

        # 비선형 활성 함수(non-linear activation function)입니다.
        self.activation = nn.ReLU()

        # 두 번째 선형 변환(linear transformation): 원래 hidden_size로 투영(project)합니다.
        self.linear_proj = nn.Linear(intermediate_size, hidden_size)

        # 드롭아웃(dropout)은 과적합(overfitting)을 줄이는 정규화(regularization)입니다.
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        # 순전파(forward pass): expand -> activation -> project -> dropout 순서입니다.
        x = self.linear_expand(x)
        x = self.activation(x)
        x = self.linear_proj(x)
        x = self.dropout(x)
        return x


class EncoderBlock(nn.Module):
    """Transformer의 단일 인코더 블록(Encoder Block)을 구현합니다."""

    def __init__(self, hidden_size, num_heads, intermediate_size, dropout=0.1, layer_norm_eps=1e-12):
        super().__init__()

        # 멀티 헤드 셀프 어텐션(Multi-Head Self-Attention) 계층입니다.
        self.mha = MultiHeadAttention(hidden_size=hidden_size, num_heads=num_heads, dropout=dropout)
        self.mha_norm = nn.LayerNorm(hidden_size, eps=layer_norm_eps)
        self.dropout = nn.Dropout(dropout)

        # 위치별 피드포워드 네트워크(Position-Wise Feed-Forward Network, FFN)입니다.
        self.pff = PositionWiseFeedForward(
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            dropout=dropout,
        )
        self.pff_norm = nn.LayerNorm(hidden_size, eps=layer_norm_eps)

    def forward(self, x, pad_mask=None):
        # Attention 서브레이어(sub-layer): residual connection 후 layer normalization을 적용합니다.
        attn_output, _, _ = self.mha(x, pad_mask)
        x = self.mha_norm(x + self.dropout(attn_output))

        # FFN 서브레이어(sub-layer): residual connection 후 layer normalization을 적용합니다.
        ffn_output = self.pff(x)
        x = self.pff_norm(x + ffn_output)
        return x


class PositionEncoding(nn.Module):
    """사인/코사인 위치 인코딩(sinusoidal positional encoding)을 구현합니다."""

    def __init__(self, hidden_size, max_len=5000, dropout=0.1):
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)

        pe = torch.zeros(max_len, hidden_size)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        denominator = torch.exp(torch.arange(0, hidden_size, 2, dtype=torch.float) * (-math.log(10000.0) / hidden_size))

        # 짝수 차원에는 sine, 홀수 차원에는 cosine 값을 넣습니다.
        pe[:, 0::2] = torch.sin(position * denominator)
        pe[:, 1::2] = torch.cos(position * denominator[: pe[:, 1::2].shape[1]])

        # register_buffer는 학습 파라미터(parameter)는 아니지만 모델과 함께 저장되는 tensor입니다.
        self.register_buffer("pe", pe.unsqueeze(0))

    def forward(self, x):
        # 입력 임베딩(embedding)에 위치 정보(positional information)를 더합니다.
        x = x + self.pe[:, : x.size(1)]
        return self.dropout(x)


class TransformerEncoder(PreTrainedModel):
    """토큰 임베딩(token embedding)과 EncoderBlock 스택으로 구성된 Transformer Encoder입니다."""

    config_class = CustomConfig
    base_model_prefix = "transformer_encoder"

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

        self.embedding = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
        self.pos_encoder = PositionEncoding(
            hidden_size=config.hidden_size,
            max_len=config.max_position_embeddings,
            dropout=config.hidden_dropout_prob,
        )
        self.layers = nn.ModuleList(
            [
                EncoderBlock(
                    hidden_size=config.hidden_size,
                    num_heads=config.num_attention_heads,
                    intermediate_size=config.intermediate_size,
                    dropout=config.hidden_dropout_prob,
                    layer_norm_eps=config.layer_norm_eps,
                )
                for _ in range(config.num_hidden_layers)
            ]
        )
        self.post_init()

    def _init_weights(self, module):
        # Hugging Face 초기화(initialization) 흐름에서 사용하는 가중치 초기화입니다.
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)

    def forward(self, input_ids, pad_mask=None):
        # input_ids를 임베딩(embedding)한 뒤 위치 인코딩(positional encoding)을 더합니다.
        x = self.embedding(input_ids)
        x = self.pos_encoder(x)

        # 여러 EncoderBlock을 순서대로 통과시킵니다.
        for layer in self.layers:
            x = layer(x, pad_mask)

        return x


class CustomTransformerEncoderModel(TransformerEncoder):
    """Hugging Face AutoModel에서 사용할 수 있는 TransformerEncoder 별칭(alias)입니다."""

    pass