//! PPO actor (policy) + critic (value) heads for the integrated RL trainer.
//!
//! Sit on the shared encoder's `h_t` output. The policy head produces a
//! categorical distribution over the 9-action grid (logits + softmax fused
//! into the surrogate loss kernel). The value head produces a scalar V(s)
//! for the baseline. Per `pearl_controller_anchors_isv_driven`, the PPO
//! clip ε and entropy bonus coefficient are read from
//! `ISV[RL_PPO_CLIP_INDEX=402]` / `ISV[RL_ENTROPY_COEF_INDEX=403]` at
//! loss-kernel time, not from hardcoded constants.
//!
//! Phase D (this commit) lands the heads + cubin handles. Phase E wires the
//! surrogate loss into the integrated trainer's joint-loss combine path.
//!
//! ## Architecture
//!
//! ```text
//!   logits[b, a] = b_pi[a] + Σ_c W_pi[a, c] * h_t[b, c]      (PolicyHead)
//!   V[b]         = b_v     + Σ_c W_v[c]     * h_t[b, c]      (ValueHead)
//! ```
//!
//! where `c` ranges over `HIDDEN_DIM`, `a` over `N_ACTIONS=9`. Softmax is
//! NOT applied here — the `ppo_clipped_surrogate_fwd` kernel computes the
//! numerically-stable softmax once, reusing it for log π, entropy and
//! the clipped surrogate.
//!
//! ## Initialisation
//!
//! Per `pearl_scoped_init_seed_for_reproducibility`, both heads draw under
//! a `scoped_init_seed` guard. Xavier uniform scaled by `0.01` so initial
//! logits sit near zero → softmax ≈ uniform → policy uninformed at step 0,
//! learning driven by gradient rather than init noise. Matches the
//! Phase C `DqnHead` initialisation pattern.

use std::sync::Arc;

use anyhow::{Context, Result};
use cudarc::driver::{
    CudaFunction, CudaModule, CudaSlice, CudaStream, DevicePtrMut,
};
use cudarc::driver::sys::CUstream;
use ml_core::cuda_autograd::init::scoped_init_seed;
use ml_core::device::MlDevice;
use rand::{Rng, SeedableRng};
use rand_chacha::ChaCha8Rng;

use crate::heads::HIDDEN_DIM;
use crate::pinned_mem::MappedF32Buffer;
use crate::rl::common::{Action, N_ACTIONS};
use crate::trainer::raw_launch::{RawArgs, raw_launch};

const PPO_SURR_CUBIN: &[u8] = include_bytes!(concat!(
    env!("OUT_DIR"),
    "/ppo_clipped_surrogate.cubin"
));
const PPO_DIAGNOSTIC_REDUCE_CUBIN: &[u8] = include_bytes!(concat!(
    env!("OUT_DIR"),
    "/rl_ppo_diagnostic_stats_reduce.cubin"
));
const PPO_LOSS_REDUCE_CUBIN: &[u8] = include_bytes!(concat!(
    env!("OUT_DIR"),
    "/ppo_loss_reduce_b.cubin"
));
const V_HEAD_CUBIN: &[u8] =
    include_bytes!(concat!(env!("OUT_DIR"), "/v_head_fwd_bwd.cubin"));

/// Construction config for [`PolicyHead`] and [`ValueHead`].
#[derive(Clone, Debug)]
pub struct PpoHeadsConfig {
    /// Encoder hidden dimension feeding both heads. Matches the production
    /// CfC trunk's output width (currently 128) so the heads can shape-check
    /// later launches without re-deriving from the encoder config.
    pub hidden_dim: usize,
    /// Random seed for Xavier init. Reproducibility-critical per
    /// `pearl_scoped_init_seed_for_reproducibility`.
    pub seed: u64,
}

/// Policy head: linear projection `h_t [B, HIDDEN_DIM] → logits [B, N_ACTIONS]`.
/// Softmax is fused into the surrogate loss kernel. The policy is categorical
/// over the same 9-action grid that [`crate::rl::dqn::DqnHead`] uses
/// (see [`crate::rl::common::Action`]).
///
/// Like Phase C's `DqnHead`, the `cfg` field is kept under
/// `#[allow(dead_code)]` because production reads happen at the Phase E
/// training-loop site (and via checkpoint round-trip), not from this struct
/// directly.
pub struct PolicyHead {
    #[allow(dead_code)]
    cfg: PpoHeadsConfig,
    stream: Arc<CudaStream>,
    raw_stream: CUstream,

