from typing import Optional, Tuple

import numpy as np
import torch
import torch.distributed as dist
import transformers

from realhf.base import constants, logging
from realhf.impl.model.utils.padding import pad_input, unpad_input

logger = logging.getLogger("Modeling Functional Utils")


@torch.jit.script
def upcast_masked_softmax(
    x: torch.Tensor,
    mask: torch.Tensor,
    mask_value: torch.Tensor,
    scale: float,
    softmax_dtype: torch.dtype,
):
    input_dtype = x.dtype
    x = x.to(softmax_dtype) * scale
    x = torch.where(mask, x, mask_value)
    x = torch.nn.functional.softmax(x, dim=-1).to(input_dtype)
    return x


@torch.jit.script
def upcast_softmax(x: torch.Tensor, scale: float, softmax_dtype: torch.dtype):
    input_dtype = x.dtype
    x = x.to(softmax_dtype) * scale
    x = torch.nn.functional.softmax(x, dim=-1).to(input_dtype)
    return x


@torch.jit.script
def masked_softmax(
    x: torch.Tensor, mask: torch.Tensor, mask_value: torch.Tensor
):
    x = torch.where(mask, x, mask_value)
    x = torch.nn.functional.softmax(x, dim=-1)
    return x


def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
    """torch.repeat_interleave(x, dim=2, repeats=n_rep)"""
    total_seqlen, n_kv_heads, head_dim = x.shape
    if n_rep == 1:
        return x
    return (
        x[:, :, None, :]
        .expand(total_seqlen, n_kv_heads, n_rep, head_dim)
        .reshape(total_seqlen, n_kv_heads * n_rep, head_dim)
    )


def mask_eos_token(
    logits: torch.Tensor,
    eos_token_id: Optional[int] = None,
) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
    # for min_new_tokens
    if eos_token_id is not None:
        logits[..., eos_token_id] = torch.finfo(logits.dtype).min
    return logits


def gather_shifted_log_probs(
    logits: torch.FloatTensor, labels: torch.LongTensor
) -> torch.FloatTensor:
    """Gather log probs of shifted labels from logits.

    Args:
        logits (torch.FloatTensor): Non-shifted logits with shape [bs, seqlen].
            The final value at [:, seqlen -1] is not used.
        labels (torch.LongTensor): Non-shifted labels/input_ids with shape [bs, seqlen].
            The first value at [:, 0] has no corresponding log prob.

    Returns:
        torch.FloatTensor: Shifted log probability with shape [bs, seqlen -1].
    """
    logits = logits[:, :-1]
    labels = labels[:, 1:]
    log_probs = torch.nn.functional.log_softmax(logits, dim=-1)
    log_probs_labels = log_probs.gather(dim=-1, index=labels.unsqueeze(-1))
    return log_probs_labels.squeeze(-1)


def build_shift_one_indices(
    x: torch.HalfTensor, cu_seqlens: torch.IntTensor
) -> torch.IntTensor:
    """Build indices for shifting labels/input_ids one step to the left.

    Equivalent to:
    ```
    shift_one_indices = torch.cat([
        torch.arange(cu_seqlens[i] + 1, cu_seqlens[i + 1], dtype=torch.long, device=cu_seqlens.device)
        for i in range(cu_seqlens.shape[0] - 1)
    ])
    ```
    but the above implementaion will implicitly convert a tensor (cu_seqlens[i]) to an integer,
    which will cause a cuda device sync and slow down performance.

    Args:
        x (torch.HalfTensor): Shape [total_seqlen]. This tensor is required to get
            total_seqlen from its shape. Computing total_seqlen from cu_seqlens will implicitly cause
            a cuda device sync.
        cu_seqlens (torch.IntTensor): Shape [bs + 1]. Indices marking the start
            and end of each sequences.

    Returns:
        torch.IntTensor: Shape [tot_seqlen - bs]. Indices for shifting labels/input_ids
            one step to the left.
    """
    total_seqlen = x.shape[0]
    bs = cu_seqlens.shape[0] - 1
    short1lens = cu_seqlens[1:] - cu_seqlens[:-1] - 1
    short1cu_seqlens = torch.nn.functional.pad(
        short1lens.cumsum(0), (1, 0), value=0
    )
    indexing_t = torch.arange(
        total_seqlen - bs, dtype=torch.long, device=cu_seqlens.device
    )
    return indexing_t + (
        indexing_t.unsqueeze(0) >= short1cu_seqlens[:-1].unsqueeze(1)
    ).sum(0)


