import torch
import torch.nn as nn
from transformers import AutoModel
import os

from config import *

    

class XLMRBiGRU(nn.Module):
    def __init__(self, xlmr_model_name, gru_hidden_size, gru_num_layers, dropout, freeze_xlmr_layers):
        super().__init__()
        self.xlmr = AutoModel.from_pretrained(xlmr_model_name)
        
        # freeze embedding layer
        for param in self.xlmr.embeddings.parameters():
            param.requires_grad = False

        # freeze first 6 encoder layers
        for num_layer in range(freeze_xlmr_layers):
            for param in self.xlmr.encoder.layer[num_layer].parameters():
                param.requires_grad = False
        
        self.gru = nn.GRU(
            input_size=self.xlmr.config.hidden_size,
            hidden_size=gru_hidden_size,
            num_layers=gru_num_layers,
            batch_first=True,
            bidirectional=True,
            dropout=dropout if gru_num_layers > 1 else 0
        )

        self.attention = nn.Linear(gru_hidden_size * 2, 1)  # For attention scoring

        self.classifier = nn.Linear(gru_hidden_size * 2, 1)  # Binary classification

        self.dropout = nn.Dropout(dropout)


    def forward(self, input_ids, attention_mask, conv_length):
        batch_size, max_conv_len, max_msg_len = input_ids.size()

        # Reshape for XLM-RoBERTa
        input_ids = input_ids.view(-1, max_msg_len)  # (batch_size * max_conv_len, max_msg_len)
        attention_mask = attention_mask.view(-1, max_msg_len)  # (batch_size * max_conv_len, max_msg_len)

        # XLM-RoBERTa encoding
        outputs = self.xlmr(input_ids=input_ids, attention_mask=attention_mask)
        cls_embeddings = outputs.last_hidden_state[:, 0, :]  # (batch_size * max_conv_len, hidden_size)

        # Reshape back to conversation format
        cls_embeddings = cls_embeddings.view(batch_size, max_conv_len, -1)  # (batch_size, max_conv_len, hidden_size)

        # GRU encoding
        gru_output, _ = self.gru(cls_embeddings)  # (batch_size, max_conv_len, gru_hidden_size*2)

        # Attention mechanism
        attn_scores = self.attention(gru_output).squeeze(-1)  # (batch_size, max_conv_len)
        
        # Mask out padding messages
        mask = torch.arange(max_conv_len).unsqueeze(0).to(conv_length.device) < conv_length.unsqueeze(1)
        attn_scores[~mask] = float('-inf')  # Set scores of padding messages to -inf

        attn_weights = torch.softmax(attn_scores, dim=1).unsqueeze(-1)  # (batch_size, max_conv_len, 1)
        context_vector = torch.sum(attn_weights * gru_output, dim=1)  # (batch_size, gru_hidden_size*2)

        context_vector = self.dropout(context_vector)
        
        logits = self.classifier(context_vector).squeeze(-1)  # (batch_size,)
        
        return logits
    


if __name__ == "__main__":

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    model = XLMRBiGRU(
        xlmr_model_name=XLMR_MODEL_NAME,
        gru_hidden_size=GRU_HIDDEN_SIZE,
        gru_num_layers=GRU_NUM_LAYERS,
        dropout=DROPOUT,
        freeze_xlmr_layers=FREEZE_XLMR_LAYERS
    )
    
    # dummy input
    input_ids     = torch.randint(0, 100, (2, 50, 64))
    attention_mask = torch.ones(2, 50, 64, dtype=torch.long)
    conv_lengths   = torch.tensor([50, 30])
    
    output = model(input_ids, attention_mask, conv_lengths)
    print(output.shape)  # should be (2,)