# Copyright 2023 PKU-Alignment Team. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

from __future__ import annotations

from typing import ClassVar

import torch
import torch.nn as nn
from transformers import GPTNeoModel, GPTNeoPreTrainedModel, PretrainedConfig, PreTrainedModel
from transformers.models.gpt_neo.modeling_gpt_neo import (
    GPT_NEO_INPUTS_DOCSTRING,
    GPT_NEO_START_DOCSTRING,
)
from transformers.utils.doc import add_start_docstrings, add_start_docstrings_to_model_forward

from safe_rlhf.models.score_model import ScoreModelOutput


@add_start_docstrings(
    """
    The GPT Neo Model transformer with a language modeling head on top (linear layer with weights tied to the input
    embeddings).
    """,
    GPT_NEO_START_DOCSTRING,
)
class GPTNeoForScore(GPTNeoPreTrainedModel):
    _keys_to_ignore_on_load_missing: ClassVar[list[str]] = [
        r'h\.\d+\.attn\.masked_bias',
        r'lm_head.weight',
        r'h\.\d+\.attn\.attention\.bias',
    ]
    _keys_to_ignore_on_save: ClassVar[list[str]] = [r'lm_head.weight']

    def __init__(self, config: PretrainedConfig) -> None:
        super().__init__(config)
        self.transformer = GPTNeoModel(config)

        config.score_dim = getattr(config, 'score_dim', 1)
        config.bias = getattr(config, 'bias', False)
        config.architectures = [self.__class__.__name__]
        self.score_head = nn.Linear(config.hidden_size, config.score_dim, bias=config.bias)

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

    def get_output_embeddings(self) -> None:
        return None

    def set_decoder(self, decoder: PreTrainedModel) -> None:
        self.transformer = decoder

    def get_decoder(self) -> PreTrainedModel:
        return self.transformer

    @add_start_docstrings_to_model_forward(GPT_NEO_INPUTS_DOCSTRING)
    def forward(  # pylint: disable=too-many-arguments
        self,
        input_ids: torch.Tensor | None = None,
        past_key_values: tuple[torch.FloatTensor] | None = None,
        attention_mask: torch.Tensor | None = None,
        token_type_ids: torch.Tensor | None = None,
        position_ids: torch.Tensor | None = None,
        head_mask: torch.Tensor | None = None,
        inputs_embeds: torch.Tensor | None = None,
        use_cache: bool | None = None,
        output_attentions: bool | None = None,
        output_hidden_states: bool | None = None,
        return_dict: bool | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor] | ScoreModelOutput:
        r"""
        Args:

        Returns:

        Examples:

        ```python
        >>> from safe_rlhf.models.llama.modeling_llama import LlamaModelForScore
        >>> from transformers import LlamaTokenizer

        >>> model = LlamaForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
        >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)

        >>> prompt = "Hey, are you conscious? Can you talk to me?"
        >>> inputs = tokenizer(prompt, return_tensors="pt")

        # got score
        >>> outputs = model(**inputs)
        >>> scores = outputs.scores
        >>> scores
        tensor([[[0.0000]]])
        ```
        """
        assert attention_mask is not None
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        transformer_outputs = self.transformer(
            input_ids,
            past_key_values=past_key_values,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        hidden_states = transformer_outputs[0]
        scores = self.score_head(hidden_states)  # size = (B, L, D)

        end_scores = []
        for i in range(input_ids.size(0)):
            end_index = attention_mask[i].nonzero()[-1].item()
            end_scores.append(scores[i, end_index])  # size = (D,)
        end_scores = torch.stack(end_scores, dim=0)  # size = (B, D)

        if not return_dict:
            return scores, end_scores

        return ScoreModelOutput(
            scores=scores,  # size = (B, L, D)
            end_scores=end_scores,  # size = (B, D)
        )
