//! Weight extraction and GPU upload for DQN experience collection kernel.
//!
//! Extracts neural network weights from `GpuVarStore` objects and copies
//! them as flat `CudaSlice<f32>` GPU buffers for the CUDA experience
//! collection kernels (`experience_kernels.cu` + cuBLAS).
//!
//! Two weight sets are managed:
//! - **Dueling Q-Network** (online + target): shared layers, value head, advantage head (12 tensors each)
//! - **Curiosity Forward Model**: two fully-connected layers (4 tensors)
//!
//! Weights are stored row-major `[out_features, in_features]`. The CUDA
//! kernel reads them in the same layout.
//!
//! All weights are `CudaSlice<f32>` on GPU. Extraction and sync use
//! device-to-device copy (zero CPU roundtrip).
// CUDA FFI module — memcpy_dtod_async requires unsafe by design.
#![allow(unsafe_code)]

use std::sync::Arc;

use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DeviceRepr, PushKernelArg};
use ml_core::cuda_autograd::GpuVarStore;
use tracing::info;

use crate::MLError;

// ---------------------------------------------------------------------------
// Weight name constants
// ---------------------------------------------------------------------------

/// Dueling Q-Network weight tensor names (12 total).
const DUELING_WEIGHT_NAMES: [&str; 12] = [
    "shared_0.weight",
    "shared_0.bias",
    "shared_1.weight",
    "shared_1.bias",
    "value_fc.weight",
    "value_fc.bias",
    "value_out.weight",
    "value_out.bias",
    "advantage_fc.weight",
    "advantage_fc.bias",
    "advantage_out.weight",
    "advantage_out.bias",
];

/// Curiosity forward model weight tensor names (4 total).
const CURIOSITY_WEIGHT_NAMES: [&str; 4] = [
    "fc1.weight",
    "fc1.bias",
    "fc2.weight",
    "fc2.bias",
];

/// Branching DQN extra head weight tensor names (8 total).
///
/// In branching mode the standard advantage head becomes the exposure head (branch 0).
/// These names cover the order head (branch 1) and urgency head (branch 2).
const BRANCHING_WEIGHT_NAMES: [&str; 8] = [
    "branch_1_fc.weight",
    "branch_1_fc.bias",
    "branch_1_out.weight",
    "branch_1_out.bias",
    "branch_2_fc.weight",
    "branch_2_fc.bias",
    "branch_2_out.weight",
    "branch_2_out.bias",
];

/// PPO Actor weight tensor names (6 total).
const PPO_ACTOR_WEIGHT_NAMES: [&str; 6] = [
    "policy_layer_0.weight",
    "policy_layer_0.bias",
    "policy_layer_1.weight",
    "policy_layer_1.bias",
    "policy_output.weight",
    "policy_output.bias",
];

/// PPO Critic weight tensor names (12 total).
const PPO_CRITIC_WEIGHT_NAMES: [&str; 12] = [
    "value_layer_0.weight",
    "value_layer_0.bias",
    "value_layer_1.weight",
    "value_layer_1.bias",
    "value_layer_2.weight",
    "value_layer_2.bias",
    "value_layer_3.weight",
    "value_layer_3.bias",
    "value_layer_4.weight",
    "value_layer_4.bias",
    "value_output.weight",
    "value_output.bias",
];

// ---------------------------------------------------------------------------
// Structs
// ---------------------------------------------------------------------------

/// GPU buffers holding all 12 weight tensors of a Dueling Q-Network.
///
/// When used with `BranchingDuelingQNetwork` (distributional C51 always enabled):
/// - Shared layers: `[SHARED_H1, STATE_DIM]`, `[SHARED_H1]`, `[SHARED_H2, SHARED_H1]`, `[SHARED_H2]`
/// - Value head:    `[VALUE_H, SHARED_H2]`, `[VALUE_H]`, `[NUM_ATOMS, VALUE_H]`, `[NUM_ATOMS]`
/// - Exposure head: `[ADV_H, SHARED_H2]`, `[ADV_H]`, `[BRANCH_0_SIZE*NUM_ATOMS, ADV_H]`, `[BRANCH_0_SIZE*NUM_ATOMS]`
///
/// Defaults: SHARED_H1=SHARED_H2=256, VALUE_H=ADV_H=128, NUM_ATOMS=51, BRANCH_0_SIZE=5.
#[allow(missing_debug_implementations)] // CudaSlice does not implement Debug
pub struct DuelingWeightSet {
    // Shared layers
    pub w_s1: CudaSlice<f32>, // shared_0.weight [SHARED_H1, STATE_DIM]
    pub b_s1: CudaSlice<f32>, // shared_0.bias   [SHARED_H1]
    pub w_s2: CudaSlice<f32>, // shared_1.weight [SHARED_H2, SHARED_H1]
    pub b_s2: CudaSlice<f32>, // shared_1.bias   [SHARED_H2]
    // Value head (distributional: NUM_ATOMS outputs, not 1)
    pub w_v1: CudaSlice<f32>, // value_fc.weight  [VALUE_H, SHARED_H2]
    pub b_v1: CudaSlice<f32>, // value_fc.bias    [VALUE_H]
    pub w_v2: CudaSlice<f32>, // value_out.weight [NUM_ATOMS, VALUE_H]
    pub b_v2: CudaSlice<f32>, // value_out.bias   [NUM_ATOMS]
    // Exposure head / branch 0 (distributional: BRANCH_0_SIZE*NUM_ATOMS outputs)
    pub w_a1: CudaSlice<f32>, // branch_0_fc.weight  [ADV_H, SHARED_H2]
    pub b_a1: CudaSlice<f32>, // branch_0_fc.bias    [ADV_H]
    pub w_a2: CudaSlice<f32>, // branch_0_out.weight [BRANCH_0_SIZE*NUM_ATOMS, ADV_H]
    pub b_a2: CudaSlice<f32>, // branch_0_out.bias   [BRANCH_0_SIZE*NUM_ATOMS]
}

/// GPU buffers holding the 2 extra advantage head weight sets for Branching DQN.
///
/// In branching mode, the standard advantage head (w_a1/b_a1/w_a2/b_a2 in `DuelingWeightSet`)
/// becomes the exposure head (branch 0). This struct holds the order head (branch 1)
/// and urgency head (branch 2).
///
/// Layout matches `BranchingDuelingQNetwork` in `crates/ml-dqn/src/branching.rs`:
/// - Order FC:    `[ADV_H, SHARED_H2]`, `[ADV_H]`
/// - Order Out:   `[BRANCH_1_SIZE*NUM_ATOMS, ADV_H]`, `[BRANCH_1_SIZE*NUM_ATOMS]`
/// - Urgency FC:  `[ADV_H, SHARED_H2]`, `[ADV_H]`
/// - Urgency Out: `[BRANCH_2_SIZE*NUM_ATOMS, ADV_H]`, `[BRANCH_2_SIZE*NUM_ATOMS]`
///
/// Defaults: ADV_H=128, SHARED_H2=256, BRANCH_1_SIZE=BRANCH_2_SIZE=3, NUM_ATOMS=51.
#[allow(missing_debug_implementations)] // CudaSlice does not implement Debug
pub struct BranchingWeightSet {
    // Order head (branch 1)
    pub w_bo1: CudaSlice<f32>, // branch_1_fc.weight  [ADV_H, SHARED_H2]
    pub b_bo1: CudaSlice<f32>, // branch_1_fc.bias    [ADV_H]
    pub w_bo2: CudaSlice<f32>, // branch_1_out.weight [BRANCH_1_SIZE*NUM_ATOMS, ADV_H]
    pub b_bo2: CudaSlice<f32>, // branch_1_out.bias   [BRANCH_1_SIZE*NUM_ATOMS]
    // Urgency head (branch 2)
    pub w_bu1: CudaSlice<f32>, // branch_2_fc.weight  [ADV_H, SHARED_H2]
    pub b_bu1: CudaSlice<f32>, // branch_2_fc.bias    [ADV_H]
    pub w_bu2: CudaSlice<f32>, // branch_2_out.weight [BRANCH_2_SIZE*NUM_ATOMS, ADV_H]
    pub b_bu2: CudaSlice<f32>, // branch_2_out.bias   [BRANCH_2_SIZE*NUM_ATOMS]
}

/// GPU buffers holding all 4 weight tensors of the Curiosity Forward Model.
///
/// Layout matches `ForwardDynamicsModel` in `crates/ml/src/dqn/curiosity.rs`:
/// - Layer 1: `[64, 35]`, `[64]`
/// - Layer 2: `[32, 64]`, `[32]`
#[allow(missing_debug_implementations)] // CudaSlice does not implement Debug
/// Curiosity forward model dimensions (must match CUDA defines).
/// MARKET_DIM=42, CUR_HIDDEN=128, CUR_INPUT=45, CUR_OUTPUT=42
const CUR_MARKET_DIM: usize = 42;
const CUR_HIDDEN_DIM: usize = 128;
const CUR_INPUT_DIM: usize = CUR_MARKET_DIM + 3; // 45
const CUR_OUTPUT_DIM: usize = CUR_MARKET_DIM;     // 42

