import torch
import numpy as np
import sys
import os
import random
import pandas as pd
import copy
import logging
import csv
import socket
import traceback as tb_module
from datetime import datetime

from aiter.ops.triton.flash_attention_forward import _attention
from aiter.test_common import perftest
from aiter.ops.triton.utils import arch_info

class QuantizationUtils:
    @staticmethod
    def calculate_fp8_scale(tensor, quantile: float = 0.99):
        """Calculate FP8 scale based on tensor distribution.
        Use max-abs as in device reference to avoid outliers issues.
        """
        abs_tensor = torch.abs(tensor.float())
        max_val = torch.max(abs_tensor).item()
        max_val = max(max_val, 1e-6)
        FP8_MAX = 448.0
        scale = max_val / FP8_MAX
        scale = max(scale, 1e-4)
        scale = min(scale, 1.0)
        return scale

    @staticmethod
    def quantize_fp8(tensor: torch.Tensor, scale: float) -> torch.Tensor:
        descale = 1.0 / scale
        quantized = (tensor * descale).clamp(min=-448.0, max=448.0)
        return quantized.to(torch.float8_e4m3fn)

    @staticmethod
    def dequantize_fp8(tensor: torch.Tensor, scale: float) -> torch.Tensor:
        return tensor.to(torch.float32) * scale

def seed_everything(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)

# Format: (batch_size, q_heads, k_heads, seq_lens_q, seq_lens_k, head_dim)
attention_test_cases = [
    # Ali case
    (1, 16, 16, [8192], [8192], 128),
    (1, 32, 32, [8192], [8192], 128),
    (1, 32, 4, [8192], [8192], 128),
    (1, 52, 4, [8192], [8192], 128),
    (1, 16, 2, [8192], [8192], 128),
    (1, 26, 2, [8192], [8192], 128),
    (1,  8, 1, [8192], [8192], 128),
    (1, 13, 1, [8192], [8192], 128),
    # llama2-7b
    (1, 32, 32, [4096], [4096], 128),
    (1, 16, 16, [4096], [4096], 128),
    (1,  8,  8, [4096], [4096], 128),
    (1,  4,  4, [4096], [4096], 128),
    # llama2-13b
    (1, 40, 40, [4096], [4096], 128),
    (1, 20, 20, [4096], [4096], 128),
    (1, 10, 10, [4096], [4096], 128),
    (1,  5,  5, [4096], [4096], 128),
    # llama3-8b / qwen3-8b
    (1, 32,  8, [4096], [4096], 128),
    (1, 16,  4, [4096], [4096], 128),
    (1,  8,  2, [4096], [4096], 128),
    (1,  4,  1, [4096], [4096], 128),
    # qwen2-7b
    (1, 28,  4, [4096], [4096], 128),
    (1, 14,  2, [4096], [4096], 128),
    (1,  7,  1, [4096], [4096], 128),
    # corner cases
    (1, 28,  4, [4096], [4096], 124),
]

def pytorch_flash_attention_reference(q, k, v, cu_seqlens_q, cu_seqlens_k, causal=False, sm_scale=1.0, bias=None):
    """
    PyTorch reference implementation of flash attention for comparison
    Using standard softmax (no log2_e scaling needed)
    """
    batch_size = len(cu_seqlens_q) - 1
    nheads_q, head_dim = q.shape[1], q.shape[2]
    nheads_k = k.shape[1]

    # Initialize output tensor
    o = torch.empty_like(q) # [total_seq_len_q, nheads_q, head_dim]

    # Process each batch
    for b in range(batch_size):
        start_q = cu_seqlens_q[b]
        end_q = cu_seqlens_q[b + 1]
        start_k = cu_seqlens_k[b]
        end_k = cu_seqlens_k[b + 1]

        seq_len_q = end_q - start_q
        seq_len_k = end_k - start_k

        # Extract Q, K, V for this batch
        q_batch = q[start_q:end_q]  # [seq_len_q, nheads_q, head_dim]
        k_batch = k[start_k:end_k]  # [seq_len_k, nheads_k, head_dim]
        v_batch = v[start_k:end_k]  # [seq_len_k, nheads_k, head_dim]

        # Reshape for attention computation
        q_batch = q_batch.transpose(0, 1)  # [nheads_q, seq_len_q, head_dim]
        k_batch = k_batch.transpose(0, 1)  # [nheads_k, seq_len_k, head_dim]
        v_batch = v_batch.transpose(0, 1)  # [nheads_k, seq_len_k, head_dim]

        # Handle multi-query attention (MQA/GQA)
        if nheads_q != nheads_k:
            # Repeat K and V for multi-query attention
            group_size = nheads_q // nheads_k
            k_batch = k_batch.repeat_interleave(group_size, dim=0)
            v_batch = v_batch.repeat_interleave(group_size, dim=0)

        # Compute attention scores: Q @ K^T
        scores = torch.matmul(q_batch, k_batch.transpose(-2, -1))  # [nheads_q, seq_len_q, seq_len_k]

        # Apply scaling
        scores = scores * sm_scale

        # Apply causal mask if needed
        if causal:
            # Create causal mask
            mask = torch.triu(torch.ones(seq_len_q, seq_len_k, device=q.device), diagonal=seq_len_k - seq_len_q + 1)
            mask = mask.bool()
            scores = scores.masked_fill(mask, float('-inf'))

        # Apply bias if provided
        if bias is not None:
            # bias shape: [1, nheads_q, seq_len_q, seq_len_k]
            bias_batch = bias[:, :, :seq_len_q, :seq_len_k]
            scores = scores + bias_batch

        # Apply softmax
        attn_weights = torch.softmax(scores, dim=-1)

        # Apply attention to values: attn_weights @ V
        out = torch.matmul(attn_weights, v_batch)  # [nheads_q, seq_len_q, head_dim]

        # Transpose back and store in output
        out = out.transpose(0, 1)  # [seq_len_q, nheads_q, head_dim]
        o[start_q:end_q, :, :] = out

    return o

