//! Forward-Return-Distribution (FRD) head — supervised forecaster over
//! 3 horizons × 21 return-bucket atoms (SP20 P3).
//!
//! Architecture (per SP20 spec §3 P3 + §4.3):
//!
//! ```text
//!   hidden [B, 64]  = ReLU(h_t [B, HIDDEN_DIM=128] @ W1 [HIDDEN_DIM, 64] + b1)
//!   logits [B, 63]  = hidden @ W2 [64, 63] + b2          // 63 = 3 × 21
//! ```
//!
//! Each row of 63 logits decomposes into 3 horizons × 21 return-bucket
//! atoms. Softmax + cross-entropy happen in the backward kernel (next
//! commit). Atom span is ISV-driven (`RL_FRD_BUCKET_RANGE_SIGMA_INDEX`,
//! default ±3σ); only the 21-atom count is structural compile-time
//! per the SP20 §0.1 audit.
//!
//! ## Why supervised forecaster (not survivor-biased classifier)
//!
//! Critical-review CRIT-1 fix: a checklist-style head trained on
//! "was this past trade profitable?" labels reads market state THROUGH
//! the surviving-policy filter. The FRD head replaces it with a forecast
//! of the FORWARD return distribution at fixed horizons — no policy
//! involvement, no survivor bias.
//!
//! ## Initialisation
//!
//! Per [`pearl_scoped_init_seed_for_reproducibility`]:
//! [`FrdHead::new`] installs a [`scoped_init_seed`] guard so Xavier
//! draws are deterministic given the seed.
//!
//!   * `w1`: Xavier uniform × 0.1 — `scale = 0.1 × sqrt(6 / (HIDDEN_DIM + FRD_HIDDEN_DIM))`
//!   * `w2`: same Xavier × 0.1 — `scale = 0.1 × sqrt(6 / (FRD_HIDDEN_DIM + FRD_OUT_DIM))`
//!   * `b1`, `b2`: zeros (keeps initial logits near 0 → softmax uniform)
//!
//! ## Constraints honoured
//!
//!   * `feedback_no_atomicadd.md` — forward writes per-batch slots; bwd
//!     will follow the same per-batch scratch + `reduce_axis0` pattern.
//!   * `feedback_no_htod_htoh_only_mapped_pinned.md` — weight uploads
//!     stage through `write_slice_f32_d_pub`.
//!   * `feedback_no_nvrtc.md` — pre-compiled cubin via `build.rs`.
//!   * `feedback_isv_for_adaptive_bounds.md` — bucket-range σ lives in
//!     ISV (slot 503), not in the kernel source.

use std::sync::Arc;

use anyhow::{Context, Result};
use cudarc::driver::{
    CudaFunction, CudaModule, CudaSlice, CudaStream,
};
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::rl::common::{FRD_HIDDEN_DIM, FRD_N_ATOMS, FRD_N_HORIZONS};
use crate::trainer::integrated::write_slice_f32_d_pub;
use crate::trainer::raw_launch::{RawArgs, raw_launch};

const FRD_FWD_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/rl_frd_fwd.cubin"));
const FRD_SOFTMAX_CE_GRAD_CUBIN: &[u8] =
    include_bytes!(concat!(env!("OUT_DIR"), "/rl_frd_softmax_ce_grad.cubin"));
const FRD_LAYER2_BWD_CUBIN: &[u8] =
    include_bytes!(concat!(env!("OUT_DIR"), "/rl_frd_layer2_bwd.cubin"));
const FRD_LAYER1_BWD_CUBIN: &[u8] =
    include_bytes!(concat!(env!("OUT_DIR"), "/rl_frd_layer1_bwd.cubin"));

/// Per-batch output width: `FRD_N_HORIZONS × FRD_N_ATOMS`. Each batch
/// row contains 3 contiguous horizon blocks of 21 atom logits.
pub const FRD_OUT_DIM: usize = FRD_N_HORIZONS * FRD_N_ATOMS;

