/**
 * Batch statistics via single-block parallel reduction.
 *
 * Outputs raw aggregates (sums, min, max, counts) — the host computes
 * derived statistics (mean, variance, Sharpe, win rate) from these.
 * This avoids the multi-block atomicAdd-of-means correctness bug.
 *
 * Output layout: 10 floats
 *   [0] reward_sum
 *   [1] reward_sq_sum
 *   [2] global_min_pv       (min portfolio value — for drawdown)
 *   [3] global_max_pv       (max portfolio value — for drawdown)
 *   [4] win_count            (rewards > 0)
 *   [5] total_count          (= N, for convenience)
 *   [6] total_pnl            (= reward_sum, alias for clarity)
 *   [7] position_sum         (sum of abs positions — for Kelly)
 *   [8] position_sq_sum      (sum of position^2 — for Kelly variance)
 *   [9] reserved             (zero)
 *
 * Launch: grid=(1,1,1), block=(256,1,1).
 * Grid-stride loop handles N > 256.
 */
extern "C" __global__ void batch_statistics(
    const float* __restrict__ rewards,
    const float* __restrict__ portfolio_values,
    const float* __restrict__ positions,
    float* output_stats,
    int N
) {
    __shared__ float s_sum[256];
    __shared__ float s_sq_sum[256];
    __shared__ float s_min[256];
    __shared__ float s_max[256];
    __shared__ int   s_win[256];
    __shared__ float s_pos_sum[256];
    __shared__ float s_pos_sq_sum[256];

    int tid = threadIdx.x;

    /* Load phase — grid-stride loop */
    float local_sum = 0.0f;
    float local_sq_sum = 0.0f;
    float local_min_val = 1e30f;
    float local_max_val = -1e30f;
    int   local_wins = 0;
    float local_pos_sum = 0.0f;
    float local_pos_sq_sum = 0.0f;

    for (int i = tid; i < N; i += blockDim.x) {
        float r = rewards[i];
        local_sum += r;
        local_sq_sum += r * r;
        if (r > 0.0f) local_wins++;

        float pv = portfolio_values[i];
        if (pv < local_min_val) local_min_val = pv;
        if (pv > local_max_val) local_max_val = pv;

        float pos = fabsf(positions[i]);
        local_pos_sum += pos;
        local_pos_sq_sum += pos * pos;
    }

    s_sum[tid] = local_sum;
    s_sq_sum[tid] = local_sq_sum;
    s_min[tid] = local_min_val;
    s_max[tid] = local_max_val;
    s_win[tid] = local_wins;
    s_pos_sum[tid] = local_pos_sum;
    s_pos_sq_sum[tid] = local_pos_sq_sum;
    __syncthreads();

    /* Reduction phase 1: shared memory reduction down to 32 threads */
    for (int s = blockDim.x / 2; s >= 32; s >>= 1) {
        if (tid < s) {
            s_sum[tid] += s_sum[tid + s];
            s_sq_sum[tid] += s_sq_sum[tid + s];
            s_min[tid] = fminf(s_min[tid], s_min[tid + s]);
            s_max[tid] = fmaxf(s_max[tid], s_max[tid + s]);
            s_win[tid] += s_win[tid + s];
            s_pos_sum[tid] += s_pos_sum[tid + s];
            s_pos_sq_sum[tid] += s_pos_sq_sum[tid + s];
        }
        __syncthreads();
    }

    /* Reduction phase 2: warp-level reduction (no __syncthreads needed) */
    if (tid < 32) {
        float my_sum        = s_sum[tid];
        float my_sq_sum     = s_sq_sum[tid];
        float my_min        = s_min[tid];
        float my_max        = s_max[tid];
        int   my_win        = s_win[tid];
        float my_pos_sum    = s_pos_sum[tid];
        float my_pos_sq_sum = s_pos_sq_sum[tid];
        for (int mask = 16; mask >= 1; mask >>= 1) {
            my_sum        += __shfl_xor_sync(0xFFFFFFFF, my_sum, mask);
            my_sq_sum     += __shfl_xor_sync(0xFFFFFFFF, my_sq_sum, mask);
            my_min         = fminf(my_min, __shfl_xor_sync(0xFFFFFFFF, my_min, mask));
            my_max         = fmaxf(my_max, __shfl_xor_sync(0xFFFFFFFF, my_max, mask));
            my_win        += __shfl_xor_sync(0xFFFFFFFF, my_win, mask);
            my_pos_sum    += __shfl_xor_sync(0xFFFFFFFF, my_pos_sum, mask);
            my_pos_sq_sum += __shfl_xor_sync(0xFFFFFFFF, my_pos_sq_sum, mask);
        }
        if (tid == 0) {
            s_sum[0]        = my_sum;
            s_sq_sum[0]     = my_sq_sum;
            s_min[0]        = my_min;
            s_max[0]        = my_max;
            s_win[0]        = my_win;
            s_pos_sum[0]    = my_pos_sum;
            s_pos_sq_sum[0] = my_pos_sq_sum;
        }
    }
    __syncthreads();

    /* Write results (thread 0 only — single block, no atomics needed) */
    if (tid == 0) {
        output_stats[0] = s_sum[0];                 /* reward_sum */
        output_stats[1] = s_sq_sum[0];              /* reward_sq_sum */
        output_stats[2] = s_min[0];                 /* global_min_pv */
        output_stats[3] = s_max[0];                 /* global_max_pv */
        output_stats[4] = (float)s_win[0];          /* win_count */
        output_stats[5] = (float)N;                 /* total_count */
        output_stats[6] = s_sum[0];                 /* total_pnl (alias) */
        output_stats[7] = s_pos_sum[0];             /* position_sum */
        output_stats[8] = s_pos_sq_sum[0];          /* position_sq_sum */
        output_stats[9] = 0.0f;                     /* reserved */
    }
}
