"""
This repo is forked from [Boyuan Chen](https://boyuan.space/)'s research 
template [repo](https://github.com/buoyancy99/research-template). 
By its MIT license, you must keep the above sentence in `README.md` 
and the `LICENSE` file to credit the author.
"""

from typing import Callable
import torch
import torch.distributed as dist
from lightning.pytorch.utilities.rank_zero import rank_zero_only
from utils.print_utils import print_once

is_rank_zero = rank_zero_only.rank == 0

rank_zero_print = rank_zero_only(print)
rank_zero_print_once = rank_zero_only(print_once)


def broadcast_from_zero(
    init_func: Callable[[], torch.Tensor], func: Callable[[], torch.Tensor]
) -> torch.Tensor:
    """
    Broadcasts the result of func() from rank 0 to all other processes.
    If distributed is not initialized, returns the result of func() directly.
    Args:
        init_func: function that initializes a tensor (will be overwritten)
        func: function that generates the tensor to broadcast from rank 0
    Returns:
        torch.Tensor: the tensor generated by func() on rank 0
    """
    if not dist.is_initialized():
        return func()
    tensor = init_func()
    if is_rank_zero:
        tensor = func()
    dist.broadcast(tensor, src=0)
    return tensor