    _module: Arc<CudaModule>,
    /// F4.1 reducer cubin module — kept alive while `loss_reduce_fn`
    /// references it.
    _loss_reduce_module: Arc<CudaModule>,
    /// Forward kernel: `ppo_clipped_surrogate_fwd`. Consumes new-policy
    /// logits, old log π, advantages + ISV[402,403,440] and writes
    /// per-batch log π / entropy + per-batch loss slots (no atomicAdd).
    pub surrogate_fwd_fn: CudaFunction,
    /// F4.1 reducer kernel: `ppo_loss_reduce_b`. Takes the per-batch
    /// `[B]` loss buffers written by `surrogate_fwd_fn` and produces
    /// the two scalar diagnostics the trainer / LR controller reads.
    pub loss_reduce_fn: CudaFunction,
    /// Backward kernel: `ppo_clipped_surrogate_bwd`. Computes per-logit
    /// gradient with the clip-mask applied (zero outside the clip band).
    pub surrogate_bwd_fn: CudaFunction,
    /// Phase E.2 logits-projection kernel: `ppo_clipped_logits_fwd`. The
    /// policy head uses the existing `ppo_clipped_surrogate_fwd` for the
    /// loss path; for the standalone logits forward (used in the rollout
    /// path AND as the input to the surrogate kernel) we run a separate
    /// linear projection. Since PPO's surrogate kernel takes pre-computed
    /// logits as input, the trainer first runs a small linear forward
    /// (`policy_linear_fwd_fn`) that produces `[B × N_ACTIONS]` logits
    /// from `h_t`. The surrogate kernel then consumes them.
    pub policy_linear_fwd_fn: CudaFunction,
    /// Phase E.2 backward chain stage 2: `ppo_grad_w_b_h_t`. Maps the
    /// `grad_logits` output of `surrogate_bwd_fn` into per-batch
    /// `grad_w` / `grad_b` scratch buffers + an OVERWRITE `grad_h_t`
    /// slot for the integrated trainer to fold into the encoder grad
    /// accumulator. Per-batch scratch — cross-batch reduction is the
    /// caller's responsibility via `reduce_axis0`.
    pub grad_w_b_h_t_fn: CudaFunction,

    /// B-10 (2026-06-01) G3+G4: cross-batch tree-reduce of per-batch
    /// PPO advantage + ratio + surrogate observability scratches written
    /// by `surrogate_fwd_fn`. Reads slot 612 (`var_pre_norm`) to compute
    /// `σ_used = sqrt(var)` for the |A_unnorm| reconstruction. Writes
    /// ISV slots 735-742. Pure observability — does not feed any
    /// controller, does not perturb training dynamics.
    pub ppo_diagnostic_reduce_fn: CudaFunction,
    _ppo_diagnostic_reduce_module: Arc<CudaModule>,

    /// Online weights `[N_ACTIONS, HIDDEN_DIM]`, row-major. Each row is
    /// one action's projection vector.
    pub w_d: CudaSlice<f32>,
    /// Online biases `[N_ACTIONS]`.
    pub b_d: CudaSlice<f32>,
}

