from dataclasses import dataclass
from typing import Any
import torch
from vidrial.jit.static.types import Shape, Layout, Int
from vidrial.jit.static.util import layout_from_shape_stride, torch_dtype_to_c
from vidrial.jit.timingcache import ConfigTimingCache
from vidrial.jit.decorator import pickbest
from vidrial.jit.jit import jit, render
from vidrial.kernels.flash.configurator import base_configurator, advanced_configurator


# ------------------- Source Code -------------------

@dataclass
class SourceCode:
    T: str
    LT: str
    Atom1: str
    Atom2: str
    MNKTileShape1: Shape
    NTile2: Int
    MNKAtomPlacement1: Shape
    GQSlab: Layout
    GKSlab: Layout
    GVSlab: Layout
    GOSlab: Layout
    GLSlab: Layout
    smempipe1: int
    smempipe2: int
    regpipe1: int
    regpipe2: int
    use_ldsm1: bool
    use_ldsm2: bool
    swizzle1: int
    swizzle2: int
    q_in_reg: bool
    @property
    def template(self):
        return """
#include <cuda.h>
#include <cuda_runtime.h>
#include <iostream>
#include "cute/tensor.hpp"
#include "kernels/flash/kernel.cuh"

using namespace cute;
using namespace vidrial;

extern "C" int launch(void* __raw_Q, void* __raw_K, void* __raw_V, void* __raw_O, void* __raw_l) {
    using T = {T};
    using LT = {LT};
    using Atom1 = MMA_Atom<{Atom1}>;
    using Atom2 = MMA_Atom<{Atom2}>;
    using MNKTileShape_1 = decltype(static_tree_cast<int64_t>({MNKTileShape1}{}));
    using NTile2 = decltype(static_tree_cast<int64_t>({NTile2}{}));
    using MNKAtomPlacement_1 = decltype(static_tree_cast<int64_t>({MNKAtomPlacement1}{}));
    using GQSlab = decltype(static_tree_cast<int64_t>({GQSlab}{}));
    using GKSlab = decltype(static_tree_cast<int64_t>({GKSlab}{}));
    using GVSlab = decltype(static_tree_cast<int64_t>({GVSlab}{}));
    using GOSlab = decltype(static_tree_cast<int64_t>({GOSlab}{}));
    using GLSlab = decltype(static_tree_cast<int64_t>({GLSlab}{}));
    using PerfCfg1 = FlashPerfCfg<{smempipe1}, {regpipe1}, {use_ldsm1}, {swizzle1}, {q_in_reg}>;
    using PerfCfg2 = PerfCfg<{smempipe2}, {regpipe2}, {use_ldsm2}, {swizzle2}>;
    auto cfg = make_FlashKernelCfg<T, LT, Atom1, MNKTileShape_1, MNKAtomPlacement_1, Atom2, NTile2, GQSlab, GKSlab, GVSlab, GOSlab, GLSlab, PerfCfg1, PerfCfg2>();
    auto Q = reinterpret_cast<T*>(__raw_Q);
    auto K = reinterpret_cast<T*>(__raw_K);
    auto V = reinterpret_cast<T*>(__raw_V);
    auto O = reinterpret_cast<T*>(__raw_O);
    auto l = reinterpret_cast<LT*>(__raw_l);
    return launch_flash_attention_kernel(cfg, Q, K, V, O, l);
}"""
    def __str__(self):
        return render(self.template, self.__dict__)


# ------------------- Binding -------------------

