//! SP4 ISV slot indices for signal-driven magnitude bounds.
//!
//! Allocated at indices 131..173 (extending ISV_TOTAL_DIM 131 → 173).
//! Per param-group families use 8 contiguous slots.
//!
//! Layer C #260 (2026-05-01) added two new SP4 slots:
//!   - 171 = `BW_D_H_S2_BOUND_INDEX` (cuBLAS backward-grad sanitiser bound for
//!           `bw_d_h_s2`; replaces SP1-Phase-C `1e6 × ISV[H_S2_RMS_EMA].max(1.0)`).
//!   - 172 = `Q_DIR_GRAD_BOUND_INDEX` (cuBLAS backward-grad sanitiser bound for
//!           `d_value_logits ∪ d_adv_logits`; replaces SP1-Phase-C
//!           `1e6 × ISV[Q_DIR_ABS_REF].max(1.0)`).

/// Number of distinct param groups (DQN trunk/value/branch/IQN/IQL-hi/IQL-lo/attn/curiosity).
pub const SP4_PARAM_GROUP_COUNT: usize = 8;

/// Number of distinct branches for ATOM_POS_BOUND (DQN's 4-branch architecture:
/// direction / magnitude / order / urgency).
pub const SP4_BRANCH_COUNT: usize = 4;

/// p99(|target_q|) — Mech 1 clamp bound.
pub const TARGET_Q_BOUND_INDEX: usize = 131;

/// p99(|atom_positions[branch]|) — Mech 2 clamp bound (per-branch).
/// Indexed [132..136): ATOM_POS_BOUND[branch] = ATOM_POS_BOUND_BASE + branch.
pub const ATOM_POS_BOUND_BASE: usize = 132;

/// p99(|params|) per param group — Mech 9 clamp bound.
/// Indexed [136..144): WEIGHT_BOUND[group] = WEIGHT_BOUND_BASE + group.
pub const WEIGHT_BOUND_BASE: usize = 136;

/// p99(|adam_m|) per param group — Mech 5 diagnostic.
/// Indexed [144..152).
pub const ADAM_M_BOUND_BASE: usize = 144;

/// p99(|adam_v|) per param group — Mech 5 diagnostic.
/// Indexed [152..160).
pub const ADAM_V_BOUND_BASE: usize = 152;

/// EMA of |w·g|/||w||² per param group — AdamW weight_decay arg.
/// Indexed [160..168).
pub const WD_RATE_BASE: usize = 160;

/// p99(grad_norm) — Mech 6 adaptive_clip upper bound.
pub const GRAD_CLIP_BOUND_INDEX: usize = 168;

/// p99(|h_s2|) — Mech 10 clamp bound.
pub const H_S2_BOUND_INDEX: usize = 169;

/// (mean|g|/mean|w|) × entropy_deficit — trunk L1 lambda. Group 0 only.
pub const L1_LAMBDA_TRUNK_INDEX: usize = 170;

/// SP4 Layer C #260 (2026-05-01): backward-grad sanitiser bound for
/// `bw_d_h_s2` buffer (slot 35 NaN-check accomplice).
///
/// Migrated from SP1-Phase-C hardcoded formula
/// `1e6 × ISV[H_S2_RMS_EMA_INDEX].max(1.0)` per
/// `feedback_isv_for_adaptive_bounds`. Producer: `bw_d_h_s2_p99_kernel`
/// (single-buffer 256-thread histogram p99 over `bw_d_h_s2 [B × SH2]`)
/// — see `GpuDqnTrainer::launch_sp4_bw_d_h_s2_p99`. Pearls A+D applied
/// via `apply_pearls_ad_kernel` chained on the same stream. Consumer:
/// `gpu_dqn_trainer.rs::apply_iqn_trunk_gradient` reads
/// `ISV[BW_D_H_S2_BOUND_INDEX].max(EPS_CLAMP_FLOOR)` directly.
pub const BW_D_H_S2_BOUND_INDEX: usize = 171;

/// SP4 Layer C #260 (2026-05-01): backward-grad sanitiser bound for
/// `d_value_logits ∪ d_adv_logits` buffers (slot 24/25 NaN-check
/// accomplices).
///
/// Migrated from SP1-Phase-C hardcoded formula
/// `1e6 × ISV[Q_DIR_ABS_REF_INDEX].max(1.0)` per
/// `feedback_isv_for_adaptive_bounds`. Producer: `q_dir_grad_p99_kernel`
/// (multi-sub-buffer 256-thread histogram p99 with `n_sub=2`, treating
/// the union of `d_value_logits` + `d_adv_logits` as one logical
/// distribution) — see `GpuDqnTrainer::launch_sp4_q_dir_grad_p99`.
/// Pearls A+D applied via `apply_pearls_ad_kernel` chained on the same
/// stream. Consumer: `gpu_dqn_trainer.rs::launch_cublas_backward_to`
/// reads `ISV[Q_DIR_GRAD_BOUND_INDEX].max(EPS_CLAMP_FLOOR)` directly.
pub const Q_DIR_GRAD_BOUND_INDEX: usize = 172;

