# This file is part of the R package "aifeducation".
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 3 as published by
# the Free Software Foundation.
#
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>

"""
BERT            MPNet
0   - [PAD];    1     - <pad>
100 - [UNK];    3     - <unk>
101 - [CLS];    0     - <s>
102 - [SEP];    2     - </s>
103 - [MASK];   30526 - <mask>
"""

# ======================================================================================
# Imports
# ======================================================================================

import torch
from torch import nn

from transformers import DataCollatorForLanguageModeling, \
                         DataCollatorForWholeWordMask, \
                         MPNetForMaskedLM

# ======================================================================================
# LossFunction for MLM and PLM tasks
# ======================================================================================

class MPLMLoss_PT(nn.Module):
    def __init__(self, ignore_index = -100):
        super().__init__()
        """
        The combination of nn.LogSoftmax and nn.NLLLoss is equivalent to using nn.CrossEntropyLoss
        """
        self.criterion_mlm = nn.CrossEntropyLoss(ignore_index = ignore_index)
        self.criterion_plm = nn.CrossEntropyLoss(ignore_index = ignore_index)

    def forward(self, mlm_logits, plm_logits, mlm_labels, plm_labels):
        # forward(y_pred, y_true)
        mlm_loss = self.criterion_mlm(mlm_logits, mlm_labels)
        plm_loss = self.criterion_plm(plm_logits, plm_labels)

        return mlm_loss + plm_loss

# ======================================================================================
# Head of model for MLM and PLM tasks
# ======================================================================================

class MPNetForMPLM_PT(MPNetForMaskedLM):
    def __init__(self, config):
        super().__init__(config)
        self.plm_head = nn.Linear(config.hidden_size, config.vocab_size)  # Adding PLM output layer

        # Initialize weights and apply final processing
        self.post_init()

    def forward(self, 
                input_ids = None, 
                attention_mask = None, 
                mlm_labels = None, 
                plm_labels = None,
                return_dict = None,
                output_hidden_states = None):
        """
        outputs is BaseModelOutputWithPooling with args:
        - last_hidden_state (torch.FloatTensor of shape (batch_size, sequence_length, hidden_size))
        - pooler_output
        - hidden_states
        - attentions

        hidden_size = 768
        """

        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        outputs = self.mpnet(input_ids, 
                             attention_mask = attention_mask,
                             output_hidden_states = output_hidden_states)
        
        """
        mlm_logits, plm_logits (torch.FloatTensor of shape (batch_size, sequence_length, vocab_size))

        vocab_size = 30527
        """

        sequence_output = outputs[0]

        mlm_logits = self.lm_head(sequence_output)
        plm_logits = self.plm_head(sequence_output)  # Computing PLM logits

        """
        Logits (torch.FloatTensor of shape (batch_size x sequence_length, vocab_size)):
            mlm_logits.view(-1, mlm_logits.size(-1))
            plm_logits.view(-1, plm_logits.size(-1))
        Targets (torch.FloatTensor of shape (batch_size x sequence_length)):
            mlm_labels.view(-1)
            plm_labels.view(-1)

            [0 <= targets[i] < vocab_size]
        """

        loss = None
        if mlm_labels is not None and plm_labels is not None:
            loss_fct = MPLMLoss_PT()
            loss = loss_fct(mlm_logits.view(-1, mlm_logits.size(-1)),
                            plm_logits.view(-1, plm_logits.size(-1)), 
                            mlm_labels.view(-1),
                            plm_labels.view(-1))
        
        if not return_dict: 
            output = (mlm_logits, plm_logits) + outputs[2:]
            return ((loss,) + output) if loss is not None else output
        
        return {
            'loss': loss,
            'logits': mlm_logits,
            'hidden_states': outputs.hidden_states,
            'attentions': outputs.attentions
        }
