//! `CfcTrunk` — owns CfC weights, multi-horizon heads, projection, and a
//! ping-pong hidden buffer. `forward_snapshot` dispatches the four
//! perception kernels sequentially. The Mamba2 prefix (per the
//! 2026-05-16 spec amendment) lands in a follow-up task before Graph A
//! capture; this task validates the CfC + heads + projection forward
//! path standalone.
//!
//! Per the addendum API patterns: cache `Arc<CudaStream>` and the
//! `CudaFunction` handles in the struct so the hot path avoids module
//! reload.

use std::sync::Arc;

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

use crate::cfc::bucket_routing::MAX_BUCKET_DIM;
use crate::heads::{HEAD_MID_DIM, HIDDEN_DIM, N_HORIZONS};
use crate::mamba2_block::{Mamba2Block, Mamba2BlockConfig};
use serde::{Deserialize, Serialize};

/// On-disk checkpoint envelope for CfcTrunk weights. Covers the full
/// inference graph (VSN, Mamba2 ×2, LN ×2, attn-pool, CfC, GRN heads).
/// Greenfield: no backward-compat, no version field — bincode either
/// deserialises a current envelope or errors.
#[derive(Clone, Debug, Serialize, Deserialize)]
struct Checkpoint {
    n_in: usize,                   // snap feature dim (for sanity check)
    n_hid: usize,                  // HIDDEN_DIM
    cfc_n_in: usize,               // CfC input dim (HIDDEN_DIM)
    mamba2_state_dim: usize,
    mamba2_seq_len: usize,         // captured at save; load validates
    mamba2_l1_in_dim: usize,       // FEATURE_DIM
    mamba2_l2_in_dim: usize,       // HIDDEN_DIM
    // VSN
    vsn_w: Vec<f32>,
    vsn_b: Vec<f32>,
    // Mamba2 stack 1 weights
    m1_w_in_w: Vec<f32>,  m1_w_in_b: Vec<f32>,
    m1_w_a_w: Vec<f32>,   m1_w_a_b: Vec<f32>,
    m1_w_b_w: Vec<f32>,   m1_w_b_b: Vec<f32>,
    m1_w_c: Vec<f32>,
    m1_w_out_w: Vec<f32>, m1_w_out_b: Vec<f32>,
    // LN_a
    ln_a_gain: Vec<f32>, ln_a_bias: Vec<f32>,
    // Mamba2 stack 2 weights
    m2_w_in_w: Vec<f32>,  m2_w_in_b: Vec<f32>,
    m2_w_a_w: Vec<f32>,   m2_w_a_b: Vec<f32>,
    m2_w_b_w: Vec<f32>,   m2_w_b_b: Vec<f32>,
    m2_w_c: Vec<f32>,
    m2_w_out_w: Vec<f32>, m2_w_out_b: Vec<f32>,
    // LN_b
    ln_b_gain: Vec<f32>, ln_b_bias: Vec<f32>,
    // Attention pool
    attn_q: Vec<f32>,
    // CfC (ORIGINAL channel layout — ALPHA fix 2026-05-21: no reorder)
    w_in: Vec<f32>, w_rec: Vec<f32>, b: Vec<f32>, tau_all: Vec<f32>,
    // Bucket metadata (populated at Phase 1→2 transition; zero-initialized in new_random)
    bucket_id_per_channel: Vec<u8>,    // [HIDDEN_DIM]
    bucket_channel_offset: Vec<u32>,   // [N_HORIZONS + 1]
    bucket_dim_k: Vec<u32>,            // [N_HORIZONS]
    channels_in_bucket: Vec<u32>,      // [N_HORIZONS × MAX_BUCKET_DIM] (ALPHA fix)
    heads_w_skip_offset: Vec<u32>,     // [N_HORIZONS + 1]
    // Heads (GRN structure × 5 horizons)
    heads_w1: Vec<f32>, heads_b1: Vec<f32>,
    heads_w2: Vec<f32>, heads_b2: Vec<f32>,
    heads_w_gate: Vec<f32>, heads_b_gate: Vec<f32>,
    heads_w_main: Vec<f32>, heads_b_main: Vec<f32>,
    heads_w_skip: Vec<f32>, heads_b_skip: Vec<f32>,
}

const SNAP_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/snap_feature_assemble.cubin"));
const STEP_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/cfc_step.cubin"));
const HEADS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/multi_horizon_heads.cubin"));
// X10: forward kernel cubins — loaded into the trunk so it owns
// every kernel handle the forward chain needs. Trainer will read
// these handles via `self.trunk.<fn>` in a subsequent commit (X10b).
const LAYER_NORM_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/layer_norm.cubin"));
const VARIABLE_SELECTION_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/variable_selection.cubin"));
const ATTENTION_POOL_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/attention_pool.cubin"));
// Per-horizon CfC Phase 2 dispatch (spec §5.4 points 1–2). Modules
// stay alive on the trunk so the cached `CudaFunction` handles below
// remain valid through the trainer's lifetime.
const CFC_STEP_PER_BRANCH_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/cfc_step_per_branch.cubin"));
const HEADS_BLOCK_DIAGONAL_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/heads_block_diagonal_fwd.cubin"));

