//! SDD-3 Layer B3 — auxiliary CfC trunk for outcome-supervision (D-labels).
//!
//! A SECOND, smaller CfC trunk that operates on the same encoder output
//! as the main BCE trunk, but with INDEPENDENT parameters and a smaller
//! hidden dim. Lives next to the main per-branch trunk in `cfc::trunk`.
//!
//! ## Why separate trunk (not shared)
//!
//! Per `pearl_separate_aux_trunk_when_shared_starves`: aux supervision
//! plateauing at ln(K) on a shared trunk signals starvation. Separate
//! aux trunk + own optimiser group prevents the BCE direction signal
//! from dominating gradient flow toward the aux head.
//!
//! Per the E3 design memo: asymmetric stop-grad at encoder boundary —
//! BCE drives encoder; aux gets gradient only into its own params at
//! first; lifted conditionally once aux CE < 0.4 + dir_acc > 0.85
//! within 200 steps.
//!
//! ## Algorithm
//!
//! Identical to `cfc_step_per_branch` but stripped of the per-bucket
//! channel split (single bucket, single launch grid).
//!
//!   pre = b[c] + Σ_k W_in[c, k] * x[batch, k] + Σ_k W_rec[c, k] * h_old[batch, k]
//!   decay = exp(-dt / max(tau[c], 1e-6))
//!   h_new[batch, c] = h_old[batch, c] * decay + (1 - decay) * tanh(pre)
//!
//! ## Initialisation
//!
//! Per `pearl_scoped_init_seed_for_reproducibility`: `AuxTrunk::new`
//! installs a `scoped_init_seed(seed)` guard so the Xavier draws are
//! deterministic given `seed`.
//!
//!  * `w_in`: Xavier uniform `[-sqrt(6/(feat_dim+AUX_HIDDEN)), +]`
//!  * `w_rec`: Xavier uniform `[-sqrt(6/(2*AUX_HIDDEN)), +]`
//!  * `b`: zeros
//!  * `tau`: log-uniform over `[2 s, 200 s]` (single-bucket so a wide
//!    span across all channels is fine; no per-horizon assignment).
//!
//! ## Constraints honoured
//!
//!   * `feedback_no_atomicadd.md` — no atomicAdd in fwd or bwd kernels;
//!     bwd writes per-batch grad scratch and the caller reduces axis 0.
//!   * `feedback_no_htod_htoh_only_mapped_pinned.md` — uploads go through
//!     mapped-pinned staging like the rest of `cfc/step.rs`.
//!   * `feedback_no_nvrtc.md` — pre-compiled cubin via build.rs.
//!   * `pearl_no_host_branches_in_captured_graph` — kernel uses no host
//!     branching; the `feat_dim` runtime arg is read once and used as
//!     loop bound.

use std::sync::Arc;

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

use crate::pinned_mem::MappedF32Buffer;
use crate::trainer::raw_launch::{RawArgs, raw_launch};

const KERNEL_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/aux_trunk.cubin"));

/// Aux trunk hidden dimension. Half the main trunk's HIDDEN_DIM=128.
/// Kept constant (not configurable) so cubin compile-time `#define
/// AUX_HIDDEN` stays in lockstep with the host wrapper.
pub const AUX_HIDDEN: usize = 64;

/// Upper bound on the input feature dim. Matches the cubin's
/// `MAX_FEAT_DIM = 128`. Covers both raw-snap input (FEATURE_DIM=40)
/// and post-encoder input (HIDDEN_DIM=128).
pub const MAX_AUX_FEAT_DIM: usize = 128;

/// Tau initialisation range (seconds). Single-bucket trunk gets a wide
/// log-uniform span so the aux head sees diverse memory horizons within
/// one trunk. Matches the spirit of the main trunk's [10 ms, 1000 s]
/// span but narrowed since aux supervision targets ~1k-step D-label
/// outcomes — too-long τ wastes capacity on horizons the head doesn't
/// supervise.
const AUX_TAU_LO_S: f32 = 2.0;
const AUX_TAU_HI_S: f32 = 200.0;

/// Host-side weight storage for `AuxTrunk` (used by checkpoint round-trip).
#[derive(Clone, Debug)]
pub struct AuxTrunkWeights {
    pub w_in: Vec<f32>,  // [AUX_HIDDEN × feat_dim]
    pub w_rec: Vec<f32>, // [AUX_HIDDEN × AUX_HIDDEN]
    pub b: Vec<f32>,     // [AUX_HIDDEN]
    pub tau: Vec<f32>,   // [AUX_HIDDEN]
    pub feat_dim: usize,
}

