//! Factored noisy linear layer (Fortunato et al. 2017) for
//! state-dependent exploration in the C51 / IQN final projections.
//!
//! Replaces the ε-greedy probability floor (`P_MIN` in
//! `rl_pi_action_kernel.cu`) with learned, parameter-space noise.
//! During training, noise is resampled each forward pass; at eval
//! time, noise buffers are zeroed for deterministic action selection.
//!
//! ## Architecture
//!
//! ```text
//!   y[b,o] = sum_i (mu_w[o,i] + sigma_w[o,i] × f(eps_i) × f(eps_o)) × x[b,i]
//!          + mu_b[o] + sigma_b[o] × f(eps_o)
//!
//!   where f(x) = sign(x) × sqrt(|x|)   (factored noise transform)
//! ```
//!
//! Only `in_dim + out_dim` noise samples per forward (not `in_dim × out_dim`).
//!
//! ## Initialisation
//!
//!   * `mu_w`:    uniform `[-1/√in_dim, +1/√in_dim]`
//!   * `sigma_w`: constant `sigma_init / √in_dim`
//!   * `mu_b`:    uniform `[-1/√in_dim, +1/√in_dim]`
//!   * `sigma_b`: constant `sigma_init / √in_dim`
//!
//! `sigma_init` is read from `ISV[RL_NOISY_SIGMA_INIT_INDEX=547]`
//! at construction time (default 0.5, canonical NoisyNet value).
//!
//! ## Constraints honoured
//!
//!   * `feedback_no_atomicadd.md` — per-thread writes only in both kernels.
//!   * `feedback_no_htod_htoh_only_mapped_pinned.md` — weight uploads
//!     stage through `MappedF32Buffer`. Noise is device-side PRNG.
//!   * `feedback_no_nvrtc.md` — pre-compiled cubins via `build.rs`.
//!   * `feedback_isv_for_adaptive_bounds.md` — `sigma_init` in ISV slot 547.
//!   * `feedback_cpu_is_read_only.md` — forward/backward/resample are
//!     GPU kernels; host only touches init (Xavier draws + PRNG seed).
//!   * `pearl_scoped_init_seed_for_reproducibility` — seeded RNG for init.

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::pinned_mem::MappedF32Buffer;
use crate::trainer::raw_launch::{RawArgs, raw_launch};

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

const NOISY_FWD_CUBIN: &[u8] = include_bytes!(concat!(
    env!("OUT_DIR"),
    "/rl_noisy_linear_forward.cubin"
));
const NOISY_BWD_CUBIN: &[u8] = include_bytes!(concat!(
    env!("OUT_DIR"),
    "/rl_noisy_linear_backward.cubin"
));

/// Default σ_init when the ISV slot is sentinel-zero or not yet wired.
const DEFAULT_SIGMA_INIT: f32 = 0.5;

/// Construction config for [`NoisyLinear`].
#[derive(Clone, Debug)]
pub struct NoisyLinearConfig {
    /// Input dimension (columns of the weight matrix).
    pub in_dim: usize,
    /// Output dimension (rows of the weight matrix).
    pub out_dim: usize,
    /// Initial noise scale — `sigma_w[o,i] = sigma_init / sqrt(in_dim)`.
    /// Read from `ISV[RL_NOISY_SIGMA_INIT_INDEX]` at construction; falls
    /// back to `DEFAULT_SIGMA_INIT` (0.5) if the caller passes 0.0.
    pub sigma_init: f32,
    /// Random seed for Xavier init + noise RNG. Reproducibility-critical
    /// per `pearl_scoped_init_seed_for_reproducibility`.
    pub seed: u64,
}

impl Default for NoisyLinearConfig {
    fn default() -> Self {
        Self {
            in_dim: 128,
            out_dim: 11,
            sigma_init: DEFAULT_SIGMA_INIT,
            seed: 0x4E01,
        }
    }
}