@dataclass
class BindingCfg:
    P: int # batch dimension
    M: int
    N: int
    D: int
    E: int
    Q_shape: tuple[int, ...]
    K_shape: tuple[int, ...]
    V_shape: tuple[int, ...]
    O_shape: tuple[int, ...]
    l_shape: tuple[int, ...]
    Q_stride: tuple[int, ...]
    K_stride: tuple[int, ...]
    V_stride: tuple[int, ...]
    O_stride: tuple[int, ...]
    l_stride: tuple[int, ...]
    Atom1: str
    Atom2: str
    MNKTileShape1: tuple[int, ...]
    NTile2: int
    MNKAtomPlacement1: tuple[int, ...]
    smempipe1: int
    smempipe2: int
    regpipe1: int
    regpipe2: int
    use_ldsm1: bool
    use_ldsm2: bool
    swizzle1: int
    swizzle2: int
    q_in_reg: bool
    dtype: torch.dtype
    acc_dtype: torch.dtype
    @classmethod
    def from_args(cls, Q, K, V, O, l, Atom1, Atom2, MNKTileShape1, NTile2, MNKAtomPlacement1, smempipe1, smempipe2, regpipe1, regpipe2, use_ldsm1, use_ldsm2, swizzle1, swizzle2, q_in_reg):
        assert Q.device.type == K.device.type == V.device.type == O.device.type == l.device.type == 'cuda', f"Invalid {Q.device}, {K.device}, {V.device}, {O.device}, {l.device}."
        assert Q.dtype == K.dtype == V.dtype == O.dtype, f"Invalid {Q.dtype}, {K.dtype}, {V.dtype}, {O.dtype},  Q, K, V, O must have the same dtype"
        P, M, N, D, E = Q.shape[0], Q.shape[1], K.shape[1], K.shape[2], V.shape[2]
        assert O.shape == (P, M, E), f"Invalid {O.shape}, O must be of shape ({P}, {M}, {E})"
        assert l.shape == (P, M), f"Invalid {l.shape}, l must be of shape ({P}, {M})"

        return cls(P, M, N, D, E, Q.shape, K.shape, V.shape, O.shape, l.shape, Q.stride(), K.stride(), V.stride(), O.stride(), l.stride(), Atom1, Atom2, MNKTileShape1, NTile2, MNKAtomPlacement1, smempipe1, smempipe2, regpipe1, regpipe2, use_ldsm1, use_ldsm2, swizzle1, swizzle2, q_in_reg, Q.dtype, l.dtype)

    @property
    def source(self):
        """ Conversion between the different variable name and shape conventions
                Python                        c++
            Q_shape=[P,M1,K1]      ->      GQSlabShape=[M1,K1,P]
            K_shape=[P,N1,K1]      ->      GKSlabShape=[N1,K1,P]
            V_shape=[P,K2,N2]      ->      GVSlabShape=[N2,K2,P]
            O_shape=[P,M2,N2]      ->      GO_slabShape=[M2,N2,P]
            l_shape=[P,M2]      ->      GL_slabShape=[M2,P]
        """
        GQslab = layout_from_shape_stride(self.Q_shape, self.Q_stride, (1, 2, 0))
        GKslab = layout_from_shape_stride(self.K_shape, self.K_stride, (1, 2, 0))
        GVslab = layout_from_shape_stride(self.V_shape, self.V_stride, (2, 1, 0))
        GOslab = layout_from_shape_stride(self.O_shape, self.O_stride, (1, 2, 0))
        GLslab = layout_from_shape_stride(self.l_shape, self.l_stride, (1, 0))
        return SourceCode(
            T=torch_dtype_to_c(self.dtype),
            LT=torch_dtype_to_c(self.acc_dtype),
            Atom1=self.Atom1,
            Atom2=self.Atom2,
            MNKTileShape1=Shape(*[Int(x) for x in self.MNKTileShape1]),
            NTile2=Int(self.NTile2),
            MNKAtomPlacement1=Shape(*[Int(x) for x in self.MNKAtomPlacement1]),
            GQSlab=GQslab,
            GKSlab=GKslab,
            GVSlab=GVslab,
            GOSlab=GOslab,
            GLSlab=GLslab,
            smempipe1=self.smempipe1,
            smempipe2=self.smempipe2,
            regpipe1=self.regpipe1,
            regpipe2=self.regpipe2,
            use_ldsm1=self.use_ldsm1,
            use_ldsm2=self.use_ldsm2,
            swizzle1=self.swizzle1,
            swizzle2=self.swizzle2,
            q_in_reg=self.q_in_reg,
        )

# ---------------------- Autotune --------------------------------

def make_mma_configurator(smempipe=None, regpipe=None, use_ldsm=None, swizzle=None):
    def mma_configurator(args: dict) -> list[dict[str, Any]]:
        Q, K, V = args['Q'], args['K'], args['V']
        P, M, N, D, E = Q.shape[0], Q.shape[1], K.shape[1], K.shape[2], V.shape[2]
        return advanced_configurator(MNDEP=(M, N, D, E, P), dtype=Q.dtype, acc_dtype=torch.float32, smempipe=smempipe, regpipe=regpipe, use_ldsm=use_ldsm, swizzle=1, random_seed=42)
    return mma_configurator