/// Construction config for `AuxTrunk`.
#[derive(Clone, Debug)]
pub struct AuxTrunkConfig {
    /// Input feature dimensionality. Must be ≤ `MAX_AUX_FEAT_DIM` so the
    /// cubin's shared-memory staging buffer fits.
    pub feat_dim: usize,
    /// Random seed for Xavier + tau init. Reproducibility-critical per
    /// `pearl_scoped_init_seed_for_reproducibility`.
    pub seed: u64,
}

/// Auxiliary CfC trunk. Owns device weights + a cached cubin module
/// with `aux_trunk_fwd` / `aux_trunk_bwd` function handles.
pub struct AuxTrunk {
    cfg: AuxTrunkConfig,
    stream: Arc<CudaStream>,

    // Cubin module kept alive so the cached `CudaFunction` handles stay
    // valid for the trunk's lifetime.
    _module: Arc<CudaModule>,
    pub fwd_fn: CudaFunction,
    pub bwd_fn: CudaFunction,

    // Device weights. Public so the trainer (B5) can mutate via AdamW
    // and the checkpoint envelope (B5) can memcpy_htod from a loaded
    // file.
    pub w_in_d: CudaSlice<f32>,  // [AUX_HIDDEN × feat_dim]
    pub w_rec_d: CudaSlice<f32>, // [AUX_HIDDEN × AUX_HIDDEN]
    pub b_d: CudaSlice<f32>,     // [AUX_HIDDEN]
    pub tau_d: CudaSlice<f32>,   // [AUX_HIDDEN]
}

impl AuxTrunk {
    /// Construct a fresh `AuxTrunk` with seeded Xavier + log-uniform τ
    /// initialisation. Loads the `aux_trunk.cubin` and resolves the
    /// `aux_trunk_fwd` / `aux_trunk_bwd` function handles once.
    pub fn new(dev: &MlDevice, cfg: AuxTrunkConfig) -> Result<Self> {
        anyhow::ensure!(
            cfg.feat_dim > 0 && cfg.feat_dim <= MAX_AUX_FEAT_DIM,
            "AuxTrunk feat_dim={} out of (0, {}]",
            cfg.feat_dim,
            MAX_AUX_FEAT_DIM,
        );

        let stream: Arc<CudaStream> = dev.cuda_stream().context("aux_trunk stream")?.clone();
        let ctx = dev.cuda_context().context("aux_trunk ctx")?;
        let module = ctx
            .load_cubin(KERNEL_CUBIN.to_vec())
            .context("load aux_trunk cubin")?;
        let fwd_fn = module
            .load_function("aux_trunk_fwd")
            .context("load aux_trunk_fwd")?;
        let bwd_fn = module
            .load_function("aux_trunk_bwd")
            .context("load aux_trunk_bwd")?;

        // Deterministic init guard — every random draw below funnels
        // through `scoped_init_seed`'s thread-local RNG so the same
        // cfg.seed produces the same weights bit-for-bit.
        let _seed_guard = scoped_init_seed(cfg.seed);
        let mut r = ChaCha8Rng::seed_from_u64(cfg.seed);
        let scale_in = (6.0_f32 / (cfg.feat_dim + AUX_HIDDEN) as f32).sqrt();
        let scale_rec = (6.0_f32 / (2 * AUX_HIDDEN) as f32).sqrt();
        let w_in: Vec<f32> = (0..AUX_HIDDEN * cfg.feat_dim)
            .map(|_| r.gen_range(-scale_in..scale_in))
            .collect();
        let w_rec: Vec<f32> = (0..AUX_HIDDEN * AUX_HIDDEN)
            .map(|_| r.gen_range(-scale_rec..scale_rec))
            .collect();
        let b: Vec<f32> = vec![0.0; AUX_HIDDEN];
        // tau: log-uniform over [AUX_TAU_LO_S, AUX_TAU_HI_S].
        let lo = AUX_TAU_LO_S.ln();
        let hi = AUX_TAU_HI_S.ln();
        let tau: Vec<f32> = (0..AUX_HIDDEN)
            .map(|_| {
                let u: f32 = r.gen_range(0.0..1.0);
                (lo + u * (hi - lo)).exp()
            })
            .collect();

        let w_in_d = upload(&stream, &w_in)?;
        let w_rec_d = upload(&stream, &w_rec)?;
        let b_d = upload(&stream, &b)?;
        let tau_d = upload(&stream, &tau)?;

        Ok(Self {
            cfg,
            stream,
            _module: module,
            fwd_fn,
            bwd_fn,
            w_in_d,
            w_rec_d,
            b_d,
            tau_d,
        })
    }