impl PolicyHead {
    /// Allocate device weights, load the cubin, cache kernel handles.
    /// Xavier-uniform-scaled-by-0.01 init keeps initial logits near zero
    /// so softmax ≈ uniform; the head is uninformed at step 0 and the
    /// policy-gradient signal drives learning.
    pub fn new(dev: &MlDevice, cfg: PpoHeadsConfig) -> Result<Self> {
        let stream: Arc<CudaStream> = dev
            .cuda_stream()
            .context("policy_head stream")?
            .clone();
        let ctx = dev.cuda_context().context("policy_head ctx")?;
        let module = ctx
            .load_cubin(PPO_SURR_CUBIN.to_vec())
            .context("load ppo_clipped_surrogate cubin")?;
        let surrogate_fwd_fn = module
            .load_function("ppo_clipped_surrogate_fwd")
            .context("load ppo_clipped_surrogate_fwd")?;
        let surrogate_bwd_fn = module
            .load_function("ppo_clipped_surrogate_bwd")
            .context("load ppo_clipped_surrogate_bwd")?;
        let loss_reduce_module = ctx
            .load_cubin(PPO_LOSS_REDUCE_CUBIN.to_vec())
            .context("load ppo_loss_reduce_b cubin")?;
        let loss_reduce_fn = loss_reduce_module
            .load_function("ppo_loss_reduce_b")
            .context("load ppo_loss_reduce_b")?;
        let policy_linear_fwd_fn = module
            .load_function("ppo_policy_logits_fwd")
            .context("load ppo_policy_logits_fwd")?;
        let grad_w_b_h_t_fn = module
            .load_function("ppo_grad_w_b_h_t")
            .context("load ppo_grad_w_b_h_t")?;

        // B-10 (2026-06-01): PPO diagnostic stats reducer.
        let ppo_diagnostic_reduce_module = ctx
            .load_cubin(PPO_DIAGNOSTIC_REDUCE_CUBIN.to_vec())
            .context("load rl_ppo_diagnostic_stats_reduce cubin")?;
        let ppo_diagnostic_reduce_fn = ppo_diagnostic_reduce_module
            .load_function("rl_ppo_diagnostic_stats_reduce")
            .context("load rl_ppo_diagnostic_stats_reduce")?;

        // Per pearl_scoped_init_seed_for_reproducibility: install the
        // scoped seed guard BEFORE drawing any Xavier samples.
        let _seed_guard = scoped_init_seed(cfg.seed);
        let mut rng = ChaCha8Rng::seed_from_u64(cfg.seed);

        let n_out = N_ACTIONS;
        let n_in = cfg.hidden_dim;
        let scale = 0.01_f32 * (6.0_f32 / (n_in + n_out) as f32).sqrt();
        let w_host: Vec<f32> = (0..n_out * n_in)
            .map(|_| rng.gen_range(-scale..scale))
            .collect();
        // Phase 3D-A (2026-06-03): Hold-action logit bias.
        //
        // The 11-action grid (see `crate::rl::common::Action`) has 4 open
        // variants (0=ShortLarge, 1=ShortSmall, 5=LongSmall, 6=LongLarge)
        // vs 1 hold (2=Hold). At uniform random init the structural prior
        // is P(open)=4/11≈36% and P(hold)=1/11≈9% BEFORE any learning —
        // a 4:1 bias toward opening. Counter-balance by adding +log(4)≈1.386
        // to the Hold action bias so the asymmetry-break probabilities
        // are P(hold)=4/(4+10)≈28.6% and P(any specific open)=1/14≈7.1%,
        // restoring near-uniform open-vs-hold mass at init.
        //
        // Per Cao et al. 2026 + Goodhart-Skalse 2024 literature on
        // action-space prior bias, this is one of three concurrent fixes
        // (B: quadratic impact cost; C: PPO surrogate restoration) for
        // foxhunt's structural overtrading pathology.
        let hold_bias: f32 = (4.0_f32).ln();
        let mut b_host: Vec<f32> = vec![0.0_f32; n_out];
        b_host[Action::Hold as usize] += hold_bias;

        let w_d = upload(&stream, &w_host)?;
        let b_d = upload(&stream, &b_host)?;
        let raw_stream = stream.cu_stream();

        Ok(Self {
            cfg,
            stream,
            raw_stream,
            _module: module,
            _loss_reduce_module: loss_reduce_module,
            surrogate_fwd_fn,
            loss_reduce_fn,
            surrogate_bwd_fn,
            policy_linear_fwd_fn,
            grad_w_b_h_t_fn,
            ppo_diagnostic_reduce_fn,
            _ppo_diagnostic_reduce_module: ppo_diagnostic_reduce_module,
            w_d,
            b_d,
        })
    }

