# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
# 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.
"""Utilities for score models."""

from __future__ import annotations

from abc import abstractmethod
from dataclasses import dataclass
from typing import Any, Literal

import paddle
import paddle.distributed as dist
import paddle.nn as nn

from ...transformers.configuration_utils import PretrainedConfig
from ...transformers.model_outputs import ModelOutput


@dataclass
class ScoreModelOutput(ModelOutput):
    """
    Output of the score model.

    Args:
        scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, score_dim, sequence_length)`):
            Prediction scores of the score model.
        end_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, score_dim)`):
            Prediction scores of the end of the sequence.
    """

    scores: paddle.Tensor | None = None  # size = (B, L, D)
    end_scores: paddle.Tensor | None = None  # size = (B, D)


class ScoreModelMixin:
    """Base class for score models."""

    score_head: nn.Linear
    # normalizer: Normalizer
    do_normalize: bool = False
    normalize_function: NormalizeFunction = "affine"
    _initialized: bool = False

    @classmethod
    def init_score_head(self, config: PretrainedConfig, hidden_size: int, **kwargs: Any) -> None:
        """Initialize the score head."""
        if self._initialized:
            return

        config.score_dim = kwargs.pop("score_dim", getattr(config, "score_dim", 1))
        config.bias = kwargs.pop("bias", getattr(config, "bias", False))

        config.score_type = kwargs.pop("score_type", getattr(config, "score_type", "reward"))
        if config.score_type == "reward":
            self.normalize_function = "affine"
        elif config.score_type == "cost":
            self.normalize_function = "scale"
        elif config.score_type == "critic":
            self.normalize_function = "identity"
        else:
            raise ValueError(
                f"Invalid score type: {config.score_type}. Expected one of 'reward', 'cost', or 'critic'.",
            )

        config.do_normalize = kwargs.pop(
            "do_normalize",
            getattr(config, "do_normalize", False),
        )
        self.do_normalize = config.do_normalize

        config.normalizer_type = kwargs.pop(
            "normalizer_type",
            getattr(config, "normalizer_type", None),
        )
        if config.normalizer_type not in {
            "RunningMeanStd",
            "ExponentialMovingAverage",
            None,
        }:
            raise ValueError(
                f"Invalid norm type: {config.normalizer_type}."
                "Expected one of 'RunningMeadStd', 'ExponentialMovingAverage', or None.",
            )
        if config.normalizer_type == "ExponentialMovingAverage":
            config.momentum = kwargs.pop("momentum", getattr(config, "momentum", None))
        momentum = getattr(config, "momentum", None)

        self.score_head = nn.Linear(hidden_size, config.score_dim, bias_attr=config.bias)
        self.normalizer = Normalizer.instantiate(
            normalizer_type=config.normalizer_type,
            normalize_function=self.normalize_function,
            shape=(config.score_dim,),
            momentum=momentum,
        )

        mean = getattr(config, "mean", None)
        var = getattr(config, "var", None)
        self.normalizer.set_mean_var(mean, var)

        self._initialized = True

    def get_score(
        self,
        hidden_state: paddle.Tensor,  # size = (B, L, E)
        position_ids: paddle.Tensor | None = None,  # size = (B, L)
        attn_mask_startend_row_indices: paddle.Tensor | None = None,  # size = (B, 1), (B, 2), (B, 3) or (B, 4)
        return_dict: bool | None = None,
        attention_mask: paddle.Tensor | None = None,  # size = (B, L)
    ) -> ScoreModelOutput:
        """Forward pass of the score model."""
        hidden_state = hidden_state.cast(paddle.float32)
        scores = self.score_head(hidden_state)  # size = (B, L, D)
        if scores.dtype != hidden_state.dtype:  # EB rm cast to float32
            scores = scores.cast(hidden_state.dtype)

        if position_ids is not None:
            first_pos = paddle.arange(hidden_state.shape[0]).unsqueeze(-1)
            # Take left padding into account, which has 0s in left and max_len
            # in right.
            left_pad_mask = position_ids == 0
            # position_ids = paddle.where(
            #     left_pad_mask, position_ids, position_ids + left_pad_mask.sum(-1, keepdim=True) - 1
            # )
            # the above limits right padding must not be 0s, the following suits
            # to both left and right padding with 0s
            left_pad_num = (
                paddle.where(left_pad_mask, position_ids.shape[-1] + 100, position_ids).argmin(axis=-1, keepdim=True)
                - 1
            )
            position_ids = left_pad_num + position_ids
            second_pos = paddle.max(position_ids, axis=-1, keepdim=True)
            end_pos = paddle.stack([first_pos, second_pos], axis=-1).squeeze(1)
            end_score = scores.gather_nd(end_pos)
        else:
            # attention_mask passed from pipeline pre-stage is shaped (bs, 1, seq_len, seq_len)
            assert attention_mask is not None and len(attention_mask.shape) == 2
            end_score = []
            end_pos = []
            for i in range(hidden_state.shape[0]):
                end_index = attention_mask[i].nonzero()[-1].item()
                end_pos.append((i, end_index))
                end_score.append(scores[i, end_index])  # size = (D,)
            end_score = paddle.stack(end_score, axis=0)  # size = (B, D)

        if self.training and self.do_normalize:
            if dist.is_initialized():
                gathered_end_score_list = []
                try:
                    # gather among data parallel group
                    hcg = dist.fleet.get_hybrid_communicate_group()
                    group = hcg.get_sharding_parallel_group()
                    dist.all_gather(gathered_end_score_list, end_score, group)
                except:
                    dist.all_gather(gathered_end_score_list, end_score)
                gathered_end_score = paddle.concat(gathered_end_score_list, axis=0)
                self.normalizer.update(gathered_end_score)
            else:
                self.normalizer.update(end_score)
            self.config.mean = self.normalizer.mean.tolist()
            self.config.var = self.normalizer.var.tolist()

        if self.do_normalize:
            scores = self.normalizer.normalize(scores)

        if not return_dict:
            return scores, end_score

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

        return scores

    def set_normalize(self, mode: bool = True) -> None:
        """
        Set whether to normalize the input data, default is True.
        If mode is True, normalize the input data; if mode is False, do not normalize the input data.

        Args:
            mode (bool, optional): Whether to normalize the input data, default is True. Defaults to True.

        Returns:
            None: No return value, directly modifies the instance's do_normalize attribute and the do_normalize attribute in config.

        Raises:
            None: No exceptions are raised.
        """
        if self.do_normalize == mode:
            return

        self.do_normalize = self.config.do_normalize = mode