    pub fn config(&self) -> &AuxTrunkConfig {
        &self.cfg
    }

    pub fn stream(&self) -> &Arc<CudaStream> {
        &self.stream
    }

    /// Download the trunk's current weights into host vectors. Used by
    /// checkpoint save and unit tests.
    pub fn download_weights(&self) -> Result<AuxTrunkWeights> {
        let mut w_in = vec![0.0_f32; self.w_in_d.len()];
        let mut w_rec = vec![0.0_f32; self.w_rec_d.len()];
        let mut b = vec![0.0_f32; self.b_d.len()];
        let mut tau = vec![0.0_f32; self.tau_d.len()];
        self.stream
            .memcpy_dtoh(&self.w_in_d, w_in.as_mut_slice())
            .context("aux_trunk dtoh w_in")?;
        self.stream
            .memcpy_dtoh(&self.w_rec_d, w_rec.as_mut_slice())
            .context("aux_trunk dtoh w_rec")?;
        self.stream
            .memcpy_dtoh(&self.b_d, b.as_mut_slice())
            .context("aux_trunk dtoh b")?;
        self.stream
            .memcpy_dtoh(&self.tau_d, tau.as_mut_slice())
            .context("aux_trunk dtoh tau")?;
        Ok(AuxTrunkWeights {
            w_in,
            w_rec,
            b,
            tau,
            feat_dim: self.cfg.feat_dim,
        })
    }
}

/// Launch the fused aux-trunk forward kernel.
///
/// Buffer contract:
///   * `x_d`     : `[B × feat_dim]`        per-batch encoder output
///   * `h_old_d` : `[B × AUX_HIDDEN]`      per-batch prior aux-trunk state
///   * `h_new_d` : `[B × AUX_HIDDEN]`      overwrite — new aux-trunk state
///
/// Shared memory layout: `(feat_dim + AUX_HIDDEN) * sizeof(float)`
/// covering `x_local` + `h_old_local` cooperative staging.
#[allow(clippy::too_many_arguments)]
pub fn aux_trunk_fwd_gpu(
    stream: &Arc<CudaStream>,
    func: &CudaFunction,
    feat_dim: usize,
    w_in_d: &CudaSlice<f32>,
    w_rec_d: &CudaSlice<f32>,
    b_d: &CudaSlice<f32>,
    tau_d: &CudaSlice<f32>,
    x_d: &CudaSlice<f32>,
    h_old_d: &CudaSlice<f32>,
    dt_s: f32,
    b_sz: i32,
    h_new_d: &mut CudaSlice<f32>,
) -> Result<()> {
    debug_assert_eq!(w_in_d.len(), AUX_HIDDEN * feat_dim);
    debug_assert_eq!(w_rec_d.len(), AUX_HIDDEN * AUX_HIDDEN);
    debug_assert_eq!(b_d.len(), AUX_HIDDEN);
    debug_assert_eq!(tau_d.len(), AUX_HIDDEN);
    debug_assert_eq!(x_d.len(), (b_sz as usize) * feat_dim);
    debug_assert_eq!(h_old_d.len(), (b_sz as usize) * AUX_HIDDEN);
    debug_assert_eq!(h_new_d.len(), (b_sz as usize) * AUX_HIDDEN);

    let feat_dim_i: i32 = feat_dim as i32;
    let smem = ((feat_dim + AUX_HIDDEN) * std::mem::size_of::<f32>()) as u32;
    let rs = stream.cu_stream();
    {
        let mut args = RawArgs::new();
        args.push_ptr(w_in_d.raw_ptr());
        args.push_ptr(w_rec_d.raw_ptr());
        args.push_ptr(b_d.raw_ptr());
        args.push_ptr(tau_d.raw_ptr());
        args.push_ptr(x_d.raw_ptr());
        args.push_ptr(h_old_d.raw_ptr());
        args.push_f32(dt_s);
        args.push_i32(b_sz);
        args.push_i32(feat_dim_i);
        args.push_ptr(h_new_d.raw_ptr());
        let mut ptrs = args.build_arg_ptrs();
        unsafe {
            raw_launch(
                func.cu_function(),
                (b_sz as u32, 1, 1),
                (AUX_HIDDEN as u32, 1, 1),
                smem,
                rs,
                &mut ptrs[..args.len()],
            )
            .map_err(|e| anyhow::anyhow!("aux_trunk_fwd: {:?}", e))?;
        }
    }
    Ok(())
}

