from functools import partial
import math
import torch
import torch.nn as nn
import torch.nn.functional as F

from einops import rearrange

try:
    import cudnn
except ImportError:
    cudnn = None


# from flash_attn_interface import flash_attn_func as flash_attn_func_v3
# from flash_attn_interface import flash_attn_varlen_func as flash_attn_varlen_func_v3
# from flash_attn_interface import flash_attn_with_kvcache

from flash_attn.cute.interface import flash_attn_func as flash_attn_func_v3_python

def convert_to_cudnn_type(torch_type):
    if torch_type == torch.float16:
        return cudnn.data_type.HALF
    elif torch_type == torch.bfloat16:
        return cudnn.data_type.BFLOAT16
    elif torch_type == torch.float32:
        return cudnn.data_type.FLOAT
    elif torch_type == torch.int32:
        return cudnn.data_type.INT32
    elif torch_type == torch.int64:
        return cudnn.data_type.INT64
    else:
        raise ValueError("Unsupported tensor data type.")


def cudnn_spda_setup(q, k, v, causal=False):
    b, nheads, seqlen_q, headdim = q.shape
    _, _, seqlen_k, _ = k.shape
    assert v.shape == (b, nheads, seqlen_k, headdim)
    assert cudnn is not None, 'CUDNN is not available'
    q_gpu, k_gpu, v_gpu = q, k, v
    o_gpu = torch.empty_like(q_gpu)
    stats_gpu = torch.empty(b, nheads, seqlen_q, 1, dtype=torch.float32, device=q.device)
    graph = cudnn.pygraph(
        io_data_type=convert_to_cudnn_type(q.dtype),
        intermediate_data_type=cudnn.data_type.FLOAT,
        compute_data_type=cudnn.data_type.FLOAT,
    )
    q = graph.tensor_like(q_gpu.detach())
    k = graph.tensor_like(k_gpu.detach())
    v = graph.tensor_like(v_gpu.detach())

    o, stats = graph.sdpa(
        name="sdpa",
        q=q,
        k=k,
        v=v,
        is_inference=False,
        attn_scale=1.0 / math.sqrt(headdim),
        use_causal_mask=causal,
    )

    o.set_output(True).set_dim(o_gpu.shape).set_stride(o_gpu.stride())
    stats.set_output(True).set_data_type(cudnn.data_type.FLOAT)

    graph.validate()
    graph.build_operation_graph()
    graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK])
    graph.check_support()
    graph.build_plans()

    variant_pack = {
        q: q_gpu,
        k: k_gpu,
        v: v_gpu,
        o: o_gpu,
        stats: stats_gpu,
    }

    workspace = torch.empty(graph.get_workspace_size(), device="cuda", dtype=torch.uint8)

    def cudnn_spda_fwd(*args, **kwargs):
        graph.execute(variant_pack, workspace)
        return o_gpu

    return cudnn_spda_fwd


def cudnn_spda_bwd_setup(q, k, v, o, g, lse, causal=False):
    b, nheads, seqlen_q, headdim = q.shape
    _, _, seqlen_k, _ = k.shape
    assert v.shape == (b, nheads, seqlen_k, headdim)
    assert g.shape == (b, nheads, seqlen_q, headdim)
    assert o.shape == (b, nheads, seqlen_q, headdim)
    assert lse.shape == (b, nheads, seqlen_q, 1)
    assert cudnn is not None, 'CUDNN is not available'
    q_gpu, k_gpu, v_gpu, o_gpu, g_gpu = q, k, v, o, g
    dq_gpu = torch.empty_like(q_gpu)
    dk_gpu = torch.empty_like(k_gpu)
    dv_gpu = torch.empty_like(v_gpu)
    graph = cudnn.pygraph(
        io_data_type=convert_to_cudnn_type(q.dtype),
        intermediate_data_type=cudnn.data_type.FLOAT,
        compute_data_type=cudnn.data_type.FLOAT,
    )
    q = graph.tensor_like(q_gpu.detach())
    k = graph.tensor_like(k_gpu.detach())
    v = graph.tensor_like(v_gpu.detach())
    o = graph.tensor_like(o_gpu.detach())
    g = graph.tensor_like(g_gpu.detach())
    stats = graph.tensor_like(lse.detach())

    dq, dk, dv = graph.sdpa_backward(
        name="sdpa_backward",
        q=q,
        k=k,
        v=v,
        o=o,
        dO=g,
        stats=stats,
        attn_scale=1.0 / math.sqrt(headdim),
        use_causal_mask=causal,
    )

    dq.set_output(True).set_dim(dq_gpu.shape).set_stride(dq_gpu.stride())
    dk.set_output(True).set_dim(dk_gpu.shape).set_stride(dk_gpu.stride())
    dv.set_output(True).set_dim(dv_gpu.shape).set_stride(dv_gpu.stride())

    graph.validate()
    graph.build_operation_graph()
    graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK])
    graph.check_support()
    graph.build_plans()

    variant_pack = {
        q: q_gpu,
        k: k_gpu,
        v: v_gpu,
        o: o_gpu,
        g: g_gpu,
        stats: lse,
        dq: dq_gpu,
        dk: dk_gpu,
        dv: dv_gpu,
    }

    workspace = torch.empty(graph.get_workspace_size(), device="cuda", dtype=torch.uint8)

    def cudnn_spda_fwd(*args, **kwargs):
        graph.execute(variant_pack, workspace)
        return dq_gpu, dk_gpu, dv_gpu

    return cudnn_spda_fwd


