#!/usr/bin/env python3
"""
VAD Transformer Model Definition

This module defines the custom RoBERTa-based model for VAD prediction
that will be uploaded to Hugging Face.

Author: AI Assistant
"""

import torch
from transformers import RobertaModel, RobertaPreTrainedModel

class RobertaForVADRegression(RobertaPreTrainedModel):
    """
    RoBERTa model for predicting Valence, Arousal, and Dominance values.
    This model extends RobertaPreTrainedModel to be compatible with the Hugging Face ecosystem.
    """
    
    def __init__(self, config):
        super().__init__(config)
        self.roberta = RobertaModel(config)
        self.dropout = torch.nn.Dropout(config.hidden_dropout_prob)
        self.valence_head = torch.nn.Linear(config.hidden_size, 1)
        self.arousal_head = torch.nn.Linear(config.hidden_size, 1)
        self.dominance_head = torch.nn.Linear(config.hidden_size, 1)
        
        # Initialize weights
        self.init_weights()
    
    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
        labels=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
    ):
        """
        Forward pass of the model.
        
        Returns:
            tuple: Tuple containing valence, arousal, and dominance predictions
            or
            dict: Dictionary containing loss and predictions if labels are provided
        """
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        
        outputs = self.roberta(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        
        sequence_output = outputs[0]
        pooled_output = sequence_output[:, 0, :]  # Take CLS token representation
        pooled_output = self.dropout(pooled_output)
        
        valence = self.valence_head(pooled_output)
        arousal = self.arousal_head(pooled_output)
        dominance = self.dominance_head(pooled_output)
        
        loss = None
        if labels is not None:
            # If labels are provided, calculate loss
            # Assuming labels is a tensor of shape [batch_size, 3] with VAD values
            loss_fct = torch.nn.MSELoss()
            v_loss = loss_fct(valence.squeeze(), labels[:, 0])
            a_loss = loss_fct(arousal.squeeze(), labels[:, 1])
            d_loss = loss_fct(dominance.squeeze(), labels[:, 2])
            loss = v_loss + a_loss + d_loss
        
        if not return_dict:
            output = (valence.squeeze(), arousal.squeeze(), dominance.squeeze()) + outputs[2:]
            return ((loss,) + output) if loss is not None else output
        
        return {
            "loss": loss,
            "valence": valence.squeeze(),
            "arousal": arousal.squeeze(),
            "dominance": dominance.squeeze(),
            "hidden_states": outputs.hidden_states,
            "attentions": outputs.attentions,
        }