// snap_feature_assemble.cu
//
// Reads one MBP-10 snapshot in struct-of-arrays layout, emits the 40-dim
// feature vector documented in spec Section 2 (bumped from 32 → 40 in
// Phase 2C to add TGN-style Δt Fourier features).
//
// Layout enforced by tests:
//   out[0]     = (mid - prev_mid) / tick_size       [O(1), tick-normalized return]
//   out[1]     = (ask[0] - bid[0]) / tick_size      [O(1), spread in ticks]
//   out[2..7]  = log1p(bid_sz[0..5])                [0..~7, log size]
//   out[7..12] = log1p(ask_sz[0..5])                [0..~7]
//   out[12..17]= sgn(d) * log1p(|d|), d = bid_delta - ask_delta
//                                                   [signed-log OFI, ±~7]
//   out[17]    = log1p(trade_count)                 [0..~5]
//   out[18]    = sgn(v) * log1p(|v|), v = trade_signed_vol
//                                                   [signed-log, ±~7]
//   out[19]    = log1p(max(dt_ms, 0))               [0..~9]
//   out[20..26]= regime[0..6] (loader-precomputed EMA cascade — gives
//                              model multi-minute context unreachable
//                              inside the K-snapshot BPTT window)
//   out[26..32]= reserved, zero (future macro context)
//   out[32..40]= TGN-style Δt Fourier features:
//                  (cos(ω_k · Δt_ns), sin(ω_k · Δt_ns))_{k=0..3}
//                with log-spaced periods [60s, 6s, 600ms, 60ms].
//                Gives the model explicit Δt encoding that's
//                discriminative across scales — critical with
//                decision-stride > 1 where the gap between K-positions
//                becomes irregular. Range: [-1, 1] for each slot.
//
// All transforms are monotone, bounded-rate, dimensionally clean —
// tick_size is a market constant; log1p / signed-log carry no tuned
// parameters; Fourier features are intrinsically [-1, 1]. Brings every
// feature to comparable O(1)-O(10) scale so Mamba2's W_in projection
// doesn't saturate on the unbounded features (raw bid/ask size deltas
// were ±1e2, dt_ms could exceed 1e4).
//
// Single-thread per-snapshot kernel; the call site launches with one
// block of one thread per snapshot.

#define SNAP_FEATURE_DIM 40

__device__ __forceinline__ float signed_log1p(float x) {
    return copysignf(log1pf(fabsf(x)), x);
}

// Angular frequencies for TGN Δt Fourier encoding (rad / ns).
// Periods are log-spaced: 60s, 6s, 600ms, 60ms — span sub-second
// microstructure to minute-scale macro context. 2π / period_ns.
__device__ __constant__ float SNAP_DT_OMEGAS[4] = {
    1.0471975511965977e-10f,   // 2π / 60e9 ns (60s)
    1.0471975511965977e-9f,    // 2π / 6e9 ns  (6s)
    1.0471975511965977e-8f,    // 2π / 600e6 ns (600ms)
    1.0471975511965977e-7f,    // 2π / 60e6 ns  (60ms)
};

__device__ __forceinline__ void emit_dt_fourier(float* o_32, long dt_ns) {
    // Convert to float once — at the longest period (60s) the angle for
    // Δt = 60s is exactly 2π, so float precision is fine within ±100s.
    const float dt = (float)dt_ns;
    #pragma unroll
    for (int k = 0; k < 4; ++k) {
        const float theta = SNAP_DT_OMEGAS[k] * dt;
        o_32[2 * k]     = cosf(theta);
        o_32[2 * k + 1] = sinf(theta);
    }
}