def build_inputs(batch_size, q_heads, k_heads, seq_lens_q, seq_lens_k, head_dim,
                        bias=False, use_fp8=False):
    total_seq_len_q = sum(seq_lens_q)
    total_seq_len_k = sum(seq_lens_k)
    base_dtype = torch.float32 if use_fp8 else torch.float16
    q_base = torch.randn(total_seq_len_q, q_heads, head_dim, dtype=base_dtype, device="cuda").normal_(mean=0., std=0.5)
    k_base = torch.randn(total_seq_len_k, k_heads, head_dim, dtype=base_dtype, device="cuda").normal_(mean=0., std=0.5)
    v_base = torch.randn(total_seq_len_k, k_heads, head_dim, dtype=base_dtype, device="cuda").normal_(mean=0., std=0.5)

    cu_seqlens_q = torch.zeros(batch_size + 1, dtype=torch.int32, device="cuda")
    cu_seqlens_k = torch.zeros(batch_size + 1, dtype=torch.int32, device="cuda")
    for i in range(batch_size):
        cu_seqlens_q[i + 1] = cu_seqlens_q[i] + seq_lens_q[i]
        cu_seqlens_k[i + 1] = cu_seqlens_k[i] + seq_lens_k[i]
    max_seq_len_q = max(seq_lens_q)
    max_seq_len_k = max(seq_lens_k)
    bias_tensor = torch.randn(1, q_heads, max_seq_len_q, max_seq_len_k, dtype=base_dtype, device="cuda") if bias else None
    if use_fp8:
        q_scale = QuantizationUtils.calculate_fp8_scale(q_base)
        k_scale = QuantizationUtils.calculate_fp8_scale(k_base)
        v_scale = QuantizationUtils.calculate_fp8_scale(v_base)
        q_kernel = QuantizationUtils.quantize_fp8(q_base, q_scale)
        k_kernel = QuantizationUtils.quantize_fp8(k_base, k_scale)
        v_kernel = QuantizationUtils.quantize_fp8(v_base, v_scale)
        q_ref = QuantizationUtils.dequantize_fp8(q_kernel, q_scale)
        k_ref = QuantizationUtils.dequantize_fp8(k_kernel, k_scale)
        v_ref = QuantizationUtils.dequantize_fp8(v_kernel, v_scale)
        scales = {"q": q_scale, "k": k_scale, "v": v_scale}
        return (q_kernel, k_kernel, v_kernel, q_ref, k_ref, v_ref,
                cu_seqlens_q, cu_seqlens_k, max_seq_len_q, max_seq_len_k, bias_tensor, scales)
    else:
        return (q_base, k_base, v_base, q_base, k_base, v_base,
                cu_seqlens_q, cu_seqlens_k, max_seq_len_q, max_seq_len_k, bias_tensor, None)

