/**
 * rl_hindsight_forward.cu — Forward HER: post-exit continuation evaluation
 *
 * Bidirectional HER Phase 3: after a trade closes, monitors whether the
 * market continued in the position's direction. If forward_pnl >> actual_pnl
 * after a lookahead window, injects a synthetic "should-have-held-longer"
 * transition into the replay buffer.
 *
 * Operates on the closed-trades ring buffer (CLOSED_RING_SIZE entries),
 * populated by the host/experience kernel when trades close.
 *
 * Coordination: Grid=(1), Block=(CLOSED_RING_SIZE=256). Single-block
 * prefix-sum assigns contiguous replay write slots.
 *
 * ISV slots read:
 *   549 — HINDSIGHT_THRESHOLD          (forward_pnl / actual_pnl threshold)
 *   550 — HINDSIGHT_PRIORITY_BOOST     (multiplier on max_priority)
 *   552 — HINDSIGHT_FORWARD_LOOKAHEAD  (steps to wait before evaluating)
 *   548 — STEP_COUNTER                 (current training step)
 *   406 — REWARD_SCALE                 (for scaled reward computation)
 *
 * No atomicAdd — block prefix-sum only.
 */

#define CLOSED_RING_SIZE   256
#define CLOSED_RING_FIELDS 5    /* entry_price, exit_step, direction, actual_pnl, batch_idx */
#define HIDDEN_DIM         128
#define SCALARS_PER_TRANSITION 7

/* ISV slot indices. */
#define RL_HINDSIGHT_THRESHOLD_INDEX         549
#define RL_HINDSIGHT_PRIORITY_BOOST_INDEX    550
#define RL_HINDSIGHT_FORWARD_LOOKAHEAD_INDEX 552
#define RL_STEP_COUNTER_ISV_INDEX            548
#define RL_REWARD_SCALE_INDEX                406

extern "C" __global__ void rl_hindsight_forward(
    const float* __restrict__ bid_px,             /* [BOOK_LEVELS]              */
    const float* __restrict__ ask_px,             /* [BOOK_LEVELS]              */
    const float* __restrict__ isv,
    /* Closed trades ring */
    float* __restrict__ closed_ring,              /* [CLOSED_RING_SIZE * 5]     */
    const float* __restrict__ closed_h_t,         /* [CLOSED_RING_SIZE * 128]   */
    const unsigned int* __restrict__ closed_step,  /* [CLOSED_RING_SIZE]        */
    /* Replay write targets */
    float* __restrict__ replay_h_t,
    float* __restrict__ replay_h_tp1,
    float* __restrict__ replay_scalars,
    float* __restrict__ priority_tree,
    unsigned int* __restrict__ write_head,
    unsigned int* __restrict__ replay_len,
    float* __restrict__ max_priority,
    int capacity
)
{
    int i = threadIdx.x;
    if (i >= CLOSED_RING_SIZE) return;

    /* ── Phase 1: Evaluate each closed-trade entry ─────────────── */
    int wants_inject = 0;
    float forward_pnl = 0.0f;
    float actual_pnl = 0.0f;
    float direction = 0.0f;
    float entry_price = 0.0f;
    int should_clear = 0;

    entry_price = closed_ring[i * CLOSED_RING_FIELDS + 0];

    if (entry_price != 0.0f) {
        int current_step = (int)isv[RL_STEP_COUNTER_ISV_INDEX];
        int age = current_step - (int)closed_step[i];
        int lookahead = (int)isv[RL_HINDSIGHT_FORWARD_LOOKAHEAD_INDEX];

        if (age < lookahead) {
            /* Too early to evaluate — skip. */
        } else if (age > 2 * lookahead) {
            /* Expired — clear the slot. */
            should_clear = 1;
        } else {
            /* Evaluate forward PnL. */
            direction  = closed_ring[i * CLOSED_RING_FIELDS + 2];
            actual_pnl = closed_ring[i * CLOSED_RING_FIELDS + 3];

            float mid = 0.5f * (bid_px[0] + ask_px[0]);
            forward_pnl = direction * (mid - entry_price);

            float threshold = isv[RL_HINDSIGHT_THRESHOLD_INDEX];

            if (forward_pnl > 0.0f && forward_pnl > actual_pnl * threshold) {
                wants_inject = 1;
            }

            /* Mark for clearing after evaluation (whether or not we inject). */
            should_clear = 1;
        }
    }

    /* ── Phase 2: Prefix-sum for contiguous write slots ──────── */
    __shared__ unsigned int smem[CLOSED_RING_SIZE + 2];

    smem[i] = (unsigned int)wants_inject;
    __syncthreads();

    if (i == 0) {
        unsigned int running = 0;
        for (int j = 0; j < CLOSED_RING_SIZE; j++) {
            unsigned int val = smem[j];
            smem[j] = running;
            running += val;
        }
        smem[CLOSED_RING_SIZE]     = running;     /* total count */
        smem[CLOSED_RING_SIZE + 1] = write_head[0]; /* base offset */

        /* Advance write head and replay length. */
        if (running > 0) {
            write_head[0] = (write_head[0] + running) % (unsigned int)capacity;
            unsigned int new_len = replay_len[0] + running;
            if (new_len > (unsigned int)capacity) new_len = (unsigned int)capacity;
            replay_len[0] = new_len;
        }
    }
    __syncthreads();

    /* ── Phase 3: Clear expired / evaluated slots ────────────── */
    if (should_clear) {
        closed_ring[i * CLOSED_RING_FIELDS + 0] = 0.0f;
    }

    /* ── Phase 4: Write synthetic transitions ────────────────── */
    if (!wants_inject) return;

    unsigned int my_offset = smem[i];
    unsigned int base_head = smem[CLOSED_RING_SIZE + 1];
    unsigned int slot = (base_head + my_offset) % (unsigned int)capacity;

    float reward_scale = isv[RL_REWARD_SCALE_INDEX];

    /* h_t: copy the hidden state captured at trade close. */
    for (int d = 0; d < HIDDEN_DIM; d++) {
        float h = closed_h_t[(long long)i * HIDDEN_DIM + d];
        replay_h_t[(long long)slot * HIDDEN_DIM + d] = h;
        /* h_tp1 = self-loop (synthetic terminal). */
        replay_h_tp1[(long long)slot * HIDDEN_DIM + d] = h;
    }

    /* Scalars: [action, reward, raw_reward, done, log_pi, n_step_gamma, return] */
    long long sbase = (long long)slot * SCALARS_PER_TRANSITION;
    replay_scalars[sbase + 0] = 2.0f;                           /* Hold */
    replay_scalars[sbase + 1] = forward_pnl * reward_scale;     /* reward (scaled) */
    replay_scalars[sbase + 2] = forward_pnl;                    /* raw_reward */
    replay_scalars[sbase + 3] = 1.0f;                           /* done */
    replay_scalars[sbase + 4] = -1.0f;                          /* log_pi (synthetic) */
    replay_scalars[sbase + 5] = 0.0f;                           /* n_step_gamma (terminal) */
    replay_scalars[sbase + 6] = forward_pnl;                    /* return */

    /* Priority: boosted. Tree leaves at [cap..2*cap). */
    float boost = isv[RL_HINDSIGHT_PRIORITY_BOOST_INDEX];
    float prio = boost * max_priority[0];
    if (prio < 1e-6f) prio = 1e-6f;
    priority_tree[capacity + (int)slot] = prio;
}