/// Construction config for [`FrdHead`].
#[derive(Clone, Debug)]
pub struct FrdHeadConfig {
    /// Random seed for Xavier init. Reproducibility-critical per
    /// `pearl_scoped_init_seed_for_reproducibility`.
    pub seed: u64,
}

impl Default for FrdHeadConfig {
    fn default() -> Self {
        Self { seed: 0xF8D }
    }
}

/// FRD head — owns device weights for the 2-layer MLP + cached cubin
/// function handles. Forward (F.2) ships separately from backward
/// (F.3a/b/c) so each stage has its own GPU-oracle test before
/// composition.
pub struct FrdHead {
    stream: Arc<CudaStream>,
    raw_stream: CUstream,
    _module: Arc<CudaModule>,
    fwd_fn: CudaFunction,
    _softmax_ce_grad_module: Arc<CudaModule>,
    softmax_ce_grad_fn: CudaFunction,
    _layer2_bwd_module: Arc<CudaModule>,
    layer2_bwd_fn: CudaFunction,
    _layer1_bwd_module: Arc<CudaModule>,
    layer1_bwd_fn: CudaFunction,

    pub w1_d: CudaSlice<f32>, // [HIDDEN_DIM, FRD_HIDDEN_DIM]
    pub b1_d: CudaSlice<f32>, // [FRD_HIDDEN_DIM]
    pub w2_d: CudaSlice<f32>, // [FRD_HIDDEN_DIM, FRD_OUT_DIM]
    pub b2_d: CudaSlice<f32>, // [FRD_OUT_DIM]
}

impl FrdHead {
    pub fn new(dev: &MlDevice, cfg: FrdHeadConfig) -> Result<Self> {
        let stream: Arc<CudaStream> = dev.cuda_stream().context("frd stream")?.clone();
        let ctx = dev.cuda_context().context("frd ctx")?;
        let module = ctx
            .load_cubin(FRD_FWD_CUBIN.to_vec())
            .context("load rl_frd_fwd cubin")?;
        let fwd_fn = module
            .load_function("rl_frd_fwd")
            .context("load rl_frd_fwd fn")?;
        let softmax_ce_grad_module = ctx
            .load_cubin(FRD_SOFTMAX_CE_GRAD_CUBIN.to_vec())
            .context("load rl_frd_softmax_ce_grad cubin")?;
        let softmax_ce_grad_fn = softmax_ce_grad_module
            .load_function("rl_frd_softmax_ce_grad")
            .context("load rl_frd_softmax_ce_grad fn")?;
        let layer2_bwd_module = ctx
            .load_cubin(FRD_LAYER2_BWD_CUBIN.to_vec())
            .context("load rl_frd_layer2_bwd cubin")?;
        let layer2_bwd_fn = layer2_bwd_module
            .load_function("rl_frd_layer2_bwd")
            .context("load rl_frd_layer2_bwd fn")?;
        let layer1_bwd_module = ctx
            .load_cubin(FRD_LAYER1_BWD_CUBIN.to_vec())
            .context("load rl_frd_layer1_bwd cubin")?;
        let layer1_bwd_fn = layer1_bwd_module
            .load_function("rl_frd_layer1_bwd")
            .context("load rl_frd_layer1_bwd fn")?;

        // Per pearl_scoped_init_seed_for_reproducibility — guard around
        // Xavier draws so GPU init helpers downstream see the same RNG.
        let _seed_guard = scoped_init_seed(cfg.seed);
        let mut r = ChaCha8Rng::seed_from_u64(cfg.seed);

        // Xavier × 0.1 — keeps initial logits near zero so the softmax
        // is near-uniform (target entropy ≈ ln(FRD_N_ATOMS) = 3.04) and
        // CE loss is dominated by the label, not by random logit noise.
        let scale_w1 = 0.1_f32 * (6.0_f32 / (HIDDEN_DIM + FRD_HIDDEN_DIM) as f32).sqrt();
        let w1: Vec<f32> = (0..HIDDEN_DIM * FRD_HIDDEN_DIM)
            .map(|_| r.gen_range(-scale_w1..scale_w1))
            .collect();
        let scale_w2 = 0.1_f32 * (6.0_f32 / (FRD_HIDDEN_DIM + FRD_OUT_DIM) as f32).sqrt();
        let w2: Vec<f32> = (0..FRD_HIDDEN_DIM * FRD_OUT_DIM)
            .map(|_| r.gen_range(-scale_w2..scale_w2))
            .collect();
        let b1 = vec![0.0_f32; FRD_HIDDEN_DIM];
        let b2 = vec![0.0_f32; FRD_OUT_DIM];

        let mut w1_d = stream.alloc_zeros::<f32>(w1.len())?;
        write_slice_f32_d_pub(&stream, &w1, &mut w1_d)?;
        let mut b1_d = stream.alloc_zeros::<f32>(b1.len())?;
        write_slice_f32_d_pub(&stream, &b1, &mut b1_d)?;
        let mut w2_d = stream.alloc_zeros::<f32>(w2.len())?;
        write_slice_f32_d_pub(&stream, &w2, &mut w2_d)?;
        let mut b2_d = stream.alloc_zeros::<f32>(b2.len())?;
        write_slice_f32_d_pub(&stream, &b2, &mut b2_d)?;

        let raw_stream = stream.cu_stream();
        Ok(Self {
            stream,
            raw_stream,
            _module: module,
            fwd_fn,
            _softmax_ce_grad_module: softmax_ce_grad_module,
            softmax_ce_grad_fn,
            _layer2_bwd_module: layer2_bwd_module,
            layer2_bwd_fn,
            _layer1_bwd_module: layer1_bwd_module,
            layer1_bwd_fn,
            w1_d,
            b1_d,
            w2_d,
            b2_d,
        })
    }

