/**
 * HER (Hindsight Experience Replay) goal relabeling kernel.
 *
 * Operates on GPU replay buffer tensors. For each experience to relabel:
 *   1. Copies original (state, next_state, action, done) from source index
 *   2. Reads achieved goal from a donor experience
 *   3. Replaces goal portion (first GOAL_DIM elements) in state and next_state
 *   4. Recomputes reward (sparse: +1 if goal achieved, -0.01 otherwise)
 *   5. Writes relabeled experience to output staging buffer
 *
 * Launch: grid=(her_batch_size, 1, 1), block=(32, 1, 1)
 * One warp per relabeled experience.
 *
 * Compile-time defines (injected via NVRTC dim_overrides):
 *   STATE_DIM      -- tensor-core-aligned state width (e.g. 48 or 56)
 *   GOAL_DIM       -- number of goal elements at front of state vector
 *   GOAL_THRESHOLD -- L2 distance threshold for goal achievement (default 0.01)
 */

#ifndef GOAL_DIM
#define GOAL_DIM 1
#endif

#ifndef GOAL_THRESHOLD
#define GOAL_THRESHOLD 0.01f
#endif

extern "C" __global__
void her_relabel_kernel(
    const float* __restrict__ states,        /* [capacity, STATE_DIM] */
    const float* __restrict__ next_states,   /* [capacity, STATE_DIM] */
    const int*   __restrict__ actions,       /* [capacity]            */
    const float* __restrict__ rewards,       /* [capacity]            */
    const float* __restrict__ dones,         /* [capacity]            */
    const int*   __restrict__ source_indices,/* [her_batch_size]      */
    const int*   __restrict__ donor_indices, /* [her_batch_size]      */
    int buffer_size,
    float* __restrict__ out_states,          /* [her_batch_size, STATE_DIM] */
    float* __restrict__ out_next_states,     /* [her_batch_size, STATE_DIM] */
    int*   __restrict__ out_actions,         /* [her_batch_size]            */
    float* __restrict__ out_rewards,         /* [her_batch_size]            */
    float* __restrict__ out_dones            /* [her_batch_size]            */
)
{
    int sample_idx = blockIdx.x;
    int lane_id = threadIdx.x;  /* 0-31 within the warp */

    int src_idx = source_indices[sample_idx];
    int donor_idx = donor_indices[sample_idx];

    /* 1. Warp-cooperative coalesced copy of full state and next_state vectors.
     *    Each lane handles a stride-32 subset of STATE_DIM elements. */
    for (int d = lane_id; d < STATE_DIM; d += 32) {
        out_states[sample_idx * STATE_DIM + d] = states[src_idx * STATE_DIM + d];
        out_next_states[sample_idx * STATE_DIM + d] = next_states[src_idx * STATE_DIM + d];
    }
    __syncwarp();

    /* 2. Replace goal portion (first GOAL_DIM elements) with donor's achieved goal.
     *    The donor's achieved goal is extracted from the donor's next_state. */
    for (int d = lane_id; d < GOAL_DIM; d += 32) {
        float achieved = next_states[donor_idx * STATE_DIM + d];
        out_states[sample_idx * STATE_DIM + d] = achieved;
        out_next_states[sample_idx * STATE_DIM + d] = achieved;
    }

    /* 3. Scalar work: copy action/done, recompute sparse reward (lane 0 only). */
    if (lane_id == 0) {
        out_actions[sample_idx] = actions[src_idx];
        out_dones[sample_idx] = dones[src_idx];

        /* Goal distance: L2 between source's achieved goal and donor's achieved goal.
         * Source achieved = next_states[src_idx, 0..GOAL_DIM]
         * Donor achieved  = next_states[donor_idx, 0..GOAL_DIM] */
        float dist_sq = 0.0f;
        for (int d = 0; d < GOAL_DIM; d++) {
            float src_achieved = next_states[src_idx * STATE_DIM + d];
            float donor_achieved = next_states[donor_idx * STATE_DIM + d];
            float diff = src_achieved - donor_achieved;
            dist_sq += diff * diff;
        }
        float dist = sqrtf(dist_sq);

        /* Sparse reward: +1.0 if goal achieved, -0.01 otherwise */
        out_rewards[sample_idx] = (dist < GOAL_THRESHOLD) ? 1.0f : -0.01f;
    }
}
