//! `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};
use ml_core::device::MlDevice;

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

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;

/// Encoder input dimension: raw market features (40) + per-batch state
/// (4 trade_context + 12 multires = 16). The encoder's first layer and
/// VSN gate size from this, not FEATURE_DIM.
pub const ENCODER_INPUT_DIM: usize = FEATURE_DIM + 16;

/// Number of book-side levels in one MBP-10 snapshot (bid + ask).
/// Canonical structural dim — every Rust array sized for book depth
/// (`[f32; BOOK_LEVELS]`) and every loop iterating over book depth
/// (`0..BOOK_LEVELS`) MUST reference this constant, not the literal `10`.
/// Per feedback_use_consts_not_literals_for_structural_dims.
///
/// Re-exported by `ml_backtesting::lob::BOOK_LEVELS` for downstream use
/// in the LOB simulator (ml-backtesting depends on ml-alpha already via
/// `Mbp10RawInput` import).
pub const BOOK_LEVELS: usize = 10;

/// 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 microstructure-burst (~2.6ms) and 1-2-second-clustering
/// (~100ms) trend / volatility / liquidity context that is structurally
/// unreachable inside the K-snapshot BPTT window. (NOTE: these are NOT
/// macro-minute regimes — see `loader.rs` ALPHA_* constants for the
/// empirically-derived half-lives at the actual 74µs median Δt.)
///
/// Slot assignments (loader computes via EMA cascade):
///   regime[0] = (mid - ema_mid_fast) / sqrt(ema_rv_fast + eps)  z-score, ~2.6ms half-life
///   regime[1] = (mid - ema_mid_med)  / sqrt(ema_rv_med  + eps)  z-score, ~100ms half-life
///   regime[2] = (ema_mid_fast - ema_mid_med) / sqrt(ema_rv_med + eps)  trend signal
///   regime[3] = log1p(ema_rv_med * 1e4)                         log-compressed med vol
///   regime[4] = log1p(ema_spread_in_ticks_fast)                 liquidity regime
///   regime[5] = log1p(ema_trades_per_snap_fast * 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.
///
/// `#[repr(C)]` guarantees stable C-ABI layout so the GPU
/// `snapshot_aos_to_soa` kernel can read this struct directly from
/// a mapped-pinned buffer without field-by-field staging.
#[repr(C)]
#[derive(Clone, Copy, Debug, Default)]
pub struct Mbp10RawInput {
    pub bid_px: [f32; BOOK_LEVELS],
    pub bid_sz: [f32; BOOK_LEVELS],
    pub ask_px: [f32; BOOK_LEVELS],
    pub ask_sz: [f32; BOOK_LEVELS],
    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],
}

/// Byte size of one `Mbp10RawInput` in C ABI layout. The GPU
/// `snapshot_aos_to_soa` kernel's `Mbp10Raw` struct must match this
/// exactly — a mismatch silently corrupts every SoA field.
pub const MBP10_RAW_INPUT_BYTES: usize = std::mem::size_of::<Mbp10RawInput>();
const _: () = assert!(MBP10_RAW_INPUT_BYTES == 216, "Mbp10RawInput size drift — update snapshot_aos_to_soa.cu");

/// 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 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 rs = stream.cu_stream();
        let mut args = RawArgs::new();
        args.push_ptr(bid_px.raw_ptr());
        args.push_ptr(bid_sz.raw_ptr());
        args.push_ptr(ask_px.raw_ptr());
        args.push_ptr(ask_sz.raw_ptr());
        args.push_ptr(prev_bid_sz.raw_ptr());
        args.push_ptr(prev_ask_sz.raw_ptr());
        args.push_ptr(regime.raw_ptr());
        args.push_f32(prev_mid);
        args.push_f32(trade_signed_vol);
        args.push_i32(trade_count);
        args.push_u64(ts_ns as u64);
        args.push_u64(prev_ts_ns as u64);
        args.push_f32(tick_size);
        args.push_ptr(out_d.raw_ptr());
        let mut ptrs = args.build_arg_ptrs();
        unsafe {
            raw_launch(
                func.cu_function(), (1,1,1), (1,1,1),
                0, rs, &mut ptrs[..args.len()],
            ).map_err(|e| anyhow::anyhow!("snap_feature_assemble: {:?}", e))?;
        }
    }

    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 dst = stream.alloc_zeros::<f32>(n).context("upload alloc")?;
    if n > 0 {
        let nbytes = n * std::mem::size_of::<f32>();
        unsafe {
            let dst_ptr = dst.raw_ptr();
            raw_memcpy_dtod_async(dst_ptr, staging.dev_ptr, nbytes, stream.cu_stream())
                .map_err(|e| anyhow::anyhow!("upload DtoD: {e:?}"))?;
        }
    }
    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 = src.raw_ptr();
        raw_memcpy_dtod_async(staging.dev_ptr, src_ptr, nbytes, stream.cu_stream())
            .map_err(|e| anyhow::anyhow!("download DtoD: {e:?}"))?;
        raw_stream_sync(stream.cu_stream())
            .map_err(|e| anyhow::anyhow!("download sync: {e:?}"))?;
    }
    Ok(staging.read_all())
}
