/**
 * branch_grad_balance_kernel — ISV-driven adaptive per-branch gradient-norm
 *                              balancer.
 *
 * Scales each of the 4 factored-action branches' gradient toward a
 * per-branch target norm maintained as an adaptive EMA in the ISV signal
 * bus (slots GRAD_NORM_TARGET_{DIR,MAG,ORD,URG}). The scale factor is
 * clamped to `[1/limit, limit]` where `limit` is itself an ISV-bus
 * adaptive EMA of observed `max_norm / min_norm` spread
 * (GRAD_SCALE_LIMIT_INDEX).
 *
 * Rationale — what was wrong with the pre-spec design:
 *   The previous kernel capped branches whose norm exceeded `4 × median`.
 *   On L40S train-5wb4n (2026-04-23) with grad distribution
 *     [dir=130, mag=20000, ord=18000, urg=16000]
 *   median = 17000, cap = 68000 → no branch triggered the cap, the
 *   balancer was a geometric no-op for 11 consecutive epochs. The
 *   algorithm only handles outlier-reduction; it cannot address
 *   starvation ("one branch far below the pack"). The pattern here is
 *   starvation: direction is 150× below magnitude, not magnitude above
 *   the rest.
 *
 * Why ISV-driven (architecturally, NOT tuned):
 *   * Per-branch target tracking (EMA → median) causes every branch to
 *     equalise over time, regardless of which side is the outlier.
 *     Starved branches are boosted; elevated branches are suppressed.
 *   * The scale clamp `limit` is not a fixed constant. It tracks the
 *     observed max/min norm spread via adaptive EMA — grows just enough
 *     to correct the actual asymmetry, bounded by an absolute safety
 *     cap (1e4) written on the Rust side.
 *   * All state lives in ISV slots. The producer `grad_balance_isv_update`
 *     reads observed branch norms and writes updated EMAs. Consumers
 *     (this kernel) read ISV at rescale time. A single producer, one
 *     source of truth — the ISV invariant applies everywhere.
 *   * See `feedback_isv_for_adaptive_bounds.md` (2026-04-23) for the
 *     project-wide rule this satisfies.
 *
 * Pipeline (three kernels, three launches):
 *   branch_grad_norm_reduce:    grid=(4,1,1), block=(256,1,1). Unchanged
 *                               from the pre-spec version. Writes per-
 *                               branch L2 norms to branch_norms_dev[4].
 *
 *   grad_balance_isv_update:    grid=(1,1,1), block=(32,1,1). Thread 0
 *                               reads branch_norms_dev, computes median
 *                               + max/min spread, applies adaptive EMAs
 *                               to ISV slots. Single-threaded work;
 *                               extra threads idle-exit. No atomicAdd.
 *                               Writes: ISV[31..35]. Safe in graph
 *                               capture — no host syncs, no dynamic
 *                               alloc.
 *
 *   branch_grad_rescale:        grid=(max_blocks, 4, 1), block=(256,1,1).
 *                               blockIdx.y = branch id; blockIdx.x =
 *                               256-element sub-chunk. Thread 0 reads
 *                               ISV target + limit, computes scale,
 *                               writes diag slot. Each thread multiplies
 *                               one element by its branch's scale. No
 *                               atomicAdd.
 *
 * CUDA Graph safety:
 *   All three launches have fixed grid/block. No host syncs, no dynamic
 *   allocs, no atomicAdd. Safe to capture inside `adam_grad_child`
 *   between the post_aux phase and `compute_grad_norm_for_adam`.
 */

#include "state_layout.cuh"