    /// Backward stage 1 (F.3a): per-(batch, horizon) softmax + CE loss
    /// + gradient w.r.t. logits. Pairs with the forward kernel's
    /// `logits_out_d`. The `labels_d` buffer is `[B, FRD_N_HORIZONS]`
    /// of i32 — pass `-1` as the sentinel for missing horizons (the
    /// kernel zeros that row's gradient + loss).
    ///
    /// Outputs:
    ///   * `grad_logits_d [B, FRD_OUT_DIM]` — mean-reduced over batch
    ///     (multiplied by `1/B` at the source per v_head_bwd convention)
    ///   * `loss_per_batch_h_d [B, FRD_N_HORIZONS]` — raw CE per row;
    ///     caller sums + scales by `λ_frd` before backward continues
    pub fn softmax_ce_grad(
        &self,
        logits_d: &CudaSlice<f32>,
        labels_d: &CudaSlice<i32>,
        grad_logits_d: &mut CudaSlice<f32>,
        loss_per_batch_h_dev_ptr: &u64,
        b_size: usize,
    ) -> Result<()> {
        debug_assert_eq!(logits_d.len(), b_size * FRD_OUT_DIM);
        debug_assert_eq!(labels_d.len(), b_size * FRD_N_HORIZONS);
        debug_assert_eq!(grad_logits_d.len(), b_size * FRD_OUT_DIM);
        let b_size_i = b_size as i32;
        {
            let mut args = RawArgs::new();
            args.push_ptr(logits_d.raw_ptr());
            args.push_ptr(labels_d.raw_ptr());
            args.push_i32(b_size_i);
            args.push_ptr(grad_logits_d.raw_ptr());
            args.push_ptr(*loss_per_batch_h_dev_ptr);
            let mut ptrs = args.build_arg_ptrs();
            unsafe {
                raw_launch(
                    self.softmax_ce_grad_fn.cu_function(),
                    (b_size as u32, FRD_N_HORIZONS as u32, 1),
                    (FRD_N_ATOMS as u32, 1, 1),
                    0,
                    self.raw_stream,
                    &mut ptrs[..args.len()],
                )
                .map_err(|e| anyhow::anyhow!("rl_frd_softmax_ce_grad: {:?}", e))?;
            }
        }
        Ok(())
    }

