/* =====================================================================
 * rl_outcome_ce.cu -- softmax cross-entropy loss + gradient for K=3
 * trade-outcome classifier.
 *
 * Grid  = (B, 1, 1)   -- one block per batch element.
 * Block = (1, 1, 1)    -- single thread per batch (K=3 is tiny).
 *
 * Per batch element:
 *   - If label == -1 (no label): loss = 0, grad = 0 for all 3 classes.
 *   - Otherwise: numerically stable softmax, CE = -log(p[label]),
 *     grad[k] = p[k] - 1(k == label).
 *
 * No atomicAdd -- single-thread per batch writes its own loss + grad.
 * No shared memory -- K=3 fits trivially in registers.
 *
 * Layout:
 *   logits      : [B, 3]   IN
 *   labels      : [B]      IN  (0=Loss, 1=Timeout, 2=Profit, -1=skip)
 *   loss_pb     : [B]      OUT (per-batch CE loss)
 *   grad_logits : [B, 3]   OUT (d_loss / d_logits)
 * ===================================================================== */

#define K_CLASSES 3

extern "C" __global__ void rl_outcome_ce(
    const float* __restrict__ logits,       /* [B, 3]  */
    const int*   __restrict__ labels,       /* [B]     */
    float*       __restrict__ loss_pb,      /* [B]     */
    float*       __restrict__ grad_logits,  /* [B, 3]  */
    int b_size
) {
    int batch = blockIdx.x;
    if (batch >= b_size) return;

    int label = labels[batch];

    int base = batch * K_CLASSES;

    /* Masked sample: no label available. */
    if (label < 0 || label >= K_CLASSES) {
        loss_pb[batch] = 0.0f;
        grad_logits[base + 0] = 0.0f;
        grad_logits[base + 1] = 0.0f;
        grad_logits[base + 2] = 0.0f;
        return;
    }

    /* Load logits into registers. */
    float l0 = logits[base + 0];
    float l1 = logits[base + 1];
    float l2 = logits[base + 2];

    /* Numerically stable softmax: subtract max. */
    float mx = fmaxf(l0, fmaxf(l1, l2));
    float e0 = expf(l0 - mx);
    float e1 = expf(l1 - mx);
    float e2 = expf(l2 - mx);
    float sum_exp = e0 + e1 + e2;

    /* Softmax probabilities. */
    float p0 = e0 / sum_exp;
    float p1 = e1 / sum_exp;
    float p2 = e2 / sum_exp;

    /* CE loss = -log(p[label]). */
    float p_label = (label == 0) ? p0 : ((label == 1) ? p1 : p2);
    loss_pb[batch] = -logf(fmaxf(p_label, 1e-12f));

    /* Gradient: p[k] - 1(k == label). */
    grad_logits[base + 0] = p0 - ((label == 0) ? 1.0f : 0.0f);
    grad_logits[base + 1] = p1 - ((label == 1) ? 1.0f : 0.0f);
    grad_logits[base + 2] = p2 - ((label == 2) ? 1.0f : 0.0f);
}
