#![allow(unsafe_code)] // Required for CUDA kernel launches.

//! `GpuAuxHeads` — Plan 4 Task 6 Commit A.
//!
//! Multi-task auxiliary heads orchestrator (E.6). Owns the per-head
//! `CudaFunction` handles for both forward / backward / loss-reduce kernels
//! plus the regime-label builder, and exposes thin `pub(crate)` wrappers
//! that mirror the kernel ABIs in row-major device-pointer form.
//!
//! NO PRODUCTION CALLERS in this commit — Commit B's training-loop wire-up
//! invokes these methods. The orchestrator is built as additive scaffolding
//! per the same convention used for `gpu_grn.rs` (1B-i / 2c.2 pattern):
//! the kernels and orchestrator land first, the wire-up commit lights them
//! up.
//!
//! # Architecture
//!
//! Both heads share the `Linear(SH2 → 32) → ELU → Linear(32 → K)` shape:
//!
//!   - Next-bar regression head: K = 1, MSE loss.
//!   - Regime classification head: K = 5, cross-entropy loss.
//!
//! Param tensors live in the trainer's flat `params_buf` at indices
//! `[119..127)` (8 tensors total — w1/b1/w2/b2 per head). See
//! `compute_param_sizes()` in `gpu_dqn_trainer.rs` for the exact layout.
//!
//! # Pearls applied
//!
//! - `feedback_no_atomicadd.md` — backward kernels emit per-sample partials;
//!   the trailing `aux_param_grad_reduce` kernel collapses along the batch
//!   dim with shmem-tree reduction (no atomicAdd anywhere).
//! - `pearl_cold_path_no_exception_to_gpu_drives.md` — every reduction is
//!   GPU-side; the orchestrator does zero CPU/DtoH work.
//! - `feedback_wire_everything_up.md` — this orchestrator is intentionally
//!   built as a pre-wire scaffold for Commit B, mirroring the same
//!   commit-by-commit ordering used for VSN (1B-ii orchestrator → 1B-iii
//!   wire-up). The Commit B wire-up retires the orphan-by-design status.

use std::sync::Arc;

use cudarc::driver::{CudaFunction, CudaStream, LaunchConfig, PushKernelArg};

use crate::MLError;

// ── Cubin embedding (kernels live in aux_heads_kernel.cu / aux_heads_loss_ema_kernel.cu) ──

use super::gpu_dqn_trainer::{AUX_HEADS_CUBIN, AUX_HEADS_LOSS_EMA_CUBIN};

/// Hidden dim shared by both auxiliary heads (Linear_1 output width).
/// Kept in sync with the `AUX_HIDDEN_DIM` macro at the top of
/// `aux_heads_kernel.cu`.
///
/// **Enlarged 32 → 128 (2026-05-13)** as part of the aux-head capacity
/// upgrade after Path C investigation revealed the aux head's 28%
/// accuracy at H=60 bars. The prior 32-unit bottleneck was severely
/// under-capacity for a 256-dim input → 2-way classification task. The
/// new 128 hidden dim gives 4× capacity while keeping memory footprint
/// under control. Architecture:
///   256 (h_s2_aux input) → Linear → 128 (hidden) → ELU → Linear → 2 (logits)
///
/// **Memory note**: initial attempt at 256 (8×) OOM'd L40S. Aux backward
/// stores `dW1_partial [B, H, SH2]` = batch_size × H × 256 floats =
/// 16384 × 256 × 256 × 4 bytes = 4GB per head — total 8GB for both heads,
/// which pushed total VRAM over the 48GB L40S budget when combined with
/// the rollout's `next_states[~4GB]` and other DQN buffers. At H=128,
/// the partial buffer is 2GB per head = 4GB total, fitting comfortably.
///
/// Side effects:
/// - Param buffer grows: aux_nb_w1/aux_rg_w1 went from [32, 256] to
///   [128, 256] each → +24K params per head → +48K params total.
/// - Fingerprint changes → checkpoint incompatibility (intentional,
///   fresh-start training).
/// - Thread utilisation improves: AUX_BLOCK=256 threads × H=128 hidden
///   = 2 threads per hidden unit (vs 8 threads/unit at H=32).
pub(crate) const AUX_HIDDEN_DIM: usize = 128;

/// Output cardinality of the next-bar direction classification head
/// (SP13 B1.1a, 2026-05-05). Flipped 1 → 2: the head was a K=1 MSE
/// regression head; B1.1a flips it to a K=2 softmax classifier reading
/// i32 direction labels in {-1, 0, 1} where -1 = mask/skip. The
/// per-class softmax is structurally bounded in [0, 1] so downstream
/// ISV consumers (slot 375) see naturally-bounded inputs without a
/// runtime tanh squash per `pearl_bounded_modifier_outputs_require_structural_activation`.
pub(crate) const AUX_NEXT_BAR_K: usize = 2;

/// SP22 H6 vNext Phase A3+ (2026-05-13): output cardinality of the
/// trade-outcome aux head. K=3 softmax over `{Profit=0, Stop=1,
/// Timeout=2}`. Replaces the per-bar next-bar direction prediction
/// (K=2) with a per-trade-close 3-way outcome prediction that ALIGNS
/// with the system's actual decision horizon (multi-bar trade
/// lifecycle). The architectural mismatch — per-bar prediction vs
/// multi-bar decisions — was the root cause of the H6 Phase 3
/// falsification (smoke train-xrkb7: WR=0.4346 with K=2 head fully
/// active vs baseline 0.4338-0.4346). See
/// `docs/plans/2026-05-14-sp22-h6-vNext-trade-outcome-aux.md`.
pub(crate) const AUX_OUTCOME_K: usize = 3;

/// Output cardinality of the regime classification head (5-way softmax).
pub(crate) const AUX_REGIME_K: usize = 5;

/// Block dim used by the per-sample forward / backward kernels.
const AUX_BLOCK: u32 = 256;

/// Forward orchestrator — owns the 5 forward / loss-reduce / label-builder
/// kernel handles for both auxiliary heads.
///
/// Dropping this struct releases the underlying `CudaFunction` handles via
/// cudarc's RAII; nothing else needs explicit teardown.
#[allow(missing_debug_implementations)]
pub(crate) struct AuxHeadsForwardOps {
    next_bar_forward_kernel: CudaFunction,
    regime_forward_kernel: CudaFunction,
    next_bar_loss_reduce_kernel: CudaFunction,
    regime_loss_reduce_kernel: CudaFunction,
    regime_label_kernel: CudaFunction,
    /// Aux loss EMA producer kernel (single thread, single block; mirrors
    /// `h_s2_rms_ema_update`'s shape). Loaded from `AUX_HEADS_LOSS_EMA_CUBIN`.
    ///
    /// SP13 B1.0 (2026-05-05): the sibling label-scale EMA kernel was
    /// retired together with the ISV[117] label-scale divisor in
    /// `aux_next_bar_loss_reduce` / `aux_next_bar_backward`. Labels are
    /// already z-normalised at the data layer; B1.1 will replace MSE
    /// with CE entirely.
    loss_ema_kernel: CudaFunction,
}