#[allow(missing_debug_implementations)] // CudaSlice does not implement Debug
pub struct CuriosityWeightSet {
    pub w1: CudaSlice<f32>, // fc1.weight [CUR_HIDDEN, CUR_INPUT] = [128, 45]
    pub b1: CudaSlice<f32>, // fc1.bias   [CUR_HIDDEN] = [128]
    pub w2: CudaSlice<f32>, // fc2.weight [CUR_OUTPUT, CUR_HIDDEN] = [42, 128]
    pub b2: CudaSlice<f32>, // fc2.bias   [CUR_OUTPUT] = [42]
}

impl CuriosityWeightSet {
    /// Create zero-filled curiosity weights on GPU.
    ///
    /// Used when the curiosity module is disabled (`curiosity_weight = 0`).
    /// The CUDA kernel still runs the curiosity forward pass but produces
    /// zero intrinsic reward (all-zero weights → zero output, and
    /// `curiosity_scale` is set to 0.0 in the config).
    pub fn zeros(stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        Ok(Self {
            w1: stream.alloc_zeros::<f32>(CUR_HIDDEN_DIM * CUR_INPUT_DIM).map_err(|e| {
                MLError::ModelError(format!("Failed to alloc curiosity w1 zeros: {e}"))
            })?,
            b1: stream.alloc_zeros::<f32>(CUR_HIDDEN_DIM).map_err(|e| {
                MLError::ModelError(format!("Failed to alloc curiosity b1 zeros: {e}"))
            })?,
            w2: stream.alloc_zeros::<f32>(CUR_OUTPUT_DIM * CUR_HIDDEN_DIM).map_err(|e| {
                MLError::ModelError(format!("Failed to alloc curiosity w2 zeros: {e}"))
            })?,
            b2: stream.alloc_zeros::<f32>(CUR_OUTPUT_DIM).map_err(|e| {
                MLError::ModelError(format!("Failed to alloc curiosity b2 zeros: {e}"))
            })?,
        })
    }
}

/// GPU buffers holding all 6 weight tensors of the PPO Actor (PolicyNetwork).
///
/// Layout matches `PolicyNetwork` in `crates/ml/src/ppo/ppo.rs`:
/// - Layer 0: `[128, 54]`, `[128]`
/// - Layer 1: `[64, 128]`, `[64]`
/// - Output:  `[45, 64]`,  `[45]`
#[allow(missing_debug_implementations)] // CudaSlice does not implement Debug
pub struct PpoActorWeightSet {
    pub pw1: CudaSlice<f32>, // policy_layer_0.weight [128, 54]
    pub pb1: CudaSlice<f32>, // policy_layer_0.bias   [128]
    pub pw2: CudaSlice<f32>, // policy_layer_1.weight [64, 128]
    pub pb2: CudaSlice<f32>, // policy_layer_1.bias   [64]
    pub pw3: CudaSlice<f32>, // policy_output.weight  [45, 64]
    pub pb3: CudaSlice<f32>, // policy_output.bias    [45]
}

/// GPU buffers holding all 12 weight tensors of the PPO Critic (ValueNetwork).
///
/// Layout matches `ValueNetwork` in `crates/ml/src/ppo/ppo.rs`:
/// - Layer 0: `[512, 54]`, `[512]`
/// - Layer 1: `[384, 512]`, `[384]`
/// - Layer 2: `[256, 384]`, `[256]`
/// - Layer 3: `[128, 256]`, `[128]`
/// - Layer 4: `[64, 128]`, `[64]`
/// - Output:  `[1, 64]`, `[1]`
#[allow(missing_debug_implementations)] // CudaSlice does not implement Debug
pub struct PpoCriticWeightSet {
    pub vw1: CudaSlice<f32>,  // value_layer_0.weight [512, 54]
    pub vb1: CudaSlice<f32>,  // value_layer_0.bias   [512]
    pub vw2: CudaSlice<f32>,  // value_layer_1.weight [384, 512]
    pub vb2: CudaSlice<f32>,  // value_layer_1.bias   [384]
    pub vw3: CudaSlice<f32>,  // value_layer_2.weight [256, 384]
    pub vb3: CudaSlice<f32>,  // value_layer_2.bias   [256]
    pub vw4: CudaSlice<f32>,  // value_layer_3.weight [128, 256]
    pub vb4: CudaSlice<f32>,  // value_layer_3.bias   [128]
    pub vw5: CudaSlice<f32>,  // value_layer_4.weight [64, 128]
    pub vb5: CudaSlice<f32>,  // value_layer_4.bias   [64]
    pub vw6: CudaSlice<f32>,  // value_output.weight  [1, 64]
    pub vb6: CudaSlice<f32>,  // value_output.bias    [1]
}

/// RMSNorm gamma weight names for distributional dueling network (4 total).
const RMSNORM_WEIGHT_NAMES: [&str; 4] = [
    "shared_rmsnorm_0.weight",
    "shared_rmsnorm_1.weight",
    "value_rmsnorm.weight",
    "advantage_rmsnorm.weight",
];

/// GPU buffers for RMSNorm gamma weights used by DistributionalDuelingQNetwork.
///
/// Each gamma vector is the same dimension as the corresponding layer output:
/// - shared_rmsnorm_0: [256] (after first shared layer)
/// - shared_rmsnorm_1: [256] (after second shared layer)
/// - value_rmsnorm:    [128] (after value FC)
/// - advantage_rmsnorm:[128] (after advantage FC)
#[allow(missing_debug_implementations)]
pub struct RmsNormWeightSet {
    pub gamma_s0: CudaSlice<f32>, // shared_rmsnorm_0.weight [shared_h1]
    pub gamma_s1: CudaSlice<f32>, // shared_rmsnorm_1.weight [shared_h2]
    pub gamma_v:  CudaSlice<f32>, // value_rmsnorm.weight    [value_h]
    pub gamma_a:  CudaSlice<f32>, // advantage_rmsnorm.weight [adv_h]
}

impl RmsNormWeightSet {
    /// Create all-ones RMSNorm weights (no-op scaling, used when distributional is disabled).
    ///
    /// Dimensions must match the actual network layer sizes to avoid GPU OOB reads.
    /// `network_dims`: `(shared_h1, shared_h2, value_h, adv_h)`.
    pub fn ones(
        stream: &Arc<CudaStream>,
        network_dims: (usize, usize, usize, usize),
    ) -> Result<Self, MLError> {
        let (shared_h1, shared_h2, value_h, adv_h) = network_dims;
        let alloc_ones = |size: usize| -> Result<CudaSlice<f32>, MLError> {
            let data = vec![1.0_f32; size];
            let mut buf = stream
                .alloc_zeros::<f32>(size)
                .map_err(|e| MLError::ModelError(format!("Alloc RMSNorm ones: {e}")))?;
            stream
                .memcpy_htod(&data, &mut buf)
                .map_err(|e| MLError::ModelError(format!("Upload RMSNorm ones: {e}")))?;
            Ok(buf)
        };
        Ok(Self {
            gamma_s0: alloc_ones(shared_h1)?,
            gamma_s1: alloc_ones(shared_h2)?,
            gamma_v:  alloc_ones(value_h)?,
            gamma_a:  alloc_ones(adv_h)?,
        })
    }
}

// ---------------------------------------------------------------------------
// KernelWeightPack — packed device pointers for kernel launch
// ---------------------------------------------------------------------------

/// Dueling Q-network weight device pointers (12 × u64).
///
/// Layout matches `struct DuelingWeightPtrs` in `common_device_functions.cuh`.
#[repr(C)]
#[derive(Clone, Copy)]
pub(super) struct GpuDuelingPtrs {
    pub w_s1: u64, pub b_s1: u64,
    pub w_s2: u64, pub b_s2: u64,
    pub w_v1: u64, pub b_v1: u64,
    pub w_v2: u64, pub b_v2: u64,
    pub w_a1: u64, pub b_a1: u64,
    pub w_a2: u64, pub b_a2: u64,
}

/// Curiosity forward model weight device pointers (4 × u64).
///
/// Layout matches `struct CuriosityWeightPtrs` in `common_device_functions.cuh`.
#[repr(C)]
#[derive(Clone, Copy)]
pub(super) struct GpuCuriosityPtrs {
    pub w1: u64, pub b1: u64,
    pub w2: u64, pub b2: u64,
}

