/* atoms_update — C51 adaptive atom position recompute, GPU-driven (Plan 1 Task 9).
 *
 * Replaces the CPU-side `recompute_atom_positions` loop which launched the
 * existing `adaptive_atom_positions` kernel once per branch.  Here, all 4
 * branches are dispatched in a single launch: grid = (4, 1, 1),
 * block = (256, 1, 1), shared memory = 2 * num_atoms * sizeof(float).
 * blockIdx.x selects the branch (0=dir, 1=mag, 2=ord, 3=urg).
 *
 * Reads:
 *   spacing_raw[b]  — learnable softmax spacing params for branch b,
 *                     accessed via pre-resolved device pointers passed as args.
 *   isv_signals     — ISV bus:
 *     SP5 Layer B (Pearl 1): per-branch v-range at
 *       ATOM_V_CENTER_BASE=174: v_center[b] = isv_signals[174 + b]
 *       ATOM_V_HALF_BASE=178:   v_half[b]   = isv_signals[178 + b]
 *     (SP4 slots 23+2*b / 24+2*b are superseded by SP5 Pearl 1.)
 *
 * Writes:
 *   positions_out   — atom grid [4 * num_atoms], branch b at offset b*num_atoms.
 *
 * Formula (same as adaptive_atom_positions in experience_kernels.cu):
 *   spacing = softmax(spacing_raw)
 *   positions[j] = v_min + (v_max - v_min) * cumsum(spacing)[j]
 *   with v_min = isv[174+b] - isv[178+b], v_max = isv[174+b] + isv[178+b].
 */

/* SP4 Layer B: per-branch atom-position bound from ISV[ATOM_POS_BOUND_BASE + branch].
 * Mirrors `crates/ml/src/cuda_pipeline/sp4_isv_slots.rs::ATOM_POS_BOUND_BASE = 132`.
 */
#define ATOM_POS_BOUND_BASE 132

/* SP5 Layer B (Pearl 1): per-branch v_center and v_half from ISV.
 * Mirrors sp5_isv_slots.rs: ATOM_V_CENTER_BASE=174, ATOM_V_HALF_BASE=178. */
#define ATOM_V_CENTER_BASE 174
#define ATOM_V_HALF_BASE   178

/* SP5 Layer B (Pearl 1-ext): per-branch num_atoms from ISV.
 * Mirrors sp5_isv_slots.rs: ATOM_NUM_ATOMS_BASE=274.
 * Valid discrete values: {16, 32, 64}. Rounds ISV float to nearest valid
 * count. Cold-start floor: if ISV=0 (pre-observation), fall back to the
 * global `num_atoms` parameter (Invariant-1 carve-out). */
#define ATOM_NUM_ATOMS_BASE 274

/* Round an ISV-driven f32 num_atoms signal to the nearest valid discrete C51
 * atom count {16, 32, 64}, bounded to [16, num_atoms_max]. */
__device__ __forceinline__ int round_to_valid_num_atoms(float isv_val, int num_atoms_max) {
    if (isv_val < 1.0f) return num_atoms_max;  /* cold-start floor */
    int v = (int)(isv_val + 0.5f);
    /* Snap to nearest of {16, 32, 64}: */
    int d16 = abs(v - 16);
    int d32 = abs(v - 32);
    int d64 = abs(v - 64);
    int best = (d16 <= d32 && d16 <= d64) ? 16 : (d32 <= d64) ? 32 : 64;
    /* Clamp to [16, num_atoms_max]: never exceed the buffer's allocated stride */
    if (best > num_atoms_max) best = num_atoms_max;
    if (best < 16)            best = 16;
    return best;
}

