// crates/ml/src/cuda_pipeline/pearl_1_atom_kernel.cu
//
// SP5 Task A1: Pearl 1 — per-branch C51 atom span via Q-stats + headroom controller.
// Reads q_branch_stats outputs from scratch_buf, computes:
//   v_half = max_abs_dev × headroom
//   headroom updated via clip-rate controller targeting 1% atom-clip rate.
// Writes per-branch (v_center, v_half, headroom, clip_rate) to scratch_out.
// apply_pearls_ad_kernel then smooths these via Pearls A+D into ISV.
//
// 4-thread single-block (one thread per branch: Direction, Magnitude, Order, Urgency).
// No atomicAdd (feedback_no_atomicadd) — one thread per branch, independent writes.
// __threadfence_system() after all writes ensures visibility to the chained
// apply_pearls_ad_kernel on the same stream.
//
// TARGET_CLIP_RATE = 0.01 and HEADROOM_LR = 0.01 are Invariant 1 anchors
// (structural 1% target rate, slow controller rate) — not tuned constants.

#include <cuda_runtime.h>

extern "C" __global__ void pearl_1_atom_update(
    const float* __restrict__ scratch_buf,      // q_branch_stats outputs in scratch
    int   q_stats_idx_base,                     // scratch_buf base for q_branch_stats outputs
    const int*   __restrict__ atoms_clip_count, // [4] per-branch clip counts from atoms_update_kernel
    const int*   __restrict__ action_counts,    // [4] per-branch Q-values: {4, 3, 3, 3}
    int   batch_size,
    float* __restrict__ scratch_out,            // output scratch for apply_pearls_ad_kernel
    int   v_center_idx_base,                    // scratch_out base for v_center[4]
    int   v_half_idx_base,                      // scratch_out base for v_half[4]
    int   headroom_idx_base,                    // scratch_out base for headroom[4]
    int   clip_rate_idx_base,                   // scratch_out base for clip_rate[4]
    const float* __restrict__ isv_signals,      // read current HEADROOM from ISV
    int   isv_headroom_base                     // ISV index of ATOM_HEADROOM_BASE
) {
    int b = threadIdx.x;
    if (blockIdx.x != 0 || b >= 4) return;

    // Read Q-stats for this branch from q_branch_stats output:
    //   offset + b*4 + 0 = mean (v_center)
    //   offset + b*4 + 1 = max_abs_dev
    float mean        = scratch_buf[q_stats_idx_base + b * 4 + 0];
    float max_abs_dev = scratch_buf[q_stats_idx_base + b * 4 + 1];

    // Read current HEADROOM from ISV (Pearl A sentinel 0.0 → bootstrap to 6.0).
    float headroom = isv_signals[isv_headroom_base + b];
    if (headroom <= 0.0f) headroom = 6.0f;

    // Per-branch clip rate: atoms_clip_count[b] / (batch_size × action_counts[b]).
    // The denominator is the per-branch Q-value count this step (NOT the total
    // 13 Q-values across all branches), so the controller's TARGET_CLIP_RATE=0.01
    // is calibrated against the same rate the count is measured at.
    // clip_count=0 before Layer B wires atoms_update_kernel → clip_rate=0 → headroom holds.
    float clip_rate = (float)atoms_clip_count[b] / (float)(batch_size * action_counts[b]);

    // Headroom controller: HEADROOM += HEADROOM_LR × (clip_rate - TARGET_CLIP_RATE)
    // Structural envelope [2.0, 20.0] prevents runaway in either direction.
    const float TARGET_CLIP_RATE = 0.01f;
    const float HEADROOM_LR      = 0.01f;
    float new_headroom = headroom + HEADROOM_LR * (clip_rate - TARGET_CLIP_RATE);
    new_headroom = fmaxf(2.0f, fminf(20.0f, new_headroom));

    float v_half = max_abs_dev * new_headroom;

    // Write 4 outputs per branch to scratch_out for apply_pearls_ad_kernel:
    scratch_out[v_center_idx_base  + b] = mean;
    scratch_out[v_half_idx_base    + b] = v_half;
    scratch_out[headroom_idx_base  + b] = new_headroom;
    scratch_out[clip_rate_idx_base + b] = clip_rate;
    __threadfence_system();
}