extern "C" __global__ void branch_grad_norm_reduce(
    const float* __restrict__ grad_buf,
    const int* __restrict__ branch_starts,   /* [4] element offsets into grad_buf */
    const int* __restrict__ branch_lens,     /* [4] element counts (padded, zero fill) */
    float* __restrict__ branch_norms_dev     /* [4] output L2 norms */
) {
    const int branch = blockIdx.x;            /* 0..3 */
    if (branch >= 4) return;
    const int tid = threadIdx.x;
    const int start = branch_starts[branch];
    const int len   = branch_lens[branch];

    __shared__ float sh[256];
    float acc = 0.0f;
    for (int i = tid; i < len; i += blockDim.x) {
        float g = grad_buf[start + i];
        acc += g * g;
    }
    sh[tid] = acc;
    __syncthreads();

    /* Block tree reduction 256→1 */
    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) {
            sh[tid] += sh[tid + s];
        }
        __syncthreads();
    }
    if (tid == 0) {
        branch_norms_dev[branch] = sqrtf(sh[0]);
    }
}

/* Device helper: median of 4 non-negative floats (sort-based, fully
 * branch-deterministic).  Returns the average of the middle two after
 * sorting — the standard even-length median.
 */
__device__ __forceinline__ float branch_grad_balance_median4(
    float a, float b, float c, float d
) {
    /* 5-comparator sorting network for 4 elements (Bose-Nelson). */
    #define SWAP_MIN_MAX(x, y) do { \
        float mn = fminf((x), (y)); \
        float mx = fmaxf((x), (y)); \
        (x) = mn; \
        (y) = mx; \
    } while (0)
    SWAP_MIN_MAX(a, c);
    SWAP_MIN_MAX(b, d);
    SWAP_MIN_MAX(a, b);
    SWAP_MIN_MAX(c, d);
    SWAP_MIN_MAX(b, c);
    #undef SWAP_MIN_MAX
    /* a ≤ b ≤ c ≤ d; median = (b + c) / 2 */
    return 0.5f * (b + c);
}

/* grad_balance_isv_update — producer kernel for the ISV-driven balancer.
 *
 * Reads the per-branch L2 norms (written by branch_grad_norm_reduce) and
 * updates the ISV signal bus:
 *   * target[b] (slots 31..35): adaptive-rate EMA toward the current
 *     step's median norm. All four targets track the same median, so
 *     the rescale pass equalises branches.
 *   * limit (slot 35): adaptive-rate EMA toward `max/min` spread.
 *     Hard-clamped to [2.0, 1e4] as a numerical-safety bound only.
 *
 * Adaptive-rate EMA: alpha = clamp(err / (err + baseline), 0.01, 0.3).
 * High alpha when error is large relative to current value (fast
 * response); low alpha when error is small (slow drift). No fixed-rate
 * constant — the rate adapts to observed error size.
 *
 * Single-threaded: the work is a trivial 4-float sort + EMA update.
 * Launched with block=(32,1,1), grid=(1,1,1). Threads 1..31 idle-exit.
 * Kept at 32 threads to match a warp (hardware scheduling alignment).
 */