def build_leave_one_indices(
    x: torch.HalfTensor, cu_seqlens: torch.IntTensor
) -> torch.IntTensor:
    """Build indices for leaving one token out at the end of each sequence.

    Equivalent to:
    ```
    leave_one_indices = torch.cat([
        torch.arange(cu_seqlens[i], cu_seqlens[i + 1] - 1, dtype=torch.long, device=cu_seqlens.device)
        for i in range(cu_seqlens.shape[0] - 1)
    ])
    ```
    but the above implementaion will implicitly convert a tensor (cu_seqlens[i]) to an integer,
    which will cause a cuda device sync and slow down performance.

    Args:
        x (torch.HalfTensor): Shape [total_seqlen]. This tensor is required to get
            total_seqlen from its shape. Computing total_seqlen from cu_seqlens will implicitly cause
            a cuda device sync.
        cu_seqlens (torch.IntTensor): Shape [bs + 1]. Indices marking the start
            and end of each sequences.

    Returns:
        torch.IntTensor: Shape [tot_seqlen - bs]. Indices for shifting labels/input_ids
            one step to the left.
    """
    total_seqlen = x.shape[0]
    bs = cu_seqlens.shape[0] - 1
    short1lens = cu_seqlens[1:] - cu_seqlens[:-1] - 1
    short1cu_seqlens = torch.nn.functional.pad(
        short1lens.cumsum(0), (1, 0), value=0
    )
    indexing_t = torch.arange(
        total_seqlen - bs, dtype=torch.long, device=cu_seqlens.device
    )
    return (
        indexing_t
        + (indexing_t.unsqueeze(0) >= short1cu_seqlens[:-1].unsqueeze(1)).sum(0)
        - 1
    )


def gather_packed_shifted_log_probs(
    logits: torch.FloatTensor,
    cu_seqlens: torch.Tensor,
    labels: torch.LongTensor,
) -> torch.FloatTensor:
    """Gather log probs from packed input_ids and logits.

    Args:
        logits_ (torch.FloatTensor): Shape [tot_seqlen]. The final value at the end of
            each sequence is not used.
        cu_seqlens (torch.Tensor): Shape [#seqs + 1]. Indices marking the start
            and end of each sequences.
        labels_ (torch.LongTensor): Labels or input_ids with shape [tot_seqlen].
            The first value at the beginning of each sequence has no corresponding log prob.

    Returns:
        torch.FloatTensor: Log probability with shape [tot_seqlen - #seqs].
    """
    labels = torch.nn.functional.pad(labels[1:], (0, 1), value=0)
    leave_one_indices = build_leave_one_indices(logits, cu_seqlens)
    if constants.model_parallel_world_size() > 1:
        # NOTE: logprobs is freaking sensitive to input_ids. If the input sequence is a natural sequence, everything will be fine.
        # However, if we input random token IDs, parallel cross entropy can produce VERY different results than the normal
        # torch.gather based version (e.g., the maximum absolute different can reach ~50).
        from realhf.impl.model.parallelism.model_parallel.modules import (
            vocab_parallel_cross_entropy,
        )

        logprobs = -vocab_parallel_cross_entropy(logits, labels)[
            leave_one_indices
        ]
        ########### sanity check ###########
        # world_size = constants.model_parallel_world_size()
        # dim_size = [logits.shape[1] * world_size, logits.shape[0]]
        # all_gather_buffer = torch.zeros(*dim_size, dtype=logits.dtype, device=logits.device)
        # torch.distributed._all_gather_base(
        #     all_gather_buffer,
        #     logits.transpose(0, 1).contiguous(),
        #     group=constants.model_parallel_group(),
        # )
        # logits2 = all_gather_buffer.transpose(0, 1).contiguous()
        # logprobs2 = gather_packed_shifted_log_probs(logits2, cu_seqlens, packed_input_ids).float()
        # assert torch.allclose(logprobs, logprobs2, atol=2e-2), (
        #     (logprobs - logprobs2).abs().max(),
        #     logprobs,
        #     logprobs2,
        # )
        ########### sanity check ###########
        return logprobs
    logits_shape = logits.shape
    # shift_one_indices = torch.cat([
    #     torch.arange(cu_seqlens[i] + 1 , cu_seqlens[i + 1], dtype=torch.long, device=cu_seqlens.device)
    #     for i in range(cu_seqlens.shape[0] - 1)
    # ])
    # shift labels one step to the left and pad it to match the shape of logits
    log_probs = torch.nn.functional.log_softmax(logits, dim=-1)
    log_probs_labels = log_probs.gather(
        dim=-1, index=labels.unsqueeze(-1)
    ).squeeze(-1)
    log_probs_labels = log_probs_labels[leave_one_indices]
    assert (
        log_probs_labels.shape[0] == logits_shape[0] - cu_seqlens.shape[0] + 1
    ), (
        log_probs_labels.shape,
        logits_shape,
        cu_seqlens.shape,
        cu_seqlens,
        # shift_one_indices,
    )
    return log_probs_labels