/// First slot used by SP4 (for sanity-check assertions).
pub const SP4_SLOT_BASE: usize = 131;

/// One past last slot used by SP4. Must equal `ISV_TOTAL_DIM` post-bump.
///
/// Layer C #260 (2026-05-01) extended this from 171 → 173 to add
/// `BW_D_H_S2_BOUND_INDEX=171` + `Q_DIR_GRAD_BOUND_INDEX=172` for the
/// SP1-Phase-C `1e6 × isv` migration.
pub const SP4_SLOT_END: usize = 173;

// ── SP4 Task A14: Pearl C engagement-counter buffer layout ───────────

/// Maximum block grid size per Adam-kernel launch — the engagement
/// counter buffer reserves `MAX_BLOCKS_PER_ADAM` slots per param group,
/// indexed by `blockIdx.x`. 4096 covers up to ≈1M params per group at
/// blockDim=256 threads.
///
/// **Layer B fix-up (2026-05-01)**: bumped 256 → 4096 to accommodate
/// the trunk-extras tail sub-launch (tensors `[33..163)` co-located
/// under DqnTrunk). Production cfg `shared_h2=256, num_atoms=51,
/// adv_h=128, MOE_NUM_EXPERTS=8, MOE_EXPERT_BOTTLENECK=64` puts
/// trunk-extras at ≈600k floats → ≈2400 blocks at blockDim=256, so
/// 4096 leaves a comfortable safety margin (1.7×) for future config
/// growth without wasted memory (each engage slot is 4 bytes × 11
/// groups × 4096 = 180 KiB total — negligible vs. params buffer).
/// Other groups (DqnValue ≈ 4k params, DqnBranches ≈ 80k, IQN/IQL/
/// Attn/Curiosity bounded similarly) easily fit within 4096 blocks.
pub const MAX_BLOCKS_PER_ADAM: usize = 4096;

/// Total length of the Pearl C engagement-counter mapped-pinned buffer
/// (`clamp_engage_per_block_buf`). The 8 main param groups occupy
/// offsets `[0..8 × MAX_BLOCKS_PER_ADAM)`. Curiosity has 4 sub-launches
/// (W1/b1/W2/b2) sharing the same kernel; to avoid block-index
/// collisions between sub-launches, each gets its own offset slot.
/// Sub-launch 0 reuses Curiosity's main group offset (group_idx=7).
/// Sub-launches 1/2/3 occupy extra slots — 3 extra logical "groups"
/// beyond the canonical 8. Total length: (8 + 3) × MAX_BLOCKS_PER_ADAM
/// = 11 × 4096 = 45056 i32s. Host-side
/// `pearl_c_post_adam_engagement_check` for `ParamGroup::Curiosity`
/// sums all 4 sub-launch ranges.
pub const SP4_ENGAGE_EXTRA_CURIOSITY_SUBLAUNCHES: usize = 3;
pub const SP4_ENGAGE_BUF_LEN: usize =
    (SP4_PARAM_GROUP_COUNT + SP4_ENGAGE_EXTRA_CURIOSITY_SUBLAUNCHES) * MAX_BLOCKS_PER_ADAM;

/// Curiosity sub-launch offsets in `clamp_engage_per_block_buf`.
/// Sub-launch 0 (W1) reuses the Curiosity main-group offset (7×256=1792).
/// Sub-launches 1/2/3 (b1/W2/b2) get extra slots 8/9/10 × 256
/// (= 2048/2304/2560).
pub const SP4_ENGAGE_OFFSET_CURIOSITY_W1: i32 = 7 * (MAX_BLOCKS_PER_ADAM as i32);
pub const SP4_ENGAGE_OFFSET_CURIOSITY_B1: i32 = 8 * (MAX_BLOCKS_PER_ADAM as i32);
pub const SP4_ENGAGE_OFFSET_CURIOSITY_W2: i32 = 9 * (MAX_BLOCKS_PER_ADAM as i32);
pub const SP4_ENGAGE_OFFSET_CURIOSITY_B2: i32 = 10 * (MAX_BLOCKS_PER_ADAM as i32);