torch.manual_seed(0)
repeats = 100
batch_size = 4
# seqlen = 2048
seqlen = 8192
# seqlen = 4096
# seqlen_q = 256
# seqlen_q = 1
seqlen_q = seqlen
# seqlen = 2047
dim = 2048
headdim = 128
# headdim = 64
nheads = dim // headdim
# nheads = 24
# headdim = 64
# batch_size = 64
# seqlen = 512
# nheads = 8
# headdim = 128
# nheads_kv = nheads // 4
nheads_kv = nheads
dropout_p = 0.0
causal = False
varlen = False
page_size = None
dtype = torch.bfloat16
# dtype = torch.float8_e4m3fn
dtype_gen = torch.bfloat16 if dtype == torch.float8_e4m3fn else dtype
device = 'cuda'

q = torch.randn(batch_size, seqlen_q, nheads, headdim, device=device, dtype=dtype_gen).to(dtype).requires_grad_()
k = torch.randn(batch_size, seqlen, nheads, headdim, device=device, dtype=dtype_gen).to(dtype).requires_grad_()
v = torch.randn(batch_size, seqlen, nheads, headdim, device=device, dtype=dtype_gen).to(dtype).requires_grad_()
# knew = torch.randn(batch_size, seqlen_q, nheads, headdim, device=device, dtype=dtype_gen).to(dtype)
# vnew = torch.randn(batch_size, seqlen_q, nheads, headdim, device=device, dtype=dtype_gen).to(dtype)
# cache_seqlens = torch.tensor([seqlen - seqlen_q] * batch_size, device=device, dtype=torch.int32)
# cos = torch.randn(seqlen, headdim // 2, device=device, dtype=dtype_gen).to(dtype)
# sin = torch.randn(seqlen, headdim // 2, device=device, dtype=dtype_gen).to(dtype)

# if page_size is not None:
#     assert seqlen % page_size == 0
#     k_paged, v_paged = [rearrange(x, "b (n p) h d -> (b n) p h d", p=page_size) for x in [k, v]]
#     page_table = rearrange(torch.arange(batch_size * seqlen // page_size, device=device, dtype=torch.int32),
#                             "(b s) -> b s", s=seqlen // page_size)
# else:
#     page_table = None

if not varlen:
    out, lse = flash_attn_func_v3_python(q, k, v, causal=causal)
    # out, lse, *rest = flash_attn_with_kvcache(q, k if page_size is None else k_paged, v if page_size is None else v_paged,
    #                                           knew, vnew, rotary_cos=cos, rotary_sin=sin, rotary_interleaved=True, cache_seqlens=cache_seqlens, page_table=page_table, causal=causal, num_splits=8)
    # out, lse, *rest = flash_attn_with_kvcache(q, k if page_size is None else k_paged, v if page_size is None else v_paged,
    #                                           knew, vnew, rotary_cos=cos, rotary_sin=sin, rotary_interleaved=False, cache_seqlens=cache_seqlens, page_table=page_table, causal=causal, num_splits=8)
else:
    q_unpad, k_unpad, v_unpad = [rearrange(x.detach(), "b s h d -> (b s) h d").requires_grad_() for x in [q, k, v]]
    cu_seqlens = torch.arange(batch_size + 1, device=device, dtype=torch.int32) * seqlen
    out, lse = flash_attn_varlen_func_v3(q_unpad, k_unpad, v_unpad, cu_seqlens, cu_seqlens, seqlen, seqlen, causal=causal)

if dtype != torch.float8_e4m3fn:
    g = torch.randn_like(out).requires_grad_()
    out.backward(g)

if cudnn is not None:
    g = torch.randn(batch_size, seqlen_q, nheads, headdim, device=device, dtype=dtype_gen).to(dtype).requires_grad_()
    o = torch.randn(batch_size, seqlen_q, nheads, headdim, device=device, dtype=dtype_gen).to(dtype).requires_grad_()
    cudnn_spda_fwd = cudnn_spda_setup(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), causal=causal)
    cudnn_spda_fwd()
    stats = torch.randn(batch_size, seqlen_q, nheads, 1, device=device, dtype=torch.float32)
    cudnn_spda_bwd = cudnn_spda_bwd_setup(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), o.transpose(1, 2), g.transpose(1, 2), stats.transpose(1, 2), causal=causal)
    cudnn_spda_bwd()

# num_splits = 128
# seqlen = 15
# nheads = 32
# d = 128
# out_partial = torch.randn(num_splits, batch_size, seqlen, nheads, d, device=device, dtype=torch.float32)
# lse_partial = torch.randn(num_splits, batch_size, nheads, seqlen, device=device, dtype=torch.float32).transpose(-1, -2)
# from flash_attn_interface import flash_attn_combine
# out, lse = flash_attn_combine(out_partial, lse_partial)