impl AuxHeadsForwardOps {
    /// Load all forward / reduce / EMA kernel handles from the precompiled
    /// cubins. Mirrors `GrnBlock::new`'s loader pattern.
    pub(crate) fn new(stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let context = stream.context();

        let aux_module = context
            .load_cubin(AUX_HEADS_CUBIN.to_vec())
            .map_err(|e| MLError::ModelError(format!("aux_heads cubin load: {e}")))?;
        let next_bar_forward_kernel = aux_module
            .load_function("aux_next_bar_forward")
            .map_err(|e| MLError::ModelError(format!("aux_next_bar_forward load: {e}")))?;
        let regime_forward_kernel = aux_module
            .load_function("aux_regime_forward")
            .map_err(|e| MLError::ModelError(format!("aux_regime_forward load: {e}")))?;
        let next_bar_loss_reduce_kernel = aux_module
            .load_function("aux_next_bar_loss_reduce")
            .map_err(|e| MLError::ModelError(format!("aux_next_bar_loss_reduce load: {e}")))?;
        let regime_loss_reduce_kernel = aux_module
            .load_function("aux_regime_loss_reduce")
            .map_err(|e| MLError::ModelError(format!("aux_regime_loss_reduce load: {e}")))?;
        let regime_label_kernel = aux_module
            .load_function("aux_regime_label_from_states")
            .map_err(|e| MLError::ModelError(format!("aux_regime_label_from_states load: {e}")))?;

        let ema_module = context
            .load_cubin(AUX_HEADS_LOSS_EMA_CUBIN.to_vec())
            .map_err(|e| MLError::ModelError(format!("aux_heads_loss_ema cubin load: {e}")))?;
        let loss_ema_kernel = ema_module
            .load_function("aux_heads_loss_ema_update")
            .map_err(|e| MLError::ModelError(format!("aux_heads_loss_ema_update load: {e}")))?;

        Ok(Self {
            next_bar_forward_kernel,
            regime_forward_kernel,
            next_bar_loss_reduce_kernel,
            regime_loss_reduce_kernel,
            regime_label_kernel,
            loss_ema_kernel,
        })
    }

    /// Launch `aux_next_bar_forward`: `Linear(h_s2) → ELU → Linear → softmax`.
    ///
    /// SP13 B1.1a (2026-05-05): K-flipped 1 → 2 with softmax tile output.
    /// Forward writes BOTH the `[B, K]` logits buffer (saved for backward
    /// re-read; the previous regression-mode design only saved the scalar
    /// prediction) AND the `[B, K]` softmax tile (the new structurally-
    /// bounded prediction surface read by 3 downstream consumers: loss
    /// reduce, backward, dir-acc / isv-tanh producers).
    ///
    /// Caller-owned buffers (raw `u64` device pointers for graph-capture
    /// safety; mirrors `gpu_grn.rs::forward_raw`):
    ///   * `h_s2_aux_ptr`    — `[B, SH2]` row-major (SP14-C.5b: aux trunk
    ///                         output, NOT Q's GRN trunk `save_h_s2`).
    ///   * `w1_ptr`/`b1_ptr` — `[H, SH2]` / `[H]`. Slice from `params_buf[119]`/`[120]`.
    ///   * `w2_ptr`/`b2_ptr` — `[K, H]` / `[K]`.   Slice from `params_buf[121]`/`[122]`.
    ///   * `hidden_out_ptr`  — `[B, H]` SAVED post-ELU buffer (consumed by backward).
    ///   * `logits_out_ptr`  — `[B, K]` SAVED logits (consumed by backward
    ///                         only; the consumers below read the softmax tile).
    ///   * `softmax_out_ptr` — `[B, K]` SAVED softmax tile (consumed by
    ///                         loss reduce + backward + dir-acc + isv-tanh).
    ///
    /// `k` is `AUX_NEXT_BAR_K = 2` in B1.1a; passed as runtime arg for
    /// kernel-signature stability (mirrors `forward_regime`'s K-runtime
    /// pattern — a future K bump won't require an ABI change).
    ///
    /// Block: `AUX_BLOCK` threads. Grid: `B` blocks.
    /// Shared mem: `(H + K) * sizeof(f32)` bytes (h cache + softmax temp).
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn forward_next_bar(
        &self,
        stream: &Arc<CudaStream>,
        h_s2_aux_ptr: u64,
        w1_ptr: u64,
        b1_ptr: u64,
        w2_ptr: u64,
        b2_ptr: u64,
        b: usize,
        sh2: usize,
        k: usize,
        hidden_out_ptr: u64,
        logits_out_ptr: u64,
        softmax_out_ptr: u64,
    ) -> Result<(), MLError> {
        let b_i32 = b as i32;
        let sh2_i32 = sh2 as i32;
        let k_i32 = k as i32;
        let smem_bytes = (AUX_HIDDEN_DIM as u32 + k as u32) * std::mem::size_of::<f32>() as u32;
        unsafe {
            stream
                .launch_builder(&self.next_bar_forward_kernel)
                .arg(&h_s2_aux_ptr)
                .arg(&w1_ptr)
                .arg(&b1_ptr)
                .arg(&w2_ptr)
                .arg(&b2_ptr)
                .arg(&b_i32)
                .arg(&sh2_i32)
                .arg(&k_i32)
                .arg(&hidden_out_ptr)
                .arg(&logits_out_ptr)
                .arg(&softmax_out_ptr)
                .launch(LaunchConfig {
                    grid_dim: (b as u32, 1, 1),
                    block_dim: (AUX_BLOCK, 1, 1),
                    shared_mem_bytes: smem_bytes,
                })
                .map_err(|e| MLError::ModelError(format!("aux_next_bar_forward: {e}")))?;
        }
        Ok(())
    }