/// RMSNorm gamma weight device pointers (4 × u64).
///
/// Layout matches `struct RMSNormWeightPtrs` in `common_device_functions.cuh`.
#[repr(C)]
#[derive(Clone, Copy)]
pub(super) struct GpuRMSNormPtrs {
    pub s0: u64,
    pub s1: u64,
    pub v: u64,
    pub a: u64,
}

/// Branching DQN extra head weight device pointers (8 × u64).
///
/// Layout matches `struct BranchWeightPtrs` in `common_device_functions.cuh`.
#[repr(C)]
#[derive(Clone, Copy)]
pub(super) struct GpuBranchPtrs {
    pub w_o1: u64, pub b_o1: u64,
    pub w_o2: u64, pub b_o2: u64,
    pub w_u1: u64, pub b_u1: u64,
    pub w_u2: u64, pub b_u2: u64,
}

/// Packed weight pointers for the DQN experience collection kernel.
///
/// Passed by value as a single kernel argument via `DeviceRepr`.
/// Total size: 48 × 8 = 384 bytes (fits within CUDA's 4 KB parameter limit).
///
/// Layout matches `struct KernelWeightPack` in `common_device_functions.cuh` EXACTLY.
/// Both use `#[repr(C)]` / C struct layout with identical field ordering.
#[repr(C)]
#[derive(Clone, Copy)]
pub(super) struct KernelWeightPack {
    pub online: GpuDuelingPtrs,
    pub target: GpuDuelingPtrs,
    pub curiosity: GpuCuriosityPtrs,
    pub rmsnorm: GpuRMSNormPtrs,
    pub online_br: GpuBranchPtrs,
    pub target_br: GpuBranchPtrs,
}

// Safety: KernelWeightPack is #[repr(C)] with only u64 fields.
// All u64 values are valid CUDA device pointers (CUdeviceptr).
#[allow(unsafe_code)]
unsafe impl DeviceRepr for KernelWeightPack {}

/// Extract raw CUDA device pointer (CUdeviceptr = u64) from a CudaSlice.
///
/// Uses `DevicePtr::device_ptr()` and wraps the `SyncOnDrop` guard in
/// `ManuallyDrop` to skip recording a read event. This is safe because:
/// - The CudaSlice is owned by the caller and outlives the kernel launch
/// - The kernel launches on the same stream, so CUDA ordering is guaranteed
/// - `SyncOnDrop::Record` only records events for cross-stream sync (not needed)
fn raw_device_ptr(slice: &CudaSlice<f32>, stream: &CudaStream) -> u64 {
    let (ptr, guard) = slice.device_ptr(stream);
    // Keep guard alive without running its Drop (skip read event recording).
    // Named binding prevents `let_underscore_must_use`; ManuallyDrop prevents drop.
    let _no_drop = std::mem::ManuallyDrop::new(guard);
    ptr
}

impl KernelWeightPack {
    /// Build a `KernelWeightPack` from the GPU weight sets.
    ///
    /// Extracts raw device pointers from each `CudaSlice` and packs them
    /// into the struct for a single kernel argument pass.
    pub(super) fn build(
        online: &DuelingWeightSet,
        target: &DuelingWeightSet,
        curiosity: &CuriosityWeightSet,
        rmsnorm: &RmsNormWeightSet,
        online_br: &BranchingWeightSet,
        target_br: &BranchingWeightSet,
        stream: &CudaStream,
    ) -> Self {
        Self {
            online: GpuDuelingPtrs {
                w_s1: raw_device_ptr(&online.w_s1, stream),
                b_s1: raw_device_ptr(&online.b_s1, stream),
                w_s2: raw_device_ptr(&online.w_s2, stream),
                b_s2: raw_device_ptr(&online.b_s2, stream),
                w_v1: raw_device_ptr(&online.w_v1, stream),
                b_v1: raw_device_ptr(&online.b_v1, stream),
                w_v2: raw_device_ptr(&online.w_v2, stream),
                b_v2: raw_device_ptr(&online.b_v2, stream),
                w_a1: raw_device_ptr(&online.w_a1, stream),
                b_a1: raw_device_ptr(&online.b_a1, stream),
                w_a2: raw_device_ptr(&online.w_a2, stream),
                b_a2: raw_device_ptr(&online.b_a2, stream),
            },
            target: GpuDuelingPtrs {
                w_s1: raw_device_ptr(&target.w_s1, stream),
                b_s1: raw_device_ptr(&target.b_s1, stream),
                w_s2: raw_device_ptr(&target.w_s2, stream),
                b_s2: raw_device_ptr(&target.b_s2, stream),
                w_v1: raw_device_ptr(&target.w_v1, stream),
                b_v1: raw_device_ptr(&target.b_v1, stream),
                w_v2: raw_device_ptr(&target.w_v2, stream),
                b_v2: raw_device_ptr(&target.b_v2, stream),
                w_a1: raw_device_ptr(&target.w_a1, stream),
                b_a1: raw_device_ptr(&target.b_a1, stream),
                w_a2: raw_device_ptr(&target.w_a2, stream),
                b_a2: raw_device_ptr(&target.b_a2, stream),
            },
            curiosity: GpuCuriosityPtrs {
                w1: raw_device_ptr(&curiosity.w1, stream),
                b1: raw_device_ptr(&curiosity.b1, stream),
                w2: raw_device_ptr(&curiosity.w2, stream),
                b2: raw_device_ptr(&curiosity.b2, stream),
            },
            rmsnorm: GpuRMSNormPtrs {
                s0: raw_device_ptr(&rmsnorm.gamma_s0, stream),
                s1: raw_device_ptr(&rmsnorm.gamma_s1, stream),
                v: raw_device_ptr(&rmsnorm.gamma_v, stream),
                a: raw_device_ptr(&rmsnorm.gamma_a, stream),
            },
            online_br: GpuBranchPtrs {
                w_o1: raw_device_ptr(&online_br.w_bo1, stream),
                b_o1: raw_device_ptr(&online_br.b_bo1, stream),
                w_o2: raw_device_ptr(&online_br.w_bo2, stream),
                b_o2: raw_device_ptr(&online_br.b_bo2, stream),
                w_u1: raw_device_ptr(&online_br.w_bu1, stream),
                b_u1: raw_device_ptr(&online_br.b_bu1, stream),
                w_u2: raw_device_ptr(&online_br.w_bu2, stream),
                b_u2: raw_device_ptr(&online_br.b_bu2, stream),
            },
            target_br: GpuBranchPtrs {
                w_o1: raw_device_ptr(&target_br.w_bo1, stream),
                b_o1: raw_device_ptr(&target_br.b_bo1, stream),
                w_o2: raw_device_ptr(&target_br.w_bo2, stream),
                b_o2: raw_device_ptr(&target_br.b_bo2, stream),
                w_u1: raw_device_ptr(&target_br.w_bu1, stream),
                b_u1: raw_device_ptr(&target_br.b_bu1, stream),
                w_u2: raw_device_ptr(&target_br.w_bu2, stream),
                b_u2: raw_device_ptr(&target_br.b_bu2, stream),
            },
        }
    }
}

// ---------------------------------------------------------------------------
// BF16 Weight Mirror — parallel BF16 copies of all F32 weight buffers
// ---------------------------------------------------------------------------

/// BF16 mirror of a DuelingWeightSet. Stored as CudaSlice<u16> where each u16
/// holds the bit pattern of an __nv_bfloat16 value.
///
/// Updated automatically after each F32 weight sync via the `f32_to_bf16_kernel`
/// CUDA kernel. The training forward pass reads from these BF16 buffers;
/// backward/Adam operate on the F32 originals.
#[allow(missing_debug_implementations)]
pub struct DuelingWeightSetBf16 {
    pub w_s1: CudaSlice<u16>,
    pub b_s1: CudaSlice<u16>,
    pub w_s2: CudaSlice<u16>,
    pub b_s2: CudaSlice<u16>,
    pub w_v1: CudaSlice<u16>,
    pub b_v1: CudaSlice<u16>,
    pub w_v2: CudaSlice<u16>,
    pub b_v2: CudaSlice<u16>,
    pub w_a1: CudaSlice<u16>,
    pub b_a1: CudaSlice<u16>,
    pub w_a2: CudaSlice<u16>,
    pub b_a2: CudaSlice<u16>,
}

/// BF16 mirror of BranchingWeightSet (order + urgency heads).
#[allow(missing_debug_implementations)]
pub struct BranchingWeightSetBf16 {
    pub w_bo1: CudaSlice<u16>,
    pub b_bo1: CudaSlice<u16>,
    pub w_bo2: CudaSlice<u16>,
    pub b_bo2: CudaSlice<u16>,
    pub w_bu1: CudaSlice<u16>,
    pub b_bu1: CudaSlice<u16>,
    pub w_bu2: CudaSlice<u16>,
    pub b_bu2: CudaSlice<u16>,
}

