/**
 * Test multi-head FMHA kernel (6-warp, grid launch).
 * Compile with -DHD_VAL=64 etc.
 *
 * Tests:
 * 1. Multi-head MHA: n_h independent heads, each with its own K/V
 * 2. Multi-head MQA: n_h Q heads sharing 1 K/V head
 * 3. Batched: batch_size > 1
 * 4. LSE output correctness
 */

#include <cuda_runtime.h>
#include <cstdio>
#include <cmath>
#include <cstdlib>
#include <cstring>

#ifndef HD_VAL
#define HD_VAL 64
#endif

#include "dsv4/kernels/attention/fmha_common.cuh"
#include "dsv4/kernels/attention/fmha_umma_desc.cuh"

using namespace dsv4::kernels::attention;

static bf16_t f32_to_bf16_host(float f) { uint32_t u; memcpy(&u,&f,4); return (uint16_t)(u>>16); }
static float bf16_to_f32_host(bf16_t h) { uint32_t u=(uint32_t)h<<16; float f; memcpy(&f,&u,4); return f; }

constexpr int HD = HD_VAL;
constexpr int SK = 128;
constexpr int TILE_SZ = 128 * MMA_K_BF16;
constexpr int V_SUB_SZ = 256;
constexpr int TMEM_N = (HD <= 128) ? 128 : 256;

#include "dsv4/kernels/attention/fmha_6warp_multihead.cuh"

// Reference: compute attention for one head
static void reference_attention(
    const bf16_t* q, const bf16_t* k, const bf16_t* v,
    float* o_ref, float* lse_ref,
    int hd, int s_k, float scale
) {
    float s[512];  // max s_k
    for (int j = 0; j < s_k; j++) {
        float dot = 0.0f;
        for (int d = 0; d < hd; d++) dot += bf16_to_f32_host(q[d]) * bf16_to_f32_host(k[j * hd + d]);
        s[j] = dot * scale;
    }
    float mx = -INFINITY;
    for (int j = 0; j < s_k; j++) mx = fmaxf(mx, s[j]);
    float sm = 0.0f;
    for (int j = 0; j < s_k; j++) { s[j] = expf(s[j] - mx); sm += s[j]; }
    for (int j = 0; j < s_k; j++) s[j] /= sm;
    for (int d = 0; d < hd; d++) {
        float ov = 0.0f;
        for (int int_j = 0; int_j < s_k; int_j++) ov += s[int_j] * bf16_to_f32_host(v[d * s_k + int_j]);
        o_ref[d] = ov;
    }
    if (lse_ref) *lse_ref = logf(sm) + mx;
}