def apply_logits_mask(logits: torch.HalfTensor, mask: torch.BoolTensor):
    assert (
        mask.shape[-1]
        == logits.shape[-1] * constants.model_parallel_world_size()
    ), (
        constants.model_parallel_world_size(),
        logits.shape,
        mask.shape,
    )
    parallel_vocab_size = logits.shape[-1]
    mp_rank = constants.model_parallel_rank()
    mask = mask[
        :, mp_rank * parallel_vocab_size : (mp_rank + 1) * parallel_vocab_size
    ]
    logits.masked_fill_(mask, torch.finfo(logits.dtype).min)


@torch.no_grad()
def masked_normalization(
    x: torch.Tensor,
    mask: Optional[torch.BoolTensor] = None,
    dim=None,
    inplace=False,
    unbiased=False,
    eps=1e-5,
    high_precision=True,
):
    """Normalize x with a mask. Typically used in advantage normalization.

    Args:
        x (torch.Tensor):
            Tensor to be normalized.
        mask (torch.Tensor, optional):
            A mask with the same shape as x. Defaults to None.
        dim (int or tuple of ints, optional):
            Dimensions to be normalized. Defaults to None.
        inplace (bool, optional):
            Whether to perform in-place operation. Defaults to False.
        eps (torch.Tensor, optional):
            Minimal denominator. Defaults to 1e-5.

    Returns:
        torch.Tensor:
            Normalized x, with the same shape as x.
    """
    dtype = torch.float64 if high_precision else torch.float32
    x = x.to(dtype)
    if not inplace:
        x = x.clone()
    if dim is None:
        dim = tuple(range(len(x.shape)))
    if mask is None:
        factor = torch.tensor(
            np.prod([x.shape[d] for d in dim]), dtype=dtype, device=x.device
        )
    else:
        mask = mask.to(dtype)
        assert len(mask.shape) == len(x.shape), (mask.shape, x.shape, dim)
        for i in range(len(x.shape)):
            if i in dim:
                assert mask.shape[i] == x.shape[i], (mask.shape, x.shape, dim)
            else:
                assert mask.shape[i] == 1, (mask.shape, x.shape, dim)
        x = x * mask
        factor = mask.sum(dim, keepdim=True)
    x_sum = x.sum(dim=dim, keepdim=True)
    x_sum_sq = x.square().sum(dim=dim, keepdim=True)
    if dist.is_initialized():
        dist.all_reduce(
            factor, op=dist.ReduceOp.SUM, group=constants.data_parallel_group()
        )
        dist.all_reduce(
            x_sum, op=dist.ReduceOp.SUM, group=constants.data_parallel_group()
        )
        dist.all_reduce(
            x_sum_sq,
            op=dist.ReduceOp.SUM,
            group=constants.data_parallel_group(),
        )
    mean = x_sum / factor
    meansq = x_sum_sq / factor
    var = meansq - mean**2
    if unbiased:
        var *= factor / (factor - 1)
    return ((x - mean) / (var.sqrt() + eps)).float()