    /// B-10 (2026-06-01) G3+G4: launch the cross-batch reducer over the
    /// 4 per-batch observability scratches written by `surrogate_forward`.
    /// Reads slot 612 (`var_pre_norm`) inside the kernel to compute
    /// `σ_used = sqrt(var)` for the |A_unnorm| reconstruction. Writes
    /// ISV slots 735-742. Pure observability; safe to call after any
    /// `surrogate_forward` in the step.
    #[allow(clippy::too_many_arguments)]
    pub fn launch_ppo_diagnostic_stats_reduce(
        &self,
        isv_dev_ptr: &u64,
        b_size: usize,
        a_norm_pb: &CudaSlice<f32>,
        ratio_dev_pb: &CudaSlice<f32>,
        ratio_clipped_pb: &CudaSlice<f32>,
        surrogate_pb: &CudaSlice<f32>,
    ) -> Result<()> {
        debug_assert_eq!(a_norm_pb.len(), b_size);
        debug_assert_eq!(ratio_dev_pb.len(), b_size);
        debug_assert_eq!(ratio_clipped_pb.len(), b_size);
        debug_assert_eq!(surrogate_pb.len(), b_size);

        // Power-of-2 block size, capped at 256; identity-padded shared
        // mem ensures threads with no work don't pollute the reduce.
        let block_dim = (b_size as u32).next_power_of_two().clamp(1, 256);
        // 5 reductions × block_dim floats (matches kernel layout in
        // rl_ppo_diagnostic_stats_reduce.cu).
        let smem = (5 * block_dim as usize * std::mem::size_of::<f32>()) as u32;
        let b_i = b_size as i32;
        let mut args = RawArgs::new();
        args.push_ptr(*isv_dev_ptr);
        args.push_ptr(a_norm_pb.raw_ptr());
        args.push_ptr(ratio_dev_pb.raw_ptr());
        args.push_ptr(ratio_clipped_pb.raw_ptr());
        args.push_ptr(surrogate_pb.raw_ptr());
        args.push_i32(b_i);
        let mut ptrs = args.build_arg_ptrs();
        unsafe {
            raw_launch(
                self.ppo_diagnostic_reduce_fn.cu_function(),
                (1, 1, 1), (block_dim, 1, 1), smem,
                self.raw_stream,
                &mut ptrs[..args.len()],
            )
            .map_err(|e| anyhow::anyhow!("rl_ppo_diagnostic_stats_reduce: {:?}", e))?;
        }
        Ok(())
    }

    /// Stream used to launch all kernels owned by this head. Phase E's
    /// training loop reads this when scheduling the surrogate / Adam
    /// kernels.
    pub fn stream(&self) -> &Arc<CudaStream> {
        &self.stream
    }

    /// Encoder hidden dim this head was built for. Used by Phase E
    /// shape-checking when assembling the launch config.
    pub fn hidden_dim(&self) -> usize {
        self.cfg.hidden_dim
    }

    /// Phase E.2 forward — materialise per-batch action logits `[B × N_ACTIONS]`
    /// via a small linear projection. Caller then feeds into the
    /// surrogate kernel which fuses softmax + ratio + clip + loss.
    pub fn forward_logits(
        &self,
        h_t: &CudaSlice<f32>,
        b_size: usize,
        logits_out: &mut CudaSlice<f32>,
    ) -> Result<()> {
        debug_assert_eq!(h_t.len(), b_size * HIDDEN_DIM);
        debug_assert_eq!(logits_out.len(), b_size * N_ACTIONS);

        let b_i = b_size as i32;
        {
            let mut args = RawArgs::new();
            args.push_ptr(self.w_d.raw_ptr());
            args.push_ptr(self.b_d.raw_ptr());
            args.push_ptr(h_t.raw_ptr());
            args.push_i32(b_i);
            args.push_ptr(logits_out.raw_ptr());
            let mut ptrs = args.build_arg_ptrs();
            unsafe {
                raw_launch(
                    self.policy_linear_fwd_fn.cu_function(),
                    (b_size as u32, N_ACTIONS as u32, 1),
                    (HIDDEN_DIM as u32, 1, 1),
                    0,
                    self.raw_stream,
                    &mut ptrs[..args.len()],
                )
                .map_err(|e| anyhow::anyhow!("ppo_policy_logits_fwd: {:?}", e))?;
            }
        }
        Ok(())
    }