/// Factored noisy linear layer. Owns device weights (`mu_w`, `sigma_w`,
/// `mu_b`, `sigma_b`) and noise buffers (`eps_in`, `eps_out`), plus
/// cached kernel handles for the forward and backward cubins.
///
/// Caller is responsible for:
///   1. `resample_noise()` before each forward pass during training.
///   2. Zeroing `eps_in` / `eps_out` at eval time (or simply skipping
///      `resample_noise` — the buffers are zero-init'd).
pub struct NoisyLinear {
    in_dim: usize,
    out_dim: usize,
    stream: Arc<CudaStream>,
    raw_stream: CUstream,

    _fwd_module: Arc<CudaModule>,
    fwd_fn: CudaFunction,

    _bwd_module: Arc<CudaModule>,
    bwd_fn: CudaFunction,

    _sample_noise_module: Arc<CudaModule>,
    sample_noise_fn: CudaFunction,

    /// Mean weight `[out_dim, in_dim]`, row-major.
    pub mu_w: CudaSlice<f32>,
    /// Noise-scale weight `[out_dim, in_dim]`, row-major.
    pub sigma_w: CudaSlice<f32>,
    /// Mean bias `[out_dim]`.
    pub mu_b: CudaSlice<f32>,
    /// Noise-scale bias `[out_dim]`.
    pub sigma_b: CudaSlice<f32>,

    /// Pre-computed `f(eps)` for input dimension `[in_dim]`.
    /// `f(x) = sign(x) * sqrt(|x|)`.
    pub eps_in: CudaSlice<f32>,
    /// Pre-computed `f(eps)` for output dimension `[out_dim]`.
    pub eps_out: CudaSlice<f32>,

    /// Device-resident xorshift32 seed `[1]` for noise PRNG.
    noisy_prng_state_d: CudaSlice<u32>,
}

impl NoisyLinear {
    /// Allocate device weights, load cubins, and cache kernel handles.
    ///
    /// Weights are initialised per the Fortunato et al. 2017 recipe:
    ///   * `mu_w`:    uniform `[-1/sqrt(in), +1/sqrt(in)]`
    ///   * `sigma_w`: constant `sigma_init / sqrt(in_dim)`
    ///   * `mu_b`:    uniform `[-1/sqrt(in), +1/sqrt(in)]`
    ///   * `sigma_b`: constant `sigma_init / sqrt(in_dim)`
    ///
    /// Noise buffers start at zero (no exploration until first
    /// `resample_noise()` call).
    pub fn new(dev: &MlDevice, cfg: NoisyLinearConfig) -> Result<Self> {
        let stream: Arc<CudaStream> = dev.cuda_stream().context("noisy_linear stream")?.clone();
        let ctx = dev.cuda_context().context("noisy_linear ctx")?;

        // Load cubins.
        let fwd_module = ctx
            .load_cubin(NOISY_FWD_CUBIN.to_vec())
            .context("load rl_noisy_linear_forward cubin")?;
        let fwd_fn = fwd_module
            .load_function("rl_noisy_linear_forward")
            .context("load rl_noisy_linear_forward fn")?;

        let bwd_module = ctx
            .load_cubin(NOISY_BWD_CUBIN.to_vec())
            .context("load rl_noisy_linear_backward cubin")?;
        let bwd_fn = bwd_module
            .load_function("rl_noisy_linear_backward")
            .context("load rl_noisy_linear_backward fn")?;

        let sample_noise_module = ctx
            .load_cubin(SAMPLE_NOISE_CUBIN.to_vec())
            .context("load rl_sample_noise cubin")?;
        let sample_noise_fn = sample_noise_module
            .load_function("rl_sample_noise")
            .context("load rl_sample_noise fn")?;

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

        let sigma_init = if cfg.sigma_init <= 0.0 {
            DEFAULT_SIGMA_INIT
        } else {
            cfg.sigma_init
        };
        let in_dim = cfg.in_dim;
        let out_dim = cfg.out_dim;
        let in_sqrt = (in_dim as f32).sqrt();
        let mu_range = 1.0_f32 / in_sqrt;
        let sigma_const = sigma_init / in_sqrt;

        // mu_w: uniform [-1/sqrt(in), +1/sqrt(in)]
        let mu_w_host: Vec<f32> = (0..out_dim * in_dim)
            .map(|_| rng.gen_range(-mu_range..mu_range))
            .collect();
        // sigma_w: constant sigma_init / sqrt(in_dim)
        let sigma_w_host: Vec<f32> = vec![sigma_const; out_dim * in_dim];
        // mu_b: uniform [-1/sqrt(in), +1/sqrt(in)]
        let mu_b_host: Vec<f32> = (0..out_dim)
            .map(|_| rng.gen_range(-mu_range..mu_range))
            .collect();
        // sigma_b: constant sigma_init / sqrt(in_dim)
        let sigma_b_host: Vec<f32> = vec![sigma_const; out_dim];

        let mu_w = upload(&stream, &mu_w_host)?;
        let sigma_w = upload(&stream, &sigma_w_host)?;
        let mu_b = upload(&stream, &mu_b_host)?;
        let sigma_b = upload(&stream, &sigma_b_host)?;

        // Noise buffers: zero-init (no exploration until resample_noise).
        let eps_in = stream
            .alloc_zeros::<f32>(in_dim)
            .context("noisy_linear eps_in alloc")?;
        let eps_out = stream
            .alloc_zeros::<f32>(out_dim)
            .context("noisy_linear eps_out alloc")?;

        // Device-resident PRNG seed for noise generation. Starts at
        // zero; the rl_sample_noise kernel self-seeds on first call
        // (no memcpy needed).
        let noisy_prng_state_d = stream
            .alloc_zeros::<u32>(1)
            .context("noisy_linear prng_state alloc")?;

        let raw_stream = stream.cu_stream();
        Ok(Self {
            in_dim,
            out_dim,
            stream,
            raw_stream,
            _fwd_module: fwd_module,
            fwd_fn,
            _bwd_module: bwd_module,
            bwd_fn,
            _sample_noise_module: sample_noise_module,
            sample_noise_fn,
            mu_w,
            sigma_w,
            mu_b,
            sigma_b,
            eps_in,
            eps_out,
            noisy_prng_state_d,
        })
    }

