from flash_attn import flash_attn_varlen_qkvpacked_func
import torch
import torch.distributed as dist
from yunchang import zigzag_ring_flash_attn_varlen_qkvpacked_func


def log(msg, a, rank0_only=False):
    world_size = dist.get_world_size()
    rank = dist.get_rank()
    if rank0_only:
        if rank == 0:
            print(
                f"{msg}: "
                f"max {a.abs().max().item()}, "
                f"mean {a.abs().mean().item()}",
                flush=True,
            )
        return

    for i in range(world_size):
        if i == rank:
            if rank == 0:
                print(f"{msg}:")
            print(
                f"[{rank}] "
                f"max {a.abs().max().item()}, "
                f"mean {a.abs().mean().item()}",
                flush=True,
            )
        dist.barrier()


def extract_local(value, cu_seqlens, rank, world_size):
    local_values = []
    for i in range(len(cu_seqlens) - 1):
        start, end = cu_seqlens[i], cu_seqlens[i + 1]
        local_value = value[start:end].chunk(2 * world_size, dim=0)
        local_values.extend(
            [
                local_value[rank].detach().clone(),
                local_value[2 * world_size - 1 - rank].detach().clone(),
            ]
        )
    return torch.cat(local_values, dim=0).contiguous()


def extract_lse(lse, cu_seqlens):
    values = []
    for i in range(len(cu_seqlens) - 1):
        start, end = cu_seqlens[i], cu_seqlens[i + 1]
        value = lse[i, :, : end - start]
        values.append(value)
    return values


if __name__ == "__main__":
    dist.init_process_group("nccl")
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    dtype = torch.bfloat16
    device = torch.device(f"cuda:{rank}")

    batch_size = 1
    nheads = 5
    d = 128
    dropout_p = 0
    causal = True
    deterministic = False

    cu_seqlens = [0, 128, 1248, 4240]
    cu_seqlens_tensor = torch.tensor(cu_seqlens, dtype=torch.int32, device=device)
    max_seqlen = (cu_seqlens_tensor[1:] - cu_seqlens_tensor[:-1]).max().item()
    total_length = cu_seqlens[-1]
    num_seq = len(cu_seqlens) - 1

    assert torch.all(cu_seqlens_tensor % (2 * world_size) == 0)
    assert d % 8 == 0

    qkv = torch.randn(
        total_length, 3, nheads, d, device=device, dtype=dtype, requires_grad=True
    )
    dist.broadcast(qkv, src=0)

    dout = torch.randn(total_length, nheads, d, device=device, dtype=dtype)
    dist.broadcast(dout, src=0)

    local_cu_seqlens_tensor = cu_seqlens_tensor // world_size
    local_max_seqlen = max_seqlen // world_size

    local_qkv = extract_local(qkv, cu_seqlens, rank, world_size)
    local_qkv.requires_grad = True
    local_dout = extract_local(dout, cu_seqlens, rank, world_size)

    dist.barrier()
    if rank == 0:
        print("#" * 30)
        print("# forward:")
        print("#" * 30)

    out, lse, _ = flash_attn_varlen_qkvpacked_func(
        qkv,
        cu_seqlens_tensor,
        max_seqlen,
        dropout_p=dropout_p,
        causal=causal,
        window_size=(-1, -1),
        softcap=0.0,
        alibi_slopes=None,
        deterministic=deterministic,
        return_attn_probs=True,
    )

    local_out = extract_local(out, cu_seqlens, rank, world_size)
    lse_list = extract_lse(lse, cu_seqlens)

    ring_out, ring_lse, _ = zigzag_ring_flash_attn_varlen_qkvpacked_func(
        local_qkv,
        local_cu_seqlens_tensor,
        local_max_seqlen,
        dropout_p=dropout_p,
        causal=causal,
        window_size=(-1, -1),
        softcap=0.0,
        alibi_slopes=None,
        deterministic=deterministic,
        return_attn_probs=True,
    )

    ring_lse_list = extract_lse(ring_lse, local_cu_seqlens_tensor.tolist())

    log("out", out, rank0_only=True)
    log("out diff", local_out - ring_out)

    for i, (lse, ring_lse) in enumerate(zip(lse_list, ring_lse_list)):
        local_lse = lse.chunk(2 * world_size, dim=-1)
        local_lse = torch.cat(
            [local_lse[rank], local_lse[2 * world_size - 1 - rank]], dim=-1
        )
        log(f"lse {i}", lse, rank0_only=True)
        log(f"lse diff {i}", local_lse - ring_lse)

    dist.barrier()
    if rank == 0:
        print("#" * 30)
        print("# backward:")
        print("#" * 30)

    out.backward(dout)
    dqkv = qkv.grad
    local_dqkv = extract_local(dqkv, cu_seqlens, rank, world_size)

    ring_out.backward(local_dout)
    ring_dqkv = local_qkv.grad

    log("load_dq", local_dqkv[:, 0])
    log("dq diff", local_dqkv - ring_dqkv)

    log("load_dk", local_dqkv[:, 1])
    log("dk diff", local_dqkv[:, 1] - ring_dqkv[:, 1])

    log("load_dv", local_dqkv[:, 2])
    log("dv diff", local_dqkv[:, 2] - ring_dqkv[:, 2])