    /// Phase E.2 surrogate forward — runs `ppo_clipped_surrogate_fwd`
    /// then `ppo_loss_reduce_b` to produce the two scalar loss
    /// diagnostics. Reads ε from `isv[RL_PPO_CLIP_INDEX]` and entropy
    /// coef from `isv[RL_ENTROPY_COEF_INDEX]`.
    ///
    /// F4.1 (2026-05-31) — replaces the prior atomicAdd-into-scalar
    /// pattern. The fwd kernel writes per-batch loss values to the
    /// `[B]` scratch buffers; the reducer then computes the batch mean
    /// and stores it (single-writer, no accumulation) into the scalar
    /// mapped-pinned slots the trainer reads. This eliminates the
    /// step-count staleness bug that inflated `loss.pi` linearly in
    /// step count (local 1k × ~9 ≈ 9080; cluster 20k × ~1200 ≈ 24M).
    /// Per `feedback_no_atomicadd`.
    ///
    /// The V-loss path (formerly bundled with this surrogate forward) was
    /// canonicalised into the dedicated V-head kernels per
    /// `feedback_single_source_of_truth_no_duplicates`; this method no
    /// longer accepts `returns` / `v_pred` / `loss_v` (greenfield removal
    /// 2026-05-22).
    #[allow(clippy::too_many_arguments)]
    #[allow(clippy::too_many_arguments)]
    pub fn surrogate_forward(
        &self,
        logits: &CudaSlice<f32>,
        log_pi_old: &CudaSlice<f32>,
        actions: &CudaSlice<i32>,
        advantages: &CudaSlice<f32>,
        isv_dev_ptr: &u64,
        b_size: usize,
        pi_log_prob: &mut CudaSlice<f32>,
        entropy: &mut CudaSlice<f32>,
        loss_pi_per_b: &mut CudaSlice<f32>,
        loss_entropy_per_b: &mut CudaSlice<f32>,
        loss_pi_scalar_dev_ptr: &u64,
        loss_entropy_scalar_dev_ptr: &u64,
        // B-10 (2026-06-01) G3+G4: per-batch observability scratches
        // consumed downstream by `launch_ppo_diagnostic_stats_reduce`.
        // Pure observability — kernel writes alongside existing
        // loss_pi_per_b without perturbing the loss reduction path.
        ppo_a_norm_pb: &mut CudaSlice<f32>,
        ppo_ratio_dev_pb: &mut CudaSlice<f32>,
        ppo_ratio_clipped_pb: &mut CudaSlice<f32>,
        ppo_surrogate_pb: &mut CudaSlice<f32>,
    ) -> Result<()> {
        debug_assert_eq!(logits.len(), b_size * N_ACTIONS);
        debug_assert_eq!(log_pi_old.len(), b_size);
        debug_assert_eq!(actions.len(), b_size);
        debug_assert_eq!(advantages.len(), b_size);
        debug_assert_eq!(pi_log_prob.len(), b_size);
        debug_assert_eq!(entropy.len(), b_size);
        debug_assert_eq!(loss_pi_per_b.len(), b_size);
        debug_assert_eq!(loss_entropy_per_b.len(), b_size);
        debug_assert_eq!(ppo_a_norm_pb.len(), b_size);
        debug_assert_eq!(ppo_ratio_dev_pb.len(), b_size);
        debug_assert_eq!(ppo_ratio_clipped_pb.len(), b_size);
        debug_assert_eq!(ppo_surrogate_pb.len(), b_size);

        let b_i = b_size as i32;
        {
            let mut args = RawArgs::new();
            args.push_ptr(logits.raw_ptr());
            args.push_ptr(log_pi_old.raw_ptr());
            args.push_ptr(actions.raw_ptr());
            args.push_ptr(advantages.raw_ptr());
            args.push_ptr(*isv_dev_ptr);
            args.push_i32(b_i);
            args.push_ptr(pi_log_prob.raw_ptr());
            args.push_ptr(entropy.raw_ptr());
            args.push_ptr(loss_pi_per_b.raw_ptr());
            args.push_ptr(loss_entropy_per_b.raw_ptr());
            // B-10 G3+G4 scratch outputs.
            args.push_ptr(ppo_a_norm_pb.raw_ptr());
            args.push_ptr(ppo_ratio_dev_pb.raw_ptr());
            args.push_ptr(ppo_ratio_clipped_pb.raw_ptr());
            args.push_ptr(ppo_surrogate_pb.raw_ptr());
            let mut ptrs = args.build_arg_ptrs();
            unsafe {
                raw_launch(
                    self.surrogate_fwd_fn.cu_function(),
                    (b_size as u32, 1, 1),
                    (N_ACTIONS as u32, 1, 1),
                    0,
                    self.raw_stream,
                    &mut ptrs[..args.len()],
                )
                .map_err(|e| anyhow::anyhow!("ppo_clipped_surrogate_fwd: {:?}", e))?;
            }
        }
        // ── F4.1 block tree-reduce [B] → [1] for both diagnostics ──
        // Single block, 256 threads (matches BLOCK_X in
        // ppo_loss_reduce_b.cu). Grid-strided over B so the same launch
        // serves local b_size=16 and cluster b_size=1024 alike.
        {
            let mut args = RawArgs::new();
            args.push_ptr(loss_pi_per_b.raw_ptr());
            args.push_ptr(loss_entropy_per_b.raw_ptr());
            args.push_i32(b_i);
            args.push_ptr(*loss_pi_scalar_dev_ptr);
            args.push_ptr(*loss_entropy_scalar_dev_ptr);
            let mut ptrs = args.build_arg_ptrs();
            unsafe {
                raw_launch(
                    self.loss_reduce_fn.cu_function(),
                    (1, 1, 1),
                    (256, 1, 1),
                    0,
                    self.raw_stream,
                    &mut ptrs[..args.len()],
                )
                .map_err(|e| anyhow::anyhow!("ppo_loss_reduce_b: {:?}", e))?;
            }
        }
        Ok(())
    }