def hash_fn(args: dict) -> str:
    Q, K, V, O, l = args['Q'], args['K'], args['V'], args['O'], args['l']
    P, M, N, D, E = Q.shape[0], Q.shape[1], K.shape[1], K.shape[2], V.shape[2]
    Q_major = 'K_major' if Q.stride(-1) == 1 else 'M_major'
    K_major = 'N_major' if K.stride(-1) == 1 else 'K_major'
    V_major = 'N_major' if V.stride(-1) == 1 else 'K_major'
    return f"{Q.dtype}-{K.dtype}-{V.dtype}-{O.dtype}-{l.dtype}-M_mod1024_{M % 1024}-N_mod1024_{N % 1024}-D_{D}-E_{E}-P_{P}-{Q_major}-{K_major}-{V_major}"


@pickbest(cache=ConfigTimingCache('flash_attn', hash_fn),
          sweep=make_mma_configurator(smempipe=(1,2,3), regpipe=(1,2,4), use_ldsm=True, swizzle=1))
def binding(Q, K, V, O, l, Atom1, Atom2, MNKTileShape1, NTile2, MNKAtomPlacement1, smempipe1, smempipe2, regpipe1, regpipe2, use_ldsm1, use_ldsm2, swizzle1, swizzle2, q_in_reg):
    binding_cfg = BindingCfg.from_args(Q, K, V, O, l, Atom1, Atom2, MNKTileShape1, NTile2, MNKAtomPlacement1, smempipe1, smempipe2, regpipe1, regpipe2, use_ldsm1, use_ldsm2, swizzle1, swizzle2, q_in_reg)
    jit(name = "flash_attn",
        code = str(binding_cfg.source)
    )(Q, K, V, O, l)



