/* =====================================================================
 * Per-head LR scaling kernel — ml-alpha RL training stability
 *
 * Single-thread controller that adjusts per-head learning-rate scale
 * factors based on loss improvement rate. Four heads: Q, pi, V, IQN.
 *
 * Grid=(1), Block=(1). Single-thread controller (4 heads).
 * No atomicAdd — single thread, no reduction needed.
 * No #include beyond cuda_runtime.h — pre-compiled cubin, no nvrtc.
 *
 * ISV slot layout (outputs — LR scale factors):
 *   564 — LR_SCALE_Q:   Q-network LR multiplier ∈ [0.5, 2.0]
 *   565 — LR_SCALE_PI:  policy head LR multiplier ∈ [0.5, 2.0]
 *   566 — LR_SCALE_V:   value head LR multiplier ∈ [0.5, 2.0]
 *   567 — LR_SCALE_IQN: IQN head LR multiplier ∈ [0.5, 2.0]
 *
 * ISV slot layout (inputs + scratch — dual-purpose per head):
 *   568 — PREV_LOSS_Q:   previous Q loss (read prev, write current)
 *   569 — PREV_LOSS_PI:  previous pi loss (read prev, write current)
 *   570 — PREV_LOSS_V:   previous V loss (read prev, write current)
 *   571 — PREV_LOSS_IQN: previous IQN loss (read prev, write current)
 *
 * Current losses are passed as kernel arguments, not read from ISV.
 *
 * Controller logic per head:
 *   improvement_rate = (prev - current) / max(|prev|, 1e-6)
 *   If improvement > 0.1:  scale *= 0.8  (learning fast → slow down)
 *   If improvement < -0.01: scale *= 1.2  (stalled/regressing → speed up)
 *   Clamp scale to [0.5, 2.0]
 *
 * Pearl A bootstrap: when prev_loss == 0 (sentinel), skip the controller
 * update and just snapshot the current loss as prev. Scale stays at its
 * current value (1.0 if freshly initialized by host).
 * ===================================================================== */

#include <cuda_runtime.h>

#define LR_SCALE_Q_INDEX    564
#define LR_SCALE_PI_INDEX   565
#define LR_SCALE_V_INDEX    566
#define LR_SCALE_IQN_INDEX  567

/* Dual-purpose slots: host writes current loss BEFORE launch,
 * kernel reads as prev, then writes current_loss (from args) back.
 * Eliminates the separate prev_loss slot bank that collided with
 * RL_OUTCOME_AUX_LAMBDA_INDEX at 572. */
#define PREV_LOSS_Q_INDEX    568
#define PREV_LOSS_PI_INDEX   569
#define PREV_LOSS_V_INDEX    570
#define PREV_LOSS_IQN_INDEX  571

/* Static tables — 4 heads, indexed 0..3 in the loop. */
__device__ static const int SCALE_SLOTS[4] = {
    LR_SCALE_Q_INDEX,
    LR_SCALE_PI_INDEX,
    LR_SCALE_V_INDEX,
    LR_SCALE_IQN_INDEX
};

__device__ static const int PREV_LOSS_SLOTS[4] = {
    PREV_LOSS_Q_INDEX,
    PREV_LOSS_PI_INDEX,
    PREV_LOSS_V_INDEX,
    PREV_LOSS_IQN_INDEX
};

extern "C" __global__ void rl_per_branch_lr(
    float* __restrict__ isv,
    float loss_q,                   /* current Q head loss                 */
    float loss_pi,                  /* current pi head loss                */
    float loss_v,                   /* current V head loss                 */
    float loss_iqn                  /* current IQN head loss               */
) {
    if (blockIdx.x != 0 || threadIdx.x != 0) return;

    float current_losses[4] = { loss_q, loss_pi, loss_v, loss_iqn };

    #pragma unroll
    for (int h = 0; h < 4; h++) {
        float current_loss = current_losses[h];
        float prev_loss    = isv[PREV_LOSS_SLOTS[h]];
        float scale        = isv[SCALE_SLOTS[h]];

        /* Defensive: if scale not yet seeded (sentinel 0), init to 1.0. */
        if (scale < 1e-9f) {
            scale = 1.0f;
        }

        /* Pearl A bootstrap: sentinel prev_loss == 0 means no previous
         * observation. Snapshot current and skip controller update. */
        if (prev_loss > -1e-9f && prev_loss < 1e-9f) {
            isv[PREV_LOSS_SLOTS[h]] = current_loss;
            isv[SCALE_SLOTS[h]] = scale;
            continue;
        }

        /* improvement_rate: positive = loss decreased (good),
         * negative = loss increased (bad). */
        float denom = fmaxf(fabsf(prev_loss), 1e-6f);
        float improvement_rate = (prev_loss - current_loss) / denom;

        /* Bounded discrete step per
         * pearl_multiplicative_controllers_need_bounded_step_and_noise_floor:
         * — learning fast (improvement > 0.1):  slow down (× 0.8)
         * — stalled/regressing (improvement < -0.01): speed up (× 1.2)
         * — otherwise: no change (× 1.0)                                */
        if (improvement_rate > 0.1f) {
            scale *= 0.8f;
        } else if (improvement_rate < -0.01f) {
            scale *= 1.2f;
        }

        /* Bilateral clamp per pearl_symmetric_clamp_audit. */
        scale = fmaxf(0.5f, fminf(scale, 2.0f));

        /* Write outputs: updated scale + snapshot current as prev. */
        isv[SCALE_SLOTS[h]] = scale;
        isv[PREV_LOSS_SLOTS[h]] = current_loss;
    }
}