static int test_mha(int n_h) {
    printf("\n=== Test MHA: n_h=%d, HD=%d, SK=%d ===\n", n_h, HD, SK);
    const float SCALE = 1.0f / sqrtf((float)HD);
    int pass = 1;

    // Allocate host tensors: Q(n_h, hd), K(n_h, SK*hd), V(n_h, hd*SK), O(n_h, hd)
    bf16_t* h_q = (bf16_t*)malloc(n_h * HD * sizeof(bf16_t));
    bf16_t* h_k = (bf16_t*)malloc(n_h * SK * HD * sizeof(bf16_t));
    bf16_t* h_v = (bf16_t*)malloc(n_h * HD * SK * sizeof(bf16_t));
    bf16_t* h_o = (bf16_t*)calloc(n_h * HD, sizeof(bf16_t));
    float* h_lse = (float*)calloc(n_h, sizeof(float));

    srand(42);
    for (int i = 0; i < n_h * HD; i++) h_q[i] = f32_to_bf16_host((float)(rand()%100)/100.0f - 0.5f);
    for (int i = 0; i < n_h * SK * HD; i++) h_k[i] = f32_to_bf16_host((float)(rand()%100)/100.0f - 0.5f);
    for (int i = 0; i < n_h * HD * SK; i++) h_v[i] = f32_to_bf16_host((float)(rand()%100)/100.0f - 0.5f);

    bf16_t *d_q, *d_k, *d_v, *d_o;
    float *d_lse;
    cudaMalloc(&d_q, n_h * HD * sizeof(bf16_t));
    cudaMalloc(&d_k, n_h * SK * HD * sizeof(bf16_t));
    cudaMalloc(&d_v, n_h * HD * SK * sizeof(bf16_t));
    cudaMalloc(&d_o, n_h * HD * sizeof(bf16_t));
    cudaMalloc(&d_lse, n_h * sizeof(float));
    cudaMemcpy(d_q, h_q, n_h * HD * sizeof(bf16_t), cudaMemcpyHostToDevice);
    cudaMemcpy(d_k, h_k, n_h * SK * HD * sizeof(bf16_t), cudaMemcpyHostToDevice);
    cudaMemcpy(d_v, h_v, n_h * HD * SK * sizeof(bf16_t), cudaMemcpyHostToDevice);
    cudaMemset(d_o, 0, n_h * HD * sizeof(bf16_t));
    cudaMemset(d_lse, 0, n_h * sizeof(float));

    FmhaParams params;
    params.q = d_q;
    params.k = d_k;
    params.v = d_v;
    params.o = d_o;
    params.lse = d_lse;
    params.s_k = SK;
    params.scale = SCALE;
    params.head_dim = HD;
    params.q_head_stride = HD;       // T=1, stride = 1 * hd
    params.q_batch_stride = n_h * HD;
    params.k_head_stride = SK * HD;  // each head has its own K
    params.k_batch_stride = n_h * SK * HD;
    params.v_head_stride = HD * SK;  // each head has its own V
    params.v_batch_stride = n_h * HD * SK;
    params.o_head_stride = HD;
    params.o_batch_stride = n_h * HD;
    params.lse_head_stride = 1;
    params.lse_batch_stride = n_h;

    int smem = (4 + 8 + 16 + TILE_SZ*2 + TILE_SZ*2 + TILE_SZ*2 + V_SUB_SZ*2 + SK*4 + 256 + 127) & ~127;

    if (smem > 48 * 1024) {
        cudaFuncSetAttribute(fmha_6warp_multihead_kernel<HD>, cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
    }

    fmha_6warp_multihead_kernel<HD><<<dim3(1, n_h, 1), 192, smem>>>(params);

    cudaError_t launch_err = cudaGetLastError();
    cudaError_t sync_err = cudaSuccess;
    if (launch_err != cudaSuccess) {
        printf("LAUNCH ERROR: %s\n", cudaGetErrorString(launch_err));
        pass = 0; goto cleanup;
    }

    sync_err = cudaDeviceSynchronize();
    if (sync_err != cudaSuccess) {
        printf("CUDA ERROR: %s\n", cudaGetErrorString(sync_err));
        pass = 0; goto cleanup;
    }

    cudaMemcpy(h_o, d_o, n_h * HD * sizeof(bf16_t), cudaMemcpyDeviceToHost);
    cudaMemcpy(h_lse, d_lse, n_h * sizeof(float), cudaMemcpyDeviceToHost);

    // Verify each head
    for (int h = 0; h < n_h; h++) {
        float o_ref[512];  // max HD
        float lse_ref;
        reference_attention(
            h_q + h * HD, h_k + h * SK * HD, h_v + h * HD * SK,
            o_ref, &lse_ref, HD, SK, SCALE
        );

        float cs = 0, na = 0, nb = 0;
        for (int d = 0; d < HD; d++) {
            float a = bf16_to_f32_host(h_o[h * HD + d]), b = o_ref[d];
            if (fabsf(b) > 1e-4f) { cs += a*b; na += a*a; nb += b*b; }
        }
        cs /= (sqrtf(na) * sqrtf(nb) + 1e-10f);

        float lse_err = fabsf(h_lse[h] - lse_ref) / (fabsf(lse_ref) + 1e-10f);
        printf("  Head %2d: cos=%.8f  lse_err=%.6f (kernel=%.6f ref=%.6f)\n",
               h, cs, lse_err, h_lse[h], lse_ref);
        if (cs < 0.999f) {
            printf("  HEAD %d FAILED (cos=%.6f < 0.999)\n", h, cs);
            pass = 0;
        }
    }

    printf("MHA test %s\n", pass ? "PASSED" : "FAILED");

cleanup:
    cudaFree(d_q); cudaFree(d_k); cudaFree(d_v); cudaFree(d_o); cudaFree(d_lse);
    free(h_q); free(h_k); free(h_v); free(h_o); free(h_lse);
    return pass;
}

static int test_mqa(int n_q, int n_kv) {
    printf("\n=== Test MQA: n_q=%d, n_kv=%d, HD=%d, SK=%d ===\n", n_q, n_kv, HD, SK);
    const float SCALE = 1.0f / sqrtf((float)HD);
    int pass = 1;
    int q_per_kv = n_q / n_kv;

    // Q: (n_q, hd), K: (n_kv, SK*hd), V: (n_kv, hd*SK), O: (n_q, hd)
    bf16_t* h_q = (bf16_t*)malloc(n_q * HD * sizeof(bf16_t));
    bf16_t* h_k = (bf16_t*)malloc(n_kv * SK * HD * sizeof(bf16_t));
    bf16_t* h_v = (bf16_t*)malloc(n_kv * HD * SK * sizeof(bf16_t));
    bf16_t* h_o = (bf16_t*)calloc(n_q * HD, sizeof(bf16_t));

    srand(123);
    for (int i = 0; i < n_q * HD; i++) h_q[i] = f32_to_bf16_host((float)(rand()%100)/100.0f - 0.5f);
    for (int i = 0; i < n_kv * SK * HD; i++) h_k[i] = f32_to_bf16_host((float)(rand()%100)/100.0f - 0.5f);
    for (int i = 0; i < n_kv * HD * SK; i++) h_v[i] = f32_to_bf16_host((float)(rand()%100)/100.0f - 0.5f);

    bf16_t *d_q, *d_k, *d_v, *d_o;
    cudaMalloc(&d_q, n_q * HD * sizeof(bf16_t));
    cudaMalloc(&d_k, n_kv * SK * HD * sizeof(bf16_t));
    cudaMalloc(&d_v, n_kv * HD * SK * sizeof(bf16_t));
    cudaMalloc(&d_o, n_q * HD * sizeof(bf16_t));
    cudaMemcpy(d_q, h_q, n_q * HD * sizeof(bf16_t), cudaMemcpyHostToDevice);
    cudaMemcpy(d_k, h_k, n_kv * SK * HD * sizeof(bf16_t), cudaMemcpyHostToDevice);
    cudaMemcpy(d_v, h_v, n_kv * HD * SK * sizeof(bf16_t), cudaMemcpyHostToDevice);
    cudaMemset(d_o, 0, n_q * HD * sizeof(bf16_t));

    // For MQA: n_q heads each with different Q, but shared K/V
    // We launch n_q CTAs. Each CTA reads its own Q head but the SAME K/V
    // (k_head_stride=0, v_head_stride=0 for pure MQA)
    // But we need to map Q head h to KV head (h / q_per_kv)
    // The kernel doesn't know about the mapping — it just uses blockIdx.y as head_idx
    // and k_head_stride to index K. For MQA, we set k_head_stride=0 so all CTAs read the same K.
    // For GQA, we'd need a different approach (grouped launches).
    //
    // Pure MQA test: 1 KV head, all Q heads share it
    if (n_kv != 1) {
        printf("MQA test requires n_kv=1 for stride=0 trick, skipping\n");
        pass = 1; goto cleanup;
    }

    {
        FmhaParams params;
        params.q = d_q;
        params.k = d_k;
        params.v = d_v;
        params.o = d_o;
        params.lse = nullptr;  // skip LSE for MQA test
        params.s_k = SK;
        params.scale = SCALE;
        params.head_dim = HD;
        params.q_head_stride = HD;
        params.q_batch_stride = n_q * HD;
        params.k_head_stride = 0;  // MQA: all heads share same K
        params.k_batch_stride = SK * HD;
        params.v_head_stride = 0;  // MQA: all heads share same V
        params.v_batch_stride = HD * SK;
        params.o_head_stride = HD;
        params.o_batch_stride = n_q * HD;
        params.lse_head_stride = 0;
        params.lse_batch_stride = 0;

        int smem = (4 + 8 + 16 + TILE_SZ*2 + TILE_SZ*2 + TILE_SZ*2 + V_SUB_SZ*2 + SK*4 + 256 + 127) & ~127;
        if (smem > 48 * 1024) {
            cudaFuncSetAttribute(fmha_6warp_multihead_kernel<HD>, cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
        }

        fmha_6warp_multihead_kernel<HD><<<dim3(1, n_q, 1), 192, smem>>>(params);

        cudaError_t launch_err = cudaGetLastError();
        cudaError_t sync_err2 = cudaSuccess;
        if (launch_err != cudaSuccess) {
            printf("LAUNCH ERROR: %s\n", cudaGetErrorString(launch_err));
            pass = 0; goto cleanup;
        }
        sync_err2 = cudaDeviceSynchronize();
        if (sync_err2 != cudaSuccess) {
            printf("CUDA ERROR: %s\n", cudaGetErrorString(sync_err2));
            pass = 0; goto cleanup;
        }

        cudaMemcpy(h_o, d_o, n_q * HD * sizeof(bf16_t), cudaMemcpyDeviceToHost);

        // All Q heads share the same K/V (k_head[0], v_head[0])
        for (int h = 0; h < n_q; h++) {
            float o_ref[512];
            reference_attention(
                h_q + h * HD, h_k, h_v,
                o_ref, nullptr, HD, SK, SCALE
            );

            float cs = 0, na = 0, nb = 0;
            for (int d = 0; d < HD; d++) {
                float a = bf16_to_f32_host(h_o[h * HD + d]), b = o_ref[d];
                if (fabsf(b) > 1e-4f) { cs += a*b; na += a*a; nb += b*b; }
            }
            cs /= (sqrtf(na) * sqrtf(nb) + 1e-10f);
            printf("  Q-head %2d (shared KV): cos=%.8f\n", h, cs);
            if (cs < 0.999f) { printf("  HEAD %d FAILED\n", h); pass = 0; }
        }
    }

    printf("MQA test %s\n", pass ? "PASSED" : "FAILED");

cleanup:
    cudaFree(d_q); cudaFree(d_k); cudaFree(d_v); cudaFree(d_o);
    free(h_q); free(h_k); free(h_v); free(h_o);
    return pass;
}

static int test_batched(int n_h, int batch_size) {
    printf("\n=== Test Batched: n_h=%d, batch=%d, HD=%d, SK=%d ===\n", n_h, batch_size, HD, SK);
    const float SCALE = 1.0f / sqrtf((float)HD);
    int pass = 1;

    int total_q = batch_size * n_h;
    bf16_t* h_q = (bf16_t*)malloc(total_q * HD * sizeof(bf16_t));
    bf16_t* h_k = (bf16_t*)malloc(total_q * SK * HD * sizeof(bf16_t));
    bf16_t* h_v = (bf16_t*)malloc(total_q * HD * SK * sizeof(bf16_t));
    bf16_t* h_o = (bf16_t*)calloc(total_q * HD, sizeof(bf16_t));

    srand(999);
    for (int i = 0; i < total_q * HD; i++) h_q[i] = f32_to_bf16_host((float)(rand()%100)/100.0f - 0.5f);
    for (int i = 0; i < total_q * SK * HD; i++) h_k[i] = f32_to_bf16_host((float)(rand()%100)/100.0f - 0.5f);
    for (int i = 0; i < total_q * HD * SK; i++) h_v[i] = f32_to_bf16_host((float)(rand()%100)/100.0f - 0.5f);

    bf16_t *d_q, *d_k, *d_v, *d_o;
    cudaMalloc(&d_q, total_q * HD * sizeof(bf16_t));
    cudaMalloc(&d_k, total_q * SK * HD * sizeof(bf16_t));
    cudaMalloc(&d_v, total_q * HD * SK * sizeof(bf16_t));
    cudaMalloc(&d_o, total_q * HD * sizeof(bf16_t));
    cudaMemcpy(d_q, h_q, total_q * HD * sizeof(bf16_t), cudaMemcpyHostToDevice);
    cudaMemcpy(d_k, h_k, total_q * SK * HD * sizeof(bf16_t), cudaMemcpyHostToDevice);
    cudaMemcpy(d_v, h_v, total_q * HD * SK * sizeof(bf16_t), cudaMemcpyHostToDevice);
    cudaMemset(d_o, 0, total_q * HD * sizeof(bf16_t));

    FmhaParams params;
    params.q = d_q;
    params.k = d_k;
    params.v = d_v;
    params.o = d_o;
    params.lse = nullptr;
    params.s_k = SK;
    params.scale = SCALE;
    params.head_dim = HD;
    params.q_head_stride = HD;
    params.q_batch_stride = n_h * HD;
    params.k_head_stride = SK * HD;
    params.k_batch_stride = n_h * SK * HD;
    params.v_head_stride = HD * SK;
    params.v_batch_stride = n_h * HD * SK;
    params.o_head_stride = HD;
    params.o_batch_stride = n_h * HD;
    params.lse_head_stride = 0;
    params.lse_batch_stride = 0;

    int smem = (4 + 8 + 16 + TILE_SZ*2 + TILE_SZ*2 + TILE_SZ*2 + V_SUB_SZ*2 + SK*4 + 256 + 127) & ~127;
    if (smem > 48 * 1024) {
        cudaFuncSetAttribute(fmha_6warp_multihead_kernel<HD>, cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
    }

    fmha_6warp_multihead_kernel<HD><<<dim3(1, n_h, batch_size), 192, smem>>>(params);

    cudaError_t launch_err = cudaGetLastError();
    cudaError_t sync_err3 = cudaSuccess;
    int checked = 0, failed = 0;
    if (launch_err != cudaSuccess) {
        printf("LAUNCH ERROR: %s\n", cudaGetErrorString(launch_err));
        pass = 0; goto cleanup;
    }
    sync_err3 = cudaDeviceSynchronize();
    if (sync_err3 != cudaSuccess) {
        printf("CUDA ERROR: %s\n", cudaGetErrorString(sync_err3));
        pass = 0; goto cleanup;
    }

    cudaMemcpy(h_o, d_o, total_q * HD * sizeof(bf16_t), cudaMemcpyDeviceToHost);

    // Verify a sample of heads across batches
    for (int b = 0; b < batch_size; b++) {
        for (int h = 0; h < n_h; h++) {
            int idx = b * n_h + h;
            float o_ref[512];
            reference_attention(
                h_q + idx * HD,
                h_k + idx * SK * HD,
                h_v + idx * HD * SK,
                o_ref, nullptr, HD, SK, SCALE
            );

            float cs = 0, na = 0, nb = 0;
            for (int d = 0; d < HD; d++) {
                float a = bf16_to_f32_host(h_o[idx * HD + d]), b2 = o_ref[d];
                if (fabsf(b2) > 1e-4f) { cs += a*b2; na += a*a; nb += b2*b2; }
            }
            cs /= (sqrtf(na) * sqrtf(nb) + 1e-10f);
            checked++;
            if (cs < 0.999f) {
                printf("  FAIL batch=%d head=%d: cos=%.6f\n", b, h, cs);
                failed++;
            }
        }
    }
    printf("  Checked %d heads, %d failed\n", checked, failed);
    pass = (failed == 0);
    printf("Batched test %s\n", pass ? "PASSED" : "FAILED");

cleanup:
    cudaFree(d_q); cudaFree(d_k); cudaFree(d_v); cudaFree(d_o);
    free(h_q); free(h_k); free(h_v); free(h_o);
    return pass;
}

int main() {
    printf("========================================\n");
    printf("Multi-head FMHA test suite (HD=%d)\n", HD);
    printf("========================================\n");

    int all_pass = 1;

    // Test 1: MHA with 4 heads
    all_pass &= test_mha(4);

    // Test 2: MHA with 8 heads (covers Pro's hd=128 with 128 heads in smaller test)
    all_pass &= test_mha(8);

    // Test 3: MQA: 4 Q heads sharing 1 KV head
    all_pass &= test_mqa(4, 1);

    // Test 4: Batched: 4 heads × 2 batch
    all_pass &= test_batched(4, 2);

    printf("\n========================================\n");
    printf("Overall: %s\n", all_pass ? "ALL PASSED" : "SOME FAILED");
    printf("========================================\n");
    return all_pass ? 0 : 1;
}