#[derive(Clone, Debug)]
pub struct CfcConfig {
    /// Raw snap feature dimensionality (V1 forward path uses this).
    pub n_in: usize,
    pub n_hid: usize,
    /// CfC layer input dimensionality. In V1 this was conflated with
    /// `n_in` (CfC consumed snap features directly). In the trunk the CfC sits
    /// AFTER LN_b, consuming HIDDEN_DIM-dim tokens. Default matches trunk
    /// usage.
    pub cfc_n_in: usize,
    /// Mamba2 state dimension — used by the trunk weight skeleton
    /// (X1) and the per-stack Mamba2Block constructions added in X3/X5.
    /// Default matches `PerceptionTrainerConfig::default().mamba2_state_dim`.
    pub mamba2_state_dim: usize,
    /// X11: inference batch + sequence dims for the forward buffers
    /// allocated on the trunk. Backtesters typically use (1, 32);
    /// training-time consumers (PerceptionTrainer) pass their own dims.
    pub n_batch: usize,
    pub seq_len: usize,
}

impl Default for CfcConfig {
    fn default() -> Self {
        Self {
            n_in: crate::cfc::snap_features::ENCODER_INPUT_DIM,
            n_hid: HIDDEN_DIM,
            cfc_n_in: HIDDEN_DIM,
            mamba2_state_dim: 16,
            n_batch: 1,
            seq_len: 32,
        }
    }
}

pub struct CfcTrunk {
    cfg: CfcConfig,
    stream: Arc<CudaStream>,

    // Cubin modules — kept alive for the function refs below.
    _snap_module: Arc<CudaModule>,
    _step_module: Arc<CudaModule>,
    _heads_module: Arc<CudaModule>,
    _ln_module: Arc<CudaModule>,
    _vsn_module: Arc<CudaModule>,
    _attn_module: Arc<CudaModule>,
    // Per-horizon CfC Phase 2 cubin modules. Kept alive alongside the
    // cached function handles below.
    _cfc_step_per_branch_module: Arc<CudaModule>,
    _heads_block_diagonal_module: Arc<CudaModule>,
    // Forward kernel handles — read by PerceptionTrainer's evaluate_batched.
    pub vsn_fwd_fn: CudaFunction,
    pub ln_fwd_fn: CudaFunction,
    pub attn_fwd_fn: CudaFunction,
    pub snap_batched_fn: CudaFunction,
    pub step_batched_fn: CudaFunction,
    pub heads_grn_fwd_fn: CudaFunction,
    pub transpose_3d_fn: CudaFunction,
    // Per-horizon CfC Phase 2 dispatch handles (spec §5.4 points 1–2).
    // `cfc_step_per_branch_fwd_fn` is read by `dispatch_train_step`'s
    // Phase 2 branch (Task 10) and the eventual `forward_step_into`
    // Phase 2 path (Task 13). The `_bwd_fn` handle is staged here so
    // Task 11's backward wiring lands without re-touching the trunk
    // constructor. `heads_block_diagonal_fwd_fn` is loaded for the
    // future compact-w_skip heads dispatch; per Task 10's scope (see
    // perception.rs Phase 2 dispatch comment), it is NOT consumed yet
    // — heads continue through the existing GRN kernel until a later
    // task wires the compact-storage path through the GRN's skip term.
    pub cfc_step_per_branch_fwd_fn: CudaFunction,
    pub cfc_step_per_branch_bwd_fn: CudaFunction,
    pub heads_block_diagonal_fwd_fn: CudaFunction,

    // CfC weights (cfc_n_in = HIDDEN_DIM). Public so PerceptionTrainer can
    // mutate during init + AdamW.
    pub w_in_d: CudaSlice<f32>,
    pub w_rec_d: CudaSlice<f32>,
    pub b_d: CudaSlice<f32>,
    pub tau_all_d: CudaSlice<f32>,

    // Bucket metadata for per-horizon CfC routing (Phase 1→2 transition
    // populates these; zero-initialized at `new_random`). See spec §2.2.
    // ALPHA fix 2026-05-21: adds `channels_in_bucket_d` — the bucket →
    // original-channel lookup that replaces the τ-reorder approach.
    pub bucket_id_per_channel_d: CudaSlice<u8>,
    pub bucket_channel_offset_d: CudaSlice<u32>,
    pub bucket_dim_k_d: CudaSlice<u32>,
    pub channels_in_bucket_d: CudaSlice<u32>,
    pub heads_w_skip_offset_d: CudaSlice<u32>,

    // Inference weights for the full perception chain:
    //   snap → vsn → mamba2_l1 → ln_a → mamba2_l2 → ln_b → attn_pool →
    //   cfc → grn heads.
    // All public so PerceptionTrainer can mutate during init + AdamW.
    pub vsn_w_d: CudaSlice<f32>,
    pub vsn_b_d: CudaSlice<f32>,
    pub mamba2_stack_1: Option<crate::mamba2_block::Mamba2Block>,
    pub mamba2_stack_2: Option<crate::mamba2_block::Mamba2Block>,
    pub ln_a_gain_d: CudaSlice<f32>,
    pub ln_a_bias_d: CudaSlice<f32>,
    pub ln_b_gain_d: CudaSlice<f32>,
    pub ln_b_bias_d: CudaSlice<f32>,
    pub attn_q_d: CudaSlice<f32>,
    // GRN heads: heads_w1/b1 is the input projection (HIDDEN → HEAD_MID),
    // heads_w2/b2 is the mid→mid layer, gate/main/skip feed the final
    // multi-horizon probabilities.
    pub heads_w1_d: CudaSlice<f32>,
    pub heads_b1_d: CudaSlice<f32>,
    pub heads_w2_d: CudaSlice<f32>,
    pub heads_b2_d: CudaSlice<f32>,
    pub heads_w_gate_d: CudaSlice<f32>,
    pub heads_b_gate_d: CudaSlice<f32>,
    pub heads_w_main_d: CudaSlice<f32>,
    pub heads_b_main_d: CudaSlice<f32>,
    pub heads_w_skip_d: CudaSlice<f32>,
    pub heads_b_skip_d: CudaSlice<f32>,
}