    /// Launch `aux_regime_forward`: `Linear(h_s2) → ELU → Linear → logits[K]`.
    ///
    /// `K` is `AUX_REGIME_K = 5` in this commit's design; passed as runtime
    /// arg for kernel-signature stability.
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn forward_regime(
        &self,
        stream: &Arc<CudaStream>,
        h_s2_aux_ptr: u64,
        w1_ptr: u64,
        b1_ptr: u64,
        w2_ptr: u64,
        b2_ptr: u64,
        b: usize,
        sh2: usize,
        k: usize,
        hidden_out_ptr: u64,
        logits_out_ptr: u64,
    ) -> Result<(), MLError> {
        let b_i32 = b as i32;
        let sh2_i32 = sh2 as i32;
        let k_i32 = k as i32;
        let smem_bytes = (AUX_HIDDEN_DIM as u32) * std::mem::size_of::<f32>() as u32;
        unsafe {
            stream
                .launch_builder(&self.regime_forward_kernel)
                .arg(&h_s2_aux_ptr)
                .arg(&w1_ptr)
                .arg(&b1_ptr)
                .arg(&w2_ptr)
                .arg(&b2_ptr)
                .arg(&b_i32)
                .arg(&sh2_i32)
                .arg(&k_i32)
                .arg(&hidden_out_ptr)
                .arg(&logits_out_ptr)
                .launch(LaunchConfig {
                    grid_dim: (b as u32, 1, 1),
                    block_dim: (AUX_BLOCK, 1, 1),
                    shared_mem_bytes: smem_bytes,
                })
                .map_err(|e| MLError::ModelError(format!("aux_regime_forward: {e}")))?;
        }
        Ok(())
    }

    /// Launch `aux_next_bar_loss_reduce`: scalar mean cross-entropy +
    /// scalar valid-count.
    ///
    /// SP13 B1.1a (2026-05-05): MSE → softmax CE flip. Reads the `[B, K]`
    /// softmax tile (produced by `forward_next_bar`) and the `[B] i32`
    /// labels buffer (`-1 / 0..K`; -1 is the producer's mask sentinel).
    /// Writes mean CE = `numerator / max(B_valid, 1)` to `loss_out[1]`
    /// AND the `B_valid` count to `valid_count_out[1]` for the backward
    /// kernel to consume (so loss + backward share the same divisor —
    /// derivatives of the same scalar function).
    ///
    /// Two parallel partial-reduction strips (loss numerator + valid
    /// count); shmem matches the K-class `regime_loss_reduce` shape.
    pub(crate) fn next_bar_loss_reduce(
        &self,
        stream: &Arc<CudaStream>,
        softmax_ptr: u64,
        labels_ptr: u64,
        b: usize,
        k: usize,
        loss_out_ptr: u64,
        valid_count_out_ptr: u64,
    ) -> Result<(), MLError> {
        let b_i32 = b as i32;
        let k_i32 = k as i32;
        // Two parallel partial-reduction strips: loss + valid-count. See
        // `aux_next_bar_loss_reduce` shmem layout in the kernel.
        let smem_bytes = 2 * AUX_BLOCK * std::mem::size_of::<f32>() as u32;
        unsafe {
            stream
                .launch_builder(&self.next_bar_loss_reduce_kernel)
                .arg(&softmax_ptr)
                .arg(&labels_ptr)
                .arg(&b_i32)
                .arg(&k_i32)
                .arg(&loss_out_ptr)
                .arg(&valid_count_out_ptr)
                .launch(LaunchConfig {
                    grid_dim: (1, 1, 1),
                    block_dim: (AUX_BLOCK, 1, 1),
                    shared_mem_bytes: smem_bytes,
                })
                .map_err(|e| MLError::ModelError(format!("aux_next_bar_loss_reduce: {e}")))?;
        }
        Ok(())
    }

    /// Launch `aux_regime_loss_reduce`: scalar mean CE + scalar mean accuracy.
    pub(crate) fn regime_loss_reduce(
        &self,
        stream: &Arc<CudaStream>,
        logits_ptr: u64,
        label_ptr: u64,
        b: usize,
        k: usize,
        loss_out_ptr: u64,
        correct_out_ptr: u64,
    ) -> Result<(), MLError> {
        let b_i32 = b as i32;
        let k_i32 = k as i32;
        // Two parallel partial-reduction strips: loss + accuracy. See
        // `aux_regime_loss_reduce` shmem layout in the kernel.
        let smem_bytes = 2 * AUX_BLOCK * std::mem::size_of::<f32>() as u32;
        unsafe {
            stream
                .launch_builder(&self.regime_loss_reduce_kernel)
                .arg(&logits_ptr)
                .arg(&label_ptr)
                .arg(&b_i32)
                .arg(&k_i32)
                .arg(&loss_out_ptr)
                .arg(&correct_out_ptr)
                .launch(LaunchConfig {
                    grid_dim: (1, 1, 1),
                    block_dim: (AUX_BLOCK, 1, 1),
                    shared_mem_bytes: smem_bytes,
                })
                .map_err(|e| MLError::ModelError(format!("aux_regime_loss_reduce: {e}")))?;
        }
        Ok(())
    }