    /// Resample factored noise on device. Call before each training
    /// forward pass.
    ///
    /// Launches `rl_sample_noise` kernel — pure device-side xorshift32
    /// PRNG + f(x) = sign(x) × sqrt(|x|) transform. Graph-safe (no
    /// host RNG, no mapped-pinned upload in the hot path).
    pub fn resample_noise(&mut self) -> Result<()> {
        let max_dim = self.in_dim.max(self.out_dim);
        let in_dim_i = self.in_dim as i32;
        let out_dim_i = self.out_dim as i32;
        {
            let mut args = RawArgs::new();
            args.push_ptr(self.noisy_prng_state_d.raw_ptr());
            args.push_ptr(self.eps_in.raw_ptr());
            args.push_ptr(self.eps_out.raw_ptr());
            args.push_i32(in_dim_i);
            args.push_i32(out_dim_i);
            let mut ptrs = args.build_arg_ptrs();
            unsafe {
                raw_launch(
                    self.sample_noise_fn.cu_function(),
                    (1, 1, 1),
                    (max_dim as u32, 1, 1),
                    0,
                    self.raw_stream,
                    &mut ptrs[..args.len()],
                )
                .map_err(|e| anyhow::anyhow!("rl_sample_noise: {:?}", e))?;
            }
        }
        Ok(())
    }