impl CfcTrunk {
    pub fn new_random(dev: &MlDevice, cfg: &CfcConfig, seed: u64) -> Result<Self> {
        let stream: Arc<CudaStream> = dev.cuda_stream().context("trunk stream")?.clone();
        let ctx = dev.cuda_context().context("trunk ctx")?;

        // Load forward kernels onto the trunk. snap/step/heads cubins are
        // each shared between V1 single-step kernels (no longer used) and
        // the batched variants we actually launch; we keep the modules
        // alive and only resolve the batched function handles.
        let snap_module = ctx.load_cubin(SNAP_CUBIN.to_vec()).context("load snap cubin")?;
        let step_module = ctx.load_cubin(STEP_CUBIN.to_vec()).context("load step cubin")?;
        let heads_module = ctx.load_cubin(HEADS_CUBIN.to_vec()).context("load heads cubin")?;
        let ln_module = ctx.load_cubin(LAYER_NORM_CUBIN.to_vec()).context("load ln cubin")?;
        let vsn_module = ctx.load_cubin(VARIABLE_SELECTION_CUBIN.to_vec()).context("load vsn cubin")?;
        let attn_module = ctx.load_cubin(ATTENTION_POOL_CUBIN.to_vec()).context("load attn cubin")?;
        // Per-horizon CfC Phase 2 cubin modules.
        let cfc_step_per_branch_module = ctx
            .load_cubin(CFC_STEP_PER_BRANCH_CUBIN.to_vec())
            .context("load cfc_step_per_branch cubin")?;
        let heads_block_diagonal_module = ctx
            .load_cubin(HEADS_BLOCK_DIAGONAL_CUBIN.to_vec())
            .context("load heads_block_diagonal cubin")?;
        let vsn_fwd_fn = vsn_module.load_function("variable_selection_fwd").context("vsn_fwd_fn")?;
        let ln_fwd_fn = ln_module.load_function("layer_norm_fwd").context("ln_fwd_fn")?;
        let attn_fwd_fn = attn_module.load_function("attention_pool_fwd").context("attn_fwd_fn")?;
        let snap_batched_fn = snap_module.load_function("snap_feature_assemble_batched").context("snap_batched_fn")?;
        let step_batched_fn = step_module.load_function("cfc_step_batched").context("step_batched_fn")?;
        let heads_grn_fwd_fn = heads_module.load_function("multi_horizon_heads_grn_fwd_batched").context("heads_grn_fwd_fn")?;
        let transpose_3d_fn = step_module.load_function("transpose_3d_swap_01").context("transpose_3d_fn")?;
        let cfc_step_per_branch_fwd_fn = cfc_step_per_branch_module
            .load_function("cfc_step_per_branch_fwd")
            .context("cfc_step_per_branch_fwd_fn")?;
        let cfc_step_per_branch_bwd_fn = cfc_step_per_branch_module
            .load_function("cfc_step_per_branch_bwd")
            .context("cfc_step_per_branch_bwd_fn")?;
        let heads_block_diagonal_fwd_fn = heads_block_diagonal_module
            .load_function("heads_block_diagonal_fwd")
            .context("heads_block_diagonal_fwd_fn")?;

        // CfC weight init via a seeded ChaCha8Rng — the same chain feeds
        // every weight tensor so cfg.seed fully determines the initial
        // state. CfC's input dim is `cfg.cfc_n_in` (HIDDEN_DIM in the
        // current architecture), separate from the snap-feature `cfg.n_in`.
        let mut r = ChaCha8Rng::seed_from_u64(seed);
        let scale_in = (1.0_f32 / cfg.cfc_n_in as f32).sqrt();
        let scale_rec = (1.0_f32 / cfg.n_hid as f32).sqrt();
        let w_in: Vec<f32> = (0..cfg.n_hid * cfg.cfc_n_in)
            .map(|_| r.gen_range(-scale_in..scale_in)).collect();
        let w_rec: Vec<f32> = (0..cfg.n_hid * cfg.n_hid)
            .map(|_| r.gen_range(-scale_rec..scale_rec)).collect();
        let b: Vec<f32> = vec![0.0; cfg.n_hid];
        // tau log-uniform over [10ms, 1000s] per spec §2 (ln(0.01)=-4.6, ln(1000)=6.9).
        let tau: Vec<f32> = (0..cfg.n_hid)
            .map(|_| {
                let u: f32 = r.gen_range(0.0..1.0);
                (-4.6 + u * (4.6 + 6.9)).exp()
            })
            .collect();
        // PRNG consumed by CfC weights only; PerceptionTrainer's ChaCha
        // chain (also seeded from cfg.seed) drives the higher-up layers
        // (VSN, LN, attn, GRN heads) — see trainer/perception.rs.

        let upload = |host: &[f32]| -> Result<CudaSlice<f32>> {
            let mut buf = stream.alloc_zeros::<f32>(host.len()).context("upload alloc")?;
            stream.memcpy_htod(host, &mut buf).context("upload htod")?;
            Ok(buf)
        };

        // Bucket metadata: zero-initialised here; populated by the Phase
        // 1→2 transition orchestrator in `cfc/bucket_routing.rs` once
        // Controller A fires. Allocated unconditionally so that
        // checkpoint save/load round-trips include them even before the
        // transition has run.
        let bucket_id_per_channel_d = stream
            .alloc_zeros::<u8>(HIDDEN_DIM)
            .context("bucket_id_per_channel alloc")?;
        let bucket_channel_offset_d = stream
            .alloc_zeros::<u32>(N_HORIZONS + 1)
            .context("bucket_channel_offset alloc")?;
        let bucket_dim_k_d = stream
            .alloc_zeros::<u32>(N_HORIZONS)
            .context("bucket_dim_k alloc")?;
        // ALPHA fix 2026-05-21: `channels_in_bucket_d` is the bucket →
        // original-channel lookup table populated at the Phase 1→2
        // transition. Zero-initialised here so `forward_step_into` on a
        // freshly-constructed trunk sees the kernel's defensive sentinel-
        // ish behaviour (channel index 0 for every slot, which is
        // harmless — `bucket_dim_k_d` is also zero so the kernel
        // early-returns before reading). Production paths populate this
        // via `execute_transition` or checkpoint load.
        let channels_in_bucket_d = stream
            .alloc_zeros::<u32>(N_HORIZONS * MAX_BUCKET_DIM)
            .context("channels_in_bucket alloc")?;
        let heads_w_skip_offset_d = stream
            .alloc_zeros::<u32>(N_HORIZONS + 1)
            .context("heads_w_skip_offset alloc")?;

        Ok(Self {
            cfg: cfg.clone(),
            w_in_d: upload(&w_in)?,
            w_rec_d: upload(&w_rec)?,
            b_d: upload(&b)?,
            tau_all_d: upload(&tau)?,
            bucket_id_per_channel_d,
            bucket_channel_offset_d,
            bucket_dim_k_d,
            channels_in_bucket_d,
            heads_w_skip_offset_d,
            // Inference layers — zero-initialised here; PerceptionTrainer's
            // construction overwrites them with its own seeded init via
            // memcpy_htod. Standalone trunk callers (tests, fxt-backtest's
            // random-init fallback) get zero-initialised inference layers,
            // which is fine — the layers are still well-formed device
            // tensors of the correct shape, just not trained.
            vsn_w_d: stream.alloc_zeros::<f32>(cfg.n_in * cfg.n_in).context("vsn_w alloc")?,
            vsn_b_d: stream.alloc_zeros::<f32>(cfg.n_in).context("vsn_b alloc")?,
            // Mamba2 stack 1: in=FEATURE_DIM, hidden=HIDDEN_DIM. Constructed
            // here so that load_checkpoint (which goes through new_random)
            // gets populated stacks ready for weight upload. Block's own
            // random init is overwritten by either PerceptionTrainer's seeded
            // re-init during training, or load_checkpoint's file weights.
            mamba2_stack_1: Some(
                Mamba2Block::new(
                    Mamba2BlockConfig {
                        in_dim: cfg.n_in,
                        hidden_dim: HIDDEN_DIM,
                        state_dim: cfg.mamba2_state_dim,
                        seq_len: cfg.seq_len,
                    },
                    stream.clone(),
                )
                .context("mamba2_stack_1 construct")?,
            ),
            // Mamba2 stack 2: same shape, consumes LN_a output (HIDDEN_DIM).
            mamba2_stack_2: Some(
                Mamba2Block::new(
                    Mamba2BlockConfig {
                        in_dim: HIDDEN_DIM,
                        hidden_dim: HIDDEN_DIM,
                        state_dim: cfg.mamba2_state_dim,
                        seq_len: cfg.seq_len,
                    },
                    stream.clone(),
                )
                .context("mamba2_stack_2 construct")?,
            ),
            ln_a_gain_d: stream.alloc_zeros::<f32>(cfg.n_hid).context("ln_a_gain alloc")?,
            ln_a_bias_d: stream.alloc_zeros::<f32>(cfg.n_hid).context("ln_a_bias alloc")?,
            ln_b_gain_d: stream.alloc_zeros::<f32>(cfg.n_hid).context("ln_b_gain alloc")?,
            ln_b_bias_d: stream.alloc_zeros::<f32>(cfg.n_hid).context("ln_b_bias alloc")?,
            attn_q_d: stream.alloc_zeros::<f32>(cfg.n_hid).context("attn_q alloc")?,
            heads_w1_d: stream.alloc_zeros::<f32>(N_HORIZONS * HEAD_MID_DIM * cfg.n_hid).context("heads_w1 alloc")?,
            heads_b1_d: stream.alloc_zeros::<f32>(N_HORIZONS * HEAD_MID_DIM).context("heads_b1 alloc")?,
            heads_w2_d: stream.alloc_zeros::<f32>(N_HORIZONS * HEAD_MID_DIM * HEAD_MID_DIM).context("heads_w2 alloc")?,
            heads_b2_d: stream.alloc_zeros::<f32>(N_HORIZONS * HEAD_MID_DIM).context("heads_b2 alloc")?,
            heads_w_gate_d: stream.alloc_zeros::<f32>(N_HORIZONS * HEAD_MID_DIM).context("heads_w_gate alloc")?,
            heads_b_gate_d: stream.alloc_zeros::<f32>(N_HORIZONS).context("heads_b_gate alloc")?,
            heads_w_main_d: stream.alloc_zeros::<f32>(N_HORIZONS * HEAD_MID_DIM).context("heads_w_main alloc")?,
            heads_b_main_d: stream.alloc_zeros::<f32>(N_HORIZONS).context("heads_b_main alloc")?,
            heads_w_skip_d: stream.alloc_zeros::<f32>(N_HORIZONS * cfg.n_hid).context("heads_w_skip alloc")?,
            heads_b_skip_d: stream.alloc_zeros::<f32>(N_HORIZONS).context("heads_b_skip alloc")?,
            stream,
            _snap_module: snap_module,
            _step_module: step_module,
            _heads_module: heads_module,
            _ln_module: ln_module,
            _vsn_module: vsn_module,
            _attn_module: attn_module,
            _cfc_step_per_branch_module: cfc_step_per_branch_module,
            _heads_block_diagonal_module: heads_block_diagonal_module,
            vsn_fwd_fn,
            ln_fwd_fn,
            attn_fwd_fn,
            snap_batched_fn,
            step_batched_fn,
            heads_grn_fwd_fn,
            transpose_3d_fn,
            cfc_step_per_branch_fwd_fn,
            cfc_step_per_branch_bwd_fn,
            heads_block_diagonal_fwd_fn,
        })
    }


