from dataclasses import dataclass
from typing import List, Optional, Tuple, Union

import torch
import torch.utils.checkpoint
from torch import nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from transformers import (
    T5PreTrainedModel,
    T5Model
    
)

from transformers.modeling_outputs import (
    SequenceClassifierOutput,
)

def mean_pooling(encoder_outputs, attention_mask):
    """
    Perform mean pooling over encoder outputs, considering the attention mask.
    """
    hidden_states = encoder_outputs.last_hidden_state  # Shape: (batch_size, seq_length, hidden_size)
    mask = attention_mask.unsqueeze(-1)  # Shape: (batch_size, seq_length, 1)
    masked_hidden_states = hidden_states * mask  # Zero out padding tokens
    sum_hidden_states = masked_hidden_states.sum(dim=1)  # Sum over sequence length
    sum_mask = mask.sum(dim=1)  # Sum the mask (number of non-padding tokens)
    return sum_hidden_states / sum_mask  # Mean pooling


class T5EncoderForSequenceClassification(T5PreTrainedModel):

    def __init__(self, checkpoint, tokenizer, config, num_labels):
        super().__init__(config)
        self.num_labels = num_labels
        self.config = config

        # we force the loading of a pre-trained model here
        self.t5 = T5Model.from_pretrained(checkpoint)
        self.t5.resize_token_embeddings(len(tokenizer))
        classifier_dropout = (
            config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
        )
        self.dropout = nn.Dropout(classifier_dropout)
        self.classifier = nn.Linear(config.hidden_size, self.num_labels)


    def forward(
        self,
        input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        token_type_ids: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
        labels: Optional[torch.Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple[torch.Tensor], SequenceClassifierOutput]:
        r"""
        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
            Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
        """
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict


        # encoder_outputs = self.t5.encoder(
        #     input_ids,
        #     attention_mask=attention_mask,
        #     head_mask=head_mask,
        #     inputs_embeds=inputs_embeds,
        #     output_attentions=output_attentions,
        #     output_hidden_states=output_hidden_states,
        #     return_dict=return_dict,
        # )


        encoder_outputs = self.t5.encoder(input_ids, attention_mask=attention_mask)
        # last_hidden_state = encoder_outputs.last_hidden_state
        # use mean of hidden state
        # pooled_output = mean_pooling(encoder_outputs, attention_mask)

        # Use the hidden state of the first token as the sequence representation
        pooled_output = encoder_outputs.last_hidden_state[:, 0, :]  # Shape: (batch_size, hidden_size)

        # pooled_output = encoder_outputs[1]

        pooled_output = self.dropout(pooled_output)
        logits = self.classifier(pooled_output)

        loss = None
        if labels is not None:
            if self.config.problem_type is None:
                if self.num_labels == 1:
                    self.config.problem_type = "regression"
                elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
                    self.config.problem_type = "single_label_classification"
                else:
                    self.config.problem_type = "multi_label_classification"

            if self.config.problem_type == "regression":
                loss_fct = MSELoss()
                if self.num_labels == 1:
                    loss = loss_fct(logits.squeeze(), labels.squeeze())
                else:
                    loss = loss_fct(logits, labels)
            elif self.config.problem_type == "single_label_classification":
                loss_fct = CrossEntropyLoss()
                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
            elif self.config.problem_type == "multi_label_classification":
                loss_fct = BCEWithLogitsLoss()
                loss = loss_fct(logits, labels)
        if not return_dict:
            output = (logits,) + encoder_outputs[2:]
            return ((loss,) + output) if loss is not None else output

        return SequenceClassifierOutput(
            loss=loss,
            logits=logits,
            hidden_states=encoder_outputs.hidden_states,
            attentions=encoder_outputs.attentions,
        )