extern "C" __global__ void grad_balance_isv_update(
    const float* __restrict__ branch_norms_dev,  /* [4] per-branch L2 norms */
    float* __restrict__ isv_signals,              /* device-mapped pinned ISV bus */
    int target_base_idx,                          /* = GRAD_NORM_TARGET_DIR_INDEX (31) */
    int limit_idx                                 /* = GRAD_SCALE_LIMIT_INDEX (35) */
) {
    if (threadIdx.x != 0 || blockIdx.x != 0) return;

    float n0 = branch_norms_dev[BRANCH_DIR];
    float n1 = branch_norms_dev[BRANCH_MAG];
    float n2 = branch_norms_dev[BRANCH_ORD];
    float n3 = branch_norms_dev[BRANCH_URG];

    /* All-zero branch norms → quiescent training step; skip the update
     * to avoid driving targets toward zero and the limit to 1.0 (which
     * would lock the balancer into identity for the rest of the run). */
    float max_n = fmaxf(fmaxf(n0, n1), fmaxf(n2, n3));
    if (max_n < 1e-9f) return;

    float median = branch_grad_balance_median4(n0, n1, n2, n3);
    float min_n  = fminf(fminf(n0, n1), fminf(n2, n3));

    /* Spread: observed max/min ratio. Floor at 1.0 (no asymmetry) and
     * cap at 1e4 (numerical-safety — no amplification beyond 10000×). */
    float spread = (min_n > 1e-9f) ? (max_n / min_n) : 1.0f;
    if (spread < 1.0f) spread = 1.0f;
    if (spread > 1e4f) spread = 1e4f;

    /* Per-branch target: EMA toward median. Adaptive alpha. */
    for (int b = 0; b < 4; b++) {
        float cur = isv_signals[target_base_idx + b];
        float err = fabsf(median - cur);
        float baseline = fmaxf(cur, 1e-6f);
        float alpha = err / (err + baseline);
        alpha = fminf(0.3f, fmaxf(0.01f, alpha));
        float next = (1.0f - alpha) * cur + alpha * median;
        isv_signals[target_base_idx + b] = next;
    }

    /* Scale clamp limit: EMA toward observed spread, hard-bounded
     * [2.0, 1e4] as safety (not tuning). */
    float cur_limit = isv_signals[limit_idx];
    float err_limit = fabsf(spread - cur_limit);
    float baseline_limit = fmaxf(cur_limit, 1.0f);
    float alpha_limit = err_limit / (err_limit + baseline_limit);
    alpha_limit = fminf(0.3f, fmaxf(0.01f, alpha_limit));
    float next_limit = (1.0f - alpha_limit) * cur_limit + alpha_limit * spread;
    if (next_limit < 2.0f)  next_limit = 2.0f;
    if (next_limit > 1e4f)  next_limit = 1e4f;
    isv_signals[limit_idx] = next_limit;
}

extern "C" __global__ void branch_grad_rescale(
    float* __restrict__ grad_buf,
    const int* __restrict__ branch_starts,        /* [4] element offsets */
    const int* __restrict__ branch_lens,          /* [4] element counts */
    const float* __restrict__ branch_norms_dev,   /* [4] L2 norms from reduce pass */
    const float* __restrict__ isv_signals,        /* device-mapped pinned ISV bus */
    int target_base_idx,                          /* = GRAD_NORM_TARGET_DIR_INDEX (31) */
    int limit_idx,                                /* = GRAD_SCALE_LIMIT_INDEX (35) */
    float* __restrict__ branch_scales_dev         /* [4] optional diagnostic output */
) {
    const int branch = blockIdx.y;
    if (branch >= 4) return;

    __shared__ float sh_scale;

    const int tid = threadIdx.x;

    /* Thread 0 reads the ISV target + limit, computes the scale,
     * publishes to shared memory and (optional) diagnostic slot. */
    if (tid == 0) {
        float norm   = branch_norms_dev[branch];
        float target = isv_signals[target_base_idx + branch];
        float limit  = isv_signals[limit_idx];
        if (limit < 1.0f) limit = 1.0f;
        float inv_limit = 1.0f / limit;

        float scale = 1.0f;
        /* Both target and norm must be non-trivial to compute a
         * defensible scale. Quiescent training (all-zero norms) or
         * uninitialised targets (bootstrap 1.0, no EMA update yet)
         * pass through unchanged. */
        if (norm > 1e-9f && target > 1e-9f) {
            scale = target / norm;
            if (scale > limit)     scale = limit;
            if (scale < inv_limit) scale = inv_limit;
        }
        sh_scale = scale;

        /* Only block (0, branch, 0) emits the diagnostic to avoid
         * redundant stores across element-chunk blocks. */
        if (blockIdx.x == 0 && branch_scales_dev != nullptr) {
            branch_scales_dev[branch] = scale;
        }
    }
    __syncthreads();

    const float scale = sh_scale;
    /* Skip the multiply entirely when scale == 1.0f — preserves the
     * original gradient bit-identically on healthy branches. */
    if (scale == 1.0f) return;

    const int start = branch_starts[branch];
    const int len   = branch_lens[branch];
    const int idx   = blockIdx.x * blockDim.x + tid;
    if (idx < len) {
        grad_buf[start + idx] *= scale;
    }
}