    /// Forward pass: compute `y [B, out_dim]` from `x [B, in_dim]`.
    ///
    /// The noisy weight matrix is reconstructed on-the-fly inside the
    /// kernel from `mu_w`, `sigma_w`, `eps_in`, `eps_out`. Caller must
    /// have called `resample_noise()` before this if training.
    pub fn forward(
        &self,
        x: &CudaSlice<f32>,
        b_size: usize,
        y_out: &mut CudaSlice<f32>,
    ) -> Result<()> {
        debug_assert_eq!(x.len(), b_size * self.in_dim);
        debug_assert_eq!(y_out.len(), b_size * self.out_dim);

        let b_i = b_size as i32;
        let in_dim_i = self.in_dim as i32;
        let out_dim_i = self.out_dim as i32;
        {
            let mut args = RawArgs::new();
            args.push_ptr(x.raw_ptr());
            args.push_ptr(self.mu_w.raw_ptr());
            args.push_ptr(self.sigma_w.raw_ptr());
            args.push_ptr(self.mu_b.raw_ptr());
            args.push_ptr(self.sigma_b.raw_ptr());
            args.push_ptr(self.eps_in.raw_ptr());
            args.push_ptr(self.eps_out.raw_ptr());
            args.push_ptr(y_out.raw_ptr());
            args.push_i32(b_i);
            args.push_i32(in_dim_i);
            args.push_i32(out_dim_i);
            let mut ptrs = args.build_arg_ptrs();
            unsafe {
                raw_launch(
                    self.fwd_fn.cu_function(),
                    (b_size as u32, 1, 1),
                    (self.out_dim as u32, 1, 1),
                    0,
                    self.raw_stream,
                    &mut ptrs[..args.len()],
                )
                .map_err(|e| anyhow::anyhow!("rl_noisy_linear_forward: {:?}", e))?;
            }
        }
        Ok(())
    }

    /// Backward pass: given `grad_y [B, out_dim]`, produce per-batch
    /// gradients for `mu_w`, `sigma_w`, `mu_b`, `sigma_b`.
    ///
    /// Outputs are per-batch scratch — caller is responsible for
    /// `reduce_axis0` across batches and feeding the reduced gradients
    /// to Adam.
    #[allow(clippy::too_many_arguments)]
    pub fn backward(
        &self,
        x: &CudaSlice<f32>,
        grad_y: &CudaSlice<f32>,
        b_size: usize,
        grad_mu_w_pb: &mut CudaSlice<f32>,
        grad_sigma_w_pb: &mut CudaSlice<f32>,
        grad_mu_b_pb: &mut CudaSlice<f32>,
        grad_sigma_b_pb: &mut CudaSlice<f32>,
    ) -> Result<()> {
        debug_assert_eq!(x.len(), b_size * self.in_dim);
        debug_assert_eq!(grad_y.len(), b_size * self.out_dim);
        debug_assert_eq!(grad_mu_w_pb.len(), b_size * self.out_dim * self.in_dim);
        debug_assert_eq!(grad_sigma_w_pb.len(), b_size * self.out_dim * self.in_dim);
        debug_assert_eq!(grad_mu_b_pb.len(), b_size * self.out_dim);
        debug_assert_eq!(grad_sigma_b_pb.len(), b_size * self.out_dim);

        let b_i = b_size as i32;
        let in_dim_i = self.in_dim as i32;
        let out_dim_i = self.out_dim as i32;
        {
            let mut args = RawArgs::new();
            args.push_ptr(x.raw_ptr());
            args.push_ptr(grad_y.raw_ptr());
            args.push_ptr(self.eps_in.raw_ptr());
            args.push_ptr(self.eps_out.raw_ptr());
            args.push_i32(b_i);
            args.push_i32(in_dim_i);
            args.push_i32(out_dim_i);
            args.push_ptr(grad_mu_w_pb.raw_ptr());
            args.push_ptr(grad_sigma_w_pb.raw_ptr());
            args.push_ptr(grad_mu_b_pb.raw_ptr());
            args.push_ptr(grad_sigma_b_pb.raw_ptr());
            let mut ptrs = args.build_arg_ptrs();
            unsafe {
                raw_launch(
                    self.bwd_fn.cu_function(),
                    (b_size as u32, 1, 1),
                    (self.out_dim as u32, 1, 1),
                    0,
                    self.raw_stream,
                    &mut ptrs[..args.len()],
                )
                .map_err(|e| anyhow::anyhow!("rl_noisy_linear_backward: {:?}", e))?;
            }
        }
        Ok(())
    }

    /// Input dimension of this layer.
    pub fn in_dim(&self) -> usize {
        self.in_dim
    }

    /// Output dimension of this layer.
    pub fn out_dim(&self) -> usize {
        self.out_dim
    }

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

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