/// Sentinel `engage_buf_offset` value passed by aux trainers outside
/// the SP4 8-group taxonomy (DT, OFI embed, denoise, recursive_conf,
/// sel). Causes the kernel's per-block writeback to be skipped — Pearl
/// C only tracks the canonical 8 param groups.
pub const SP4_ENGAGE_OFFSET_DISABLED: i32 = -1;

// ── SP4 Task A14: per-Adam-kernel diag_slot allocations ──────────────
// `nan_flags_buf` slots 50-54 record per-Adam-kernel sticky-engagement
// flags (idempotent writes — no race). The buffer is sized 50 today
// (see gpu_dqn_trainer.rs `nan_flags_buf` alloc); A14 grows it to 55.

/// `dqn_adam_update_kernel` (main DQN trainer + 3 post_aux helpers).
pub const SP4_DQN_ADAM_DIAG_SLOT: i32 = 50;

/// `iqn_adam_kernel` (IQN dual-head trainer).
pub const SP4_IQN_ADAM_DIAG_SLOT: i32 = 51;

/// `iql_adam_kernel` (both high- and low-tau IQL trainers share the
/// slot — identical kernel code; per-group separation is via
/// `engage_buf_offset`, not the sticky flag).
pub const SP4_IQL_ADAM_DIAG_SLOT: i32 = 52;

/// `attn_adam_kernel` (GpuAttention + GpuTlob — TLOB shares the cubin).
pub const SP4_ATTN_ADAM_DIAG_SLOT: i32 = 53;

/// `curiosity_adam_step` (4 sub-buffer launches per step).
pub const SP4_CURIOSITY_ADAM_DIAG_SLOT: i32 = 54;

/// One past the last `nan_flags_buf` index used by SP4 — drives the
/// buffer allocation size in `GpuDqnTrainer::new`.
pub const SP4_NAN_FLAGS_END: usize = 55;

/// Convenience: slot index for `WEIGHT_BOUND[group]`.
#[inline]
pub const fn weight_bound(group: usize) -> usize {
    debug_assert!(group < SP4_PARAM_GROUP_COUNT, "weight_bound group out of range");
    WEIGHT_BOUND_BASE + group
}

/// Convenience: slot index for `ADAM_M_BOUND[group]`.
#[inline]
pub const fn adam_m_bound(group: usize) -> usize {
    debug_assert!(group < SP4_PARAM_GROUP_COUNT, "adam_m_bound group out of range");
    ADAM_M_BOUND_BASE + group
}

/// Convenience: slot index for `ADAM_V_BOUND[group]`.
#[inline]
pub const fn adam_v_bound(group: usize) -> usize {
    debug_assert!(group < SP4_PARAM_GROUP_COUNT, "adam_v_bound group out of range");
    ADAM_V_BOUND_BASE + group
}

/// Convenience: slot index for `WD_RATE[group]`.
#[inline]
pub const fn wd_rate(group: usize) -> usize {
    debug_assert!(group < SP4_PARAM_GROUP_COUNT, "wd_rate group out of range");
    WD_RATE_BASE + group
}

/// Convenience: slot index for `ATOM_POS_BOUND[branch]`.
#[inline]
pub const fn atom_pos_bound(branch: usize) -> usize {
    debug_assert!(branch < SP4_BRANCH_COUNT, "atom_pos_bound branch out of range");
    ATOM_POS_BOUND_BASE + branch
}

/// Param-group enum for type-safe oracle launches.
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
#[repr(u32)]
pub enum ParamGroup {
    DqnTrunk = 0,
    DqnValue = 1,
    DqnBranches = 2,
    Iqn = 3,
    IqlHigh = 4,
    IqlLow = 5,
    Attn = 6,
    Curiosity = 7,
}

impl ParamGroup {
    pub const ALL: [ParamGroup; 8] = [
        ParamGroup::DqnTrunk, ParamGroup::DqnValue, ParamGroup::DqnBranches,
        ParamGroup::Iqn, ParamGroup::IqlHigh, ParamGroup::IqlLow,
        ParamGroup::Attn, ParamGroup::Curiosity,
    ];

    #[inline] pub fn idx(self) -> usize { self as usize }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn slot_layout_is_contiguous_and_total_42() {
        // SP4 occupies [131..173), 42 slots total post-Layer-C #260.
        assert_eq!(SP4_SLOT_BASE, 131);
        assert_eq!(SP4_SLOT_END, 173);
        assert_eq!(SP4_SLOT_END - SP4_SLOT_BASE, 42);