    /// Backward stage 2 (F.3b): layer-2 weight gradients.
    ///
    /// Given `hidden` (cached from forward) and `grad_logits` (from
    /// `softmax_ce_grad`), computes:
    ///   * `grad_w2_per_batch[b, i, j] = hidden[b, i] × grad_logits[b, j]`
    ///   * `grad_b2_per_batch[b, j]    = grad_logits[b, j]`
    ///   * `grad_hidden[b, i]          = Σ_j W2[i, j] × grad_logits[b, j]`
    ///
    /// Per-batch scratch buffers feed `reduce_axis0` to produce the
    /// final accumulated weight grads. `grad_hidden` is the upstream
    /// gradient for layer-1 backward (F.3c).
    #[allow(clippy::too_many_arguments)]
    pub fn layer2_bwd(
        &self,
        hidden_d: &CudaSlice<f32>,
        grad_logits_d: &CudaSlice<f32>,
        grad_w2_per_batch_d: &mut CudaSlice<f32>,
        grad_b2_per_batch_d: &mut CudaSlice<f32>,
        grad_hidden_d: &mut CudaSlice<f32>,
        b_size: usize,
    ) -> Result<()> {
        debug_assert_eq!(hidden_d.len(), b_size * FRD_HIDDEN_DIM);
        debug_assert_eq!(grad_logits_d.len(), b_size * FRD_OUT_DIM);
        debug_assert_eq!(
            grad_w2_per_batch_d.len(),
            b_size * FRD_HIDDEN_DIM * FRD_OUT_DIM
        );
        debug_assert_eq!(grad_b2_per_batch_d.len(), b_size * FRD_OUT_DIM);
        debug_assert_eq!(grad_hidden_d.len(), b_size * FRD_HIDDEN_DIM);
        let b_size_i = b_size as i32;
        {
            let mut args = RawArgs::new();
            args.push_ptr(hidden_d.raw_ptr());
            args.push_ptr(grad_logits_d.raw_ptr());
            args.push_ptr(self.w2_d.raw_ptr());
            args.push_i32(b_size_i);
            args.push_ptr(grad_w2_per_batch_d.raw_ptr());
            args.push_ptr(grad_b2_per_batch_d.raw_ptr());
            args.push_ptr(grad_hidden_d.raw_ptr());
            let mut ptrs = args.build_arg_ptrs();
            unsafe {
                raw_launch(
                    self.layer2_bwd_fn.cu_function(),
                    (b_size as u32, 1, 1),
                    (FRD_HIDDEN_DIM as u32, 1, 1),
                    0,
                    self.raw_stream,
                    &mut ptrs[..args.len()],
                )
                .map_err(|e| anyhow::anyhow!("rl_frd_layer2_bwd: {:?}", e))?;
            }
        }
        Ok(())
    }