/// Launch the fused aux-trunk backward kernel.
///
/// Per-batch grad scratch is OVERWRITTEN by this kernel; the caller is
/// responsible for the cross-batch reduction (e.g. `reduce_axis0`).
///
/// Shared memory layout: `(feat_dim + 2 * AUX_HIDDEN) * sizeof(float)`
/// covering `x_local` + `h_old_local` + `d_pre_smem` (the per-channel
/// `d_pre` staging used for the in-block `grad_x` tree-reduce).
#[allow(clippy::too_many_arguments)]
pub fn aux_trunk_bwd_gpu(
    stream: &Arc<CudaStream>,
    func: &CudaFunction,
    feat_dim: usize,
    w_in_d: &CudaSlice<f32>,
    w_rec_d: &CudaSlice<f32>,
    b_d: &CudaSlice<f32>,
    tau_d: &CudaSlice<f32>,
    x_d: &CudaSlice<f32>,
    h_old_d: &CudaSlice<f32>,
    grad_h_new_d: &CudaSlice<f32>,
    dt_s: f32,
    b_sz: i32,
    grad_w_in_d: &mut CudaSlice<f32>,
    grad_w_rec_d: &mut CudaSlice<f32>,
    grad_b_d: &mut CudaSlice<f32>,
    grad_tau_d: &mut CudaSlice<f32>,
    grad_h_old_d: &mut CudaSlice<f32>,
    grad_x_d: &mut CudaSlice<f32>,
) -> Result<()> {
    let b_sz_u = b_sz as usize;
    debug_assert_eq!(w_in_d.len(), AUX_HIDDEN * feat_dim);
    debug_assert_eq!(w_rec_d.len(), AUX_HIDDEN * AUX_HIDDEN);
    debug_assert_eq!(b_d.len(), AUX_HIDDEN);
    debug_assert_eq!(tau_d.len(), AUX_HIDDEN);
    debug_assert_eq!(x_d.len(), b_sz_u * feat_dim);
    debug_assert_eq!(h_old_d.len(), b_sz_u * AUX_HIDDEN);
    debug_assert_eq!(grad_h_new_d.len(), b_sz_u * AUX_HIDDEN);
    debug_assert_eq!(grad_w_in_d.len(), b_sz_u * AUX_HIDDEN * feat_dim);
    debug_assert_eq!(grad_w_rec_d.len(), b_sz_u * AUX_HIDDEN * AUX_HIDDEN);
    debug_assert_eq!(grad_b_d.len(), b_sz_u * AUX_HIDDEN);
    debug_assert_eq!(grad_tau_d.len(), b_sz_u * AUX_HIDDEN);
    debug_assert_eq!(grad_h_old_d.len(), b_sz_u * AUX_HIDDEN);
    debug_assert_eq!(grad_x_d.len(), b_sz_u * feat_dim);

    let feat_dim_i: i32 = feat_dim as i32;
    let smem = ((feat_dim + 2 * AUX_HIDDEN) * std::mem::size_of::<f32>()) as u32;
    let rs = stream.cu_stream();
    {
        let mut args = RawArgs::new();
        args.push_ptr(w_in_d.raw_ptr());
        args.push_ptr(w_rec_d.raw_ptr());
        args.push_ptr(b_d.raw_ptr());
        args.push_ptr(tau_d.raw_ptr());
        args.push_ptr(x_d.raw_ptr());
        args.push_ptr(h_old_d.raw_ptr());
        args.push_ptr(grad_h_new_d.raw_ptr());
        args.push_f32(dt_s);
        args.push_i32(b_sz);
        args.push_i32(feat_dim_i);
        args.push_ptr(grad_w_in_d.raw_ptr());
        args.push_ptr(grad_w_rec_d.raw_ptr());
        args.push_ptr(grad_b_d.raw_ptr());
        args.push_ptr(grad_tau_d.raw_ptr());
        args.push_ptr(grad_h_old_d.raw_ptr());
        args.push_ptr(grad_x_d.raw_ptr());
        let mut ptrs = args.build_arg_ptrs();
        unsafe {
            raw_launch(
                func.cu_function(),
                (b_sz as u32, 1, 1),
                (AUX_HIDDEN as u32, 1, 1),
                smem,
                rs,
                &mut ptrs[..args.len()],
            )
            .map_err(|e| anyhow::anyhow!("aux_trunk_bwd: {:?}", e))?;
        }
    }
    Ok(())
}

// ── pinned-staging upload helper (mirrors cfc/step.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!("aux_trunk upload staging: {e}"))?;
    staging.write_from_slice(host);
    let mut dst = stream
        .alloc_zeros::<f32>(n)
        .context("aux_trunk 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("aux_trunk upload DtoD")?;
        }
    }
    Ok(dst)
}
