import torch
import torch.nn as nn
from transformers import PreTrainedModel, AutoModel
from .configuration_viconbert import ViConBERTConfig


class MLPBlock(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim,
                 num_layers=2, dropout=0.3, activation=nn.GELU, use_residual=True):
        super().__init__()
        self.use_residual = use_residual
        self.activation_fn = activation()

        self.input_layer = nn.Linear(input_dim, hidden_dim)
        self.hidden_layers = nn.ModuleList()
        self.norms = nn.ModuleList()
        self.dropouts = nn.ModuleList()
        for _ in range(num_layers):
            self.hidden_layers.append(nn.Linear(hidden_dim, hidden_dim))
            self.norms.append(nn.LayerNorm(hidden_dim))
            self.dropouts.append(nn.Dropout(dropout))
        self.output_layer = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        x = self.input_layer(x)
        for layer, norm, dropout in zip(self.hidden_layers, self.norms, self.dropouts):
            residual = x
            x = layer(x)
            x = norm(x)
            x = dropout(x)
            x = self.activation_fn(x)
            if self.use_residual:
                x = x + residual
        x = self.output_layer(x)
        return x


class ViConBERT(PreTrainedModel):
    config_class = ViConBERTConfig

    def __init__(self, config):
        super().__init__(config)
        self.context_encoder = AutoModel.from_pretrained(
            config.base_model, cache_dir=config.base_model_cache_dir
        )
        self.context_projection = MLPBlock(
            self.context_encoder.config.hidden_size,
            config.hidden_dim,
            config.out_dim,
            dropout=config.dropout,
            num_layers=config.num_layers
        )
        self.context_attention = nn.MultiheadAttention(
            self.context_encoder.config.hidden_size,
            num_heads=config.num_head,
            dropout=config.dropout
        )
        self.context_window_size = config.context_window_size
        self.context_layer_weights = nn.Parameter(
            torch.zeros(self.context_encoder.config.num_hidden_layers)
        )
        self.post_init()

    def _encode_context_attentive(self, text, target_span):
        outputs = self.context_encoder(**text)
        hidden_states = outputs[0]
        start_pos, end_pos = target_span[:, 0], target_span[:, 1]

        positions = torch.arange(hidden_states.size(1), device=hidden_states.device)
        mask = (positions >= start_pos.unsqueeze(1)) & (positions <= end_pos.unsqueeze(1))
        masked_states = hidden_states * mask.unsqueeze(-1)
        span_lengths = mask.sum(dim=1, keepdim=True).clamp(min=1)
        pooled_embeddings = masked_states.sum(dim=1) / span_lengths

        Q_value = pooled_embeddings.unsqueeze(0)
        KV_value = hidden_states.permute(1, 0, 2)
        context_emb, _ = self.context_attention(Q_value, KV_value, KV_value)
        return context_emb

    def forward(self, context, target_span):
        context_emb = self._encode_context_attentive(context, target_span)
        return self.context_projection(context_emb.squeeze(0))