    /// Serialise the trunk weights + config to a bincode file. Used by
    /// downstream backtest tooling (fxt-backtest --checkpoint) to run
    /// inference with trained weights instead of new_random init.
    ///
    /// Reads each device tensor back via memcpy_dtoh and packs into a
    /// CheckpointV1 envelope. Currently writes ~22-25k f32s = ~90KB
    /// per CfcTrunk.
    /// Accessor for Mamba2 stack 1. Panics if X3 migration hasn't run.
    pub fn mamba2_l1_mut(&mut self) -> &mut crate::mamba2_block::Mamba2Block {
        self.mamba2_stack_1.as_mut()
            .expect("Mamba2 stack 1 not populated (X3 migration must have run)")
    }
    pub fn mamba2_l1(&self) -> &crate::mamba2_block::Mamba2Block {
        self.mamba2_stack_1.as_ref()
            .expect("Mamba2 stack 1 not populated (X3 migration must have run)")
    }
    /// Accessor for Mamba2 stack 2. Panics if X5 migration hasn't run.
    pub fn mamba2_l2_mut(&mut self) -> &mut crate::mamba2_block::Mamba2Block {
        self.mamba2_stack_2.as_mut()
            .expect("Mamba2 stack 2 not populated (X5 migration must have run)")
    }
    pub fn mamba2_l2(&self) -> &crate::mamba2_block::Mamba2Block {
        self.mamba2_stack_2.as_ref()
            .expect("Mamba2 stack 2 not populated (X5 migration must have run)")
    }