NormalizeFunction = Literal["affine", "scale", "translate", "identity"]
NormalizerType = Literal["RunningMeanStd", "ExponentialMovingAverage"]


class Normalizer(nn.Layer):
    """Normalize input to have zero mean and unit variance."""

    mean: paddle.Tensor
    var: paddle.Tensor
    count: paddle.Tensor
    normalize_function: NormalizeFunction

    def __init__(
        self,
        normalize_function: NormalizeFunction,
        shape: tuple[int, ...],
        device: str | None = None,
    ) -> None:
        """Initialize."""
        super().__init__()
        if normalize_function not in {"affine", "scale", "translate", "identity"}:
            raise ValueError(
                f"Invalid normalization function type: {normalize_function}. ",
                'Expected one of "affine", "scale", "translate", "identity".',
            )
        self.normalize_function = normalize_function
        self.register_buffer("mean", paddle.zeros(shape, dtype=paddle.get_default_dtype()))  # align torch zeros/ones
        self.register_buffer("var", paddle.ones(shape, dtype=paddle.get_default_dtype()))  # align torch zeros/ones
        self.register_buffer("count", paddle.zeros(1, dtype=paddle.int64))

    @abstractmethod
    def update(self, data: paddle.Tensor) -> None:
        """Update mean and variance."""
        raise NotImplementedError

    @property
    def std(self) -> paddle.Tensor:
        """Return standard deviation."""
        return self.var.sqrt()

    def set_mean_var(
        self,
        mean: paddle.Tensor | list[float] | tuple[float, ...] | None,
        var: paddle.Tensor | list[float] | tuple[float, ...] | None,
    ) -> None:
        """Set mean and variance."""
        mean = paddle.to_tensor(mean, dtype=self.mean.dtype, place=self.mean.place) if mean is not None else self.mean
        var = paddle.to_tensor(var, dtype=self.var.dtype, place=self.var.place) if var is not None else self.var

        assert mean.shape == self.mean.shape
        assert var.shape == self.var.shape

        self.mean = mean
        self.var = var

    def forward(
        self,
        data: paddle.Tensor,
        epsilon=1e-8,
    ) -> paddle.Tensor:
        """Update and normalize input."""
        if self.training:
            self.update(data)
        return self.normalize(data, epsilon=epsilon)

    def normalize(
        self,
        data: paddle.Tensor,
        epsilon=1e-8,
    ) -> paddle.Tensor:
        """Normalize input."""
        if self.normalize_function == "affine":
            return (data - self.mean.detach()) / (self.std.detach() + epsilon)
        if self.normalize_function == "scale":
            return data / (self.std.detach() + epsilon)
        if self.normalize_function == "translate":
            return data - self.mean.detach()
        if self.normalize_function == "identity":
            return data
        raise ValueError(
            f"Invalid normalization function type: {self.normalize_function}. ",
            'Expected one of "affine", "scale", "translate", "identity".',
        )

    @classmethod
    def instantiate(
        cls,
        normalizer_type: NormalizerType | None,
        normalize_function: NormalizeFunction,
        shape: tuple[int, ...],
        device: str | None = None,
        **kwargs: Any,
    ):
        """Get a normalizer."""
        if normalizer_type == "RunningMeanStd":
            return RunningMeanStd(
                normalize_function,
                shape=shape,
                device=device,
            )
        if normalizer_type == "ExponentialMovingAverage":
            return ExponentialMovingAverage(
                normalize_function,
                shape=shape,
                device=device,
                **kwargs,
            )
        if normalizer_type is None:
            return IdentityNormalizer(
                normalize_function,
                shape=shape,
                device=device,
            )
        raise ValueError(
            f"Invalid normalization function type: {normalizer_type}. "
            'Expected one of "RunningMeanStd", "ExponentialMovingAverage".',
        )