    /// Launch `aux_regime_label_from_states`: discretise the regime score
    /// channel of the state vector into `[B] uint8` regime labels in `[0, K)`.
    ///
    /// `regime_score_idx` is a feature offset into the row-major `[B, F]`
    /// state buffer (typically `ml_core::state_layout::REGIME_STABILITY_INDEX`,
    /// resolved by the caller).
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn regime_label_from_states(
        &self,
        stream: &Arc<CudaStream>,
        states_ptr: u64,
        b: usize,
        f: usize,
        regime_score_idx: usize,
        k: usize,
        label_out_ptr: u64,
    ) -> Result<(), MLError> {
        let b_i32 = b as i32;
        let f_i32 = f as i32;
        let r_i32 = regime_score_idx as i32;
        let k_i32 = k as i32;
        let blocks = ((b as u32) + AUX_BLOCK - 1) / AUX_BLOCK;
        unsafe {
            stream
                .launch_builder(&self.regime_label_kernel)
                .arg(&states_ptr)
                .arg(&b_i32)
                .arg(&f_i32)
                .arg(&r_i32)
                .arg(&k_i32)
                .arg(&label_out_ptr)
                .launch(LaunchConfig {
                    grid_dim: (blocks, 1, 1),
                    block_dim: (AUX_BLOCK, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("aux_regime_label_from_states: {e}")))?;
        }
        Ok(())
    }

    /// Launch `aux_heads_loss_ema_update`: cold-path-cadence step
    /// observation for both scalar loss buffers into
    /// `producer_step_scratch_buf[scratch_idx_next_bar..+1)` /
    /// `[scratch_idx_regime]`. Single-thread, single-block — mirrors
    /// `h_s2_rms_ema_update`'s retrofit footprint.
    ///
    /// SP4 Layer A Task A13.1 retrofit (2026-05-01): kernel writes step
    /// observations only; the host launcher (`launch_aux_heads_loss_ema`
    /// on `GpuDqnTrainer`) syncs the stream and applies Pearls A+D
    /// host-side via `pearls_ad_update`.
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn launch_loss_ema(
        &self,
        stream: &Arc<CudaStream>,
        next_bar_loss_scalar_ptr: u64,
        regime_loss_scalar_ptr: u64,
        scratch_dev_ptr: u64,
        scratch_idx_next_bar: i32,
        scratch_idx_regime: i32,
    ) -> Result<(), MLError> {
        unsafe {
            stream
                .launch_builder(&self.loss_ema_kernel)
                .arg(&next_bar_loss_scalar_ptr)
                .arg(&regime_loss_scalar_ptr)
                .arg(&scratch_dev_ptr)
                .arg(&scratch_idx_next_bar)
                .arg(&scratch_idx_regime)
                .launch(LaunchConfig {
                    grid_dim: (1, 1, 1),
                    block_dim: (1, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("aux_heads_loss_ema_update: {e}")))?;
        }
        Ok(())
    }

}

/// Backward orchestrator — owns the 3 backward / final-reduce kernel
/// handles. Held separate from the forward ops because Commit B's wire-up
/// will dispatch forward + backward at distinct points in the training
/// step (forward runs alongside the main encoder forward; backward runs
/// after the main loss backward and SAXPY-accumulates `dh_s2` into the
/// trunk's main-path gradient).
#[allow(missing_debug_implementations)]
pub(crate) struct AuxHeadsBackwardOps {
    next_bar_backward_kernel: CudaFunction,
    regime_backward_kernel: CudaFunction,
    /// Final-reduce kernel for the per-sample param-grad partials emitted
    /// by both backward kernels. `aux_param_grad_reduce(partials [B, P], P,
    /// final [P])`.
    param_grad_reduce_kernel: CudaFunction,
}

impl AuxHeadsBackwardOps {
    /// Load the 3 backward kernel handles from the precompiled aux-heads
    /// cubin.
    pub(crate) fn new(stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let context = stream.context();
        let aux_module = context
            .load_cubin(AUX_HEADS_CUBIN.to_vec())
            .map_err(|e| MLError::ModelError(format!("aux_heads cubin load (bwd): {e}")))?;

        let next_bar_backward_kernel = aux_module
            .load_function("aux_next_bar_backward")
            .map_err(|e| MLError::ModelError(format!("aux_next_bar_backward load: {e}")))?;
        let regime_backward_kernel = aux_module
            .load_function("aux_regime_backward")
            .map_err(|e| MLError::ModelError(format!("aux_regime_backward load: {e}")))?;
        let param_grad_reduce_kernel = aux_module
            .load_function("aux_param_grad_reduce")
            .map_err(|e| MLError::ModelError(format!("aux_param_grad_reduce load: {e}")))?;

        Ok(Self {
            next_bar_backward_kernel,
            regime_backward_kernel,
            param_grad_reduce_kernel,
        })
    }

    /// Launch `aux_next_bar_backward` — emits per-sample param-grad
    /// partials + per-sample `dh_s2 [B, SH2]`. Caller follows up with
    /// `param_grad_reduce` for each of {`dW1`, `db1`, `dW2`, `db2`} to
    /// collapse along the batch dim.
    ///
    /// SP13 B1.1a (2026-05-05): MSE → softmax CE flip. Reads the saved
    /// `[B, K]` softmax tile (produced by `forward_next_bar`) plus the
    /// `[B] i32` labels buffer and the `[1]` `valid_count` scalar
    /// (written by `next_bar_loss_reduce`). Mirrors `aux_regime_backward`'s
    /// K-fanout structure but with mean-over-valid-rows derivative
    /// `d_logits[b, kc] = (softmax[b, kc] - one_hot(labels[b], kc)) / B_valid`
    /// for valid rows, zeroed across the K-vector for masked rows
    /// (label == -1) so partial-grad writes are no-ops.
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn backward_next_bar(
        &self,
        stream: &Arc<CudaStream>,
        h_s2_aux_ptr: u64,
        w1_ptr: u64,
        w2_ptr: u64,
        hidden_post_ptr: u64,
        softmax_ptr: u64,
        labels_ptr: u64,
        valid_count_ptr: u64,
        b: usize,
        sh2: usize,
        k: usize,
        dw1_partial_ptr: u64,
        db1_partial_ptr: u64,
        dw2_partial_ptr: u64,
        db2_partial_ptr: u64,
        dh_s2_aux_out_ptr: u64,
    ) -> Result<(), MLError> {
        let b_i32 = b as i32;
        let sh2_i32 = sh2 as i32;
        let k_i32 = k as i32;
        // Shared mem: 2*H (d_h_pre + h_post caches) + K (d_logits cache).
        let smem_bytes =
            ((2 * AUX_HIDDEN_DIM as u32) + (k as u32)) * std::mem::size_of::<f32>() as u32;
        unsafe {
            stream
                .launch_builder(&self.next_bar_backward_kernel)
                .arg(&h_s2_aux_ptr)
                .arg(&w1_ptr)
                .arg(&w2_ptr)
                .arg(&hidden_post_ptr)
                .arg(&softmax_ptr)
                .arg(&labels_ptr)
                .arg(&valid_count_ptr)
                .arg(&b_i32)
                .arg(&sh2_i32)
                .arg(&k_i32)
                .arg(&dw1_partial_ptr)
                .arg(&db1_partial_ptr)
                .arg(&dw2_partial_ptr)
                .arg(&db2_partial_ptr)
                .arg(&dh_s2_aux_out_ptr)
                .launch(LaunchConfig {
                    grid_dim: (b as u32, 1, 1),
                    block_dim: (AUX_BLOCK, 1, 1),
                    shared_mem_bytes: smem_bytes,
                })
                .map_err(|e| MLError::ModelError(format!("aux_next_bar_backward: {e}")))?;
        }
        Ok(())
    }

    /// Launch `aux_regime_backward` — emits per-sample param-grad partials
    /// + per-sample `dh_s2 [B, SH2]`. Caller follows up with
    /// `param_grad_reduce` for each of {`dW1`, `db1`, `dW2`, `db2`}.
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn backward_regime(
        &self,
        stream: &Arc<CudaStream>,
        h_s2_aux_ptr: u64,
        w1_ptr: u64,
        w2_ptr: u64,
        hidden_post_ptr: u64,
        logits_ptr: u64,
        label_ptr: u64,
        b: usize,
        sh2: usize,
        k: usize,
        dw1_partial_ptr: u64,
        db1_partial_ptr: u64,
        dw2_partial_ptr: u64,
        db2_partial_ptr: u64,
        dh_s2_aux_out_ptr: u64,
    ) -> Result<(), MLError> {
        let b_i32 = b as i32;
        let sh2_i32 = sh2 as i32;
        let k_i32 = k as i32;
        let smem_bytes =
            ((2 * AUX_HIDDEN_DIM as u32) + (k as u32)) * std::mem::size_of::<f32>() as u32;
        unsafe {
            stream
                .launch_builder(&self.regime_backward_kernel)
                .arg(&h_s2_aux_ptr)
                .arg(&w1_ptr)
                .arg(&w2_ptr)
                .arg(&hidden_post_ptr)
                .arg(&logits_ptr)
                .arg(&label_ptr)
                .arg(&b_i32)
                .arg(&sh2_i32)
                .arg(&k_i32)
                .arg(&dw1_partial_ptr)
                .arg(&db1_partial_ptr)
                .arg(&dw2_partial_ptr)
                .arg(&db2_partial_ptr)
                .arg(&dh_s2_aux_out_ptr)
                .launch(LaunchConfig {
                    grid_dim: (b as u32, 1, 1),
                    block_dim: (AUX_BLOCK, 1, 1),
                    shared_mem_bytes: smem_bytes,
                })
                .map_err(|e| MLError::ModelError(format!("aux_regime_backward: {e}")))?;
        }
        Ok(())
    }

    /// Launch `aux_param_grad_reduce`: collapse `partials [B, P]` along
    /// the batch dim into `final_out [P]`. One block per output element
    /// (`P` blocks); each block does a shmem-tree reduction over `B`
    /// samples. No atomicAdd.
    pub(crate) fn param_grad_reduce(
        &self,
        stream: &Arc<CudaStream>,
        partials_ptr: u64,
        b: usize,
        p: usize,
        final_out_ptr: u64,
    ) -> Result<(), MLError> {
        let b_i32 = b as i32;
        let p_i32 = p as i32;
        let smem_bytes = AUX_BLOCK * std::mem::size_of::<f32>() as u32;
        unsafe {
            stream
                .launch_builder(&self.param_grad_reduce_kernel)
                .arg(&partials_ptr)
                .arg(&b_i32)
                .arg(&p_i32)
                .arg(&final_out_ptr)
                .launch(LaunchConfig {
                    grid_dim: (p as u32, 1, 1),
                    block_dim: (AUX_BLOCK, 1, 1),
                    shared_mem_bytes: smem_bytes,
                })
                .map_err(|e| MLError::ModelError(format!("aux_param_grad_reduce: {e}")))?;
        }
        Ok(())
    }
}

// ────────────────────────────────────────────────────────────────────────────
// SP22 H6 vNext: Trade-outcome aux head ops (Phase A2-A5 GPU side wired here)
// ────────────────────────────────────────────────────────────────────────────
//
// The vNext aux head replaces per-bar direction prediction (K=2) with
// per-trade-close outcome prediction (K=3: {Profit, Stop, Timeout}) — a
// signal that aligns with the system's actual multi-bar decision horizon.
// See `docs/plans/2026-05-14-sp22-h6-vNext-trade-outcome-aux.md` for the
// full architecture spec.
//
// Phase B0 (this commit): orchestrator scaffolding only — additive,
// no production callers yet. Mirrors the AuxHeadsForwardOps /
// AuxHeadsBackwardOps shape so subsequent phases (trainer struct fields,
// `collect_experiences_gpu` wireup) plug in with minimal contract drift.
//
// Reuses existing K-generic `aux_param_grad_reduce` from
// `AuxHeadsBackwardOps` — no separate reducer needed.

use super::gpu_dqn_trainer::{
    AUX_TO_INPUT_CONCAT_CUBIN, AUX_TRADE_OUTCOME_BACKWARD_CUBIN,
    AUX_TRADE_OUTCOME_FORWARD_CUBIN, AUX_TRADE_OUTCOME_LOSS_REDUCE_CUBIN,
    STRIDED_ROW_SAXPY_CUBIN, TRADE_OUTCOME_LABEL_CUBIN,
};

/// SP22 H6 vNext Phase B5b (2026-05-14): trade-plan parameter dimension.
/// The trade plan head's output [B, P=6] gets concatenated with h_s2_aux
/// [B, 256] to form the K=3 trade-outcome head's plan-conditioned input
/// [B, SH2_TOTAL = 256 + 6 = 262]. Six features per the trade plan ABI:
/// {target_bars, profit_target, stop_loss, conviction, asymmetry,
/// trail_distance}.
pub(crate) const PLAN_PARAM_DIM: usize = 6;

/// Forward + loss-reduce + label-producer orchestrator for the SP22 H6
/// vNext trade-outcome aux head. Held separate from `AuxHeadsForwardOps`
/// because the trade-outcome head is structurally distinct — different
/// label semantics (per-trade-close vs per-bar), different K (3 vs 2),
/// different input (Phase B switches to `[h_s2_aux || plan_params]` =
/// 262-dim concat vs the K=2 head's 256-dim `h_s2_aux`-only input).
///
/// Owns three kernel handles:
///   - `aux_trade_outcome_forward`         (Phase A3 cubin)
///   - `aux_trade_outcome_loss_reduce`     (Phase A4 cubin)
///   - `trade_outcome_label_kernel`        (Phase A2 cubin)
///
/// `pub(crate)` so the collector / trainer wireup commits can construct
/// it via `new()` and call the per-launch methods.
#[allow(missing_debug_implementations)]
pub(crate) struct AuxTradeOutcomeForwardOps {
    forward_kernel: CudaFunction,
    loss_reduce_kernel: CudaFunction,
    label_kernel: CudaFunction,
    /// SP22 H6 vNext Phase B5b (2026-05-14): plan-conditioning concat
    /// kernel. Writes `aux_to_input [B, SH2+P=262]` from `h_s2_aux
    /// [B, SH2=256]` || `plan_params [B, P=6]`. NULL-tolerant on
    /// plan_params (zeros for trailing P cols).
    concat_kernel: CudaFunction,
}

impl AuxTradeOutcomeForwardOps {
    /// Load all three forward-side kernel handles from their precompiled
    /// cubins. Mirrors `AuxHeadsForwardOps::new`'s loader pattern but
    /// reads three SEPARATE cubins (one per kernel) because the vNext
    /// kernels are in distinct `.cu` files for diagnostic isolation.
    pub(crate) fn new(stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let context = stream.context();

        let fwd_module = context
            .load_cubin(AUX_TRADE_OUTCOME_FORWARD_CUBIN.to_vec())
            .map_err(|e| MLError::ModelError(format!(
                "aux_trade_outcome_forward cubin load: {e}"
            )))?;
        let forward_kernel = fwd_module
            .load_function("aux_trade_outcome_forward")
            .map_err(|e| MLError::ModelError(format!(
                "aux_trade_outcome_forward load: {e}"
            )))?;

        let loss_module = context
            .load_cubin(AUX_TRADE_OUTCOME_LOSS_REDUCE_CUBIN.to_vec())
            .map_err(|e| MLError::ModelError(format!(
                "aux_trade_outcome_loss_reduce cubin load: {e}"
            )))?;
        let loss_reduce_kernel = loss_module
            .load_function("aux_trade_outcome_loss_reduce")
            .map_err(|e| MLError::ModelError(format!(
                "aux_trade_outcome_loss_reduce load: {e}"
            )))?;

        let label_module = context
            .load_cubin(TRADE_OUTCOME_LABEL_CUBIN.to_vec())
            .map_err(|e| MLError::ModelError(format!(
                "trade_outcome_label_kernel cubin load: {e}"
            )))?;
        let label_kernel = label_module
            .load_function("trade_outcome_label_kernel")
            .map_err(|e| MLError::ModelError(format!(
                "trade_outcome_label_kernel load: {e}"
            )))?;

        // SP22 H6 vNext Phase B5b (2026-05-14): plan-conditioning concat
        // kernel handle. Loaded once at ops construction so production
        // launches stay graph-capture-friendly.
        let concat_module = context
            .load_cubin(AUX_TO_INPUT_CONCAT_CUBIN.to_vec())
            .map_err(|e| MLError::ModelError(format!(
                "aux_to_input_concat cubin load: {e}"
            )))?;
        let concat_kernel = concat_module
            .load_function("aux_to_input_concat")
            .map_err(|e| MLError::ModelError(format!(
                "aux_to_input_concat load: {e}"
            )))?;

        Ok(Self {
            forward_kernel,
            loss_reduce_kernel,
            label_kernel,
            concat_kernel,
        })
    }

    /// SP22 H6 vNext Phase B5b (2026-05-14): launch the input concat
    /// kernel to build `aux_to_input [B, SH2+P]` from `h_s2_aux
    /// [B, SH2]` || `plan_params [B, P]`. Used by both the trainer's
    /// `aux_heads_forward` (real plan_params from the trade plan head)
    /// and the collector's rollout-step forward (plan_params_ptr = 0
    /// → zero-fill, the train/inference asymmetry documented in the
    /// kernel doc).
    ///
    /// Block: 256 threads. Grid: `ceil((B * (SH2 + P)) / 256)` blocks
    /// (one thread per output element).
    pub(crate) fn launch_concat(
        &self,
        stream: &Arc<CudaStream>,
        h_s2_aux_ptr: u64,
        plan_params_ptr: u64,
        b: usize,
        sh2: usize,
        p: usize,
        out_ptr: u64,
    ) -> Result<(), MLError> {
        let b_i32 = b as i32;
        let sh2_i32 = sh2 as i32;
        let p_i32 = p as i32;
        let total_cols = sh2 + p;
        let total = b * total_cols;
        let block: u32 = 256;
        let grid = (total as u32).div_ceil(block);
        unsafe {
            stream
                .launch_builder(&self.concat_kernel)
                .arg(&h_s2_aux_ptr)
                .arg(&plan_params_ptr)
                .arg(&b_i32)
                .arg(&sh2_i32)
                .arg(&p_i32)
                .arg(&out_ptr)
                .launch(LaunchConfig {
                    grid_dim: (grid.max(1), 1, 1),
                    block_dim: (block, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!(
                    "aux_to_input_concat: {e}"
                )))?;
        }
        Ok(())
    }

    /// Launch `aux_trade_outcome_forward`:
    /// `Linear(h_s2_aux) → ELU → Linear → softmax` over K=3 trade outcomes.
    ///
    /// Caller-owned buffers (raw `u64` device pointers for graph-capture
    /// safety; mirrors `forward_next_bar`):
    ///   * `h_s2_aux_ptr`    — `[B, SH2]` row-major. Phase A3-B0: SH2 =
    ///                         256 (aux trunk output only). Phase B
    ///                         switches to SH2 = 262 (`h_s2_aux ||
    ///                         plan_params`) — kernel arg unchanged,
    ///                         only the caller's W1 shape and the buffer
    ///                         contents change.
    ///   * `w1_ptr` / `b1_ptr` — `[H, SH2]` / `[H]`. Trainer-owned;
    ///                         Phase B1 lands the allocation.
    ///   * `w2_ptr` / `b2_ptr` — `[K=3, H]` / `[K=3]`.
    ///   * `hidden_out_ptr`   — `[B, H]` SAVED post-ELU buffer
    ///                         (consumed by backward).
    ///   * `logits_out_ptr`   — `[B, K=3]` SAVED logits (consumed by
    ///                         backward only — Phase A5 backward reads
    ///                         this; current backward reads the softmax
    ///                         tile, but the buffer is allocated for
    ///                         contract symmetry with the K=2 head).
    ///   * `softmax_out_ptr`  — `[B, K=3]` SAVED softmax tile (consumed
    ///                         by loss reduce + backward + Phase C
    ///                         3-slot state assembly producer).
    ///
    /// Block: `AUX_BLOCK` threads. Grid: `B` blocks.
    /// Shared mem: `(H + K) * sizeof(f32)` bytes.
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn forward(
        &self,
        stream: &Arc<CudaStream>,
        h_s2_aux_ptr: u64,
        w1_ptr: u64,
        b1_ptr: u64,
        w2_ptr: u64,
        b2_ptr: u64,
        b: usize,
        sh2: usize,
        k: usize,
        hidden_out_ptr: u64,
        logits_out_ptr: u64,
        softmax_out_ptr: u64,
    ) -> Result<(), MLError> {
        let b_i32 = b as i32;
        let sh2_i32 = sh2 as i32;
        let k_i32 = k as i32;
        let smem_bytes =
            (AUX_HIDDEN_DIM as u32 + k as u32) * std::mem::size_of::<f32>() as u32;
        unsafe {
            stream
                .launch_builder(&self.forward_kernel)
                .arg(&h_s2_aux_ptr)
                .arg(&w1_ptr)
                .arg(&b1_ptr)
                .arg(&w2_ptr)
                .arg(&b2_ptr)
                .arg(&b_i32)
                .arg(&sh2_i32)
                .arg(&k_i32)
                .arg(&hidden_out_ptr)
                .arg(&logits_out_ptr)
                .arg(&softmax_out_ptr)
                .launch(LaunchConfig {
                    grid_dim: (b as u32, 1, 1),
                    block_dim: (AUX_BLOCK, 1, 1),
                    shared_mem_bytes: smem_bytes,
                })
                .map_err(|e| MLError::ModelError(format!(
                    "aux_trade_outcome_forward: {e}"
                )))?;
        }
        Ok(())
    }

    /// Launch `aux_trade_outcome_loss_reduce`: mean CE over the valid
    /// (label != -1) rows of the K=3 softmax tile.
    ///
    /// Sparse-label semantic: most bars (~95-99%) have label=-1; the
    /// kernel divides by `fmaxf(B_valid, 1.0)` so an all-skip batch
    /// produces loss=0 + valid_count=0 (the backward sees valid_count=0
    /// and emits zero gradients across the board, no NaN).
    pub(crate) fn loss_reduce(
        &self,
        stream: &Arc<CudaStream>,
        softmax_ptr: u64,
        labels_ptr: u64,
        b: usize,
        k: usize,
        loss_out_ptr: u64,
        valid_count_out_ptr: u64,
    ) -> Result<(), MLError> {
        let b_i32 = b as i32;
        let k_i32 = k as i32;
        // Two parallel partial-reduction strips (loss + valid-count).
        let smem_bytes = 2 * AUX_BLOCK * std::mem::size_of::<f32>() as u32;
        unsafe {
            stream
                .launch_builder(&self.loss_reduce_kernel)
                .arg(&softmax_ptr)
                .arg(&labels_ptr)
                .arg(&b_i32)
                .arg(&k_i32)
                .arg(&loss_out_ptr)
                .arg(&valid_count_out_ptr)
                .launch(LaunchConfig {
                    grid_dim: (1, 1, 1),
                    block_dim: (AUX_BLOCK, 1, 1),
                    shared_mem_bytes: smem_bytes,
                })
                .map_err(|e| MLError::ModelError(format!(
                    "aux_trade_outcome_loss_reduce: {e}"
                )))?;
        }
        Ok(())
    }

    /// Launch `trade_outcome_label_kernel`: per-env K=3 outcome
    /// classification at trade-close.
    ///
    /// Inputs:
    ///   * `trade_close_per_sample_ptr` — `[N*L]` i32 flag (read at
    ///                         offset `env*L + t` for the current step).
    ///                         Produced by `experience_env_step` at
    ///                         line ~3093 (`is_close`).
    ///   * `pnl_vs_target_at_close_ptr` — `[n_envs]` f32 save-for-backward
    ///                         from `experience_env_step` segment_complete
    ///                         (Phase A3 wireup — `pnl_vs_target_at_close
    ///                         _per_env`).
    ///   * `pnl_vs_stop_at_close_ptr`   — `[n_envs]` f32 save-for-backward
    ///                         from `experience_env_step` segment_complete
    ///                         (Phase A3 wireup — `pnl_vs_stop_at_close
    ///                         _per_env`).
    ///   * `t` — current rollout step (0..timesteps); used to compute
    ///                         the offset `env*L + t` into the per-sample
    ///                         trade-close flag.
    ///   * `n_envs` — number of envs in the rollout slice.
    ///   * `l_timesteps` — buffer stride (timesteps per env).
    ///
    /// Output: `out_labels_ptr [n_envs]` i32 in `{-1, 0, 1, 2}`. `-1`
    /// = no trade close at this bar (the common case ~95-99% of bars).
    ///
    /// Block: 256 threads. Grid: `ceil(n_envs / 256)` blocks (per-env
    /// parallelism, no reduction, no atomicAdd).
    /// `out_labels_per_sample_ptr` (Phase B4b 2026-05-14) is the per-
    /// (env, t) i32 column the replay buffer scatter consumes. Pass 0
    /// (NULL) to skip the per-(i, t) write — the per-env `out_labels`
    /// output is still produced. The kernel-side NULL guard is on the
    /// pointer-read site, mirroring the standard NULL-tolerant arg
    /// pattern used by `experience_env_step`'s
    /// `pnl_vs_target_at_close_per_env` etc. (Phase A3).
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn compute_label(
        &self,
        stream: &Arc<CudaStream>,
        trade_close_per_sample_ptr: u64,
        pnl_vs_target_at_close_ptr: u64,
        pnl_vs_stop_at_close_ptr: u64,
        t: i32,
        n_envs: usize,
        l_timesteps: usize,
        out_labels_ptr: u64,
        out_labels_per_sample_ptr: u64,
    ) -> Result<(), MLError> {
        let n_envs_i32 = n_envs as i32;
        let l_i32 = l_timesteps as i32;
        let block: u32 = 256;
        let grid = (n_envs as u32).div_ceil(block);
        unsafe {
            stream
                .launch_builder(&self.label_kernel)
                .arg(&trade_close_per_sample_ptr)
                .arg(&pnl_vs_target_at_close_ptr)
                .arg(&pnl_vs_stop_at_close_ptr)
                .arg(&t)
                .arg(&n_envs_i32)
                .arg(&l_i32)
                .arg(&out_labels_ptr)
                .arg(&out_labels_per_sample_ptr)
                .launch(LaunchConfig {
                    grid_dim: (grid, 1, 1),
                    block_dim: (block, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!(
                    "trade_outcome_label_kernel: {e}"
                )))?;
        }
        Ok(())
    }
}

/// Backward orchestrator for the SP22 H6 vNext trade-outcome aux head.
/// Held separate from `AuxHeadsBackwardOps` for the same reasons the
/// forward ops are separate. Owns ONE kernel handle (the K=3 backward) —
/// the per-sample → final-grad batch-reduce step reuses the existing
/// K-generic `aux_param_grad_reduce` from `AuxHeadsBackwardOps`.
///
/// Wireup commits will hold BOTH `AuxHeadsBackwardOps` (for the
/// K-generic `param_grad_reduce`) AND this struct (for the K=3
/// backward) on the trainer.
#[allow(missing_debug_implementations)]
pub(crate) struct AuxTradeOutcomeBackwardOps {
    backward_kernel: CudaFunction,
    /// SP22 H6 vNext Phase B5b (2026-05-14): row-truncating SAXPY for
    /// dh_s2_aux accumulate when backward emits `[B, SH2_TOTAL=262]`
    /// but the upstream consumer `dh_s2_aux_accum [B, SH2=256]` only
    /// reads the first 256 cols (h_s2_aux gradient; plan_params
    /// gradient is stop-grad).
    strided_saxpy_kernel: CudaFunction,
}

impl AuxTradeOutcomeBackwardOps {
    /// Load the K=3 backward kernel handle from the Phase A5 cubin.
    pub(crate) fn new(stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let context = stream.context();
        let bwd_module = context
            .load_cubin(AUX_TRADE_OUTCOME_BACKWARD_CUBIN.to_vec())
            .map_err(|e| MLError::ModelError(format!(
                "aux_trade_outcome_backward cubin load: {e}"
            )))?;
        let backward_kernel = bwd_module
            .load_function("aux_trade_outcome_backward")
            .map_err(|e| MLError::ModelError(format!(
                "aux_trade_outcome_backward load: {e}"
            )))?;

        // SP22 H6 vNext Phase B5b (2026-05-14): strided row SAXPY kernel
        // handle. Used by `aux_heads_backward` to feed only the first
        // SH2=256 cols of `dh_s2_aux_to_buf [B, 262]` into
        // `dh_s2_aux_accum [B, 256]` — plan_params gradient (trailing
        // 6 cols) is dropped (stop-grad).
        let strided_module = context
            .load_cubin(STRIDED_ROW_SAXPY_CUBIN.to_vec())
            .map_err(|e| MLError::ModelError(format!(
                "strided_row_saxpy cubin load: {e}"
            )))?;
        let strided_saxpy_kernel = strided_module
            .load_function("strided_row_saxpy_f32")
            .map_err(|e| MLError::ModelError(format!(
                "strided_row_saxpy_f32 load: {e}"
            )))?;

        Ok(Self { backward_kernel, strided_saxpy_kernel })
    }

    /// SP22 H6 vNext Phase B5b (2026-05-14): launch strided row SAXPY
    /// to accumulate first `n_cols_copy` columns of `src [B, src_cols]`
    /// into `dst [B, dst_cols]` scaled by `alpha`. Different row
    /// strides on src and dst — element-wise SAXPY won't work when
    /// src_cols != dst_cols.
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn launch_strided_row_saxpy(
        &self,
        stream: &Arc<CudaStream>,
        dst_ptr: u64,
        src_ptr: u64,
        alpha: f32,
        b: usize,
        src_cols: usize,
        dst_cols: usize,
        n_cols_copy: usize,
    ) -> Result<(), MLError> {
        let b_i32 = b as i32;
        let src_i32 = src_cols as i32;
        let dst_i32 = dst_cols as i32;
        let n_i32 = n_cols_copy as i32;
        let total = b * n_cols_copy;
        let block: u32 = 256;
        let grid = (total as u32).div_ceil(block);
        unsafe {
            stream
                .launch_builder(&self.strided_saxpy_kernel)
                .arg(&dst_ptr)
                .arg(&src_ptr)
                .arg(&alpha)
                .arg(&b_i32)
                .arg(&src_i32)
                .arg(&dst_i32)
                .arg(&n_i32)
                .launch(LaunchConfig {
                    grid_dim: (grid.max(1), 1, 1),
                    block_dim: (block, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!(
                    "strided_row_saxpy_f32: {e}"
                )))?;
        }
        Ok(())
    }

    /// Launch `aux_trade_outcome_backward` — emits per-sample param-grad
    /// partials + per-sample `dh_s2_aux [B, SH2]`. Caller follows up
    /// with `param_grad_reduce` (from `AuxHeadsBackwardOps`) for each
    /// of {`dW1`, `db1`, `dW2`, `db2`} to collapse along the batch dim.
    ///
    /// Sparse-label gradient amplification: `valid_count` from the loss
    /// kernel is typically small (~1-5% of nominal B) so `inv_B =
    /// 1/B_valid` is much larger than the K=2 sibling's `inv_B =
    /// 1/(~B)`. Per-trade-close gradients have proportionally higher
    /// magnitude — correct credit assignment (rare signal speaks louder)
    /// but Phase E's Adam may need class-weighted CE or per-group LR
    /// tuning to keep the trade-outcome head's update magnitude
    /// comparable to the K=2 sibling.
    ///
    /// `dh_s2_aux_out_ptr` SAXPYs into the aux trunk's
    /// `dh_s2_aux_accum` at the call site — the aux trunk's backward
    /// kernel reads this as its upstream gradient. Encoder stop-grad
    /// is enforced STRUCTURALLY by `aux_trunk_backward` (no `dx_in`
    /// output param), so Q's encoder is protected from aux
    /// contamination regardless of gradient magnitude.
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn backward(
        &self,
        stream: &Arc<CudaStream>,
        h_s2_aux_ptr: u64,
        w1_ptr: u64,
        w2_ptr: u64,
        hidden_post_ptr: u64,
        softmax_ptr: u64,
        labels_ptr: u64,
        valid_count_ptr: u64,
        b: usize,
        sh2: usize,
        k: usize,
        dw1_partial_ptr: u64,
        db1_partial_ptr: u64,
        dw2_partial_ptr: u64,
        db2_partial_ptr: u64,
        dh_s2_aux_out_ptr: u64,
    ) -> Result<(), MLError> {
        let b_i32 = b as i32;
        let sh2_i32 = sh2 as i32;
        let k_i32 = k as i32;
        // Shared mem: 2*H (d_h_pre + h_post caches) + K (d_logits cache).
        let smem_bytes =
            ((2 * AUX_HIDDEN_DIM as u32) + (k as u32)) * std::mem::size_of::<f32>() as u32;
        unsafe {
            stream
                .launch_builder(&self.backward_kernel)
                .arg(&h_s2_aux_ptr)
                .arg(&w1_ptr)
                .arg(&w2_ptr)
                .arg(&hidden_post_ptr)
                .arg(&softmax_ptr)
                .arg(&labels_ptr)
                .arg(&valid_count_ptr)
                .arg(&b_i32)
                .arg(&sh2_i32)
                .arg(&k_i32)
                .arg(&dw1_partial_ptr)
                .arg(&db1_partial_ptr)
                .arg(&dw2_partial_ptr)
                .arg(&db2_partial_ptr)
                .arg(&dh_s2_aux_out_ptr)
                .launch(LaunchConfig {
                    grid_dim: (b as u32, 1, 1),
                    block_dim: (AUX_BLOCK, 1, 1),
                    shared_mem_bytes: smem_bytes,
                })
                .map_err(|e| MLError::ModelError(format!(
                    "aux_trade_outcome_backward: {e}"
                )))?;
        }
        Ok(())
    }
}