    /// Phase E.2 surrogate backward — emits `grad_logits [B × N_ACTIONS]`.
    /// Subsequent stage is `backward_to_w_b_h`.
    #[allow(clippy::too_many_arguments)]
    pub fn surrogate_backward_logits(
        &self,
        logits: &CudaSlice<f32>,
        log_pi_old: &CudaSlice<f32>,
        actions: &CudaSlice<i32>,
        advantages: &CudaSlice<f32>,
        isv_dev_ptr: &u64,
        b_size: usize,
        grad_logits: &mut CudaSlice<f32>,
    ) -> Result<()> {
        debug_assert_eq!(logits.len(), b_size * N_ACTIONS);
        debug_assert_eq!(log_pi_old.len(), b_size);
        debug_assert_eq!(actions.len(), b_size);
        debug_assert_eq!(advantages.len(), b_size);
        debug_assert_eq!(grad_logits.len(), b_size * N_ACTIONS);

        let b_i = b_size as i32;
        {
            let mut args = RawArgs::new();
            args.push_ptr(logits.raw_ptr());
            args.push_ptr(log_pi_old.raw_ptr());
            args.push_ptr(actions.raw_ptr());
            args.push_ptr(advantages.raw_ptr());
            args.push_ptr(*isv_dev_ptr);
            args.push_i32(b_i);
            args.push_ptr(grad_logits.raw_ptr());
            let mut ptrs = args.build_arg_ptrs();
            unsafe {
                raw_launch(
                    self.surrogate_bwd_fn.cu_function(),
                    (b_size as u32, 1, 1),
                    (N_ACTIONS as u32, 1, 1),
                    0,
                    self.raw_stream,
                    &mut ptrs[..args.len()],
                )
                .map_err(|e| anyhow::anyhow!("ppo_clipped_surrogate_bwd: {:?}", e))?;
            }
        }
        Ok(())
    }

    /// Phase E.2 backward chain stage 2 — propagate `grad_logits` into
    /// `grad_w_per_batch [B × N_ACTIONS × HIDDEN_DIM]`,
    /// `grad_b_per_batch [B × N_ACTIONS]`, and
    /// `grad_h_t [B × HIDDEN_DIM]` (OVERWRITE). Caller reduces axis 0 via
    /// `reduce_axis0`.
    #[allow(clippy::too_many_arguments)]
    pub fn backward_to_w_b_h(
        &self,
        h_t: &CudaSlice<f32>,
        grad_logits: &CudaSlice<f32>,
        b_size: usize,
        grad_w_per_batch: &mut CudaSlice<f32>,
        grad_b_per_batch: &mut CudaSlice<f32>,
        grad_h_t: &mut CudaSlice<f32>,
    ) -> Result<()> {
        debug_assert_eq!(h_t.len(), b_size * HIDDEN_DIM);
        debug_assert_eq!(grad_logits.len(), b_size * N_ACTIONS);
        debug_assert_eq!(grad_w_per_batch.len(), b_size * N_ACTIONS * HIDDEN_DIM);
        debug_assert_eq!(grad_b_per_batch.len(), b_size * N_ACTIONS);
        debug_assert_eq!(grad_h_t.len(), b_size * HIDDEN_DIM);

        let b_i = b_size as i32;
        let smem = (N_ACTIONS * std::mem::size_of::<f32>()) as u32;
        {
            let mut args = RawArgs::new();
            args.push_ptr(self.w_d.raw_ptr());
            args.push_ptr(h_t.raw_ptr());
            args.push_ptr(grad_logits.raw_ptr());
            args.push_i32(b_i);
            args.push_ptr(grad_w_per_batch.raw_ptr());
            args.push_ptr(grad_b_per_batch.raw_ptr());
            args.push_ptr(grad_h_t.raw_ptr());
            let mut ptrs = args.build_arg_ptrs();
            unsafe {
                raw_launch(
                    self.grad_w_b_h_t_fn.cu_function(),
                    (b_size as u32, 1, 1),
                    (HIDDEN_DIM as u32, 1, 1),
                    smem,
                    self.raw_stream,
                    &mut ptrs[..args.len()],
                )
                .map_err(|e| anyhow::anyhow!("ppo_grad_w_b_h_t: {:?}", e))?;
            }
        }
        Ok(())
    }
}