        // Per-group families are contiguous and disjoint.
        assert_eq!(WEIGHT_BOUND_BASE - ATOM_POS_BOUND_BASE, 4);
        assert_eq!(ADAM_M_BOUND_BASE - WEIGHT_BOUND_BASE, 8);
        assert_eq!(ADAM_V_BOUND_BASE - ADAM_M_BOUND_BASE, 8);
        assert_eq!(WD_RATE_BASE - ADAM_V_BOUND_BASE, 8);
        assert_eq!(GRAD_CLIP_BOUND_INDEX - WD_RATE_BASE, 8);

        // Final 5 single-output slots: GRAD_CLIP, H_S2, L1_LAMBDA,
        // BW_D_H_S2_BOUND, Q_DIR_GRAD_BOUND.
        assert_eq!(GRAD_CLIP_BOUND_INDEX, 168);
        assert_eq!(H_S2_BOUND_INDEX, 169);
        assert_eq!(L1_LAMBDA_TRUNK_INDEX, 170);
        assert_eq!(BW_D_H_S2_BOUND_INDEX, 171);
        assert_eq!(Q_DIR_GRAD_BOUND_INDEX, 172);

        // Convenience accessors.
        assert_eq!(weight_bound(0), 136);
        assert_eq!(weight_bound(7), 143);
        assert_eq!(atom_pos_bound(0), 132);
        assert_eq!(atom_pos_bound(3), 135);

        // Branch-count constant + non-aliasing guard for atom_pos_bound max.
        assert_eq!(SP4_BRANCH_COUNT, 4);
        assert!(atom_pos_bound(SP4_BRANCH_COUNT - 1) < WEIGHT_BOUND_BASE,
            "atom_pos_bound max must not alias into WEIGHT_BOUND family");
    }

    #[test]
    fn pearl_c_engage_buf_layout() {
        // 8 main groups + 3 extra curiosity sub-launch slots.
        // Layer B fix-up (2026-05-01): MAX_BLOCKS_PER_ADAM grown 256 → 4096
        // to fit trunk-extras tail (≈2400 blocks at production cfg).
        assert_eq!(MAX_BLOCKS_PER_ADAM, 4096);
        assert_eq!(SP4_PARAM_GROUP_COUNT, 8);
        assert_eq!(SP4_ENGAGE_EXTRA_CURIOSITY_SUBLAUNCHES, 3);
        assert_eq!(SP4_ENGAGE_BUF_LEN, 11 * MAX_BLOCKS_PER_ADAM);
        assert_eq!(SP4_ENGAGE_BUF_LEN, 11 * 4096);
        assert_eq!(SP4_ENGAGE_BUF_LEN, 45056);

        // Curiosity sub-launch offsets are distinct + contiguous starting
        // from the canonical Curiosity group offset. They're derived from
        // `MAX_BLOCKS_PER_ADAM` so they auto-track its growth.
        assert_eq!(SP4_ENGAGE_OFFSET_CURIOSITY_W1, 7 * MAX_BLOCKS_PER_ADAM as i32);
        assert_eq!(SP4_ENGAGE_OFFSET_CURIOSITY_B1, 8 * MAX_BLOCKS_PER_ADAM as i32);
        assert_eq!(SP4_ENGAGE_OFFSET_CURIOSITY_W2, 9 * MAX_BLOCKS_PER_ADAM as i32);
        assert_eq!(SP4_ENGAGE_OFFSET_CURIOSITY_B2, 10 * MAX_BLOCKS_PER_ADAM as i32);

        // Highest sub-launch end must fit within the buffer.
        let highest_end = (SP4_ENGAGE_OFFSET_CURIOSITY_B2 as usize) + MAX_BLOCKS_PER_ADAM;
        assert_eq!(highest_end, SP4_ENGAGE_BUF_LEN);

        // Sentinel disabled value.
        assert_eq!(SP4_ENGAGE_OFFSET_DISABLED, -1);

        // Per-Adam-kernel diag slot indices land within the
        // sticky-flag buffer (size SP4_NAN_FLAGS_END=55).
        assert!((SP4_DQN_ADAM_DIAG_SLOT as usize) < SP4_NAN_FLAGS_END);
        assert!((SP4_IQN_ADAM_DIAG_SLOT as usize) < SP4_NAN_FLAGS_END);
        assert!((SP4_IQL_ADAM_DIAG_SLOT as usize) < SP4_NAN_FLAGS_END);
        assert!((SP4_ATTN_ADAM_DIAG_SLOT as usize) < SP4_NAN_FLAGS_END);
        assert!((SP4_CURIOSITY_ADAM_DIAG_SLOT as usize) < SP4_NAN_FLAGS_END);
    }
}
