// rl_reward_shaping.cu — surfer-philosophy reward shaping.
//
// Three components applied AFTER extract_realized_pnl_delta, BEFORE
// apply_reward_scale:
//
// 1. ENTRY COST: subtract a fixed penalty when a new position opens.
//    "Don't paddle for ripples" — the agent must believe the trade
//    will profit more than the cost to justify entering. ISV-driven.
//
// 2. SHORT-HOLD PENALTY: multiply reward by a factor < 1 when a trade
//    closes with hold_time < min_hold. "Don't bail on the first bump."
//    Discourages churning by making quick flips less rewarding.
//
// 3. HOLD BONUS: small per-step bonus when in a profitable position
//    (unrealized_r > 0). "Ride the wave" — rewards patience when the
//    trade is working. Scales with sqrt(hold_time) so the incentive
//    grows as the ride extends.
//
// All thresholds and coefficients are ISV-driven for runtime tuning.
// Runs BEFORE reward_scale so the shaping is in raw dollar terms.
//
// Per `feedback_no_atomicadd`: per-batch element-wise.
// Per `feedback_cpu_is_read_only`: pure device-side.

#include <stdint.h>

#define RL_ENTRY_COST_INDEX              532
#define RL_SHORT_HOLD_MIN_STEPS_INDEX    533
#define RL_SHORT_HOLD_PENALTY_INDEX      534
#define RL_HOLD_BONUS_INDEX              535

extern "C" __global__ void rl_reward_shaping(
    float*               __restrict__ rewards,         // [B] IN/OUT (raw PnL delta)
    const int*           __restrict__ prev_pos_lots,   // [B] position BEFORE this step
    const unsigned char* __restrict__ pos_state,       // [B × pos_bytes] position AFTER fill
    const float*         __restrict__ dones,           // [B] (1.0 if done, 0.0 otherwise)
    const int*           __restrict__ steps_since_done, // [B] hold time counter
    const float*         __restrict__ isv,
    int                                b_size,
    int                                pos_bytes
) {
    const int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= b_size) return;

    const int prev_pos = prev_pos_lots[b];
    const int curr_pos = *(const int*)(pos_state + b * pos_bytes);
    const float entry_cost = isv[RL_ENTRY_COST_INDEX];
    const float min_hold   = isv[RL_SHORT_HOLD_MIN_STEPS_INDEX];
    const float penalty    = isv[RL_SHORT_HOLD_PENALTY_INDEX];
    const float hold_bonus = isv[RL_HOLD_BONUS_INDEX];
    const int hold_time    = steps_since_done[b];

    float r = rewards[b];

    // 1. Entry cost: fires when transitioning from flat to positioned.
    const bool was_flat = (prev_pos == 0);
    const bool is_positioned = (curr_pos != 0);
    if (was_flat && is_positioned) {
        r -= entry_cost;
    }

    // 2. Short-hold penalty: fires at trade close (done) if hold time
    //    is below minimum. Multiplicative so it scales with trade size.
    if (dones[b] > 0.5f && (float)hold_time < min_hold) {
        r *= penalty;
    }

    // 3. Long-ride bonus: at trade close, multiply profit by a
    //    hold-time factor. Longer winning rides get amplified reward.
    //    "The best wave of the session deserves the biggest cheer."
    //    Only fires on profitable closes (r > 0 at done).
    if (dones[b] > 0.5f && r > 0.0f && hold_time > 0) {
        const float ride_mult = 1.0f + hold_bonus * sqrtf((float)hold_time);
        r *= ride_mult;
    }

    // 4. Per-step hold bonus for staying in a profitable position.
    //    "Ride the wave" incentive between closes.
    if (!was_flat && is_positioned && r > 0.0f && hold_time > 0) {
        r += hold_bonus * sqrtf((float)hold_time);
    }

    rewards[b] = r;
}
