/* ── Sort-based O(N log N) reward rank normalization ───────────────────
 *
 * r_shaped[i] = sign(sharpe[i]) * rank(|sharpe[i]|)
 * where sharpe[i] = (r[i] - mean_ema) / std_ema   (G13 Sharpe contribution)
 *
 * rank(|r[i]|) = (sorted position of |r[i]| + 1) / N
 *
 * Three-kernel pipeline replaces O(N²) tiled kernel:
 *   1. reward_compute_abs_sharpe — compute |sharpe| keys + indices, pad to N_padded
 *   2. bitonic_sort_step — O(log²N) passes, each O(N) work
 *   3. reward_scatter_rank — write rank from sorted position to original index
 *
 * For N=4M: ~276 bitonic launches × ~0.05ms ≈ 14ms vs 10.6s before.
 */

/* ── Step 1: Compute |sharpe| sort keys + original indices ──────────── */
/* Grid: ceil(N_padded/256), Block: 256.
 * Handles padding: indices [N..N_padded) get +inf key and index -1. */
extern "C" __global__ void reward_compute_abs_sharpe(
    const float* __restrict__ rewards_in,  /* [N] raw rewards */
    float* __restrict__ abs_sharpe_out,    /* [N_padded] sort keys */
    int* __restrict__ indices_out,         /* [N_padded] original indices */
    int N,                                 /* actual data count */
    int N_padded,                          /* next power of 2 >= N */
    float std_ema,
    float return_mean_ema
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= N_padded) return;

    if (i < N) {
        float raw = rewards_in[i];
        float sharpe = (std_ema > 1e-6f) ? (raw - return_mean_ema) / std_ema : raw;
        abs_sharpe_out[i] = fabsf(sharpe);
        indices_out[i] = i;
    } else {
        /* Padding: sentinel sorts to end, invalid index skipped by scatter */
        abs_sharpe_out[i] = 3.4e38f;
        indices_out[i] = -1;
    }
}

/* ── Step 2: Bitonic sort — one compare-and-swap pass ─────────────── */
/* Each thread handles one compare-swap pair at distance `step`
 * within a bitonic block of size `2 * stage`.
 * Grid: ceil(N_padded/2/256), Block: 256. */
extern "C" __global__ void bitonic_sort_step(
    float* __restrict__ keys,    /* [N_padded] sort keys (in-place) */
    int* __restrict__ values,    /* [N_padded] indices (in-place) */
    int N_padded,
    int stage,                   /* current merge block half-size (1, 2, 4, ...) */
    int step                     /* compare distance within stage (stage, stage/2, ..., 1) */
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int pair_dist = step;
    int block_size = stage << 1;

    /* Map thread idx to (left, right) pair */
    int left = (idx / pair_dist) * (pair_dist << 1) + (idx % pair_dist);
    int right = left + pair_dist;

    if (right >= N_padded) return;

    /* Sort direction: ascending in even blocks, descending in odd blocks */
    int block_id = left / block_size;
    int ascending = ((block_id & 1) == 0) ? 1 : 0;

    float kl = keys[left], kr = keys[right];

    int should_swap = ascending ? (kl > kr) : (kl < kr);
    if (should_swap) {
        keys[left] = kr;
        keys[right] = kl;
        int vl = values[left], vr = values[right];
        values[left] = vr;
        values[right] = vl;
    }
}

/* ── Step 3: Scatter rank from sorted position to original index ───── */
/* After sort, sorted_indices[pos] holds the original index of the element
 * with rank (pos+1)/N. Write sign(sharpe) * rank to output.
 * Grid: ceil(N/256), Block: 256. Only iterates over real data [0..N). */
extern "C" __global__ void reward_scatter_rank(
    const float* __restrict__ rewards_in,      /* [N] raw rewards (re-read for sign) */
    const int* __restrict__ sorted_indices,    /* [N_padded] sorted original indices */
    float* __restrict__ rewards_out,           /* [N] output shaped rewards */
    int N,
    float std_ema,
    float return_mean_ema
) {
    int pos = blockIdx.x * blockDim.x + threadIdx.x;
    if (pos >= N) return;

    int orig_idx = sorted_indices[pos];
    /* Skip padding sentinels (should not appear in [0..N) after sort, but guard) */
    if (orig_idx < 0 || orig_idx >= N) return;

    float raw = rewards_in[orig_idx];
    float sharpe = (std_ema > 1e-6f) ? (raw - return_mean_ema) / std_ema : raw;

    /* Holding bars: passthrough zero, don't rank */
    int is_trade = (fabsf(sharpe) > 1e-5f) ? 1 : 0;  /* lowered for dense micro-rewards */
    if (!is_trade) {
        rewards_out[orig_idx] = 0.0f;
        return;
    }

    /* Rank in [0, 1], sign of Sharpe contribution preserved */
    float rank = (float)(pos + 1) / (float)N;
    float sign = (sharpe > 0.0f) ? 1.0f : -1.0f;
    rewards_out[orig_idx] = sign * rank;
}