/// BF16 mirror of CuriosityWeightSet.
#[allow(missing_debug_implementations)]
pub struct CuriosityWeightSetBf16 {
    pub w1: CudaSlice<u16>,
    pub b1: CudaSlice<u16>,
    pub w2: CudaSlice<u16>,
    pub b2: CudaSlice<u16>,
}

/// Allocate a BF16 buffer matching an F32 buffer's element count.
fn alloc_bf16_mirror(
    f32_buf: &CudaSlice<f32>,
    stream: &Arc<CudaStream>,
) -> Result<CudaSlice<u16>, MLError> {
    let n = f32_buf.len();
    stream.alloc_zeros::<u16>(n).map_err(|e| {
        MLError::ModelError(format!("BF16 mirror alloc ({n} elems): {e}"))
    })
}

/// Convert F32 buffer → BF16 buffer on GPU using a pre-compiled CUDA kernel.
///
/// Launches `f32_to_bf16_kernel` (defined in `common_device_functions.cuh`,
/// compiled via NVRTC as part of the training kernel module).
///
/// `func`: pre-compiled `CudaFunction` for `f32_to_bf16_kernel`.
pub fn convert_f32_to_bf16(
    src: &CudaSlice<f32>,
    dst: &mut CudaSlice<u16>,
    func: &cudarc::driver::CudaFunction,
    stream: &Arc<CudaStream>,
) -> Result<(), MLError> {
    let n = src.len();
    if n == 0 { return Ok(()); }
    let block = 256_u32;
    let grid = ((n as u32) + block - 1) / block;
    let cfg = cudarc::driver::LaunchConfig {
        grid_dim: (grid, 1, 1),
        block_dim: (block, 1, 1),
        shared_mem_bytes: 0,
    };
    unsafe {
        stream
            .launch_builder(func)
            .arg(src)
            .arg(dst)
            .arg(&(n as i32))
            .launch(cfg)
            .map_err(|e| MLError::ModelError(format!("f32_to_bf16 launch: {e}")))?;
    }
    Ok(())
}

impl DuelingWeightSetBf16 {
    /// Allocate BF16 mirrors matching an existing F32 weight set.
    pub fn alloc_from(f32_set: &DuelingWeightSet, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        Ok(Self {
            w_s1: alloc_bf16_mirror(&f32_set.w_s1, stream)?,
            b_s1: alloc_bf16_mirror(&f32_set.b_s1, stream)?,
            w_s2: alloc_bf16_mirror(&f32_set.w_s2, stream)?,
            b_s2: alloc_bf16_mirror(&f32_set.b_s2, stream)?,
            w_v1: alloc_bf16_mirror(&f32_set.w_v1, stream)?,
            b_v1: alloc_bf16_mirror(&f32_set.b_v1, stream)?,
            w_v2: alloc_bf16_mirror(&f32_set.w_v2, stream)?,
            b_v2: alloc_bf16_mirror(&f32_set.b_v2, stream)?,
            w_a1: alloc_bf16_mirror(&f32_set.w_a1, stream)?,
            b_a1: alloc_bf16_mirror(&f32_set.b_a1, stream)?,
            w_a2: alloc_bf16_mirror(&f32_set.w_a2, stream)?,
            b_a2: alloc_bf16_mirror(&f32_set.b_a2, stream)?,
        })
    }

    /// Sync all 12 BF16 mirrors from F32 originals using GPU kernel.
    pub fn sync_from_f32(
        &mut self,
        f32_set: &DuelingWeightSet,
        func: &cudarc::driver::CudaFunction,
        stream: &Arc<CudaStream>,
    ) -> Result<(), MLError> {
        convert_f32_to_bf16(&f32_set.w_s1, &mut self.w_s1, func, stream)?;
        convert_f32_to_bf16(&f32_set.b_s1, &mut self.b_s1, func, stream)?;
        convert_f32_to_bf16(&f32_set.w_s2, &mut self.w_s2, func, stream)?;
        convert_f32_to_bf16(&f32_set.b_s2, &mut self.b_s2, func, stream)?;
        convert_f32_to_bf16(&f32_set.w_v1, &mut self.w_v1, func, stream)?;
        convert_f32_to_bf16(&f32_set.b_v1, &mut self.b_v1, func, stream)?;
        convert_f32_to_bf16(&f32_set.w_v2, &mut self.w_v2, func, stream)?;
        convert_f32_to_bf16(&f32_set.b_v2, &mut self.b_v2, func, stream)?;
        convert_f32_to_bf16(&f32_set.w_a1, &mut self.w_a1, func, stream)?;
        convert_f32_to_bf16(&f32_set.b_a1, &mut self.b_a1, func, stream)?;
        convert_f32_to_bf16(&f32_set.w_a2, &mut self.w_a2, func, stream)?;
        convert_f32_to_bf16(&f32_set.b_a2, &mut self.b_a2, func, stream)?;
        Ok(())
    }
}

impl BranchingWeightSetBf16 {
    /// Allocate BF16 mirrors matching an existing F32 weight set.
    pub fn alloc_from(f32_set: &BranchingWeightSet, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        Ok(Self {
            w_bo1: alloc_bf16_mirror(&f32_set.w_bo1, stream)?,
            b_bo1: alloc_bf16_mirror(&f32_set.b_bo1, stream)?,
            w_bo2: alloc_bf16_mirror(&f32_set.w_bo2, stream)?,
            b_bo2: alloc_bf16_mirror(&f32_set.b_bo2, stream)?,
            w_bu1: alloc_bf16_mirror(&f32_set.w_bu1, stream)?,
            b_bu1: alloc_bf16_mirror(&f32_set.b_bu1, stream)?,
            w_bu2: alloc_bf16_mirror(&f32_set.w_bu2, stream)?,
            b_bu2: alloc_bf16_mirror(&f32_set.b_bu2, stream)?,
        })
    }

    /// Sync all 8 BF16 mirrors from F32 originals.
    pub fn sync_from_f32(
        &mut self,
        f32_set: &BranchingWeightSet,
        func: &cudarc::driver::CudaFunction,
        stream: &Arc<CudaStream>,
    ) -> Result<(), MLError> {
        convert_f32_to_bf16(&f32_set.w_bo1, &mut self.w_bo1, func, stream)?;
        convert_f32_to_bf16(&f32_set.b_bo1, &mut self.b_bo1, func, stream)?;
        convert_f32_to_bf16(&f32_set.w_bo2, &mut self.w_bo2, func, stream)?;
        convert_f32_to_bf16(&f32_set.b_bo2, &mut self.b_bo2, func, stream)?;
        convert_f32_to_bf16(&f32_set.w_bu1, &mut self.w_bu1, func, stream)?;
        convert_f32_to_bf16(&f32_set.b_bu1, &mut self.b_bu1, func, stream)?;
        convert_f32_to_bf16(&f32_set.w_bu2, &mut self.w_bu2, func, stream)?;
        convert_f32_to_bf16(&f32_set.b_bu2, &mut self.b_bu2, func, stream)?;
        Ok(())
    }
}

impl CuriosityWeightSetBf16 {
    /// Allocate BF16 mirrors matching an existing F32 weight set.
    pub fn alloc_from(f32_set: &CuriosityWeightSet, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        Ok(Self {
            w1: alloc_bf16_mirror(&f32_set.w1, stream)?,
            b1: alloc_bf16_mirror(&f32_set.b1, stream)?,
            w2: alloc_bf16_mirror(&f32_set.w2, stream)?,
            b2: alloc_bf16_mirror(&f32_set.b2, stream)?,
        })
    }

    /// Sync all 4 BF16 mirrors from F32 originals.
    pub fn sync_from_f32(
        &mut self,
        f32_set: &CuriosityWeightSet,
        func: &cudarc::driver::CudaFunction,
        stream: &Arc<CudaStream>,
    ) -> Result<(), MLError> {
        convert_f32_to_bf16(&f32_set.w1, &mut self.w1, func, stream)?;
        convert_f32_to_bf16(&f32_set.b1, &mut self.b1, func, stream)?;
        convert_f32_to_bf16(&f32_set.w2, &mut self.w2, func, stream)?;
        convert_f32_to_bf16(&f32_set.b2, &mut self.b2, func, stream)?;
        Ok(())
    }
}

/// Extract raw CUDA device pointer from a BF16 CudaSlice (u16 bit pattern).
fn raw_device_ptr_bf16(slice: &CudaSlice<u16>, stream: &CudaStream) -> u64 {
    let (ptr, guard) = slice.device_ptr(stream);
    let _no_drop = std::mem::ManuallyDrop::new(guard);
    ptr
}

