# 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 Callable
from typing_extensions import TypedDict  # Python 3.10+

import torch

from safe_rlhf.datasets.base import CollatorBase, RawSample, TokenizedDataset
from safe_rlhf.datasets.utils import format_prompt, right_padding


__all__ = [
    'SafetyPreferenceDataset',
    'SafetyPreferenceCollator',
    'SafetyPreferenceSample',
    'SafetyPreferenceBatch',
]


class SafetyPreferenceSample(TypedDict, total=True):
    safer_input_ids: torch.LongTensor  # size = (L,)
    # +1 for safe / -1 for unsafe
    safer_sign: torch.LongTensor  # size = (L,)

    unsafer_input_ids: torch.LongTensor  # size = (L,)
    # +1 for safe / -1 for unsafe
    unsafer_sign: torch.LongTensor  # size = (L,)


class SafetyPreferenceBatch(TypedDict, total=True):
    safer_input_ids: torch.LongTensor  # size = (B, L)
    safer_attention_mask: torch.BoolTensor  # size = (B, L)
    # +1 for safe / -1 for unsafe
    safer_safety_sign: torch.LongTensor  # size = (B,)

    unsafer_input_ids: torch.LongTensor  # size = (B, L)
    unsafer_attention_mask: torch.BoolTensor  # size = (B, L)
    # +1 for safe / -1 for unsafe
    unsafer_safety_sign: torch.LongTensor  # size = (B,)


class SafetyPreferenceDataset(TokenizedDataset):
    def preprocess(self, raw_sample: RawSample) -> SafetyPreferenceSample:
        prompt = format_prompt(input=raw_sample['input'], eos_token=self.tokenizer.eos_token)
        answer = raw_sample['answer']
        other_answer = raw_sample['other_answer']
        safer = raw_sample['safer']
        is_safe = raw_sample['is_safe']
        is_other_safe = raw_sample['is_other_safe']

        safer_answer, unsafer_answer = answer, other_answer
        safer_sign, unsafer_sign = (  # +1 for safe / -1 for unsafe
            2 * int(is_safe) - 1,
            2 * int(is_other_safe) - 1,
        )
        if not safer:
            safer_answer, unsafer_answer = unsafer_answer, safer_answer
            safer_sign, unsafer_sign = unsafer_sign, safer_sign

        if safer_sign < unsafer_sign:
            raise ValueError(
                'The safer answer is not safer than the unsafer answer.\n\n'
                f'Prompt: {prompt}\n\n'
                f'Safer answer (labeled as unsafe): {safer_answer}\n\n'
                f'Unsafer answer (labeled as safe): {unsafer_answer}',
            )

        # size = (L,)
        safer_input_ids = self.tokenize(prompt + safer_answer + self.tokenizer.eos_token)
        unsafer_input_ids = self.tokenize(prompt + unsafer_answer + self.tokenizer.eos_token)
        if (
            safer_input_ids.size() == unsafer_input_ids.size()
            and torch.all(torch.eq(safer_input_ids, unsafer_input_ids)).item()
        ):
            raise ValueError(
                'Two responses get the same `input_ids` after tokenization.\n\n'
                f'Prompt: {prompt}\n\n'
                f'Safer answer: {safer_answer}\n\n'
                f'Unsafer answer: {unsafer_answer}',
            )
        return {
            'safer_input_ids': safer_input_ids,  # size = (L,)
            'safer_sign': torch.tensor(safer_sign),  # size = ()
            'unsafer_input_ids': unsafer_input_ids,  # size = (L,)
            'unsafer_sign': torch.tensor(unsafer_sign),  # size = ()
        }

    def get_collator(self) -> Callable[[list[dict[str, torch.Tensor]]], dict[str, torch.Tensor]]:
        return SafetyPreferenceCollator(self.tokenizer.pad_token_id)


class SafetyPreferenceCollator(CollatorBase):
    def __call__(self, samples: list[SafetyPreferenceSample]) -> SafetyPreferenceBatch:
        input_ids = [sample['safer_input_ids'] for sample in samples] + [
            sample['unsafer_input_ids'] for sample in samples
        ]
        attention_mask = [
            input_id.new_ones(input_id.size(), dtype=torch.bool) for input_id in input_ids
        ]
        safety_sign = [sample['safer_sign'] for sample in samples] + [
            sample['unsafer_sign'] for sample in samples
        ]

        # size = (2 * B, L)
        input_ids = right_padding(input_ids, padding_value=self.pad_token_id)
        attention_mask = right_padding(attention_mask, padding_value=0)
        # size = (2 * B,)
        safety_sign = torch.tensor(safety_sign, dtype=torch.long)

        # size = (B, L)
        safer_input_ids, unsafer_input_ids = input_ids.chunk(chunks=2, dim=0)
        safer_attention_mask, unsafer_attention_mask = attention_mask.chunk(chunks=2, dim=0)
        # size = (B,)
        safer_safety_sign, unsafer_safety_sign = safety_sign.chunk(chunks=2, dim=0)
        return {
            'safer_input_ids': safer_input_ids,  # size = (B, L)
            'safer_attention_mask': safer_attention_mask,  # size = (B, L)
            'safer_safety_sign': safer_safety_sign,  # size = (B,)
            'unsafer_input_ids': unsafer_input_ids,  # size = (B, L)
            'unsafer_attention_mask': unsafer_attention_mask,  # size = (B, L)
            'unsafer_safety_sign': unsafer_safety_sign,  # size = (B,)
        }