    /// Backward stage 3 (F.3c): layer-1 weight gradients with ReLU mask.
    ///
    /// Given `h_t` (encoder hidden, the layer-1 input), `hidden` (post-
    /// ReLU cache from forward — used for the ReLU mask), `grad_hidden`
    /// (from `layer2_bwd`), and `self.w1_d`, computes:
    ///   * `grad_w1_per_batch[b, k, i] = h_t[b, k] × dpre_hidden[b, i]`
    ///   * `grad_b1_per_batch[b, i]    = dpre_hidden[b, i]`
    ///   * `grad_h_t[b, k]             = Σ_i W1[k, i] × dpre_hidden[b, i]`
    ///
    /// where `dpre_hidden[b, i] = grad_hidden[b, i] × 1{hidden[b, i] > 0}`
    /// is the ReLU-masked upstream gradient.
    ///
    /// `grad_h_t` is the encoder-upstream gradient — the caller's
    /// `grad_h_accumulate` kernel folds it into the encoder's grad with
    /// a `λ_frd`-scaled accumulator (same pattern as Q / π / V heads).
    #[allow(clippy::too_many_arguments)]
    pub fn layer1_bwd(
        &self,
        h_t_d: &CudaSlice<f32>,
        hidden_d: &CudaSlice<f32>,
        grad_hidden_d: &CudaSlice<f32>,
        grad_w1_per_batch_d: &mut CudaSlice<f32>,
        grad_b1_per_batch_d: &mut CudaSlice<f32>,
        grad_h_t_d: &mut CudaSlice<f32>,
        b_size: usize,
    ) -> Result<()> {
        debug_assert_eq!(h_t_d.len(), b_size * HIDDEN_DIM);
        debug_assert_eq!(hidden_d.len(), b_size * FRD_HIDDEN_DIM);
        debug_assert_eq!(grad_hidden_d.len(), b_size * FRD_HIDDEN_DIM);
        debug_assert_eq!(
            grad_w1_per_batch_d.len(),
            b_size * HIDDEN_DIM * FRD_HIDDEN_DIM
        );
        debug_assert_eq!(grad_b1_per_batch_d.len(), b_size * FRD_HIDDEN_DIM);
        debug_assert_eq!(grad_h_t_d.len(), b_size * HIDDEN_DIM);
        let b_size_i = b_size as i32;
        {
            let mut args = RawArgs::new();
            args.push_ptr(h_t_d.raw_ptr());
            args.push_ptr(hidden_d.raw_ptr());
            args.push_ptr(grad_hidden_d.raw_ptr());
            args.push_ptr(self.w1_d.raw_ptr());
            args.push_i32(b_size_i);
            args.push_ptr(grad_w1_per_batch_d.raw_ptr());
            args.push_ptr(grad_b1_per_batch_d.raw_ptr());
            args.push_ptr(grad_h_t_d.raw_ptr());
            let mut ptrs = args.build_arg_ptrs();
            unsafe {
                raw_launch(
                    self.layer1_bwd_fn.cu_function(),
                    (b_size as u32, 1, 1),
                    (HIDDEN_DIM as u32, 1, 1),
                    0,
                    self.raw_stream,
                    &mut ptrs[..args.len()],
                )
                .map_err(|e| anyhow::anyhow!("rl_frd_layer1_bwd: {:?}", e))?;
            }
        }
        Ok(())
    }

    /// Forward pass — fills `logits_out_d [B, FRD_OUT_DIM]` and caches
    /// the post-ReLU hidden activations in `hidden_out_d [B,
    /// FRD_HIDDEN_DIM]` for the backward kernel.
    pub fn forward(
        &self,
        h_t_d: &CudaSlice<f32>,
        hidden_out_d: &mut CudaSlice<f32>,
        logits_out_d: &mut CudaSlice<f32>,
        b_size: usize,
    ) -> Result<()> {
        debug_assert_eq!(h_t_d.len(), b_size * HIDDEN_DIM);
        debug_assert_eq!(hidden_out_d.len(), b_size * FRD_HIDDEN_DIM);
        debug_assert_eq!(logits_out_d.len(), b_size * FRD_OUT_DIM);
        let b_size_i = b_size as i32;
        {
            let mut args = RawArgs::new();
            args.push_ptr(h_t_d.raw_ptr());
            args.push_ptr(self.w1_d.raw_ptr());
            args.push_ptr(self.b1_d.raw_ptr());
            args.push_ptr(self.w2_d.raw_ptr());
            args.push_ptr(self.b2_d.raw_ptr());
            args.push_ptr(hidden_out_d.raw_ptr());
            args.push_ptr(logits_out_d.raw_ptr());
            args.push_i32(b_size_i);
            let mut ptrs = args.build_arg_ptrs();
            unsafe {
                raw_launch(
                    self.fwd_fn.cu_function(),
                    (b_size as u32, 1, 1),
                    (FRD_HIDDEN_DIM as u32, 1, 1), // 64 threads -- covers both hidden and output loops
                    0,
                    self.raw_stream,
                    &mut ptrs[..args.len()],
                )
                .map_err(|e| anyhow::anyhow!("rl_frd_fwd: {:?}", e))?;
            }
        }
        Ok(())
    }

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