# Copyright (c) ModelScope Contributors. All rights reserved.
"""Megatron-specific GKD utilities: TP-aware gather/topk and CP reduce."""
import torch
from megatron.core import mpu


def vocab_parallel_topk(logits: torch.Tensor, k: int) -> tuple:
    """Global top-k from vocab-parallel sharded logits. TP=1 → plain torch.topk."""
    tp_size = mpu.get_tensor_model_parallel_world_size()
    if tp_size == 1:
        return torch.topk(logits, k=k, dim=-1)

    tp_rank = mpu.get_tensor_model_parallel_rank()
    tp_group = mpu.get_tensor_model_parallel_group()
    partition_vocab_size = logits.shape[-1]

    local_topk_vals, local_topk_ids = torch.topk(logits, k=k, dim=-1)
    local_topk_ids = local_topk_ids + tp_rank * partition_vocab_size

    gathered_vals = [torch.empty_like(local_topk_vals) for _ in range(tp_size)]
    gathered_ids = [torch.empty_like(local_topk_ids) for _ in range(tp_size)]
    torch.distributed.all_gather(gathered_vals, local_topk_vals, group=tp_group)
    torch.distributed.all_gather(gathered_ids, local_topk_ids, group=tp_group)

    all_vals = torch.cat(gathered_vals, dim=-1)
    all_ids = torch.cat(gathered_ids, dim=-1)
    global_topk_vals, sel = torch.topk(all_vals, k=k, dim=-1)
    global_topk_ids = torch.gather(all_ids, dim=-1, index=sel)
    return global_topk_vals, global_topk_ids


def tp_gather_topk(logits: torch.Tensor, indices: torch.Tensor) -> torch.Tensor:
    """Gather logits at global top-k indices with TP-aware partitioning."""
    tp_size = mpu.get_tensor_model_parallel_world_size()
    if tp_size == 1:
        return torch.gather(logits, dim=-1, index=indices)

    tp_rank = mpu.get_tensor_model_parallel_rank()
    partition_vocab_size = logits.shape[-1]
    vocab_start = tp_rank * partition_vocab_size

    in_range = (indices >= vocab_start) & (indices < vocab_start + partition_vocab_size)
    local_indices = (indices - vocab_start).clamp(0, partition_vocab_size - 1)
    gathered = torch.gather(logits, dim=-1, index=local_indices)
    gathered = gathered.masked_fill(~in_range, float('-inf'))

    gathered_for_reduce = gathered.detach()
    torch.distributed.all_reduce(
        gathered_for_reduce, op=torch.distributed.ReduceOp.MAX, group=mpu.get_tensor_model_parallel_group())
    return torch.where(in_range, gathered, gathered_for_reduce)


def cp_reduce(total_loss, num_valid, *, cp_size):
    """Normalize total_loss by num_valid with CP all-reduce when cp_size > 1."""
    num_valid_f = num_valid.float() if isinstance(num_valid, torch.Tensor) else torch.tensor(
        float(num_valid), device=total_loss.device)
    if cp_size > 1:
        torch.distributed.all_reduce(
            num_valid_f, op=torch.distributed.ReduceOp.SUM, group=mpu.get_context_parallel_group())
        torch.distributed.all_reduce(
            total_loss, op=torch.distributed.ReduceOp.SUM, group=mpu.get_context_parallel_group())
    # Avoid the host-device sync from ``if num_valid_f == 0`` (num_valid_f is a GPU scalar,
    # so a Python bool() forces a .item() stream sync every step). Keep the zero-valid
    # guard on-device with torch.where: num_valid==0 -> loss 0, else total_loss/num_valid.
    safe = torch.where(num_valid_f == 0, torch.ones_like(num_valid_f), num_valid_f)
    loss = total_loss / safe
    return torch.where(num_valid_f == 0, torch.zeros_like(loss), loss)