def get_eos_indices(
    input_ids: torch.LongTensor,
    tokenizer: transformers.PreTrainedTokenizerFast,
) -> Tuple[torch.LongTensor, torch.FloatTensor]:
    if torch.any(input_ids[:, 0] == tokenizer.eos_token_id):
        indices = (
            (input_ids[:, 0] == tokenizer.eos_token_id).nonzero().flatten()
        )
        bad_input_ids = input_ids[indices]
        bad_strs = tokenizer.batch_decode(
            bad_input_ids,
            skip_special_tokens=True,
            clean_up_tokenization_spaces=True,
        )
        raise RuntimeError(
            f"Generated sequence terminates unexpectedly early: {bad_strs}"
        )
    seq_len = input_ids.shape[1]
    eos_mask = (input_ids == tokenizer.eos_token_id).float()
    seq_no_eos_mask = (eos_mask.sum(1) == 0).float()
    eos_indices = eos_mask.argmax(1)
    eos_indices = (
        eos_indices * (1 - seq_no_eos_mask) + seq_no_eos_mask * (seq_len - 1)
    ).long()
    return eos_indices, seq_no_eos_mask


def torch_attn_func(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    causal: bool,
    cu_seqlens_q: torch.IntTensor,
    max_seqlen_q: int,
    cu_seqlens_k: torch.IntTensor,
    max_seqlen_k: int,
    dropout_p: float,
    softmax_scale: float,
    upcast_unscale: float = 1.0,
) -> torch.Tensor:
    """PyTorch implementation of the attention function with a flash-attn-like realhf.api.

    We use this function to compare the output of our model and huggingface models.
    Flash-attn/float16/CUDAkernels will all more or less suffer from float point errors.
    We call this function with float32 and CPU to get the "ground truth" output.

    Args:
        q (torch.Tensor): Shape [total_seqlen, #q, head_dim].
        k (torch.Tensor): Shape [total_seqlen, #kv, head_dim].
        v (torch.Tensor): Shape [total_seqlen, #kv, head_dim].
        causal (bool): .
        dropout_p (float): .
        softmax_scale (float): .
        upcast_unscale (float, optional): Scale factor when upcastin attention scores.
            Defaults to 1.0.

    Returns:
        torch.Tensor: Attention score. Shape [bs, seqlen, #q, head_dim].
    """
    nq = q.shape[-2]
    nkv = k.shape[-2]
    n_rep = q.shape[-2] // k.shape[-2]
    bsz = cu_seqlens_q.shape[0] - 1
    # repeat k/v heads if n_kv_heads < n_heads
    k = repeat_kv(k, n_rep)  # (total_seqlen, nq, head_dim)
    v = repeat_kv(v, n_rep)  # (total_seqlen, nq, head_dim)

    input_lens_k = cu_seqlens_k[1:] - cu_seqlens_k[:-1]
    attention_mask_k = torch.arange(
        max_seqlen_k, dtype=torch.long, device="cpu"
    ).unsqueeze(0) < input_lens_k.unsqueeze(1)
    _, _pad_indices_k, _, _ = unpad_input(attention_mask_k, attention_mask_k)

    input_lens_q = cu_seqlens_q[1:] - cu_seqlens_q[:-1]
    attention_mask_q = torch.arange(
        max_seqlen_q, dtype=torch.long, device="cpu"
    ).unsqueeze(0) < input_lens_q.unsqueeze(1)
    _, _pad_indices_q, _, _ = unpad_input(attention_mask_q, attention_mask_q)

    q = pad_input(q, _pad_indices_q, bsz, max_seqlen_q)
    k = pad_input(k, _pad_indices_k, bsz, max_seqlen_k)
    v = pad_input(v, _pad_indices_k, bsz, max_seqlen_k)

    q = q.transpose(1, 2)  # (bs, nq, seqlen, head_dim)
    k = k.transpose(1, 2)
    v = v.transpose(1, 2)
    scores = torch.matmul(q, k.transpose(2, 3)) * softmax_scale

    mask = (
        attention_mask_k.unsqueeze(1)
        .unsqueeze(1)
        .repeat(1, nq, max_seqlen_q, 1)
    )  # [bs, nq, seqlen, seqlen]
    if causal:
        _ms = max(max_seqlen_q, max_seqlen_k)
        causal_mask = torch.tril(
            torch.ones(_ms, _ms, device=q.device, dtype=torch.bool)
        )[-max_seqlen_q:, -max_seqlen_k:]
        mask = mask & causal_mask

    # if mask_softmax:
    scores = upcast_masked_softmax(
        scores,
        mask,
        mask_value=torch.full(
            [],
            torch.finfo(torch.float32).min,
            device=scores.device,
            dtype=torch.float32,
        ),
        scale=upcast_unscale,
        softmax_dtype=torch.float32,
    )
    # else:
    #     scores = upcast_softmax(scores, scale=upcast_unscale, softmax_dtype=torch.float32)
    scores = torch.nn.functional.dropout(scores, p=dropout_p)
    scores = scores.to(q.dtype)
    output = torch.matmul(scores, v)  # (bs, nq, seqlen, head_dim)
    output = output.transpose(1, 2).contiguous()

    output = unpad_input(output, attention_mask_q)[0]
    return output


