// crates/ml/src/cuda_pipeline/loss_balance_controller_kernel.cu
//
// SP7 (2026-05-03): per-branch loss-balance controller for CQL and C51 budgets.
//
// Reads per-loss decomp pinned slots [mag_norm, dir_norm, trunk_norm] for
// IQN (reference), CQL (managed, RAW pre-budget — Path A 2026-05-03),
// C51 (managed, pre-budget).
//
// Path A note: prior to 2026-05-03 the kernel read CQL_SX (post-SAXPY,
// budget-scaled = budget × raw_grad) at element offset 6, which created
// a self-perpetuating deadlock at small bootstrap budgets — the controller
// saw ≈0 norm → cold-started → never updated the budget. The new
// `cql_raw_norm_compute` kernel populates element offset 3 (the `cql` slot)
// directly from `cql_grad_scratch`, giving SP7 a budget-independent CQL
// reference. Launcher updated atomically to read offset 3.
// Reads per-branch flatness from ISV[FLATNESS_BASE..+4]. Reads prior
// budgets from ISV[BUDGET_{CQL,C51}_BASE..+4] and prior Wiener state from
// ISV[LB_*_VAR_{CQL,C51}_BASE..+4]. Writes new budgets + new Wiener
// observations to the producer scratch buffer (downstream
// `apply_pearls_ad_kernel` smooths into ISV).
//
// Math (per branch b ∈ {dir, mag, ord, urg}, per head h ∈ {CQL, C51}):
//   target_ratio[CQL][b] = ANCHOR_CQL_RATIO * flatness[b]      (Fix 35, 2026-05-03)
//   target_ratio[C51][b] = ANCHOR_C51_RATIO * flatness[b]
//
// Fix 35 (formula direction flip): the original CQL formula used
// `(1 - flatness[b])`, which inverted the relationship between Q-distribution
// variance and CQL strength. With flatness HIGH meaning "Q has variance"
// (per pearl_2_budget_kernel definition `var_q / σ²`), the original formula
// pushed CQL HIGH when Q was flat — exactly when there is no overconfidence
// for CQL to penalise — creating a death spiral on the magnitude head:
// flat Q → CQL boosted → conservative pull keeps Q flat → repeat. Validated
// 2026-05-03 in T10-v3 train-multi-seed-x7sl2 logs (cql_budget saturated to
// MAX_BUDGET=1.0 on every branch by epoch 1 once IQN-mag/ord/urg woke from
// Fix 34, and val collapsed to active_frac=0). The corrected formula
// scales CQL with flatness: more variance → more overconfidence risk → more
// CQL gradient. C51 formula was already correctly aligned and unchanged.
//   slice_idx = (b == 0) ? 1 : (b == 1) ? 0 : 2  // dir, mag, trunk
//   actual_ratio = h_norm[slice] / iqn_norm[slice]  (both > EPS_DIV via outer guard)
//   correction   = target_ratio[h][b] / max(actual_ratio, EPS_DIV)
//   candidate    = old_budget * correction
//
//   diff       = candidate - old_budget
//   sample_sq  = h_norm[slice]^2
//   wiener_α   = old_diff_var / (old_diff_var + old_sample_var + EPS_DIV)
//   welford_α  = 1 / max(1, epoch_idx_in_fold)     (full update on first
//                                                    active step, falls off
//                                                    as 1/N within the fold)
//   alpha_eff  = max(welford_α, wiener_α)          (hybrid: Welford bounds α
//                                                    from below early; Wiener
//                                                    takes over once accumulated
//                                                    variance is meaningful)
//   alpha_eff  = clamp(alpha_eff, ALPHA_FLOOR, ALPHA_CEIL)
//   new_budget = clamp(old_budget + alpha_eff * diff, EPS_DIV, MAX_BUDGET)
//
// Cold-start (per "Activation flag" spec, 2026-05-03):
//   if h_norm[slice] < EPS_DIV || iqn_norm[slice] < EPS_DIV:
//       if prior was_active >= 0.5:
//           # Transient grad gate-off; hold prior budget steady.
//           new_budget    = isv[budget_isv_base + branch]
//           scratch_active = 1.0   (still active)
//       else:
//           # Genuine cold start; let consumer's bootstrap handle this step.
//           # Don't touch scratch_budget/diff_var/sample_var (apply_pearls_ad
//           # short-circuits on step_obs == 0.0).
//           scratch_active = 0.0   (no-op via short-circuit)
//           return
//
// SP7 activation-flag fix (2026-05-03): the controller writes 1.0 into
// `scratch_active` when the active path computes a real budget update OR
// when a transient cold-start hits a previously-active branch (prior
// budget held verbatim, controller stays "engaged"). It writes 0.0 in
// the genuine cold-start branch (was_active = 0). apply_pearls_ad_kernel
// short-circuits step_obs == 0.0, so writing 0 leaves the activation
// slot at its current value; activation is therefore monotonic per fold.
// FoldReset zeroes the slots so the new fold re-bootstraps cleanly.
//
// SP8 (2026-05-03): MAX_BUDGET no longer hardcoded. The cap is now
// per-(head, branch) read from `ISV[LB_MAX_BUDGET_{CQL,C51}_BASE+b]`,
// driven by `loss_balance_max_budget_compute_kernel` from the
// `train_active_frac` canary per `pearl_controller_anchors_isv_driven.md`
// and `feedback_isv_for_adaptive_bounds.md`. Replaces the prior
// `const float MAX_BUDGET = 1.0f` regime-encoded constant.
//
// Single block, 8 threads (2 heads × 4 branches). No atomicAdd (feedback_no_atomicadd).
// __threadfence_system() after writes.

