//! `snap_feature_assemble` kernel binding.
//!
//! Per spec Section 2 + Phase 2C amendment (2026-05-17), builds a 40-dim
//! feature vector from one MBP-10 snapshot. Slots [0..32] match the
//! original layout; slots [32..40] hold TGN-style Δt Fourier features
//! (4 log-spaced periods × sin+cos). Production code calls this inside
//! the captured Graph A (Task 11); the standalone helper exists for
//! bit-equiv tests.

use std::sync::Arc;

use anyhow::{Context, Result};
use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut, LaunchConfig, PushKernelArg};
use ml_core::device::MlDevice;

use crate::pinned_mem::MappedF32Buffer;

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

pub const ES_TICK_SIZE: f32 = 0.25;
pub const FEATURE_DIM: usize = 40;

/// Number of regime-context features pre-computed by the loader and
/// piped into snap_features slots `out[20..20+REGIME_DIM]`. These give
/// the model multi-minute trend / volatility / liquidity context that
/// is structurally unreachable inside the K-snapshot BPTT window.
///
/// Slot assignments (loader computes via EMA cascade):
///   regime[0] = (mid - ema_mid_med)  / sqrt(ema_rv_med + eps)   z-score, ~35s
///   regime[1] = (mid - ema_mid_slow) / sqrt(ema_rv_slow + eps)  z-score, ~6min
///   regime[2] = (ema_mid_med - ema_mid_slow) / sqrt(ema_rv_slow + eps)  trend signal
///   regime[3] = log1p(ema_rv_slow * 1e4)                        log-compressed slow vol
///   regime[4] = log1p(ema_spread_in_ticks_med)                  liquidity regime
///   regime[5] = log1p(ema_trades_per_snap_med * 100)            activity regime
pub const REGIME_DIM: usize = 6;

/// One raw MBP-10 snapshot in struct-of-arrays form. Phase A loader
/// converts predecoded sidecar rows to this shape.
#[derive(Clone, Copy, Debug, Default)]
pub struct Mbp10RawInput {
    pub bid_px: [f32; 10],
    pub bid_sz: [f32; 10],
    pub ask_px: [f32; 10],
    pub ask_sz: [f32; 10],
    pub prev_mid: f32,
    pub trade_signed_vol: f32,
    pub trade_count: u32,
    pub ts_ns: u64,
    pub prev_ts_ns: u64,
    /// Regime context (see REGIME_DIM doc). Loader fills via EMA pass
    /// across the WHOLE source file; sliced per-snapshot when packing
    /// the training sequence. Default = all zeros (sentinel) so unit
    /// tests + the standalone GPU helper still construct cleanly.
    pub regime: [f32; REGIME_DIM],
}

/// Standalone GPU run for bit-equiv tests. Production code uses the
/// captured Graph A path in `cfc::trunk` (Task 11).
pub fn snap_feature_assemble_gpu(dev: &MlDevice, input: &Mbp10RawInput) -> Result<[f32; FEATURE_DIM]> {
    let stream: &Arc<CudaStream> = dev.cuda_stream().context("CUDA stream")?;
    let ctx = dev.cuda_context().context("CUDA context")?;
    let module = ctx.load_cubin(KERNEL_CUBIN.to_vec()).context("load snap_feature_assemble cubin")?;
    let func = module.load_function("snap_feature_assemble").context("load function")?;

    // Stage raw input via mapped-pinned -> DtoD into CudaSlice. For one-shot
    // tests this is fine; production wires the same data through the
    // captured Graph A's mapped-pinned slot.
    let bid_px = upload(stream, &input.bid_px)?;
    let bid_sz = upload(stream, &input.bid_sz)?;
    let ask_px = upload(stream, &input.ask_px)?;
    let ask_sz = upload(stream, &input.ask_sz)?;
    let prev_bid_sz = stream.alloc_zeros::<f32>(10).context("prev_bid_sz alloc")?;
    let prev_ask_sz = stream.alloc_zeros::<f32>(10).context("prev_ask_sz alloc")?;
    let regime = upload(stream, &input.regime)?;
    let mut out_d = stream.alloc_zeros::<f32>(FEATURE_DIM).context("out alloc")?;

    let prev_mid = input.prev_mid;
    let trade_signed_vol = input.trade_signed_vol;
    let trade_count = input.trade_count as i32;
    let ts_ns = input.ts_ns as i64;
    let prev_ts_ns = input.prev_ts_ns as i64;
    let tick_size = ES_TICK_SIZE;

    let cfg = LaunchConfig {
        grid_dim: (1, 1, 1),
        block_dim: (1, 1, 1),
        shared_mem_bytes: 0,
    };

    let mut launch = stream.launch_builder(&func);
    launch
        .arg(&bid_px)
        .arg(&bid_sz)
        .arg(&ask_px)
        .arg(&ask_sz)
        .arg(&prev_bid_sz)
        .arg(&prev_ask_sz)
        .arg(&regime)
        .arg(&prev_mid)
        .arg(&trade_signed_vol)
        .arg(&trade_count)
        .arg(&ts_ns)
        .arg(&prev_ts_ns)
        .arg(&tick_size)
        .arg(&mut out_d);
    unsafe { launch.launch(cfg).context("snap_feature_assemble launch")?; }

    let v = download(stream, &out_d)?;
    let mut out = [0f32; FEATURE_DIM];
    out.copy_from_slice(&v);
    Ok(out)
}

fn upload(stream: &Arc<CudaStream>, host: &[f32]) -> Result<CudaSlice<f32>> {
    let n = host.len();
    let staging = unsafe { MappedF32Buffer::new(n) }
        .map_err(|e| anyhow::anyhow!("upload staging: {e}"))?;
    staging.write_from_slice(host);
    let mut dst = stream.alloc_zeros::<f32>(n).context("upload alloc")?;
    if n > 0 {
        let nbytes = n * std::mem::size_of::<f32>();
        unsafe {
            let (dst_ptr, _g) = dst.device_ptr_mut(stream);
            cudarc::driver::result::memcpy_dtod_async(dst_ptr, staging.dev_ptr, nbytes, stream.cu_stream())
                .context("upload DtoD")?;
        }
    }
    Ok(dst)
}

fn download(stream: &Arc<CudaStream>, src: &CudaSlice<f32>) -> Result<Vec<f32>> {
    let n = src.len();
    let staging = unsafe { MappedF32Buffer::new(n) }
        .map_err(|e| anyhow::anyhow!("download staging: {e}"))?;
    let nbytes = n * std::mem::size_of::<f32>();
    unsafe {
        let (src_ptr, _g) = src.device_ptr(stream);
        cudarc::driver::result::memcpy_dtod_async(staging.dev_ptr, src_ptr, nbytes, stream.cu_stream())
            .context("download DtoD")?;
    }
    stream.synchronize().context("download sync")?;
    Ok(staging.read_all())
}