if __name__ == "__main__":
    import logging
    from vidrial.py_utils.test_utils import diff
    import torch
    torch.manual_seed(43)
    logging.basicConfig(level=logging.DEBUG)
    from vidrial.jit.decorator import set_settings, PickBest, PickAny
    from vidrial.jit.timer import timeit
    from flash_attn.flash_attn_interface import flash_attn_func
    # b, tq, tk, d, e = 4, 4096, 4096, 64, 64
    # b, tq, tk, d, e = 32, 8192, 8192, 64, 64
    b, tq, tk, d, e = 32, 8192, 8192, 96, 96
    def create_input(b, tq, tk, d, e):
        Q = torch.randn(b, tq, d, device='cuda', dtype=torch.bfloat16)
        K = torch.ones(b, tk, d, device='cuda', dtype=torch.bfloat16)
        V = torch.randn(b, tk, e, device='cuda', dtype=torch.bfloat16)
        O = torch.empty(b, tq, e, device='cuda', dtype=torch.bfloat16)
        l = torch.empty(b, tq, device='cuda', dtype=torch.float32)
        return Q, K, V, O, l

    torch.set_printoptions(precision=8, sci_mode=True, linewidth=10000, edgeitems=100)

    BM, BN = 16, 32
    print_i, print_j = [7], [0, 1, 2, 3, 4, 5]

    def make_print(i, j, *args):
        if i in print_i and j in print_j:
            print(f"({i}, {j}):", *args)

    def op_reference_block(Q, K, V, O, l):
        b, tq, tk, d, e = Q.shape[0], Q.shape[1], K.shape[1], K.shape[2], V.shape[2]
        M = torch.tril(torch.ones(tq, tk, device=Q.device, dtype=torch.bool)).unsqueeze(0).expand(b, -1, -1)
        for i in range(0, tq, BM):
            Q_block = Q[:, i:i+BM, :]
            row_max, row_sum, O_block = None, None, None
            mask_start, mask_end = (i + (tk - tq)) // BN * BN, min((i + BM + (tk - tq)) + BN - 1, tk) // BN * BN
            for j in range(mask_start, mask_end, BN):
                K_block = K[:, j:j+BN, :]
                S_block = (Q_block @ K_block.transpose(-1, -2)).to(torch.float32)
                M_block = M[:, i:i+BM, j:j+BN]
                S_block = S_block.masked_fill(~M_block, float('-inf'))
                new_row_max = torch.maximum(row_max, S_block.max(dim=-1).values) if row_max is not None else S_block.max(dim=-1).values # type: ignore
                scaler = torch.exp(row_max - new_row_max) if row_max is not None else torch.ones_like(new_row_max)
                row_max = new_row_max
                S_block = S_block - new_row_max.unsqueeze(-1)
                P_block = torch.exp(S_block)
                make_print(i//BM, j//BN, f"P_block in loop 1: {P_block}")
                V_block = V[:, j:j+BN, :]
                row_sum = P_block.sum(dim=-1) if row_sum is None else row_sum * scaler + P_block.sum(dim=-1)
                make_print(i//BM, j//BN, f"row_sum in loop 1: {row_sum}")
                O_block = P_block.to(O.dtype) @ V_block if O_block is None else O_block * scaler.unsqueeze(-1) + (P_block.to(O.dtype) @ V_block)
            for j in range(0, mask_start, BN):
                K_block = K[:, j:j+BN, :]
                S_block = (Q_block @ K_block.transpose(-1, -2)).to(torch.float32)
                M_block = M[:, i:i+BM, j:j+BN]
                S_block = S_block.masked_fill(~M_block, float('-inf'))
                new_row_max = torch.maximum(row_max, S_block.max(dim=-1).values) if row_max is not None else S_block.max(dim=-1).values # type: ignore
                scaler = torch.exp(row_max - new_row_max) if row_max is not None else torch.ones_like(new_row_max)
                row_max = new_row_max
                S_block = S_block - new_row_max.unsqueeze(-1)
                P_block = torch.exp(S_block)
                V_block = V[:, j:j+BN, :]
                row_sum = P_block.sum(dim=-1) if row_sum is None else row_sum * scaler + P_block.sum(dim=-1)
                make_print(i//BM, j//BN, f"row_sum in loop 2: {row_sum}")
                O_block = P_block.to(O.dtype) @ V_block if O_block is None else O_block * scaler.unsqueeze(-1) + (P_block.to(O.dtype) @ V_block)
            O[:, i:i+BM, :] = (O_block / row_sum.unsqueeze(-1)).to(O.dtype)
            l[:, i:i+BM] = row_sum
        return O, l
    
    def op_reference(Q, K, V, O, l):
        S = (Q @ K.transpose(-1, -2)).to(torch.float32)
        b, tq, tk = S.shape
        M = torch.tril(torch.ones(tq, tk, device=S.device, dtype=torch.bool)).unsqueeze(0).expand(b, -1, -1)
        S = S.masked_fill(~M, float('-inf'))
        S = S - S.max(dim=-1, keepdim=True).values
        P = torch.exp(S)
        l = P.sum(dim=-1)
        P = (P / l.unsqueeze(-1)).to(O.dtype)
        O = P @ V
        return O, l

    # diff(l, l_ref, atol=1e-3, rtol=1e-2, assert_close=False)
    with set_settings(PickBest, max_configs=512, plot_timings=False):
        result = []
        for b, tq, tk, d, e in [(4, 4096, 4096, 64, 64), 
                                (4, 4096, 4096, 32, 32), 
                                (4, 4096, 4096, 128, 128), 
                                (4, 8192, 8192, 64, 64),
                                (4, 8192, 8192, 32, 32),
                                (4, 8192, 8192, 128, 128),
                                (32, 1024, 1024, 64, 64),
                                (32, 1024, 1024, 32, 32),
                                (32, 1024, 1024, 128, 128),
                                (32, 4096, 4096, 64, 64),
                                (32, 4096, 4096, 32, 32),
                                (32, 4096, 4096, 128, 128),
                                (32, 8192, 8192, 64, 64), 
                                (32, 8192, 8192, 32, 32), 
                                (32, 8192, 8192, 128, 128),]:
            Q, K, V, O, l = create_input(b, tq, tk, d, e)
            O_flash = flash_attn_func(Q.unsqueeze(-2), K.unsqueeze(-2), V.unsqueeze(-2), causal=True, softmax_scale=1.0)
            binding(Q, K, V, O, l) # type: ignore
            diff(O, O_flash.squeeze(-2), atol=1e-3, rtol=1e-2, assert_close=False)

            cuda_avg, cuda_std = timeit(binding, Q, K, V, O, l, num1=10, num2=30)
            flash_avg, flash_std = timeit(flash_attn_func, Q.unsqueeze(-2), K.unsqueeze(-2), V.unsqueeze(-2), causal=True, num1=10, num2=30)
            max_diff = torch.max(torch.abs(O - O_flash.squeeze(-2))).item()
            result.append({
                'b': b,
                'tq': tq,
                'tk': tk,
                'd': d,
                'e': e,
                'cuda_avg': cuda_avg,
                'cuda_std': cuda_std,
                'flash_avg': flash_avg,
                'flash_std': flash_std,
                'speedup': flash_avg / cuda_avg,
                'max_diff': max_diff
            })
        import pandas as pd
        df = pd.DataFrame(result)
        df.to_csv("flash_binding_results.csv", index=False)
        print(df)
        