/// BF16 device pointer pack for the forward pass.
///
/// Same layout as `KernelWeightPack` but pointing to BF16 buffers.
/// The CUDA forward kernel reads these for BF16 matmul;
/// backward/Adam kernel still reads the F32 `KernelWeightPack`.
#[repr(C)]
#[derive(Clone, Copy)]
pub(super) struct KernelWeightPackBf16 {
    pub online: GpuDuelingPtrs,
    pub target: GpuDuelingPtrs,
    pub curiosity: GpuCuriosityPtrs,
    pub rmsnorm: GpuRMSNormPtrs, // RMSNorm stays F32 (gamma scaling)
    pub online_br: GpuBranchPtrs,
    pub target_br: GpuBranchPtrs,
}

#[allow(unsafe_code)]
unsafe impl DeviceRepr for KernelWeightPackBf16 {}

impl KernelWeightPackBf16 {
    /// Build BF16 weight pointer pack for forward pass.
    ///
    /// Weight matrices use BF16 pointers; RMSNorm gammas stay F32.
    pub(super) fn build(
        online_bf16: &DuelingWeightSetBf16,
        target_bf16: &DuelingWeightSetBf16,
        curiosity_bf16: &CuriosityWeightSetBf16,
        rmsnorm_f32: &RmsNormWeightSet,
        online_br_bf16: &BranchingWeightSetBf16,
        target_br_bf16: &BranchingWeightSetBf16,
        stream: &CudaStream,
    ) -> Self {
        Self {
            online: GpuDuelingPtrs {
                w_s1: raw_device_ptr_bf16(&online_bf16.w_s1, stream),
                b_s1: raw_device_ptr_bf16(&online_bf16.b_s1, stream),
                w_s2: raw_device_ptr_bf16(&online_bf16.w_s2, stream),
                b_s2: raw_device_ptr_bf16(&online_bf16.b_s2, stream),
                w_v1: raw_device_ptr_bf16(&online_bf16.w_v1, stream),
                b_v1: raw_device_ptr_bf16(&online_bf16.b_v1, stream),
                w_v2: raw_device_ptr_bf16(&online_bf16.w_v2, stream),
                b_v2: raw_device_ptr_bf16(&online_bf16.b_v2, stream),
                w_a1: raw_device_ptr_bf16(&online_bf16.w_a1, stream),
                b_a1: raw_device_ptr_bf16(&online_bf16.b_a1, stream),
                w_a2: raw_device_ptr_bf16(&online_bf16.w_a2, stream),
                b_a2: raw_device_ptr_bf16(&online_bf16.b_a2, stream),
            },
            target: GpuDuelingPtrs {
                w_s1: raw_device_ptr_bf16(&target_bf16.w_s1, stream),
                b_s1: raw_device_ptr_bf16(&target_bf16.b_s1, stream),
                w_s2: raw_device_ptr_bf16(&target_bf16.w_s2, stream),
                b_s2: raw_device_ptr_bf16(&target_bf16.b_s2, stream),
                w_v1: raw_device_ptr_bf16(&target_bf16.w_v1, stream),
                b_v1: raw_device_ptr_bf16(&target_bf16.b_v1, stream),
                w_v2: raw_device_ptr_bf16(&target_bf16.w_v2, stream),
                b_v2: raw_device_ptr_bf16(&target_bf16.b_v2, stream),
                w_a1: raw_device_ptr_bf16(&target_bf16.w_a1, stream),
                b_a1: raw_device_ptr_bf16(&target_bf16.b_a1, stream),
                w_a2: raw_device_ptr_bf16(&target_bf16.w_a2, stream),
                b_a2: raw_device_ptr_bf16(&target_bf16.b_a2, stream),
            },
            curiosity: GpuCuriosityPtrs {
                w1: raw_device_ptr_bf16(&curiosity_bf16.w1, stream),
                b1: raw_device_ptr_bf16(&curiosity_bf16.b1, stream),
                w2: raw_device_ptr_bf16(&curiosity_bf16.w2, stream),
                b2: raw_device_ptr_bf16(&curiosity_bf16.b2, stream),
            },
            rmsnorm: GpuRMSNormPtrs {
                s0: raw_device_ptr(&rmsnorm_f32.gamma_s0, stream),
                s1: raw_device_ptr(&rmsnorm_f32.gamma_s1, stream),
                v: raw_device_ptr(&rmsnorm_f32.gamma_v, stream),
                a: raw_device_ptr(&rmsnorm_f32.gamma_a, stream),
            },
            online_br: GpuBranchPtrs {
                w_o1: raw_device_ptr_bf16(&online_br_bf16.w_bo1, stream),
                b_o1: raw_device_ptr_bf16(&online_br_bf16.b_bo1, stream),
                w_o2: raw_device_ptr_bf16(&online_br_bf16.w_bo2, stream),
                b_o2: raw_device_ptr_bf16(&online_br_bf16.b_bo2, stream),
                w_u1: raw_device_ptr_bf16(&online_br_bf16.w_bu1, stream),
                b_u1: raw_device_ptr_bf16(&online_br_bf16.b_bu1, stream),
                w_u2: raw_device_ptr_bf16(&online_br_bf16.w_bu2, stream),
                b_u2: raw_device_ptr_bf16(&online_br_bf16.b_bu2, stream),
            },
            target_br: GpuBranchPtrs {
                w_o1: raw_device_ptr_bf16(&target_br_bf16.w_bo1, stream),
                b_o1: raw_device_ptr_bf16(&target_br_bf16.b_bo1, stream),
                w_o2: raw_device_ptr_bf16(&target_br_bf16.w_bo2, stream),
                b_o2: raw_device_ptr_bf16(&target_br_bf16.b_bo2, stream),
                w_u1: raw_device_ptr_bf16(&target_br_bf16.w_bu1, stream),
                b_u1: raw_device_ptr_bf16(&target_br_bf16.b_bu1, stream),
                w_u2: raw_device_ptr_bf16(&target_br_bf16.w_bu2, stream),
                b_u2: raw_device_ptr_bf16(&target_br_bf16.b_bu2, stream),
            },
        }
    }
}

// ---------------------------------------------------------------------------
// Extraction helpers
// ---------------------------------------------------------------------------

/// Async device-to-device memcpy with error context.
///
/// Wraps the single unavoidable `unsafe` call (`memcpy_dtod_async`) in a
/// bounds-documented helper so callers remain safe.
fn dtod_copy_checked(
    dst: cudarc::driver::sys::CUdeviceptr,
    src: cudarc::driver::sys::CUdeviceptr,
    num_bytes: usize,
    stream: &Arc<CudaStream>,
    name: &str,
    op: &str,
) -> Result<(), MLError> {
    // Safety: caller guarantees src/dst are valid device pointers within
    // the same CUDA context, num_bytes ≤ allocation size of both buffers,
    // and stream belongs to the same context.
    unsafe {
        cudarc::driver::result::memcpy_dtod_async(dst, src, num_bytes, stream.cu_stream())
            .map_err(|e| MLError::ModelError(format!("DtoD {op} {name}: {e}")))?;
    }
    Ok(())
}

/// Extract a single named weight from a `GpuVarStore` into a new GPU buffer
/// via device-to-device copy (zero CPU roundtrip).
fn extract_one(
    vars: &GpuVarStore,
    name: &str,
    stream: &Arc<CudaStream>,
) -> Result<CudaSlice<f32>, MLError> {
    let param = vars
        .get(name)
        .ok_or_else(|| MLError::ModelError(format!("Missing weight: {name}")))?;

    let n_elems = param.data.len();
    let mut buf = stream
        .alloc_zeros::<f32>(n_elems)
        .map_err(|e| MLError::ModelError(format!("Alloc {name}: {e}")))?;

    stream.memcpy_dtod(&param.data, &mut buf)
        .map_err(|e| MLError::ModelError(format!("DtoD extract {name}: {e}")))?;

    Ok(buf)
}

/// Re-upload a single weight from `GpuVarStore` into an existing GPU buffer (no realloc).
/// Uses device-to-device copy (zero CPU roundtrip).
fn sync_one(
    vars: &GpuVarStore,
    name: &str,
    buf: &mut CudaSlice<f32>,
    stream: &Arc<CudaStream>,
) -> Result<(), MLError> {
    let param = vars
        .get(name)
        .ok_or_else(|| MLError::ModelError(format!("Missing weight: {name}")))?;

    stream.memcpy_dtod(&param.data, buf)
        .map_err(|e| MLError::ModelError(format!("DtoD sync {name}: {e}")))?;

    Ok(())
}

// ---------------------------------------------------------------------------
// Public API
// ---------------------------------------------------------------------------