    pub fn save_checkpoint(&self, path: &std::path::Path) -> Result<()> {
        use std::io::Write;
        // Helper closure: download a `CudaSlice<f32>` into a host Vec
        // sized to the slice's length.
        let download = |slice: &CudaSlice<f32>| -> Result<Vec<f32>> {
            let mut v = vec![0.0_f32; slice.len()];
            self.stream.memcpy_dtoh(slice, v.as_mut_slice())?;
            Ok(v)
        };
        let download_u8 = |slice: &CudaSlice<u8>| -> Result<Vec<u8>> {
            let mut v = vec![0u8; slice.len()];
            self.stream.memcpy_dtoh(slice, v.as_mut_slice())?;
            Ok(v)
        };
        let download_u32 = |slice: &CudaSlice<u32>| -> Result<Vec<u32>> {
            let mut v = vec![0u32; slice.len()];
            self.stream.memcpy_dtoh(slice, v.as_mut_slice())?;
            Ok(v)
        };
        let m1 = self.mamba2_l1();
        let m2 = self.mamba2_l2();
        let ckpt = Checkpoint {
            n_in: self.cfg.n_in,
            n_hid: self.cfg.n_hid,
            cfc_n_in: self.cfg.cfc_n_in,
            mamba2_state_dim: self.cfg.mamba2_state_dim,
            mamba2_seq_len: m1.config.seq_len,
            mamba2_l1_in_dim: m1.config.in_dim,
            mamba2_l2_in_dim: m2.config.in_dim,
            vsn_w: download(&self.vsn_w_d)?,
            vsn_b: download(&self.vsn_b_d)?,
            m1_w_in_w:  download(&m1.w_in.weight)?,
            m1_w_in_b:  download(&m1.w_in.bias)?,
            m1_w_a_w:   download(&m1.w_a.weight)?,
            m1_w_a_b:   download(&m1.w_a.bias)?,
            m1_w_b_w:   download(&m1.w_b.weight)?,
            m1_w_b_b:   download(&m1.w_b.bias)?,
            m1_w_c:     download(&m1.w_c)?,
            m1_w_out_w: download(&m1.w_out.weight)?,
            m1_w_out_b: download(&m1.w_out.bias)?,
            ln_a_gain: download(&self.ln_a_gain_d)?,
            ln_a_bias: download(&self.ln_a_bias_d)?,
            m2_w_in_w:  download(&m2.w_in.weight)?,
            m2_w_in_b:  download(&m2.w_in.bias)?,
            m2_w_a_w:   download(&m2.w_a.weight)?,
            m2_w_a_b:   download(&m2.w_a.bias)?,
            m2_w_b_w:   download(&m2.w_b.weight)?,
            m2_w_b_b:   download(&m2.w_b.bias)?,
            m2_w_c:     download(&m2.w_c)?,
            m2_w_out_w: download(&m2.w_out.weight)?,
            m2_w_out_b: download(&m2.w_out.bias)?,
            ln_b_gain: download(&self.ln_b_gain_d)?,
            ln_b_bias: download(&self.ln_b_bias_d)?,
            attn_q: download(&self.attn_q_d)?,
            w_in: download(&self.w_in_d)?,
            w_rec: download(&self.w_rec_d)?,
            b: download(&self.b_d)?,
            tau_all: download(&self.tau_all_d)?,
            bucket_id_per_channel: download_u8(&self.bucket_id_per_channel_d)?,
            bucket_channel_offset: download_u32(&self.bucket_channel_offset_d)?,
            bucket_dim_k: download_u32(&self.bucket_dim_k_d)?,
            channels_in_bucket: download_u32(&self.channels_in_bucket_d)?,
            heads_w_skip_offset: download_u32(&self.heads_w_skip_offset_d)?,
            heads_w1: download(&self.heads_w1_d)?,
            heads_b1: download(&self.heads_b1_d)?,
            heads_w2: download(&self.heads_w2_d)?,
            heads_b2: download(&self.heads_b2_d)?,
            heads_w_gate: download(&self.heads_w_gate_d)?,
            heads_b_gate: download(&self.heads_b_gate_d)?,
            heads_w_main: download(&self.heads_w_main_d)?,
            heads_b_main: download(&self.heads_b_main_d)?,
            heads_w_skip: download(&self.heads_w_skip_d)?,
            heads_b_skip: download(&self.heads_b_skip_d)?,
        };
        let bytes = bincode::serialize(&ckpt).context("bincode serialize Checkpoint")?;
        let mut f = std::fs::File::create(path)
            .with_context(|| format!("create {}", path.display()))?;
        f.write_all(&bytes).context("write checkpoint bytes")?;
        Ok(())
    }