/// Value head: linear projection `h_t [B, HIDDEN_DIM] → scalar V(s) [B]`.
/// Trained by MSE between `V_pred` and the bootstrap target — Phase E.2
/// computes the per-batch squared error + per-batch grad scratch via the
/// dedicated `v_head_fwd_bwd.cu` kernel pair.
pub struct ValueHead {
    #[allow(dead_code)]
    cfg: PpoHeadsConfig,
    stream: Arc<CudaStream>,
    raw_stream: CUstream,

    _module: Arc<CudaModule>,
    /// Forward kernel: `v_head_fwd`. Reads h_t and emits per-batch
    /// V_pred [B].
    pub fwd_fn: CudaFunction,
    /// Backward kernel: `v_head_bwd`. Emits per-batch loss + per-batch
    /// grad_w / grad_b scratch + OVERWRITE grad_h_t. Caller reduces
    /// grad_w / grad_b via `reduce_axis0`.
    pub bwd_fn: CudaFunction,

    /// Online weights `[1, HIDDEN_DIM]` (single output neuron).
    pub w_d: CudaSlice<f32>,
    /// Online bias `[1]`.
    pub b_d: CudaSlice<f32>,
}

impl ValueHead {
    /// Allocate device weights for the scalar value head. Uses a distinct
    /// seed offset from [`PolicyHead`] so the two heads start from
    /// independent init draws even when handed the same `cfg.seed`.
    pub fn new(dev: &MlDevice, cfg: PpoHeadsConfig) -> Result<Self> {
        let stream: Arc<CudaStream> = dev
            .cuda_stream()
            .context("value_head stream")?
            .clone();
        let ctx = dev.cuda_context().context("value_head ctx")?;
        let module = ctx
            .load_cubin(V_HEAD_CUBIN.to_vec())
            .context("load v_head_fwd_bwd cubin")?;
        let fwd_fn = module
            .load_function("v_head_fwd")
            .context("load v_head_fwd")?;
        let bwd_fn = module
            .load_function("v_head_bwd")
            .context("load v_head_bwd")?;

        // Per pearl_scoped_init_seed_for_reproducibility.
        let _seed_guard = scoped_init_seed(cfg.seed.wrapping_add(0x5EED_1234));
        let mut rng = ChaCha8Rng::seed_from_u64(cfg.seed.wrapping_add(0x5EED_1234));

        let n_out = 1;
        let n_in = cfg.hidden_dim;
        let scale = 0.01_f32 * (6.0_f32 / (n_in + n_out) as f32).sqrt();
        let w_host: Vec<f32> = (0..n_out * n_in)
            .map(|_| rng.gen_range(-scale..scale))
            .collect();
        let b_host: Vec<f32> = vec![0.0_f32; n_out];

        let w_d = upload(&stream, &w_host)?;
        let b_d = upload(&stream, &b_host)?;
        let raw_stream = stream.cu_stream();

        Ok(Self {
            cfg,
            stream,
            raw_stream,
            _module: module,
            fwd_fn,
            bwd_fn,
            w_d,
            b_d,
        })
    }

    /// Stream used to launch all kernels owned by this head.
    pub fn stream(&self) -> &Arc<CudaStream> {
        &self.stream
    }

    /// Encoder hidden dim this head was built for.
    pub fn hidden_dim(&self) -> usize {
        self.cfg.hidden_dim
    }

    /// Phase E.2 forward — V(s) = b + Σ_c w[c] × h_t[b, c], clamped to
    /// the ISV-driven C51 atom span `[V_MIN, V_MAX]` (G.1). Writes
    /// `v_pred [B]`.
    ///
    /// G.1 — `isv` is the trainer's device ISV bus; the kernel reads
    /// slots `RL_C51_V_MIN_INDEX = 485` and `RL_C51_V_MAX_INDEX = 484`
    /// to clamp the prediction. The clamp prevents stale V_pred (during
    /// reward_scale controller transitions) from amplifying advantage
    /// magnitude into the PPO surrogate + V regression loss.
    pub fn forward(
        &self,
        h_t: &CudaSlice<f32>,
        isv_dev_ptr: &u64,
        b_size: usize,
        v_pred: &mut CudaSlice<f32>,
    ) -> Result<()> {
        debug_assert_eq!(h_t.len(), b_size * HIDDEN_DIM);
        debug_assert_eq!(v_pred.len(), b_size);

        let b_i = b_size as i32;
        let smem = (HIDDEN_DIM * std::mem::size_of::<f32>()) as u32;
        {
            let mut args = RawArgs::new();
            args.push_ptr(self.w_d.raw_ptr());
            args.push_ptr(self.b_d.raw_ptr());
            args.push_ptr(h_t.raw_ptr());
            args.push_ptr(*isv_dev_ptr);
            args.push_i32(b_i);
            args.push_ptr(v_pred.raw_ptr());
            let mut ptrs = args.build_arg_ptrs();
            unsafe {
                raw_launch(
                    self.fwd_fn.cu_function(),
                    (b_size as u32, 1, 1),
                    (HIDDEN_DIM as u32, 1, 1),
                    smem,
                    self.raw_stream,
                    &mut ptrs[..args.len()],
                )
                .map_err(|e| anyhow::anyhow!("v_head_fwd: {:?}", e))?;
            }
        }
        Ok(())
    }