/// Extract all 12 Dueling Q-Network weight tensors from a `GpuVarStore` and
/// upload them to GPU as flat `CudaSlice<f32>` buffers.
///
/// The `GpuVarStore` must contain tensor keys matching `DuelingQNetwork`:
/// `shared_0.weight`, `shared_0.bias`, `shared_1.weight`, `shared_1.bias`,
/// `value_fc.weight`, `value_fc.bias`, `value_out.weight`, `value_out.bias`,
/// `advantage_fc.weight`, `advantage_fc.bias`, `advantage_out.weight`,
/// `advantage_out.bias`.
pub fn extract_dueling_weights(
    vars: &GpuVarStore,
    stream: &Arc<CudaStream>,
) -> Result<DuelingWeightSet, MLError> {

    let extract = |name: &str| -> Result<CudaSlice<f32>, MLError> {
        extract_one(vars, name, stream)
    };

    let w_s1 = extract("shared_0.weight")?;
    let b_s1 = extract("shared_0.bias")?;
    let w_s2 = extract("shared_1.weight")?;
    let b_s2 = extract("shared_1.bias")?;
    let w_v1 = extract("value_fc.weight")?;
    let b_v1 = extract("value_fc.bias")?;
    let w_v2 = extract("value_out.weight")?;
    let b_v2 = extract("value_out.bias")?;
    let w_a1 = extract("advantage_fc.weight")?;
    let b_a1 = extract("advantage_fc.bias")?;
    let w_a2 = extract("advantage_out.weight")?;
    let b_a2 = extract("advantage_out.bias")?;

    let total_params: usize = DUELING_WEIGHT_NAMES
        .iter()
        .filter_map(|name| {
            vars.get(name).map(|p| p.data.len())
        })
        .sum();

    info!(
        total_params,
        total_bytes = total_params * std::mem::size_of::<f32>(),
        "Dueling Q-Network weights extracted and uploaded to GPU"
    );

    Ok(DuelingWeightSet {
        w_s1,
        b_s1,
        w_s2,
        b_s2,
        w_v1,
        b_v1,
        w_v2,
        b_v2,
        w_a1,
        b_a1,
        w_a2,
        b_a2,
    })
}

/// Extract all 8 Branching DQN extra head weight tensors from a `GpuVarStore`
/// and upload them to GPU as flat `CudaSlice<f32>` buffers.
///
/// The `GpuVarStore` must contain tensor keys matching `BranchingDuelingQNetwork`:
/// `branch_1_fc.weight`, `branch_1_fc.bias`, `branch_1_out.weight`, `branch_1_out.bias`,
/// `branch_2_fc.weight`, `branch_2_fc.bias`, `branch_2_out.weight`, `branch_2_out.bias`.
pub fn extract_branching_weights(
    vars: &GpuVarStore,
    stream: &Arc<CudaStream>,
) -> Result<BranchingWeightSet, MLError> {

    let extract = |name: &str| -> Result<CudaSlice<f32>, MLError> {
        extract_one(vars, name, stream)
    };

    let w_bo1 = extract("branch_1_fc.weight")?;
    let b_bo1 = extract("branch_1_fc.bias")?;
    let w_bo2 = extract("branch_1_out.weight")?;
    let b_bo2 = extract("branch_1_out.bias")?;
    let w_bu1 = extract("branch_2_fc.weight")?;
    let b_bu1 = extract("branch_2_fc.bias")?;
    let w_bu2 = extract("branch_2_out.weight")?;
    let b_bu2 = extract("branch_2_out.bias")?;

    let total_params: usize = BRANCHING_WEIGHT_NAMES
        .iter()
        .filter_map(|name| {
            vars.get(name).map(|p| p.data.len())
        })
        .sum();

    info!(
        total_params,
        total_bytes = total_params * std::mem::size_of::<f32>(),
        "Branching DQN weights extracted and uploaded to GPU"
    );

    Ok(BranchingWeightSet {
        w_bo1,
        b_bo1,
        w_bo2,
        b_bo2,
        w_bu1,
        b_bu1,
        w_bu2,
        b_bu2,
    })
}

/// Extract Dueling Q-Network weights using branching GpuVarStore key names.
///
/// In `BranchingDuelingQNetwork`, the exposure head (branch 0) uses GpuVarStore keys
/// `branch_0_fc.*` / `branch_0_out.*` instead of `advantage_fc.*` / `advantage_out.*`.
/// Shared layers and value head keys are identical to standard dueling.
///
/// The returned `DuelingWeightSet` maps:
/// - `w_a1`/`b_a1` ← `branch_0_fc.weight`/`branch_0_fc.bias`
/// - `w_a2`/`b_a2` ← `branch_0_out.weight`/`branch_0_out.bias`
pub fn extract_dueling_weights_branching(
    vars: &GpuVarStore,
    stream: &Arc<CudaStream>,
) -> Result<DuelingWeightSet, MLError> {

    let extract = |name: &str| -> Result<CudaSlice<f32>, MLError> {
        extract_one(vars, name, stream)
    };

    // Shared + value heads: same key names
    let w_s1 = extract("shared_0.weight")?;
    let b_s1 = extract("shared_0.bias")?;
    let w_s2 = extract("shared_1.weight")?;
    let b_s2 = extract("shared_1.bias")?;
    let w_v1 = extract("value_fc.weight")?;
    let b_v1 = extract("value_fc.bias")?;
    let w_v2 = extract("value_out.weight")?;
    let b_v2 = extract("value_out.bias")?;
    // Advantage/exposure head: branch_0_* in branching GpuVarStore
    let w_a1 = extract("branch_0_fc.weight")?;
    let b_a1 = extract("branch_0_fc.bias")?;
    let w_a2 = extract("branch_0_out.weight")?;
    let b_a2 = extract("branch_0_out.bias")?;

    info!("Branching DQN dueling weights (branch_0=exposure) extracted and uploaded to GPU");

    Ok(DuelingWeightSet {
        w_s1, b_s1, w_s2, b_s2,
        w_v1, b_v1, w_v2, b_v2,
        w_a1, b_a1, w_a2, b_a2,
    })
}

/// Re-upload Dueling Q-Network weights using branching GpuVarStore key names.
///
/// Counterpart to [`extract_dueling_weights_branching`] for weight sync.
pub fn sync_dueling_weights_branching(
    vars: &GpuVarStore,
    weights: &mut DuelingWeightSet,
    stream: &Arc<CudaStream>,
) -> Result<(), MLError> {

    sync_one(vars, "shared_0.weight", &mut weights.w_s1, stream)?;
    sync_one(vars, "shared_0.bias", &mut weights.b_s1, stream)?;
    sync_one(vars, "shared_1.weight", &mut weights.w_s2, stream)?;
    sync_one(vars, "shared_1.bias", &mut weights.b_s2, stream)?;
    sync_one(vars, "value_fc.weight", &mut weights.w_v1, stream)?;
    sync_one(vars, "value_fc.bias", &mut weights.b_v1, stream)?;
    sync_one(vars, "value_out.weight", &mut weights.w_v2, stream)?;
    sync_one(vars, "value_out.bias", &mut weights.b_v2, stream)?;
    // Advantage/exposure: branch_0_* keys in branching GpuVarStore
    sync_one(vars, "branch_0_fc.weight", &mut weights.w_a1, stream)?;
    sync_one(vars, "branch_0_fc.bias", &mut weights.b_a1, stream)?;
    sync_one(vars, "branch_0_out.weight", &mut weights.w_a2, stream)?;
    sync_one(vars, "branch_0_out.bias", &mut weights.b_a2, stream)?;

    Ok(())
}

/// Extract all 4 Curiosity Forward Model weight tensors from a `GpuVarStore`
/// and upload them to GPU as flat `CudaSlice<f32>` buffers.
///
/// The `GpuVarStore` must contain tensor keys matching `ForwardDynamicsModel`:
/// `fc1.weight`, `fc1.bias`, `fc2.weight`, `fc2.bias`.
pub fn extract_curiosity_weights(
    vars: &GpuVarStore,
    stream: &Arc<CudaStream>,
) -> Result<CuriosityWeightSet, MLError> {

    let extract = |name: &str| -> Result<CudaSlice<f32>, MLError> {
        extract_one(vars, name, stream)
    };

    let w1 = extract("fc1.weight")?;
    let b1 = extract("fc1.bias")?;
    let w2 = extract("fc2.weight")?;
    let b2 = extract("fc2.bias")?;

    let total_params: usize = CURIOSITY_WEIGHT_NAMES
        .iter()
        .filter_map(|name| {
            vars.get(name).map(|p| p.data.len())
        })
        .sum();

    info!(
        total_params,
        total_bytes = total_params * std::mem::size_of::<f32>(),
        "Curiosity forward model weights extracted and uploaded to GPU"
    );

    Ok(CuriosityWeightSet { w1, b1, w2, b2 })
}