    /// Load a CfcTrunk from a previously-saved Checkpoint bincode file.
    /// Validates shape fields against the supplied config.
    pub fn load_checkpoint(
        dev: &MlDevice,
        cfg: &CfcConfig,
        path: &std::path::Path,
    ) -> Result<Self> {
        let bytes = std::fs::read(path)
            .with_context(|| format!("read {}", path.display()))?;
        let ckpt: Checkpoint = bincode::deserialize(&bytes)
            .context("bincode deserialize Checkpoint")?;
        anyhow::ensure!(
            ckpt.n_in == cfg.n_in && ckpt.n_hid == cfg.n_hid
                && ckpt.cfc_n_in == cfg.cfc_n_in
                && ckpt.mamba2_state_dim == cfg.mamba2_state_dim,
            "checkpoint dims (n_in={}, n_hid={}, cfc_n_in={}, mamba2_state_dim={}) \
             ≠ config (n_in={}, n_hid={}, cfc_n_in={}, mamba2_state_dim={})",
            ckpt.n_in, ckpt.n_hid, ckpt.cfc_n_in, ckpt.mamba2_state_dim,
            cfg.n_in, cfg.n_hid, cfg.cfc_n_in, cfg.mamba2_state_dim,
        );

        // Construct a fresh trunk to wire up cubins + scratch buffers,
        // then overwrite every weight tensor from the checkpoint via
        // memcpy_htod (size-checked to match the file's vectors).
        let mut trunk = Self::new_random(dev, cfg, 0)?;
        let stream = trunk.stream.clone();
        let upload = |dst: &mut CudaSlice<f32>, src: &[f32], name: &str| -> Result<()> {
            anyhow::ensure!(
                dst.len() == src.len(),
                "size mismatch loading {name}: device {} != file {}", dst.len(), src.len()
            );
            stream.memcpy_htod(src, dst)?;
            Ok(())
        };
        let upload_u8 = |dst: &mut CudaSlice<u8>, src: &[u8], name: &str| -> Result<()> {
            anyhow::ensure!(
                dst.len() == src.len(),
                "size mismatch loading {name}: device {} != file {}", dst.len(), src.len()
            );
            stream.memcpy_htod(src, dst)?;
            Ok(())
        };
        let upload_u32 = |dst: &mut CudaSlice<u32>, src: &[u32], name: &str| -> Result<()> {
            anyhow::ensure!(
                dst.len() == src.len(),
                "size mismatch loading {name}: device {} != file {}", dst.len(), src.len()
            );
            stream.memcpy_htod(src, dst)?;
            Ok(())
        };
        upload(&mut trunk.vsn_w_d, &ckpt.vsn_w, "vsn_w")?;
        upload(&mut trunk.vsn_b_d, &ckpt.vsn_b, "vsn_b")?;
        {
            let m1 = trunk.mamba2_l1_mut();
            upload(&mut m1.w_in.weight,  &ckpt.m1_w_in_w,  "m1.w_in.weight")?;
            upload(&mut m1.w_in.bias,    &ckpt.m1_w_in_b,  "m1.w_in.bias")?;
            upload(&mut m1.w_a.weight,   &ckpt.m1_w_a_w,   "m1.w_a.weight")?;
            upload(&mut m1.w_a.bias,     &ckpt.m1_w_a_b,   "m1.w_a.bias")?;
            upload(&mut m1.w_b.weight,   &ckpt.m1_w_b_w,   "m1.w_b.weight")?;
            upload(&mut m1.w_b.bias,     &ckpt.m1_w_b_b,   "m1.w_b.bias")?;
            upload(&mut m1.w_c,          &ckpt.m1_w_c,     "m1.w_c")?;
            upload(&mut m1.w_out.weight, &ckpt.m1_w_out_w, "m1.w_out.weight")?;
            upload(&mut m1.w_out.bias,   &ckpt.m1_w_out_b, "m1.w_out.bias")?;
        }
        upload(&mut trunk.ln_a_gain_d, &ckpt.ln_a_gain, "ln_a_gain")?;
        upload(&mut trunk.ln_a_bias_d, &ckpt.ln_a_bias, "ln_a_bias")?;
        {
            let m2 = trunk.mamba2_l2_mut();
            upload(&mut m2.w_in.weight,  &ckpt.m2_w_in_w,  "m2.w_in.weight")?;
            upload(&mut m2.w_in.bias,    &ckpt.m2_w_in_b,  "m2.w_in.bias")?;
            upload(&mut m2.w_a.weight,   &ckpt.m2_w_a_w,   "m2.w_a.weight")?;
            upload(&mut m2.w_a.bias,     &ckpt.m2_w_a_b,   "m2.w_a.bias")?;
            upload(&mut m2.w_b.weight,   &ckpt.m2_w_b_w,   "m2.w_b.weight")?;
            upload(&mut m2.w_b.bias,     &ckpt.m2_w_b_b,   "m2.w_b.bias")?;
            upload(&mut m2.w_c,          &ckpt.m2_w_c,     "m2.w_c")?;
            upload(&mut m2.w_out.weight, &ckpt.m2_w_out_w, "m2.w_out.weight")?;
            upload(&mut m2.w_out.bias,   &ckpt.m2_w_out_b, "m2.w_out.bias")?;
        }
        upload(&mut trunk.ln_b_gain_d, &ckpt.ln_b_gain, "ln_b_gain")?;
        upload(&mut trunk.ln_b_bias_d, &ckpt.ln_b_bias, "ln_b_bias")?;
        upload(&mut trunk.attn_q_d, &ckpt.attn_q, "attn_q")?;
        upload(&mut trunk.w_in_d, &ckpt.w_in, "cfc.w_in")?;
        upload(&mut trunk.w_rec_d, &ckpt.w_rec, "cfc.w_rec")?;
        upload(&mut trunk.b_d, &ckpt.b, "cfc.b")?;
        upload(&mut trunk.tau_all_d, &ckpt.tau_all, "cfc.tau_all")?;
        upload_u8(&mut trunk.bucket_id_per_channel_d, &ckpt.bucket_id_per_channel, "bucket_id_per_channel")?;
        upload_u32(&mut trunk.bucket_channel_offset_d, &ckpt.bucket_channel_offset, "bucket_channel_offset")?;
        upload_u32(&mut trunk.bucket_dim_k_d, &ckpt.bucket_dim_k, "bucket_dim_k")?;
        upload_u32(&mut trunk.channels_in_bucket_d, &ckpt.channels_in_bucket, "channels_in_bucket")?;
        upload_u32(&mut trunk.heads_w_skip_offset_d, &ckpt.heads_w_skip_offset, "heads_w_skip_offset")?;
        upload(&mut trunk.heads_w1_d, &ckpt.heads_w1, "heads_w1")?;
        upload(&mut trunk.heads_b1_d, &ckpt.heads_b1, "heads_b1")?;
        upload(&mut trunk.heads_w2_d, &ckpt.heads_w2, "heads_w2")?;
        upload(&mut trunk.heads_b2_d, &ckpt.heads_b2, "heads_b2")?;
        upload(&mut trunk.heads_w_gate_d, &ckpt.heads_w_gate, "heads_w_gate")?;
        upload(&mut trunk.heads_b_gate_d, &ckpt.heads_b_gate, "heads_b_gate")?;
        upload(&mut trunk.heads_w_main_d, &ckpt.heads_w_main, "heads_w_main")?;
        upload(&mut trunk.heads_b_main_d, &ckpt.heads_b_main, "heads_b_main")?;
        upload(&mut trunk.heads_w_skip_d, &ckpt.heads_w_skip, "heads_w_skip")?;
        upload(&mut trunk.heads_b_skip_d, &ckpt.heads_b_skip, "heads_b_skip")?;
        Ok(trunk)
    }

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

#[cfg(test)]
mod checkpoint_tests {
    use super::*;
    use ml_core::device::MlDevice;