class RunningMeanStd(Normalizer):
    """Running mean and standard deviation."""

    def update(self, data: paddle.Tensor) -> None:
        """Update mean and variance."""
        batch_mean = data.mean(0)
        batch_var = data.var(0)
        batch_count = data.shape[0]

        delta = batch_mean - self.mean
        total_count = self.count + batch_count

        new_mean = self.mean + delta * batch_count / total_count
        m_a = self.var * self.count
        m_b = batch_var * batch_count
        m2 = (  # pylint: disable=invalid-name
            m_a + m_b + paddle.square(delta) * (self.count * batch_count / total_count)
        )
        new_var = m2 / total_count

        self.mean = new_mean
        self.var = new_var
        self.count = total_count


class ExponentialMovingAverage(Normalizer):
    """Exponential moving average."""

    def __init__(
        self,
        normalize_function: NormalizeFunction,
        shape: tuple[int, ...],
        device: str | None = None,
        momentum: float = 0.9,
    ) -> None:
        """
        Args:
        normalize_function (NormalizeFunction): Function to normalize the input tensor.
        shape (tuple[int, ...]): Shape of the output tensor.
        device (str, optional): Device where the tensor will be allocated. Defaults to None.
        momentum (float, optional): Momentum for the moving average. Defaults to 0.9.
        """
        super().__init__(normalize_function, shape=shape, device=device)
        self.momentum = momentum

    def update(self, data: paddle.Tensor) -> None:
        """Update mean and variance."""
        batch_mean = data.mean(0)
        batch_var = data.var(0)
        batch_count = data.shape[0]

        self.mean = self.momentum * self.mean + (1.0 - self.momentum) * batch_mean
        self.var = self.momentum * self.var + (1.0 - self.momentum) * batch_var
        self.count += batch_count  # pylint: disable=no-member


class IdentityNormalizer(Normalizer):
    """Identity normalizer."""

    def update(self, data: paddle.Tensor) -> None:
        """Update mean and variance."""
        self.count += data.shape[0]  # pylint: disable=no-member