/// Re-upload Dueling Q-Network weights from a `GpuVarStore` into existing
/// GPU buffers without re-allocating.
///
/// Call this after each training step (or target-network soft update)
/// to keep the GPU weight copies in sync with the model's GpuVarStore.
pub fn sync_dueling_weights(
    vars: &GpuVarStore,
    weights: &mut DuelingWeightSet,
    stream: &Arc<CudaStream>,
) -> Result<(), MLError> {

    sync_one(vars, "shared_0.weight", &mut weights.w_s1, stream)?;
    sync_one(vars, "shared_0.bias", &mut weights.b_s1, stream)?;
    sync_one(vars, "shared_1.weight", &mut weights.w_s2, stream)?;
    sync_one(vars, "shared_1.bias", &mut weights.b_s2, stream)?;
    sync_one(vars, "value_fc.weight", &mut weights.w_v1, stream)?;
    sync_one(vars, "value_fc.bias", &mut weights.b_v1, stream)?;
    sync_one(vars, "value_out.weight", &mut weights.w_v2, stream)?;
    sync_one(vars, "value_out.bias", &mut weights.b_v2, stream)?;
    sync_one(vars, "advantage_fc.weight", &mut weights.w_a1, stream)?;
    sync_one(vars, "advantage_fc.bias", &mut weights.b_a1, stream)?;
    sync_one(vars, "advantage_out.weight", &mut weights.w_a2, stream)?;
    sync_one(vars, "advantage_out.bias", &mut weights.b_a2, stream)?;

    Ok(())
}

/// Re-upload Branching DQN extra head weights from a `GpuVarStore` into existing
/// GPU buffers without re-allocating.
///
/// Call this after each training step (or target-network soft update)
/// to keep the GPU weight copies in sync with the model's GpuVarStore.
pub fn sync_branching_weights(
    vars: &GpuVarStore,
    weights: &mut BranchingWeightSet,
    stream: &Arc<CudaStream>,
) -> Result<(), MLError> {

    sync_one(vars, "branch_1_fc.weight", &mut weights.w_bo1, stream)?;
    sync_one(vars, "branch_1_fc.bias", &mut weights.b_bo1, stream)?;
    sync_one(vars, "branch_1_out.weight", &mut weights.w_bo2, stream)?;
    sync_one(vars, "branch_1_out.bias", &mut weights.b_bo2, stream)?;
    sync_one(vars, "branch_2_fc.weight", &mut weights.w_bu1, stream)?;
    sync_one(vars, "branch_2_fc.bias", &mut weights.b_bu1, stream)?;
    sync_one(vars, "branch_2_out.weight", &mut weights.w_bu2, stream)?;
    sync_one(vars, "branch_2_out.bias", &mut weights.b_bu2, stream)?;

    Ok(())
}

// ---------------------------------------------------------------------------
// Reverse sync: CudaSlice → GpuVarStore (write back GPU-updated weights)
// ---------------------------------------------------------------------------

/// Write a single CudaSlice back to the corresponding GpuVarStore parameter.
///
/// Uses device-to-device copy (zero CPU roundtrip).
fn reverse_sync_one(
    vars: &GpuVarStore,
    name: &str,
    src: &CudaSlice<f32>,
    stream: &Arc<CudaStream>,
) -> Result<(), MLError> {
    let param = vars
        .get(name)
        .ok_or_else(|| MLError::ModelError(format!("Missing weight: {name}")))?;

    let n_elems = param.data.len();
    if src.len() != n_elems {
        return Err(MLError::ModelError(format!(
            "reverse_sync {name}: CudaSlice len {} != param elems {n_elems}",
            src.len()
        )));
    }

    let (dst_ptr, _dst_sync) = param.data.device_ptr(stream);
    let (src_ptr, _src_sync) = src.device_ptr(stream);
    let num_bytes = n_elems * std::mem::size_of::<f32>();
    dtod_copy_checked(dst_ptr, src_ptr, num_bytes, stream, name, "reverse_sync")?;

    Ok(())
}

/// Write back all 12 branching-dueling weight tensors from CudaSlice to GpuVarStore.
///
/// Reverse of `sync_dueling_weights_branching`: copies GPU-updated weights
/// from the fused training kernel's CudaSlice buffers back to the GpuVarStore
/// so that Polyak target updates can read current online weights.
pub fn reverse_sync_dueling_weights_branching(
    vars: &GpuVarStore,
    weights: &DuelingWeightSet,
    stream: &Arc<CudaStream>,
) -> Result<(), MLError> {

    reverse_sync_one(vars, "shared_0.weight", &weights.w_s1, stream)?;
    reverse_sync_one(vars, "shared_0.bias", &weights.b_s1, stream)?;
    reverse_sync_one(vars, "shared_1.weight", &weights.w_s2, stream)?;
    reverse_sync_one(vars, "shared_1.bias", &weights.b_s2, stream)?;
    reverse_sync_one(vars, "value_fc.weight", &weights.w_v1, stream)?;
    reverse_sync_one(vars, "value_fc.bias", &weights.b_v1, stream)?;
    reverse_sync_one(vars, "value_out.weight", &weights.w_v2, stream)?;
    reverse_sync_one(vars, "value_out.bias", &weights.b_v2, stream)?;
    reverse_sync_one(vars, "branch_0_fc.weight", &weights.w_a1, stream)?;
    reverse_sync_one(vars, "branch_0_fc.bias", &weights.b_a1, stream)?;
    reverse_sync_one(vars, "branch_0_out.weight", &weights.w_a2, stream)?;
    reverse_sync_one(vars, "branch_0_out.bias", &weights.b_a2, stream)?;

    Ok(())
}

/// Write back all 8 branching weight tensors (branches 1+2) from CudaSlice to GpuVarStore.
///
/// Reverse of `sync_branching_weights`.
pub fn reverse_sync_branching_weights(
    vars: &GpuVarStore,
    weights: &BranchingWeightSet,
    stream: &Arc<CudaStream>,
) -> Result<(), MLError> {

    reverse_sync_one(vars, "branch_1_fc.weight", &weights.w_bo1, stream)?;
    reverse_sync_one(vars, "branch_1_fc.bias", &weights.b_bo1, stream)?;
    reverse_sync_one(vars, "branch_1_out.weight", &weights.w_bo2, stream)?;
    reverse_sync_one(vars, "branch_1_out.bias", &weights.b_bo2, stream)?;
    reverse_sync_one(vars, "branch_2_fc.weight", &weights.w_bu1, stream)?;
    reverse_sync_one(vars, "branch_2_fc.bias", &weights.b_bu1, stream)?;
    reverse_sync_one(vars, "branch_2_out.weight", &weights.w_bu2, stream)?;
    reverse_sync_one(vars, "branch_2_out.bias", &weights.b_bu2, stream)?;

    Ok(())
}

/// Extract RMSNorm gamma weights from a distributional dueling `GpuVarStore`.
///
/// Looks for keys: `shared_rmsnorm_0.weight`, `shared_rmsnorm_1.weight`,
/// `value_rmsnorm.weight`, `advantage_rmsnorm.weight`.
///
/// Returns `Ok(Some(set))` if all 4 keys exist, `Ok(None)` if none exist
/// (i.e. this is a plain dueling network without RMSNorm).
pub fn extract_rmsnorm_weights(
    vars: &GpuVarStore,
    stream: &Arc<CudaStream>,
) -> Result<Option<RmsNormWeightSet>, MLError> {

    // Check if the GpuVarStore contains RMSNorm keys (distributional dueling)
    if vars.get(RMSNORM_WEIGHT_NAMES[0]).is_none() {
        return Ok(None);
    }

    let gamma_s0 = extract_one(vars, RMSNORM_WEIGHT_NAMES[0], stream)?;
    let gamma_s1 = extract_one(vars, RMSNORM_WEIGHT_NAMES[1], stream)?;
    let gamma_v  = extract_one(vars, RMSNORM_WEIGHT_NAMES[2], stream)?;
    let gamma_a  = extract_one(vars, RMSNORM_WEIGHT_NAMES[3], stream)?;

    let total_params: usize = RMSNORM_WEIGHT_NAMES
        .iter()
        .filter_map(|name| vars.get(name).map(|p| p.data.len()))
        .sum();

    info!(
        total_params,
        "RMSNorm gamma weights extracted and uploaded to GPU"
    );

    Ok(Some(RmsNormWeightSet { gamma_s0, gamma_s1, gamma_v, gamma_a }))
}

/// Re-upload RMSNorm gamma weights into existing GPU buffers.
pub fn sync_rmsnorm_weights(
    vars: &GpuVarStore,
    weights: &mut RmsNormWeightSet,
    stream: &Arc<CudaStream>,
) -> Result<(), MLError> {

    // Only sync if the GpuVarStore has RMSNorm keys
    if vars.get(RMSNORM_WEIGHT_NAMES[0]).is_none() {
        return Ok(());
    }

    sync_one(vars, RMSNORM_WEIGHT_NAMES[0], &mut weights.gamma_s0, stream)?;
    sync_one(vars, RMSNORM_WEIGHT_NAMES[1], &mut weights.gamma_s1, stream)?;
    sync_one(vars, RMSNORM_WEIGHT_NAMES[2], &mut weights.gamma_v, stream)?;
    sync_one(vars, RMSNORM_WEIGHT_NAMES[3], &mut weights.gamma_a, stream)?;

    Ok(())
}