extern "C" __global__ void snap_feature_assemble(
    const float* __restrict__ bid_px,        // [10]
    const float* __restrict__ bid_sz,        // [10]
    const float* __restrict__ ask_px,        // [10]
    const float* __restrict__ ask_sz,        // [10]
    const float* __restrict__ prev_bid_sz,   // [10] for OFI baseline
    const float* __restrict__ prev_ask_sz,   // [10]
    const float* __restrict__ regime,        // [6] loader-precomputed EMA cascade
    float prev_mid,
    float trade_signed_vol,
    int   trade_count,
    long  ts_ns,
    long  prev_ts_ns,
    float tick_size,
    float* __restrict__ out                  // [40]
) {
    if (threadIdx.x != 0 || blockIdx.x != 0) return;

    const float mid = 0.5f * (bid_px[0] + ask_px[0]);
    out[0] = (mid > 0.0f && prev_mid > 0.0f) ? (mid - prev_mid) / tick_size : 0.0f;
    out[1] = (ask_px[0] - bid_px[0]) / tick_size;

    for (int i = 0; i < 5; ++i) {
        out[2 + i] = log1pf(bid_sz[i]);
        out[7 + i] = log1pf(ask_sz[i]);
        const float bid_delta = bid_sz[i] - prev_bid_sz[i];
        const float ask_delta = ask_sz[i] - prev_ask_sz[i];
        out[12 + i] = signed_log1p(bid_delta - ask_delta);
    }

    out[17] = log1pf((float) trade_count);
    out[18] = signed_log1p(trade_signed_vol);
    const long dt_ns = ts_ns - prev_ts_ns;
    const float dt_ms = (float)dt_ns * 1e-6f;
    out[19] = log1pf(fmaxf(dt_ms, 0.0f));

    // Regime context (loader-precomputed EMA cascade) — six slots.
    for (int i = 0; i < 6; ++i) out[20 + i] = regime[i];

    // Reserved-zero slots [26..32] kept for future macro context.
    for (int i = 26; i < 32; ++i) out[i] = 0.0f;

    // TGN Δt Fourier features [32..40] — 4 log-spaced periods,
    // (cos, sin) pair each. Δt clamped to >= 0 to handle the rare
    // out-of-order timestamp without producing nonsense angles.
    const long dt_ns_clamped = dt_ns >= 0 ? dt_ns : 0;
    emit_dt_fourier(out + 32, dt_ns_clamped);
}


// ─── Batched variant ──────────────────────────────────────────────────
//
// Process N snapshots in a single kernel launch — one thread per
// snapshot. Replaces N separate single-snapshot launches with one
// fused launch. Used by PerceptionTrainer to compute snap_features
// for all B*K snapshots per training step in a single kernel call
// (eliminates ~768 launches per step at B=8, K=96).
//
// Layout: all inputs are [N, ...] row-major. Output is [N, 32].
// Per-snapshot semantics are IDENTICAL to the single-snapshot kernel
// above — bit-equivalent for any one snapshot in isolation.

extern "C" __global__ void snap_feature_assemble_batched(
    const float* __restrict__ bid_px,         // [N, 10]
    const float* __restrict__ bid_sz,         // [N, 10]
    const float* __restrict__ ask_px,         // [N, 10]
    const float* __restrict__ ask_sz,         // [N, 10]
    const float* __restrict__ prev_bid_sz,    // [N, 10]
    const float* __restrict__ prev_ask_sz,    // [N, 10]
    const float* __restrict__ regime,         // [N, 6]
    const float* __restrict__ prev_mid,       // [N]
    const float* __restrict__ trade_signed_vol, // [N]
    const int*   __restrict__ trade_count,    // [N]
    const long*  __restrict__ ts_ns,          // [N]
    const long*  __restrict__ prev_ts_ns,     // [N]
    float tick_size,
    int N,
    float* __restrict__ out                   // [N, 40]
) {
    int n = blockIdx.x * blockDim.x + threadIdx.x;
    if (n >= N) return;

    const float* bx = bid_px      + n * 10;
    const float* bs = bid_sz      + n * 10;
    const float* ax = ask_px      + n * 10;
    const float* as_ = ask_sz     + n * 10;
    const float* pbs = prev_bid_sz + n * 10;
    const float* pas = prev_ask_sz + n * 10;
    const float* rg = regime      + n * 6;
    float* o = out + n * SNAP_FEATURE_DIM;

    const float mid = 0.5f * (bx[0] + ax[0]);
    const float pm  = prev_mid[n];
    o[0] = (mid > 0.0f && pm > 0.0f) ? (mid - pm) / tick_size : 0.0f;
    o[1] = (ax[0] - bx[0]) / tick_size;

    for (int i = 0; i < 5; ++i) {
        o[2 + i] = log1pf(bs[i]);
        o[7 + i] = log1pf(as_[i]);
        const float bid_delta = bs[i] - pbs[i];
        const float ask_delta = as_[i] - pas[i];
        o[12 + i] = signed_log1p(bid_delta - ask_delta);
    }

    o[17] = log1pf((float) trade_count[n]);
    o[18] = signed_log1p(trade_signed_vol[n]);
    const long dt_ns = ts_ns[n] - prev_ts_ns[n];
    const float dt_ms = (float)dt_ns * 1e-6f;
    o[19] = log1pf(fmaxf(dt_ms, 0.0f));

    for (int i = 0; i < 6; ++i) o[20 + i] = rg[i];
    for (int i = 26; i < 32; ++i) o[i] = 0.0f;

    // TGN Δt Fourier features at slots [32..40].
    const long dt_ns_clamped = dt_ns >= 0 ? dt_ns : 0;
    emit_dt_fourier(o + 32, dt_ns_clamped);
}
