#include <stdio.h>

#define imin(a,b) (a<b?a:b)

const int N = 33 * 1024;  // 要进行计算的数组的长度
const int threadsPerBlock = 256;
const int blocksPerGrid = imin(32, (N + threadsPerBlock - 1) / threadsPerBlock);

__global__ void gpu_variance(const float *a, float *partial_sum, float *partial_sum2) {
    // 每个 block 计算两件事：sum(x) 和 sum(x^2)
    __shared__ float cache_x[threadsPerBlock];
    __shared__ float cache_x2[threadsPerBlock];

    int tid = threadIdx.x + blockIdx.x * blockDim.x;
    int cacheIndex = threadIdx.x;

    float temp_x = 0;
    float temp_x2 = 0;

    // 循环遍历数组，处理数据量大于线程总数的情况
    while (tid < N) {
        float val = a[tid];
        temp_x += val;
        temp_x2 += val * val; // 同时计算平方和
        tid += blockDim.x * gridDim.x;
    }

    // 将线程计算的结果存入共享内存
    cache_x[cacheIndex] = temp_x;
    cache_x2[cacheIndex] = temp_x2;

    // 同步线程，确保所有数据都已写入共享内存
    __syncthreads();

    // 归约操作，threadsPerBlock 必须是 2 的幂
    int i = blockDim.x / 2;
    while (i != 0) {
        if (cacheIndex < i) {
            cache_x[cacheIndex] += cache_x[cacheIndex + i];
            cache_x2[cacheIndex] += cache_x2[cacheIndex + i];
        }
        __syncthreads();
        i /= 2;
    }

    // 每个 block 只由 0 号线程写回结果
    if (cacheIndex == 0) {
        partial_sum[blockIdx.x] = cache_x[0];
        partial_sum2[blockIdx.x] = cache_x2[0];
    }
}

int main(void) {
    float *a, *partial_sum, *partial_sum2;
    float total_sum, total_sum2;
    float *dev_a, *dev_partial_sum, *dev_partial_sum2;
    float mean, variance;

    // 1. 分配 CPU 内存
    a = (float *)malloc(N * sizeof(float));
    partial_sum = (float *)malloc(blocksPerGrid * sizeof(float));
    partial_sum2 = (float *)malloc(blocksPerGrid * sizeof(float));

    // 2. 分配 GPU 内存
    cudaMalloc((void **)&dev_a, N * sizeof(float));
    cudaMalloc((void **)&dev_partial_sum, blocksPerGrid * sizeof(float));
    cudaMalloc((void **)&dev_partial_sum2, blocksPerGrid * sizeof(float));

    // 3. 初始化数据
    for (int i = 0; i < N; i++) {
        a[i] = i;
    }

    // 4. 数据拷贝 Host -> Device
    cudaMemcpy(dev_a, a, N * sizeof(float), cudaMemcpyHostToDevice);

    // 5. 启动内核
    gpu_variance<<<blocksPerGrid, threadsPerBlock>>>(dev_a, dev_partial_sum, dev_partial_sum2);

    // 6. 数据拷贝 Device -> Host
    cudaMemcpy(partial_sum, dev_partial_sum, blocksPerGrid * sizeof(float), cudaMemcpyDeviceToHost);
    cudaMemcpy(partial_sum2, dev_partial_sum2, blocksPerGrid * sizeof(float), cudaMemcpyDeviceToHost);

    // 7. CPU 端合并各个 Block 的结果
    total_sum = 0;
    total_sum2 = 0;
    for (int i = 0; i < blocksPerGrid; i++) {
        total_sum += partial_sum[i];
        total_sum2 += partial_sum2[i];
    }

    // 8. 计算最终的均值和方差
    // 方差公式: Var = E(x^2) - [E(x)]^2
    mean = total_sum / N;
    variance = (total_sum2 / N) - (mean * mean);

    printf("Array Size N = %d\n", N);
    printf("Calculated Mean = %f\n", mean);
    printf("Calculated Variance = %f\n", variance);

    // 验证结果
    // 数列 0..N-1 的均值 = (N-1)/2
    // 数列 0..N-1 的方差 = (N^2 - 1) / 12
    float expected_mean = (float)(N - 1) / 2.0f;
    float expected_variance = ((float)(N * N) - 1.0f) / 12.0f;
    
    printf("Expected Mean = %f\n", expected_mean);
    printf("Expected Variance = %f\n", expected_variance);

    // 9. 释放资源
    cudaFree(dev_a);
    cudaFree(dev_partial_sum);
    cudaFree(dev_partial_sum2);
    free(a);
    free(partial_sum);
    free(partial_sum2);

    return 0;
}
