// rl_kl_reference_grad.cu — KL divergence gradient from a hold-heavy
// reference policy. Proven approach for preventing policy collapse with
// sparse rewards (RLHF/InstructGPT pattern).
//
// Adds β × (π_θ(a|s) - π_ref(a)) to pi_grad_logits for every batch
// element, every step. Unlike entropy bonus (pushes toward uniform) or
// done-gated advantages (only fires on 6% of batch), this provides
// continuous gradient signal that specifically preserves Hold as the
// default action.
//
// π_ref encodes the surfer philosophy:
//   Hold = 50%, each trading action = 5% (uniform over 10 non-Hold)
//
// The gradient ∂KL(π_θ||π_ref)/∂logit_a = π_θ(a) - π_ref(a):
//   - If π(Hold) < 50%: gradient pushes Hold probability UP
//   - If π(Hold) > 50%: gradient gently pushes Hold DOWN (allows trading)
//   - If π(BuyL1) > 5%: gradient pushes it DOWN toward the prior
//
// β is ISV-driven so it can adapt with training progress.
// Grid=(B,1,1), Block=(N_ACTIONS,1,1). One thread per action.

#define N_ACTIONS              11
#define ACTION_HOLD            2
#define RL_KL_REF_BETA_INDEX   578
#define RL_HOLD_TARGET_FRAC_INDEX 575
#define RL_HOLD_FRAC_EMA_INDEX    576
#define BETA_MIN               0.001f
#define BETA_MAX               0.1f
#define BETA_ADJUST_UP         1.02f
#define BETA_DECAY_DOWN        0.999f
#define HOLD_EMA_ALPHA         0.05f

extern "C" __global__ void rl_kl_reference_grad(
    const float* __restrict__ pi_logits,       // [B × N_ACTIONS]
    float*       __restrict__ isv,             // ISV bus (RW — β, hold_ema)
    float*       __restrict__ pi_grad_logits,  // [B × N_ACTIONS] — ADD to
    int                       B,
    const int*   __restrict__ actions          // [B] post-gate actions
) {
    const int b = blockIdx.x;
    const int a = threadIdx.x;
    if (b >= B || a >= N_ACTIONS) return;

    const float beta = isv[RL_KL_REF_BETA_INDEX];
    if (beta <= 0.0f) return;

    // Reference policy: Hold=50%, rest=5% each.
    const float pi_ref = (a == ACTION_HOLD) ? 0.5f : 0.05f;

    // Compute π_θ(a|s) via softmax.
    const int base = b * N_ACTIONS;

    __shared__ float s_logits[11];
    __shared__ float s_max;
    __shared__ float s_sumexp;

    s_logits[a] = pi_logits[base + a];
    __syncthreads();

    if (a == 0) {
        float m = s_logits[0];
        for (int i = 1; i < N_ACTIONS; ++i) m = fmaxf(m, s_logits[i]);
        s_max = m;
    }
    __syncthreads();

    float exp_val = expf(s_logits[a] - s_max);
    __shared__ float s_exp[11];
    s_exp[a] = exp_val;
    __syncthreads();

    if (a == 0) {
        float sum = 0.0f;
        for (int i = 0; i < N_ACTIONS; ++i) sum += s_exp[i];
        s_sumexp = sum;
    }
    __syncthreads();

    const float pi_theta = s_exp[a] / s_sumexp;

    // ∂KL/∂logit_a = π_θ(a) - π_ref(a)
    // Additive to existing PPO + distillation gradients.
    // /B for batch-size invariance.
    pi_grad_logits[base + a] += beta * (pi_theta - pi_ref) / (float)B;

    // ── KL-based β controller (block 0, thread 0). ────────────────
    // Controls β based on ACTUAL KL divergence (zero-lag signal),
    // NOT Hold% (50-step lag). SAC auto-tuning pattern.
    //
    // KL is computed from the softmax we just did — free, instant.
    // If KL > KL_target: π drifted too far → raise β
    // If KL < KL_target: π too constrained → lower β
    //
    // Schulman bounded step with small rate (1.005) because signal
    // has zero lag — no risk of overshoot from delayed feedback.
    #define RL_KL_REF_TARGET_INDEX 580
    #define BETA_STEP_RATE         1.005f
    if (b == 0 && a == 0) {
        // Compute KL(π_θ || π_ref) from the full batch mean.
        float kl_sum = 0.0f;
        for (int bi = 0; bi < B; bi++) {
            const int bi_base = bi * N_ACTIONS;
            // Re-compute softmax for batch bi (cheap — 11 actions).
            float mx = pi_logits[bi_base];
            for (int j = 1; j < N_ACTIONS; j++) {
                float v = pi_logits[bi_base + j];
                if (v > mx) mx = v;
            }
            float se = 0.0f;
            float exps[11];
            for (int j = 0; j < N_ACTIONS; j++) {
                exps[j] = expf(pi_logits[bi_base + j] - mx);
                se += exps[j];
            }
            for (int j = 0; j < N_ACTIONS; j++) {
                float p = exps[j] / se;
                float pr = (j == ACTION_HOLD) ? 0.5f : 0.05f;
                if (p > 1e-7f) {
                    kl_sum += p * logf(p / pr);
                }
            }
        }
        float kl_mean = kl_sum / (float)B;

        const float kl_target = isv[RL_KL_REF_TARGET_INDEX];
        float new_beta = beta;
        if (kl_mean > kl_target * 1.1f) {
            new_beta = beta * BETA_STEP_RATE;
        } else if (kl_mean < kl_target * 0.9f) {
            new_beta = beta / BETA_STEP_RATE;
        }
        isv[RL_KL_REF_BETA_INDEX] = fmaxf(BETA_MIN, fminf(new_beta, BETA_MAX));
    }
}