#include <cuda_runtime.h>

extern "C" __global__ void loss_balance_controller_update(
    // Pointers into the shared `grad_decomp_result_pinned` buffer at the
    // 3-float component offsets (the launcher computes these by adding
    // (offset_elems * sizeof(float)) to grad_decomp_result_dev_ptr):
    //   iqn_decomp = base + 0  (slot iqn,         layout [mag, dir, trunk])
    //   cql_decomp = base + 12 (slot cql,    Path A 2026-05-03 — raw norm,
    //                           layout [mag, dir, trunk])
    //   c51_decomp = base + 36 (slot c51,         layout [mag, dir, trunk])
    const float* __restrict__ iqn_decomp,
    const float* __restrict__ cql_decomp,
    const float* __restrict__ c51_decomp,
    // ISV signal bus (read-only).
    const float* __restrict__ isv_signals,
    int   flatness_isv_base,                 // FLATNESS_BASE = 206
    int   budget_cql_isv_base,               // BUDGET_CQL_BASE = 198
    int   budget_c51_isv_base,               // BUDGET_C51_BASE = 190
    int   diff_var_cql_isv_base,             // LB_DIFF_VAR_CQL_BASE = 297
    int   sample_var_cql_isv_base,           // LB_SAMPLE_VAR_CQL_BASE = 301
    int   diff_var_c51_isv_base,             // LB_DIFF_VAR_C51_BASE = 305
    int   sample_var_c51_isv_base,           // LB_SAMPLE_VAR_C51_BASE = 309
    int   active_cql_isv_base,               // LB_CQL_ACTIVE_BASE = 313
    int   active_c51_isv_base,               // LB_C51_ACTIVE_BASE = 317
    // SP8 (2026-05-03): ISV-driven adaptive MAX_BUDGET cap per head/branch.
    // Replaces the prior hardcoded `const float MAX_BUDGET = 1.0f` constant
    // (regime-encoded) per `feedback_isv_for_adaptive_bounds.md` and
    // `pearl_controller_anchors_isv_driven.md`. Producer is
    // `loss_balance_max_budget_compute_kernel` (driven by the
    // `train_active_frac` canary). Layout per branch ∈ [0..4):
    //   isv_signals[max_budget_cql_isv_base + branch]
    //   isv_signals[max_budget_c51_isv_base + branch]
    int   max_budget_cql_isv_base,           // LB_MAX_BUDGET_CQL_BASE = 322
    int   max_budget_c51_isv_base,           // LB_MAX_BUDGET_C51_BASE = 326
    int   epoch_idx_isv_index,               // EPOCH_IDX_INDEX = 39 (per-fold counter)
    // Producer scratch buffer (kernel writes here; apply_pearls_ad_kernel
    // smooths these into ISV downstream).
    float* __restrict__ scratch_out,
    int   scratch_budget_cql_base,
    int   scratch_budget_c51_base,
    int   scratch_diff_var_cql_base,
    int   scratch_sample_var_cql_base,
    int   scratch_diff_var_c51_base,
    int   scratch_sample_var_c51_base,
    int   scratch_active_cql_base,
    int   scratch_active_c51_base
) {
    // ── Invariant 1 anchors ──────────────────────────────────────────
    // Modulated per-branch by flatness (spec section "Math").
    const float ANCHOR_CQL_RATIO    = 2.0f;
    const float ANCHOR_C51_RATIO    = 1.0f;
    // Numerical-stability anchors (Wiener α can mathematically span [0,1]).
    const float ALPHA_FLOOR         = 1e-4f;
    const float ALPHA_CEIL          = 0.5f;
    const float EPS_DIV             = 1e-8f;
    // SP8 (2026-05-03): MAX_BUDGET is now per-(head, branch) and signal-
    // driven via `loss_balance_max_budget_compute_kernel` →
    // ISV[LB_MAX_BUDGET_*_BASE+b]. Read inline below where `new_budget` is
    // clamped. Replaces `const float MAX_BUDGET = 1.0f` per
    // `feedback_isv_for_adaptive_bounds.md` and Fix 36.
    // Cold-start basis — must match consumer-side bootstrap in
    // fused_training.rs (CQL_BOOTSTRAP_BUDGET, C51_BOOTSTRAP_BUDGET).
    // Used only for the seed-from-cold-start branch when was_active = 0
    // transitions to "first active step" (norms populated, old_budget = 0).
    const float COLD_START_FLOOR_CQL = 0.02f;
    const float COLD_START_FLOOR_C51 = 0.05f;
    // Activation discriminator. Binary threshold (not a tuned knob); slot
    // values are written as 0.0 or 1.0 and smoothed monotonically.
    const float ACTIVE_THRESHOLD     = 0.5f;

    // ── Thread layout: head × branch ─────────────────────────────────
    int tid = threadIdx.x;
    if (blockIdx.x != 0 || tid >= 8) return;   // 8 = 2 heads × 4 branches

    int head   = tid / 4;       // 0 = CQL, 1 = C51
    int branch = tid % 4;       // 0..3 (dir, mag, ord, urg)

    // ── Slice mapping (spec "Slice→branch mapping") ──────────────────
    // grad_decomp pinned layout is [mag, dir, trunk] per component.
    // branch=0 → dir slice [1]; branch=1 → mag slice [0];
    // branch=2,3 → trunk slice [2].
    int slice_idx = (branch == 0) ? 1 : (branch == 1) ? 0 : 2;

    // ── Read the slice norms once ────────────────────────────────────
    float iqn_n   = iqn_decomp[slice_idx];
    float h_n     = (head == 0) ? cql_decomp[slice_idx] : c51_decomp[slice_idx];

    // ── Read flatness for this branch ────────────────────────────────
    float flat_b  = isv_signals[flatness_isv_base + branch];
    flat_b        = fminf(1.0f, fmaxf(0.0f, flat_b));   // safety clamp

    // ── Compute per-branch flatness-modulated target ratio ───────────
    // Fix 35 (2026-05-03): both heads now scale UP with flatness. When Q has
    // variance (flatness ≈ 1) overconfidence risk is high → both CQL and C51
    // budgets grow. When Q is flat (flatness ≈ 0) there is nothing to
    // penalise / nothing to distribute → both budgets shrink. The prior CQL
    // formula `(1 - flat_b)` created a death spiral on magnitude (flat Q →
    // boosted CQL → pulled Q flatter). See kernel header.
    float target_ratio = (head == 0)
        ? ANCHOR_CQL_RATIO * flat_b
        : ANCHOR_C51_RATIO * flat_b;

    // ── Read prior budget + per-branch Wiener state + activation flag ─
    int budget_isv_base       = (head == 0) ? budget_cql_isv_base       : budget_c51_isv_base;
    int diff_var_isv_base     = (head == 0) ? diff_var_cql_isv_base     : diff_var_c51_isv_base;
    int sample_var_isv_base   = (head == 0) ? sample_var_cql_isv_base   : sample_var_c51_isv_base;
    int active_isv_base       = (head == 0) ? active_cql_isv_base       : active_c51_isv_base;
    int max_budget_isv_base   = (head == 0) ? max_budget_cql_isv_base   : max_budget_c51_isv_base;
    int scratch_budget_base   = (head == 0) ? scratch_budget_cql_base   : scratch_budget_c51_base;
    int scratch_diff_base     = (head == 0) ? scratch_diff_var_cql_base : scratch_diff_var_c51_base;
    int scratch_sample_base   = (head == 0) ? scratch_sample_var_cql_base : scratch_sample_var_c51_base;
    int scratch_active_base   = (head == 0) ? scratch_active_cql_base   : scratch_active_c51_base;
    float cold_start_basis    = (head == 0) ? COLD_START_FLOOR_CQL      : COLD_START_FLOOR_C51;

    float old_budget     = isv_signals[budget_isv_base + branch];
    float old_diff_var   = isv_signals[diff_var_isv_base + branch];
    float old_sample_var = isv_signals[sample_var_isv_base + branch];
    float was_active     = isv_signals[active_isv_base + branch];

    // SP8 (2026-05-03): per-(head, branch) MAX_BUDGET cap from ISV. Producer
    // (`loss_balance_max_budget_compute_kernel`) writes 0.01 (collapse) ..
    // 1.0 (healthy) based on `train_active_frac` canary; Pearls A+D smooth
    // into ISV[LB_MAX_BUDGET_{CQL,C51}_BASE+b]. Defensive clamp to the
    // structural envelope covers the Pearl A bootstrap window before the
    // first observation lands (sentinel 0 → first-observation replacement
    // produces the real value next step).
    float MAX_BUDGET = isv_signals[max_budget_isv_base + branch];
    if (MAX_BUDGET < EPS_DIV) MAX_BUDGET = EPS_DIV;
    if (MAX_BUDGET > 1.0f)    MAX_BUDGET = 1.0f;

    // ── Branch on cold-start regime ──────────────────────────────────
    if (h_n < EPS_DIV || iqn_n < EPS_DIV) {
        // Warmup hasn't fired or the loss head is gated off.
        if (was_active >= ACTIVE_THRESHOLD) {
            // Transient grad gate-off on a previously-active branch.
            // Hold prior budget verbatim and re-assert active=1; don't
            // drift the budget on what is purely a measurement gap.
            scratch_out[scratch_budget_base + branch] = old_budget;
            scratch_out[scratch_active_base + branch] = 1.0f;
            // Do NOT update diff_var/sample_var: there's no new sample.
            // (Writing 0 to those slots is a no-op via apply_pearls_ad's
            //  step_obs == 0.0 short-circuit; we must still write _something_
            //  per the launcher loop, but 0 keeps Wiener state untouched.)
            scratch_out[scratch_diff_base + branch]   = 0.0f;
            scratch_out[scratch_sample_base + branch] = 0.0f;
        } else {
            // Genuine cold start. Don't write a budget (apply_pearls_ad
            // short-circuits on 0.0 → leaves slot untouched → consumer's
            // bootstrap handles it). active stays 0.
            scratch_out[scratch_budget_base + branch] = 0.0f;
            scratch_out[scratch_active_base + branch] = 0.0f;
            scratch_out[scratch_diff_base + branch]   = 0.0f;
            scratch_out[scratch_sample_base + branch] = 0.0f;
        }
        __threadfence_system();
        return;
    }

    // ── Active path: norms populated, controller engages ─────────────
    float actual_ratio = h_n / iqn_n;
    float correction   = target_ratio / fmaxf(actual_ratio, EPS_DIV);

    float new_budget;
    float diff;
    float sample_sq = h_n * h_n;

    if (old_budget < EPS_DIV) {
        // Sentinel-0 read (cold start or fold boundary). The norm we
        // observed was produced under cql/c51_budget = cold_start_basis
        // (consumer-side bootstrap). Seed the controller at the value
        // that achieves target_ratio next step.
        float seed = cold_start_basis * correction;
        new_budget = fminf(MAX_BUDGET, fmaxf(EPS_DIV, seed));
        diff = new_budget - cold_start_basis;
    } else {
        float candidate = old_budget * correction;
        diff = candidate - old_budget;

        // Welford-α: bounds α from below as 1/N within the fold, giving
        // a full update on the first active step (epoch_idx = 0 → 1.0)
        // and falling off as the fold progresses. Wiener-α takes over
        // once accumulated variance is meaningful.
        float epoch_idx     = isv_signals[epoch_idx_isv_index];
        float welford_alpha = 1.0f / fmaxf(1.0f, epoch_idx);

        float wiener_num   = old_diff_var;
        float wiener_den   = old_diff_var + old_sample_var + EPS_DIV;
        float wiener_alpha = wiener_num / wiener_den;

        float alpha_eff    = fmaxf(welford_alpha, wiener_alpha);
        alpha_eff          = fminf(ALPHA_CEIL, fmaxf(ALPHA_FLOOR, alpha_eff));

        new_budget = old_budget + alpha_eff * diff;
        new_budget = fminf(MAX_BUDGET, fmaxf(EPS_DIV, new_budget));
    }

    // ── Write the four outputs to scratch (apply_pearls_ad smooths) ───
    scratch_out[scratch_budget_base + branch] = new_budget;
    scratch_out[scratch_diff_base + branch]   = diff * diff;        // raw observation
    scratch_out[scratch_sample_base + branch] = sample_sq;          // raw observation
    scratch_out[scratch_active_base + branch] = 1.0f;               // controller active

    __threadfence_system();
}
