import torch
import torch.nn as nn
import torch.nn.functional as F
from huggingface_hub import PyTorchModelHubMixin

class CustomBiGRUMCQModel(nn.Module, PyTorchModelHubMixin):
    """
    Bidirectional GRU model with custom global attention pooling
    """
    #vocab_size, embedding_dim=256, hidden_dim=256, num_layers=2, dropout=0.1
    def __init__(self, config:dict):
        super().__init__()
        self.config = config
        self.vocab_size =  config.get('vocab_size', 30000)
        self.embedding_dim = config.get('embedding_dim',256)
        self.hidden_dim = config.get('hidden_dim',256)
        self.num_layers = config.get('num_layers',2)
        
        self.embedding = nn.Embedding(self.vocab_size, self.embedding_dim, padding_idx=0)
        self.embedding_dropout = nn.Dropout(config.get('dropout', 0.1))
        
        self.gru = nn.GRU(
            input_size = self.embedding_dim,
            hidden_size = self.hidden_dim,
            num_layers = self.num_layers,
            bidirectional=True,
            batch_first=True,
            dropout = config.get('dropout',0.1) if self.num_layers > 1 else 0
        )
        
        # Dot product attention pooling head
        self.attention_query = nn.Linear(self.hidden_dim * 2, 1, bias=False)

        #Final classification layer
        self.classifier = nn.Linear(self.hidden_dim*2, 1)
        self.dropout = nn.Dropout(config.get('dropout',0.1))
    
    def forward(self, input_ids, attention_mask=None, label=None, **kwargs):
        input_idx = input_ids
        #expected shape [batch_size, num_of_choice=5, seq_length]
        batch_size, num_of_choices, seq_length = input_idx.shape
        
        flat_input_idx = input_ids.view(batch_size * num_of_choices, seq_length)
        
        embedded = self.embedding(flat_input_idx)
        embedded = self.embedding_dropout(embedded)
        
        gru_out, _ = self.gru(embedded) # Shape: [batch_size * 5, seq_length, hidden_dim * 2]

        attentions_scores = self.attention_query(gru_out).squeeze(-1)
        
        # Adapt Attentions mask to the choice dimension
        flaten_mask = None
        if attention_mask is not None:
            flat_mask = attention_mask.view(batch_size * num_of_choices, seq_length)
            attentions_scores = attentions_scores.masked_fill(flat_mask == 0, -1e9)
        
        attention_weight = F.softmax(attentions_scores, dim=-1).unsqueeze(-1)
        pooled_output = torch.sum(gru_out * attention_weight, dim=1)
        pooled_output = self.dropout(pooled_output)
        
        raw_scores = self.classifier(pooled_output).squeeze(-1) # Shape: [batch_size * 5]
        reshape_logits = raw_scores.view(batch_size, num_of_choices) # Shape: [batch_size, 5]
        
        loss = None
        if label is not None:
            loss_fct = nn.CrossEntropyLoss()
            loss = loss_fct(reshape_logits, label)
            
        return {
            'loss': loss,
            'logits': reshape_logits
        }