    /// Regression test for the production panic at
    /// `trunk.rs:mamba2_l1()` during `PerceptionTrainer::from_checkpoint`.
    /// Reproduces save -> load -> assert all weights including Mamba2 ×2
    /// round-trip bit-identically. Originally `new_random` left
    /// `mamba2_stack_1/2 = None`, so `save_checkpoint` itself panicked
    /// before the file even got written.
    #[test]
    #[ignore = "requires CUDA"]
    fn save_load_roundtrip_preserves_all_weights() {
        let dev = match MlDevice::cuda(0) {
            Ok(d) => d,
            Err(e) => {
                eprintln!("skipping: cuda device unavailable ({e})");
                return;
            }
        };
        let cfg = CfcConfig::default();
        let trunk_a = CfcTrunk::new_random(&dev, &cfg, 0xDEAD_BEEF).expect("new_random");
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("trunk.bin");
        // Both stacks must allocate real weight buffers — zero-length slices
        // would let bit-equivalence pass trivially while shipping a useless
        // checkpoint.
        assert!(trunk_a.mamba2_l1().w_in.weight.len() > 0, "m1 w_in empty");
        assert!(trunk_a.mamba2_l1().w_a.weight.len()  > 0, "m1 w_a empty");
        assert!(trunk_a.mamba2_l2().w_in.weight.len() > 0, "m2 w_in empty");
        trunk_a.save_checkpoint(&path).expect("save_checkpoint");
        let trunk_b = CfcTrunk::load_checkpoint(&dev, &cfg, &path).expect("load_checkpoint");
        let read = |trunk: &CfcTrunk, slice: &CudaSlice<f32>| -> Vec<f32> {
            let mut v = vec![0.0; slice.len()];
            trunk.stream.memcpy_dtoh(slice, v.as_mut_slice()).unwrap();
            v
        };
        let read_u8 = |trunk: &CfcTrunk, slice: &CudaSlice<u8>| -> Vec<u8> {
            let mut v = vec![0u8; slice.len()];
            trunk.stream.memcpy_dtoh(slice, v.as_mut_slice()).unwrap();
            v
        };
        let read_u32 = |trunk: &CfcTrunk, slice: &CudaSlice<u32>| -> Vec<u32> {
            let mut v = vec![0u32; slice.len()];
            trunk.stream.memcpy_dtoh(slice, v.as_mut_slice()).unwrap();
            v
        };
        assert_eq!(read(&trunk_a, &trunk_a.w_in_d), read(&trunk_b, &trunk_b.w_in_d));
        assert_eq!(read(&trunk_a, &trunk_a.w_rec_d), read(&trunk_b, &trunk_b.w_rec_d));
        assert_eq!(read(&trunk_a, &trunk_a.b_d), read(&trunk_b, &trunk_b.b_d));
        assert_eq!(read(&trunk_a, &trunk_a.tau_all_d), read(&trunk_b, &trunk_b.tau_all_d));
        // Bucket metadata: zero-initialised in new_random; round-trips must
        // preserve them bit-identically (trivially zero == zero pre-transition,
        // but the producer/consumer contract is verified here).
        assert_eq!(read_u8(&trunk_a, &trunk_a.bucket_id_per_channel_d), read_u8(&trunk_b, &trunk_b.bucket_id_per_channel_d));
        assert_eq!(read_u32(&trunk_a, &trunk_a.bucket_channel_offset_d), read_u32(&trunk_b, &trunk_b.bucket_channel_offset_d));
        assert_eq!(read_u32(&trunk_a, &trunk_a.bucket_dim_k_d), read_u32(&trunk_b, &trunk_b.bucket_dim_k_d));
        assert_eq!(read_u32(&trunk_a, &trunk_a.channels_in_bucket_d), read_u32(&trunk_b, &trunk_b.channels_in_bucket_d));
        assert_eq!(read_u32(&trunk_a, &trunk_a.heads_w_skip_offset_d), read_u32(&trunk_b, &trunk_b.heads_w_skip_offset_d));
        assert_eq!(read(&trunk_a, &trunk_a.vsn_w_d), read(&trunk_b, &trunk_b.vsn_w_d));
        assert_eq!(read(&trunk_a, &trunk_a.ln_a_gain_d), read(&trunk_b, &trunk_b.ln_a_gain_d));
        assert_eq!(read(&trunk_a, &trunk_a.attn_q_d), read(&trunk_b, &trunk_b.attn_q_d));
        assert_eq!(read(&trunk_a, &trunk_a.heads_w_gate_d), read(&trunk_b, &trunk_b.heads_w_gate_d));

        // Mamba2 stack 1: the exact bug we shipped to production. Both
        // saved-side `mamba2_l1()` and load-side weight upload would panic
        // if mamba2_stack_1 weren't populated by new_random.
        let m1a = trunk_a.mamba2_l1();
        let m1b = trunk_b.mamba2_l1();
        assert_eq!(read(&trunk_a, &m1a.w_in.weight),  read(&trunk_b, &m1b.w_in.weight));
        assert_eq!(read(&trunk_a, &m1a.w_in.bias),    read(&trunk_b, &m1b.w_in.bias));
        assert_eq!(read(&trunk_a, &m1a.w_a.weight),   read(&trunk_b, &m1b.w_a.weight));
        assert_eq!(read(&trunk_a, &m1a.w_b.weight),   read(&trunk_b, &m1b.w_b.weight));
        assert_eq!(read(&trunk_a, &m1a.w_c),          read(&trunk_b, &m1b.w_c));
        assert_eq!(read(&trunk_a, &m1a.w_out.weight), read(&trunk_b, &m1b.w_out.weight));

        // Mamba2 stack 2: same.
        let m2a = trunk_a.mamba2_l2();
        let m2b = trunk_b.mamba2_l2();
        assert_eq!(read(&trunk_a, &m2a.w_in.weight),  read(&trunk_b, &m2b.w_in.weight));
        assert_eq!(read(&trunk_a, &m2a.w_a.weight),   read(&trunk_b, &m2b.w_a.weight));
        assert_eq!(read(&trunk_a, &m2a.w_out.weight), read(&trunk_b, &m2b.w_out.weight));
    }
}