/* ── Step 1b: Copy raw rewards to sort keys (for quantile extraction) ─
 * Unlike reward_compute_abs_sharpe which takes |sharpe|, this copies raw
 * reward values so that bitonic sort produces a properly-ordered array
 * from which empirical quantiles can be read.
 * Grid: ceil(N_padded/256), Block: 256. */
extern "C" __global__ void reward_copy_raw_for_sort(
    const float* __restrict__ rewards_in,  /* [N] raw rewards */
    float* __restrict__ sort_keys_out,     /* [N_padded] sort keys */
    int* __restrict__ indices_out,         /* [N_padded] original indices */
    int N,
    int N_padded
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= N_padded) return;

    if (i < N) {
        sort_keys_out[i] = rewards_in[i];
        indices_out[i] = i;
    } else {
        sort_keys_out[i] = 3.4e38f;  /* padding sorts to end */
        indices_out[i] = -1;
    }
}

/* ── Gather quantiles from sorted array ──────────────────────────────
 * Extracts num_atoms evenly-spaced quantile positions from a sorted array
 * into a contiguous output buffer for a single DtoH transfer.
 * Grid: ceil(num_atoms/256), Block: 256. */
extern "C" __global__ void gather_quantiles(
    const float* __restrict__ sorted,     /* [N] sorted rewards */
    float* __restrict__ quantiles,        /* [num_atoms] output */
    int N,
    int num_atoms
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= num_atoms) return;
    int idx;
    if (num_atoms <= 1) {
        idx = N / 2;  /* single atom: median */
    } else {
        idx = (int)((long long)i * (long long)(N - 1) / (long long)(num_atoms - 1));
    }
    if (idx >= N) idx = N - 1;
    if (idx < 0) idx = 0;
    quantiles[i] = sorted[idx];
}

/* ── Legacy O(N²) kernel — kept for reference, no longer called ───── */

#define TILE_SIZE 256

extern "C" __global__ void reward_rank_normalize(
    const float* __restrict__ rewards_in,  /* [N] read-only raw rewards */
    float* __restrict__ rewards_out,       /* [N] output shaped rewards */
    int N,
    float std_ema,                         /* running std of raw rewards */
    float return_mean_ema                  /* running mean of raw rewards */
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;

    /* Each thread loads its own reward from global memory */
    float raw_r = (i < N) ? rewards_in[i] : 0.0f;
    float r_i = (std_ema > 1e-6f) ? (raw_r - return_mean_ema) / std_ema : raw_r;

    /* Skip near-zero rewards (holding bars). Only rank actual trade P&L. */
    int is_trade = (fabsf(r_i) > 0.001f) ? 1 : 0;

    float abs_i = fabsf(r_i);

    __shared__ float tile[TILE_SIZE];

    int count = 0;
    int num_tiles = (N + TILE_SIZE - 1) / TILE_SIZE;

    for (int t = 0; t < num_tiles; t++) {
        int tile_idx = t * TILE_SIZE + threadIdx.x;
        float tile_raw = (tile_idx < N) ? rewards_in[tile_idx] : 0.0f;
        float tile_sharpe = (std_ema > 1e-6f) ? (tile_raw - return_mean_ema) / std_ema : tile_raw;
        tile[threadIdx.x] = (tile_idx < N) ? fabsf(tile_sharpe) : 1e30f;
        __syncthreads();

        if (i < N) {
            int tile_end = min(TILE_SIZE, N - t * TILE_SIZE);
            for (int k = 0; k < tile_end; k++) {
                if (tile[k] <= abs_i) count++;
            }
        }
        __syncthreads();
    }

    if (i >= N) return;

    float rank = (float)count / (float)N;
    float sign = (r_i > 0.0f) ? 1.0f : (r_i < 0.0f) ? -1.0f : 0.0f;

    if (!is_trade) {
        rewards_out[i] = 0.0f;
        return;
    }

    rewards_out[i] = sign * rank;
}