/// Re-upload Curiosity Forward Model weights from a `GpuVarStore` into existing
/// GPU buffers without re-allocating.
///
/// Call this after curiosity model training to keep the GPU weight copies
/// in sync with the model's GpuVarStore.
pub fn sync_curiosity_weights(
    vars: &GpuVarStore,
    weights: &mut CuriosityWeightSet,
    stream: &Arc<CudaStream>,
) -> Result<(), MLError> {

    sync_one(vars, "fc1.weight", &mut weights.w1, stream)?;
    sync_one(vars, "fc1.bias", &mut weights.b1, stream)?;
    sync_one(vars, "fc2.weight", &mut weights.w2, stream)?;
    sync_one(vars, "fc2.bias", &mut weights.b2, stream)?;

    Ok(())
}

/// Extract all 6 PPO Actor (PolicyNetwork) weight tensors from a `GpuVarStore`
/// and upload them to GPU as flat `CudaSlice<f32>` buffers.
///
/// The `GpuVarStore` must contain tensor keys matching `PolicyNetwork`:
/// `policy_layer_0.weight`, `policy_layer_0.bias`, `policy_layer_1.weight`,
/// `policy_layer_1.bias`, `policy_output.weight`, `policy_output.bias`.
pub fn extract_ppo_actor_weights(
    vars: &GpuVarStore,
    stream: &Arc<CudaStream>,
) -> Result<PpoActorWeightSet, MLError> {

    let extract = |name: &str| -> Result<CudaSlice<f32>, MLError> {
        extract_one(vars, name, stream)
    };

    let pw1 = extract("policy_layer_0.weight")?;
    let pb1 = extract("policy_layer_0.bias")?;
    let pw2 = extract("policy_layer_1.weight")?;
    let pb2 = extract("policy_layer_1.bias")?;
    let pw3 = extract("policy_output.weight")?;
    let pb3 = extract("policy_output.bias")?;

    let total_params: usize = PPO_ACTOR_WEIGHT_NAMES
        .iter()
        .filter_map(|name| {
            vars.get(name).map(|p| p.data.len())
        })
        .sum();

    info!(
        total_params,
        total_bytes = total_params * std::mem::size_of::<f32>(),
        "PPO Actor weights extracted and uploaded to GPU"
    );

    Ok(PpoActorWeightSet {
        pw1,
        pb1,
        pw2,
        pb2,
        pw3,
        pb3,
    })
}

/// Extract all 12 PPO Critic (ValueNetwork) weight tensors from a `GpuVarStore`
/// and upload them to GPU as flat `CudaSlice<f32>` buffers.
///
/// The `GpuVarStore` must contain tensor keys matching `ValueNetwork`:
/// `value_layer_0.weight`, `value_layer_0.bias`, `value_layer_1.weight`,
/// `value_layer_1.bias`, `value_layer_2.weight`, `value_layer_2.bias`,
/// `value_layer_3.weight`, `value_layer_3.bias`, `value_layer_4.weight`,
/// `value_layer_4.bias`, `value_output.weight`, `value_output.bias`.
pub fn extract_ppo_critic_weights(
    vars: &GpuVarStore,
    stream: &Arc<CudaStream>,
) -> Result<PpoCriticWeightSet, MLError> {

    let extract = |name: &str| -> Result<CudaSlice<f32>, MLError> {
        extract_one(vars, name, stream)
    };

    let vw1 = extract("value_layer_0.weight")?;
    let vb1 = extract("value_layer_0.bias")?;
    let vw2 = extract("value_layer_1.weight")?;
    let vb2 = extract("value_layer_1.bias")?;
    let vw3 = extract("value_layer_2.weight")?;
    let vb3 = extract("value_layer_2.bias")?;
    let vw4 = extract("value_layer_3.weight")?;
    let vb4 = extract("value_layer_3.bias")?;
    let vw5 = extract("value_layer_4.weight")?;
    let vb5 = extract("value_layer_4.bias")?;
    let vw6 = extract("value_output.weight")?;
    let vb6 = extract("value_output.bias")?;

    let total_params: usize = PPO_CRITIC_WEIGHT_NAMES
        .iter()
        .filter_map(|name| {
            vars.get(name).map(|p| p.data.len())
        })
        .sum();

    info!(
        total_params,
        total_bytes = total_params * std::mem::size_of::<f32>(),
        "PPO Critic weights extracted and uploaded to GPU"
    );

    Ok(PpoCriticWeightSet {
        vw1,
        vb1,
        vw2,
        vb2,
        vw3,
        vb3,
        vw4,
        vb4,
        vw5,
        vb5,
        vw6,
        vb6,
    })
}

/// Re-upload PPO Actor weights from a `GpuVarStore` into existing GPU buffers
/// without re-allocating.
///
/// Call this after each policy gradient step to keep the GPU weight copies
/// in sync with the model's GpuVarStore.
pub fn sync_ppo_actor_weights(
    vars: &GpuVarStore,
    weights: &mut PpoActorWeightSet,
    stream: &Arc<CudaStream>,
) -> Result<(), MLError> {

    sync_one(vars, "policy_layer_0.weight", &mut weights.pw1, stream)?;
    sync_one(vars, "policy_layer_0.bias", &mut weights.pb1, stream)?;
    sync_one(vars, "policy_layer_1.weight", &mut weights.pw2, stream)?;
    sync_one(vars, "policy_layer_1.bias", &mut weights.pb2, stream)?;
    sync_one(vars, "policy_output.weight", &mut weights.pw3, stream)?;
    sync_one(vars, "policy_output.bias", &mut weights.pb3, stream)?;

    Ok(())
}

/// Re-upload PPO Critic weights from a `GpuVarStore` into existing GPU buffers
/// without re-allocating.
///
/// Call this after each value function update to keep the GPU weight copies
/// in sync with the model's GpuVarStore.
pub fn sync_ppo_critic_weights(
    vars: &GpuVarStore,
    weights: &mut PpoCriticWeightSet,
    stream: &Arc<CudaStream>,
) -> Result<(), MLError> {

    sync_one(vars, "value_layer_0.weight", &mut weights.vw1, stream)?;
    sync_one(vars, "value_layer_0.bias", &mut weights.vb1, stream)?;
    sync_one(vars, "value_layer_1.weight", &mut weights.vw2, stream)?;
    sync_one(vars, "value_layer_1.bias", &mut weights.vb2, stream)?;
    sync_one(vars, "value_layer_2.weight", &mut weights.vw3, stream)?;
    sync_one(vars, "value_layer_2.bias", &mut weights.vb3, stream)?;
    sync_one(vars, "value_layer_3.weight", &mut weights.vw4, stream)?;
    sync_one(vars, "value_layer_3.bias", &mut weights.vb4, stream)?;
    sync_one(vars, "value_layer_4.weight", &mut weights.vw5, stream)?;
    sync_one(vars, "value_layer_4.bias", &mut weights.vb5, stream)?;
    sync_one(vars, "value_output.weight", &mut weights.vw6, stream)?;
    sync_one(vars, "value_output.bias", &mut weights.vb6, stream)?;

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_dueling_weight_names_count() {
        assert_eq!(DUELING_WEIGHT_NAMES.len(), 12);
    }

    #[test]
    fn test_branching_weight_names_count() {
        assert_eq!(BRANCHING_WEIGHT_NAMES.len(), 8);
    }

    #[test]
    fn test_curiosity_weight_names_count() {
        assert_eq!(CURIOSITY_WEIGHT_NAMES.len(), 4);
    }

    #[test]
    fn test_ppo_actor_weight_names_count() {
        assert_eq!(PPO_ACTOR_WEIGHT_NAMES.len(), 6);
    }

    #[test]
    fn test_ppo_critic_weight_names_count() {
        assert_eq!(PPO_CRITIC_WEIGHT_NAMES.len(), 12);
    }

    #[test]
    fn test_rmsnorm_weight_names_count() {
        assert_eq!(RMSNORM_WEIGHT_NAMES.len(), 4);
    }

    #[test]
    fn test_curiosity_zero_weights() {
        let context = cudarc::driver::CudaContext::new(0);
        if context.is_err() {
            return; // skip on non-CUDA
        }
        let stream = context.ok().map(|c| c.default_stream());
        if let Some(ref s) = stream {
            let weights = CuriosityWeightSet::zeros(s);
            assert!(weights.is_ok());
        }
    }
}
