import importlib.util

import torch
import torch.distributed as dist

try:
    # `turbox` and `paifuser` are two internal accelerators with different module layouts: a turbox install
    # exposes the stock `xfuser` package, while paifuser keeps the xfuser code inside its own
    # `paifuser.xfuser` namespace. Only the paifuser path therefore needs its own imports, and it is taken
    # only when turbox is absent -- the same precedence as in `videox_fun/__init__.py`.
    if importlib.util.find_spec("turbox") is None and importlib.util.find_spec("paifuser") is not None:
        import paifuser
        from paifuser.xfuser.core.distributed import (
            get_sequence_parallel_rank, get_sequence_parallel_world_size,
            get_sp_group, get_world_group, init_distributed_environment,
            initialize_model_parallel, model_parallel_is_initialized)
        from paifuser.xfuser.core.long_ctx_attention import \
            xFuserLongContextAttention
        print("Import PAI DiT Turbo")
    else:
        import xfuser
        from xfuser.core.distributed import (get_sequence_parallel_rank,
                                             get_sequence_parallel_world_size,
                                             get_sp_group, get_world_group,
                                             init_distributed_environment,
                                             initialize_model_parallel,
                                             model_parallel_is_initialized)
        from xfuser.core.long_ctx_attention import xFuserLongContextAttention
        print("Xfuser import sucessful")
except Exception:
    # Without an xfuser backend every helper below takes its single-GPU path, which relies on these names
    # being `None`.
    get_sequence_parallel_world_size = get_sequence_parallel_rank = None
    get_sp_group = get_world_group = None
    init_distributed_environment = initialize_model_parallel = None
    model_parallel_is_initialized = None
    xFuserLongContextAttention = None

def set_multi_gpus_devices(ulysses_degree, ring_degree, classifier_free_guidance_degree=1):
    if ulysses_degree > 1 or ring_degree > 1 or classifier_free_guidance_degree > 1:
        if get_sp_group is None:
            raise RuntimeError("xfuser is not installed.")
        dist.init_process_group("nccl")
        print('parallel inference enabled: ulysses_degree=%d ring_degree=%d classifier_free_guidance_degree=% rank=%d world_size=%d' % (
            ulysses_degree, ring_degree, classifier_free_guidance_degree, dist.get_rank(),
            dist.get_world_size()))
        assert dist.get_world_size() == ring_degree * ulysses_degree * classifier_free_guidance_degree, \
                    "number of GPUs(%d) should be equal to ring_degree * ulysses_degree * classifier_free_guidance_degree." % dist.get_world_size()
        init_distributed_environment(rank=dist.get_rank(), world_size=dist.get_world_size())
        initialize_model_parallel(sequence_parallel_degree=ring_degree * ulysses_degree,
                classifier_free_guidance_degree=classifier_free_guidance_degree,
                ring_degree=ring_degree,
                ulysses_degree=ulysses_degree)
        # device = torch.device("cuda:%d" % dist.get_rank())
        device = torch.device(f"cuda:{get_world_group().local_rank}")
        print('rank=%d device=%s' % (get_world_group().rank, str(device)))
    else:
        device = "cuda"
    return device

def sequence_parallel_chunk(x, dim=1):
    if get_sequence_parallel_world_size is None or not model_parallel_is_initialized():
        return x

    sp_world_size = get_sequence_parallel_world_size()
    if sp_world_size <= 1:
        return x

    sp_rank = get_sequence_parallel_rank()
    sp_group = get_sp_group()

    if x.size(1) % sp_world_size != 0:
        raise ValueError(f"Dim 1 of x ({x.size(1)}) not divisible by SP world size ({sp_world_size})")

    chunks = torch.chunk(x, sp_world_size, dim=1)
    x = chunks[sp_rank]

    return x

def sequence_parallel_all_gather(x, dim=1):
    if get_sequence_parallel_world_size is None or not model_parallel_is_initialized():
        return x

    sp_world_size = get_sequence_parallel_world_size()
    if sp_world_size <= 1:
        return x  # No gathering needed

    sp_group = get_sp_group()
    gathered_x = sp_group.all_gather(x, dim=dim)
    return gathered_x

def ulysses_all_to_all(x, scatter_dim, gather_dim):
    """Ulysses (head-parallel) all-to-all on a 4D ``[B, S, H, D]`` tensor.

    ``scatter_dim`` / ``gather_dim`` follow yunchang's ``all_to_all_4D`` convention:
    ``(2, 1)`` turns a sequence-split layout ``[B, S/P, H, D]`` into a head-split one
    ``[B, S, H/P, D]``; ``(1, 2)`` is the inverse. ``x`` is returned unchanged when sequence
    parallel is not initialised or the SP world size is 1, so single-GPU paths are unaffected.
    """
    if get_sequence_parallel_world_size is None or not model_parallel_is_initialized():
        return x
    if get_sequence_parallel_world_size() <= 1:
        return x
    from yunchang.comm.all_to_all import all_to_all_4D
    return all_to_all_4D(x, scatter_idx=scatter_dim, gather_idx=gather_dim,
                         group=get_sp_group().device_group)