    /// Phase E.2 backward — MSE bwd. Emits `loss_per_batch [B]` (raw
    /// squared error; caller divides by B if mean is desired),
    /// `grad_w_per_batch [B × HIDDEN_DIM]`, `grad_b_per_batch [B]`,
    /// and `grad_h_t [B × HIDDEN_DIM]` (OVERWRITE).
    #[allow(clippy::too_many_arguments)]
    pub fn backward(
        &self,
        h_t: &CudaSlice<f32>,
        v_pred: &CudaSlice<f32>,
        r_target: &CudaSlice<f32>,
        b_size: usize,
        loss_per_batch: &mut CudaSlice<f32>,
        grad_w_per_batch: &mut CudaSlice<f32>,
        grad_b_per_batch: &mut CudaSlice<f32>,
        grad_h_t: &mut CudaSlice<f32>,
    ) -> Result<()> {
        debug_assert_eq!(h_t.len(), b_size * HIDDEN_DIM);
        debug_assert_eq!(v_pred.len(), b_size);
        debug_assert_eq!(r_target.len(), b_size);
        debug_assert_eq!(loss_per_batch.len(), b_size);
        debug_assert_eq!(grad_w_per_batch.len(), b_size * HIDDEN_DIM);
        debug_assert_eq!(grad_b_per_batch.len(), b_size);
        debug_assert_eq!(grad_h_t.len(), b_size * HIDDEN_DIM);

        let b_i = b_size as i32;
        let smem = (HIDDEN_DIM * std::mem::size_of::<f32>()) as u32;
        {
            let mut args = RawArgs::new();
            args.push_ptr(self.w_d.raw_ptr());
            args.push_ptr(h_t.raw_ptr());
            args.push_ptr(v_pred.raw_ptr());
            args.push_ptr(r_target.raw_ptr());
            args.push_i32(b_i);
            args.push_ptr(loss_per_batch.raw_ptr());
            args.push_ptr(grad_w_per_batch.raw_ptr());
            args.push_ptr(grad_b_per_batch.raw_ptr());
            args.push_ptr(grad_h_t.raw_ptr());
            let mut ptrs = args.build_arg_ptrs();
            unsafe {
                raw_launch(
                    self.bwd_fn.cu_function(),
                    (b_size as u32, 1, 1),
                    (HIDDEN_DIM as u32, 1, 1),
                    smem,
                    self.raw_stream,
                    &mut ptrs[..args.len()],
                )
                .map_err(|e| anyhow::anyhow!("v_head_bwd: {:?}", e))?;
            }
        }
        Ok(())
    }
}

// ── pinned-staging upload helper (mirrors dqn.rs::upload) ──

fn upload(stream: &Arc<CudaStream>, host: &[f32]) -> Result<CudaSlice<f32>> {
    let n = host.len();
    let staging = unsafe { MappedF32Buffer::new(n) }
        .map_err(|e| anyhow::anyhow!("ppo upload staging: {e}"))?;
    staging.write_from_slice(host);
    let mut dst = stream
        .alloc_zeros::<f32>(n)
        .context("ppo upload alloc")?;
    if n > 0 {
        let nbytes = n * std::mem::size_of::<f32>();
        unsafe {
            let (dst_ptr, _g) = dst.device_ptr_mut(stream);
            cudarc::driver::result::memcpy_dtod_async(
                dst_ptr,
                staging.dev_ptr,
                nbytes,
                stream.cu_stream(),
            )
            .context("ppo upload DtoD")?;
        }
    }
    Ok(dst)
}
