# 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

import warnings
from typing import Any, ClassVar

import torch
import torch.nn as nn
from transformers import GPTJModel, GPTJPreTrainedModel, PreTrainedModel
from transformers.configuration_utils import PretrainedConfig
from transformers.models.gptj.modeling_gptj import (
    DEPARALLELIZE_DOCSTRING,
    GPTJ_INPUTS_DOCSTRING,
    GPTJ_START_DOCSTRING,
    PARALLELIZE_DOCSTRING,
)
from transformers.utils.doc import add_start_docstrings, add_start_docstrings_to_model_forward
from transformers.utils.model_parallel_utils import assert_device_map, get_device_map

from safe_rlhf.models.score_model import ScoreModelOutput


@add_start_docstrings(
    """
    The GPT-J Model transformer with a score head on top.
    """,
    GPTJ_START_DOCSTRING,
)
class GPTJForScore(GPTJPreTrainedModel):
    _keys_to_ignore_on_load_missing: ClassVar[list[str]] = [
        r'h\.\d+\.attn\.masked_bias',
        r'h\.\d+\.attn\.bias',
    ]

    def __init__(self, config: PretrainedConfig) -> None:
        super().__init__(config)
        self.transformer = GPTJModel(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.n_embd, config.score_dim, bias=config.bias)

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

        self.device_map: dict[Any, Any]
        self.model_parallel: bool

    @add_start_docstrings(PARALLELIZE_DOCSTRING)
    def parallelize(self, device_map: str | None = None) -> None:
        warnings.warn(
            '`GPTJForCausalLM.parallelize` is deprecated and will be removed in v5 of Transformers, you should load'
            " your model with `device_map='balanced'` in the call to `from_pretrained`. You can also provide your own"
            " `device_map` but it needs to be a dictionary module_name to device, so for instance {'transformer.h.0':"
            " 0, 'transformer.h.1': 1, ...}",
            FutureWarning,
            stacklevel=1,
        )
        self.device_map = (
            get_device_map(len(self.transformer.h), range(torch.cuda.device_count()))
            if device_map is None
            else device_map
        )
        assert_device_map(self.device_map, len(self.transformer.h))
        self.transformer.parallelize(self.device_map)
        self.score_head = self.score_head.to(self.transformer.first_device)
        self.model_parallel = True

    @add_start_docstrings(DEPARALLELIZE_DOCSTRING)
    def deparallelize(self) -> None:
        warnings.warn(
            'Like `parallelize`, `deparallelize` is deprecated and will be removed in v5 of Transformers.',
            FutureWarning,
            stacklevel=1,
        )
        self.transformer.deparallelize()
        self.transformer = self.transformer.to('cpu')
        self.score_head = self.score_head.to('cpu')
        self.model_parallel = False
        torch.cuda.empty_cache()

    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(
        GPTJ_INPUTS_DOCSTRING.format('batch_size, sequence_length'),
    )
    def forward(  # pylint: disable=too-many-arguments
        self,
        input_ids: torch.LongTensor,
        attention_mask: torch.FloatTensor,
        past_key_values: tuple[tuple[torch.Tensor]] | None = None,
        token_type_ids: torch.LongTensor | None = None,
        position_ids: torch.LongTensor | None = None,
        head_mask: torch.FloatTensor | None = None,
        inputs_embeds: torch.FloatTensor | 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:
        """
        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]

        # Set device for model parallelism
        if self.model_parallel:
            torch.cuda.set_device(self.transformer.first_device)
            hidden_states = hidden_states.to(self.lm_head.weight.device)

        # make sure sampling in fp16 works correctly and
        # compute loss in fp32 to match with mesh-tf version
        # https://github.com/EleutherAI/gpt-neo/blob/89ce74164da2fb16179106f54e2269b5da8db333/models/gpt2/gpt2.py#L179
        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)
        )