def test_op_fwd(batch_size, q_heads, k_heads, seq_lens_q, seq_lens_k, head_dim, bias, causal, use_fp8):
    seed_everything(20)
    sm_scale = head_dim ** -0.5
    built = build_inputs(batch_size, q_heads, k_heads, seq_lens_q, seq_lens_k, head_dim, bias, use_fp8)
    (q_kernel, k_kernel, v_kernel, q_ref, k_ref, v_ref,
     cu_seqlens_q, cu_seqlens_k, max_seq_len_q, max_seq_len_k, bias_tensor, scales) = built

    # Unified kernel + reference flow for FP8 and non-FP8
    if use_fp8:
        kernel_q, kernel_k, kernel_v = q_kernel, k_kernel, v_kernel
        ref_q, ref_k, ref_v = q_ref, k_ref, v_ref
        fp8_scales= (scales["q"], scales["k"], scales["v"], 1.0)
        atol, rtol = 1.6e-1, 1e-2
    else:
        kernel_q, kernel_k, kernel_v = q_kernel, k_kernel, v_kernel
        ref_q, ref_k, ref_v = q_kernel, k_kernel, v_kernel
        fp8_scales = None
        atol, rtol = 1e-2, 0

    o_triton = torch.empty_like(kernel_q)
    result_triton, _ = _attention.apply(
        kernel_q, kernel_k, kernel_v, o_triton,
        cu_seqlens_q, cu_seqlens_k,
        max_seq_len_q, max_seq_len_k,
        causal, sm_scale, bias_tensor, fp8_scales, None
    )

    tri = result_triton.to(torch.float32)
    ref = pytorch_flash_attention_reference(ref_q, ref_k, ref_v, cu_seqlens_q, cu_seqlens_k, causal, sm_scale, bias_tensor)
    ref = ref.to(torch.float32)

    torch.testing.assert_close(tri.to("cpu"), ref.to("cpu"), atol=atol, rtol=rtol)
    return True

@perftest(num_warmup=5, num_iters=100)
def _bench_op_fwd(*args, **kwargs):
    return _attention.apply(*args, **kwargs)

def bench_op_fwd():
    seed_everything(20)
    # Get system information
    system_name = socket.gethostname()

    # Create CSV file to record test results
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    csv_filename = f"attention_test_results_{system_name}_{timestamp}.csv"

    # CSV file headers
    csv_headers = [
        'batch_size', 'q_heads', 'k_heads', 'head_dim', 'seq_lens_q', 'seq_lens_k',
        'causal', 'use_fp8', 'test_passed', 'TFLOPS', 'timestamp'
    ]

    # Initialize CSV file
    with open(csv_filename, 'w', newline='', encoding='utf-8') as csvfile:
        writer = csv.writer(csvfile)
        writer.writerow(csv_headers)

    fp8_options = [False, True] if arch_info.is_fp8_avail() else [False]
    for use_fp8 in fp8_options:
        for batch_size, q_heads, k_heads, seq_lens_q, seq_lens_k, head_dim in attention_test_cases:
            for causal in [True]:
                print(f"\nTesting: batch_size={batch_size}, q_heads={q_heads}, k_heads={k_heads}, "
                    f"head_dim={head_dim}, seq_lens_q={seq_lens_q}, seq_lens_k={seq_lens_k}, "
                    f"causal={causal}, use_fp8={use_fp8}")
                sm_scale = head_dim ** -0.5
                try:
                    built = build_inputs(batch_size, q_heads, k_heads, seq_lens_q, seq_lens_k, head_dim, bias=False, use_fp8=use_fp8)
                    (q_kernel, k_kernel, v_kernel, _, _, _,
                        cu_seqlens_q, cu_seqlens_k, max_seq_len_q, max_seq_len_k, bias_tensor, scales) = built
                    fp8_scales= (scales["q"], scales["k"], scales["v"], 1.0) if use_fp8 else None

                    o_triton = torch.empty_like(q_kernel)
                    _, avg_us = _bench_op_fwd(
                        q_kernel, k_kernel, v_kernel, o_triton,
                        cu_seqlens_q, cu_seqlens_k,
                        max_seq_len_q, max_seq_len_k,
                        causal, sm_scale, None, fp8_scales, None
                    )
                    total_flops = 0
                    for i in range(batch_size):
                        total_flops += 2 * seq_lens_q[i] * seq_lens_k[i] * head_dim * 2 * q_heads
                    if causal:
                        total_flops *= 0.5
                    tflops = total_flops * 1e-6 / avg_us if avg_us > 0 else 0.0
                    passed = test_op_fwd(batch_size, q_heads, k_heads, seq_lens_q, seq_lens_k, head_dim, False, causal, use_fp8)
                    print(f"Performance: {tflops:.2f} TFLOPS")
                except Exception as e:
                    print(f"Run failed with error: {e}")
                    tb_module.print_exc()
                    passed, tflops = False, 0.0

                # Record test results to CSV file
                test_timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                tflops_str = f"{tflops:.2f}"
                csv_row = [
                    batch_size, q_heads, k_heads, head_dim,
                    str(seq_lens_q), str(seq_lens_k),
                    causal, use_fp8,
                    passed, tflops_str, test_timestamp
                ]

                with open(csv_filename, 'a', newline='', encoding='utf-8') as csvfile:
                    writer = csv.writer(csvfile)
                    writer.writerow(csv_row)

    print(f"\n测试完成！结果已保存到文件: {csv_filename}")

if __name__ == "__main__":
    bench_op_fwd()