def rotate_half(x: torch.HalfTensor, interleaved: bool = False):
    if not interleaved:
        x1, x2 = x.chunk(2, dim=-1)
        return torch.cat((-x2, x1), dim=-1)
    else:
        x1, x2 = x[..., ::2], x[..., 1::2]
        # return rearrange(torch.stack((-x2, x1), dim=-1), "... d two -> ... (d two)", two=2)
        return torch.stack((-x2, x1), dim=-1).flatten(start_dim=-2)


@torch.no_grad()
@torch.jit.script
def compute_varlen_position_indices(
    total_seqlen: int,
    cu_seqlens: torch.IntTensor,
    seqlen_offsets: Optional[torch.IntTensor] = None,
) -> torch.IntTensor:
    indexing_t = torch.arange(
        total_seqlen, dtype=torch.long, device=cu_seqlens.device
    ).unsqueeze_(0)
    indexing_t = (cu_seqlens[:-1].unsqueeze(1) <= indexing_t) & (
        indexing_t < cu_seqlens[1:].unsqueeze(1)
    )
    indices = indexing_t.cumsum(1) - 1
    if seqlen_offsets is not None:
        indices += seqlen_offsets.unsqueeze(1)
    return torch.where(indexing_t, indices, 0).sum(0)


# @torch.jit.script
def apply_rotary_varlen(
    x: torch.HalfTensor,
    cos: torch.HalfTensor,
    sin: torch.HalfTensor,
    cu_seqlens: torch.IntTensor,
    interleaved: bool,
    seqlen_offsets: Optional[torch.IntTensor] = None,
    rotary_indices: Optional[torch.LongTensor] = None,
) -> Tuple[torch.HalfTensor, torch.LongTensor]:
    if rotary_indices is None:
        rotary_indices = compute_varlen_position_indices(
            x.shape[0], cu_seqlens, seqlen_offsets
        )

    ro_dim = cos.shape[-1] * 2
    assert ro_dim <= x.shape[-1]
    cos = cos[rotary_indices]
    sin = sin[rotary_indices]
    if not interleaved:
        cos = cos[:, None, None, :].repeat(1, 1, 2, 1).flatten(start_dim=-2)
        sin = sin[:, None, None, :].repeat(1, 1, 2, 1).flatten(start_dim=-2)
    else:
        cos = cos[:, None, :, None].repeat(1, 1, 1, 2).flatten(start_dim=-2)
        sin = sin[:, None, :, None].repeat(1, 1, 1, 2).flatten(start_dim=-2)

    # cos = repeat(cos, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)")
    # sin = repeat(sin, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)")
    return torch.cat(
        [
            x[..., :ro_dim] * cos
            + rotate_half(x[..., :ro_dim], interleaved) * sin,
            x[..., ro_dim:],
        ],
        dim=-1,
    )


def apply_rotary(
    x: torch.HalfTensor,
    cos: torch.HalfTensor,
    sin: torch.HalfTensor,
    interleaved: bool = False,
):
    """
    x: (batch_size, seqlen, nheads, headdim)
    cos, sin: (seqlen, rotary_dim / 2) or (batch_size, seqlen, rotary_dim / 2)
    """
    ro_dim = cos.shape[-1] * 2
    assert ro_dim <= x.shape[-1]
    if not interleaved:
        cos = cos[:, None, None, :].repeat(1, 1, 2, 1).flatten(start_dim=-2)
        sin = sin[:, None, None, :].repeat(1, 1, 2, 1).flatten(start_dim=-2)
    else:
        cos = cos[:, None, :, None].repeat(1, 1, 1, 2).flatten(start_dim=-2)
        sin = sin[:, None, :, None].repeat(1, 1, 1, 2).flatten(start_dim=-2)
    return torch.cat(
        [
            x[..., :ro_dim] * cos
            + rotate_half(x[..., :ro_dim], interleaved) * sin,
            x[..., ro_dim:],
        ],
        dim=-1,
    )