extern "C" __global__ void atoms_update(
    unsigned long long spacing_ptr_b0,   /* device ptr to spacing_raw[dir] */
    unsigned long long spacing_ptr_b1,   /* device ptr to spacing_raw[mag] */
    unsigned long long spacing_ptr_b2,   /* device ptr to spacing_raw[ord] */
    unsigned long long spacing_ptr_b3,   /* device ptr to spacing_raw[urg] */
    float*             __restrict__ positions_out,   /* [4 * num_atoms] */
    int                num_atoms,
    const float*       __restrict__ isv_signals
) {
    int branch = (int)blockIdx.x;
    if (branch >= 4) return;

    int tid = threadIdx.x;
    extern __shared__ float shmem[];
    /* Shared memory is allocated for global `num_atoms` (worst case) at launch.
     * eff_na ≤ num_atoms, so both halves fit safely. */
    float* s_softmax = shmem;
    float* s_cumsum  = shmem + num_atoms;

    /* SP5 Layer B (Pearl 1): per-branch v-range from ISV.
     * v_center[b] = ISV[ATOM_V_CENTER_BASE + b] = ISV[174 + b]
     * v_half[b]   = ISV[ATOM_V_HALF_BASE   + b] = ISV[178 + b]
     * EPS_CLAMP_FLOOR = 1.0 (Invariant 1 anchor): pre-first-observation ISV
     * reads as 0; fmaxf prevents zero/negative half-width collapsing the range. */
    float v_center = isv_signals[ATOM_V_CENTER_BASE + branch];
    float v_half   = fmaxf(isv_signals[ATOM_V_HALF_BASE + branch], 1.0f);
    float v_min    = v_center - v_half;
    float v_max    = v_center + v_half;

    /* SP5 Layer B (Pearl 1-ext): per-branch effective num_atoms from ISV.
     * ISV[ATOM_NUM_ATOMS_BASE + branch] = ISV[274 + branch].
     * Rounded to valid discrete {16,32,64}; cold-start (ISV=0) falls back to
     * the global `num_atoms` parameter (Invariant-1 carve-out). The global
     * `num_atoms` is still used for shared-memory stride and output-buffer
     * stride (buffer is allocated as 4 * config.num_atoms). */
    int eff_na = round_to_valid_num_atoms(
        isv_signals[ATOM_NUM_ATOMS_BASE + branch], num_atoms);

    /* Resolve spacing_raw pointer for this branch */
    const float* spacing_raw;
    switch (branch) {
        case 0: spacing_raw = (const float*)spacing_ptr_b0; break;
        case 1: spacing_raw = (const float*)spacing_ptr_b1; break;
        case 2: spacing_raw = (const float*)spacing_ptr_b2; break;
        default: spacing_raw = (const float*)spacing_ptr_b3; break;
    }

    /* Load spacing_raw into shared memory (only eff_na elements are valid) */
    if (tid < eff_na) {
        s_softmax[tid] = spacing_raw[tid];
    }
    __syncthreads();

    /* Find max for numerical stability (single thread, eff_na <= 256 is tiny) */
    if (tid == 0) {
        float max_val = s_softmax[0];
        for (int j = 1; j < eff_na; j++)
            max_val = fmaxf(max_val, s_softmax[j]);
        s_cumsum[0] = max_val;  /* store max temporarily */
    }
    __syncthreads();
    float max_val = s_cumsum[0];
    __syncthreads();

    /* Compute softmax */
    if (tid < eff_na) {
        s_softmax[tid] = expf(s_softmax[tid] - max_val);
    }
    __syncthreads();

    /* Sum for normalization */
    if (tid == 0) {
        float sum = 0.0f;
        for (int j = 0; j < eff_na; j++)
            sum += s_softmax[j];
        float inv_sum = 1.0f / fmaxf(sum, 1e-8f);
        for (int j = 0; j < eff_na; j++)
            s_softmax[j] *= inv_sum;
    }
    __syncthreads();

    /* Inclusive prefix sum (sequential, tid=0 only) */
    if (tid == 0) {
        s_cumsum[0] = s_softmax[0];
        for (int j = 1; j < eff_na; j++)
            s_cumsum[j] = s_cumsum[j-1] + s_softmax[j];
    }
    __syncthreads();

    /* Write positions: positions[0] = v_min, positions[j>0] = v_min + range * cumsum[j-1].
     * Output stride is the global `num_atoms` (buffer layout invariant). Positions
     * beyond eff_na are left unchanged (they are not read — C51 consumers also
     * use eff_na from ISV). */
    float* out = positions_out + (long long)branch * num_atoms;
    if (tid < eff_na) {
        float range = v_max - v_min;
        float frac  = (tid == 0) ? 0.0f : s_cumsum[tid - 1];
        float new_pos = v_min + range * frac;

        /* SP4 Layer B (Mech 2): bound atom positions to ±ISV[ATOM_POS_BOUND[branch]].
         * Per-branch p99(|atom_positions|) populated by the SP4 producer
         * `launch_sp4_atom_pos_p99` (Layer A Task A6) per branch ∈ [0,4).
         * ε floor on the bound itself (`max(.., 1.0)`) per the SP1 cold-start
         * pearl — pre-first-observation ISV reads as 0 and the consumer must
         * not collapse the bound to zero. */
        const float max_atom_abs = fmaxf(isv_signals[ATOM_POS_BOUND_BASE + branch], 1.0f);
        new_pos = fminf(fmaxf(new_pos, -max_atom_abs), max_atom_abs);

        out[tid] = new_pos;
    }
}
