//! Trait for deterministic Q-value computation.
//!
//! The backtest evaluator needs Q-values from a forward pass. Instead of
//! running its own CublasForward (which creates a separate cuBLAS handle
//! with non-deterministic internal state), it calls this trait which routes
//! through the trainer's CUDA-Graphed forward pass.
//!
//! The CUDA Graph freezes the kernel launch configuration, making the
//! forward pass bit-identical across runs regardless of TF32/algo selection.

use crate::MLError;

/// Computes Q-values from states using a deterministic forward pass.
///
/// Implementors process states in their internal batch_size chunks and
/// write Q-values to the provided output pointer via DtoD copy.
pub trait QValueProvider {
    /// Maximum batch size per forward call (typically training batch_size).
    fn max_batch_size(&self) -> usize;

    /// Number of Q-value outputs per sample (total_branch_actions).
    fn total_actions(&self) -> usize;

    /// State dimension (feature count, with padding).
    fn padded_state_dim(&self) -> usize;

    /// Compute Q-values for a batch of states.
    ///
    /// `states_ptr`: device pointer to [batch_size, padded_state_dim] f32
    /// `batch_size`: number of samples (must be <= max_batch_size)
    /// `q_out_ptr`: device pointer to write [batch_size, total_actions] f32
    ///
    /// The implementation runs the cuBLAS forward + compute_expected_q
    /// through the trainer's CUDA Graph for deterministic results.
    fn compute_q_values_to(
        &mut self,
        states_ptr: u64,
        batch_size: usize,
        q_out_ptr: u64,
    ) -> Result<(), MLError>;

    /// Compute Q-values AND copy out raw branch-advantage C51 logits.
    ///
    /// Identical to `compute_q_values_to` but additionally DtoD-copies the
    /// trainer's `on_b_logits_buf` (branch-major
    /// `[batch_size, (B0+B1+B2+B3) * num_atoms]`) to `b_logits_out_ptr`. The
    /// direction-branch slice (first `batch_size * B0 * num_atoms` floats) is
    /// consumed by the Thompson direction sampler in
    /// `experience_action_select` (Plan C T2 amendment, 2026-04-29).
    ///
    /// Per-sample support and atom_positions can be retrieved via the
    /// `per_sample_support_ptr()` / `atom_positions_ptr()` accessors — those
    /// buffers are owned by the trainer and stable across calls.
    ///
    /// `b_logits_out_ptr`: device pointer to write
    /// `[batch_size, total_branch_atoms]` f32. Must be sized for the full
    /// branch-major layout (kernel reads only the direction-branch slice but
    /// the underlying DtoD copy mirrors the trainer's full buffer to keep the
    /// caller's buffer self-consistent for diagnostics).
    fn compute_q_and_b_logits_to(
        &mut self,
        states_ptr: u64,
        batch_size: usize,
        q_out_ptr: u64,
        b_logits_out_ptr: u64,
    ) -> Result<(), MLError>;

    /// Total branch-atom count = `(B0 + B1 + B2 + B3) * num_atoms`. Defines
    /// the row stride of the buffer written by `compute_q_and_b_logits_to`.
    fn total_branch_atoms(&self) -> usize;

    /// C51 atom count (`num_atoms`). Threaded into `experience_action_select`
    /// as the `n_atoms` argument so the kernel reads the right stride.
    fn num_atoms(&self) -> usize;

    /// Device pointer to the per-sample C51 support buffer
    /// (`[max_batch_size, 4, 3]` stride-12 floats `(v_min, v_max, delta_z)`).
    /// Owned by the trainer; stable across calls. Threaded into
    /// `experience_action_select` as the `per_sample_support` argument.
    fn per_sample_support_ptr(&self) -> u64;

    /// Device pointer to the C51 adaptive atom positions
    /// (`[B0, num_atoms]` floats). Owned by the trainer; stable across
    /// calls. Threaded into `experience_action_select` as the
    /// `atom_positions` argument. Pass through unchanged — the kernel
    /// internally falls back to linear support when this is NULL, but
    /// production always supplies the adaptive buffer.
    fn atom_positions_ptr(&self) -> u64;

    /// CUDA stream handle for DtoD copies.
    fn stream_handle(&self) -> u64;

    /// Compute trade-plan parameters from the trunk hidden state captured
    /// during the most recent `compute_q_values_to` call.
    ///
    /// This drives the `plan_isv[7]` features used by the network's state
    /// vector at positions [98..105). During training these are always
    /// present; during validation the `plan_isv_buf` was historically
    /// zero-filled which biased the deterministic val policy toward
    /// conservative Flat/Hold actions regardless of genuine directional
    /// edge (task #94 val-Flat-collapse audit, 2026-04-24).
    ///
    /// Implementors must:
    ///   1. Run the trade-plan MLP on the trunk hidden state produced by
    ///      the most recent forward pass.
    ///   2. Write the 6-float-per-sample plan parameters to `plan_out_ptr`.
    ///
    /// Layout of output: `[batch_size, 6]` f32 in the order
    /// `{target_bars, profit_target, stop_loss, scale_aggression,
    /// conviction, asymmetry}` — same as `plan_params_buf` in the trainer.
    ///
    /// Providers that don't support plan computation (pre-fix impls) may
    /// return `Err`. Callers should fall back to zero-fill in that case.
    fn compute_plan_params(
        &mut self,
        plan_out_ptr: u64,
        batch_size: usize,
    ) -> Result<(), MLError>;
}
