import os
os.environ['JIT_PACKAGE_MODE'] = '1'

from vidrial.jit.timer import timeit
from flash_attn.flash_attn_interface import flash_attn_func
from vidrial.jit.decorator import set_settings, PickBest
from vidrial.kernels.flash.binding import binding as flash_binding
from vidrial.py_utils.test_utils import diff
import torch
import pandas as pd
import logging
logging.basicConfig(level=logging.INFO)

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

def compare():
    with set_settings(PickBest, max_configs=512, plot_timings=False):
        result = []
        for b, tq, tk, d, e in [(4, 4096, 4096, 64, 64), 
                                (4, 8192, 8192, 64, 64)]:
            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)
            flash_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(flash_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
            })
        return pd.DataFrame(result)

if __name__ == "__main__":
    df = compare()
    print(df)