/**
 * Shared device functions for DQN/PPO CUDA experience collection kernels.
 *
 * Included via NVRTC source concatenation (not #include).
 * Both kernels prepend this file's contents before their own source.
 */

/* ------------------------------------------------------------------ */
/*  BF16 Mixed Precision Support (H100 Tensor Core Path)              */
/* ------------------------------------------------------------------ */
/* H100 SM90: 989 TFLOPS BF16 tensor cores vs 67 TFLOPS F32 scalar.
 * Strategy: BF16 storage + BF16 matmul inputs + F32 accumulation.
 * This matches NVIDIA's native tensor core accumulate mode. */
#include <cuda_bf16.h>

/** BF16 leaky ReLU — operates on raw bf16, returns bf16. */
__device__ __forceinline__ __nv_bfloat16 leaky_relu_bf16(__nv_bfloat16 x) {
    float xf = __bfloat162float(x);
    return __float2bfloat16((xf >= 0.0f) ? xf : 0.01f * xf);
}

/** Convert F32 to BF16 (truncation, matching PyTorch default). */
__device__ __forceinline__ __nv_bfloat16 f32_to_bf16(float x) {
    return __float2bfloat16(x);
}

/* ------------------------------------------------------------------ */
/*  Constants                                                         */
/* ------------------------------------------------------------------ */
/* State layout — MUST be injected via NVRTC dim_overrides (no fallback defaults).
 * STATE_DIM = tensor-core-aligned state width (e.g. 48 without OFI, 56 with OFI).
 * MARKET_DIM = raw market feature count (42: 40 base + 2 regime).
 * PORTFOLIO_DIM = portfolio feature count (cash, position_size, unrealized_pnl).
 * OFI_DIM = order flow imbalance features from MBP-10 (0 when disabled, 8 when enabled). */
#ifndef STATE_DIM
#error "STATE_DIM must be defined via dim_overrides before including this header"
#endif
#ifndef MARKET_DIM
#error "MARKET_DIM must be defined via dim_overrides before including this header"
#endif
#ifndef PORTFOLIO_DIM
#error "PORTFOLIO_DIM must be defined via dim_overrides before including this header"
#endif
#ifndef OFI_DIM
#define OFI_DIM 0  /* Default: no OFI features (safe fallback for kernels that don't use OFI) */
#endif
#define DQN_NUM_ACTIONS    9
#define DQN_ORDER_ACTIONS  3
#define DQN_URGENCY_ACTIONS 3
#define DQN_TOTAL_ACTIONS  81   /* 9 * 3 * 3 factored action space */
#define PPO_NUM_ACTIONS    45   /* PPO unchanged — separate action space */
#define NUM_ACTIONS        DQN_NUM_ACTIONS  /* default for DQN kernels */

/* Branching DQN head sizes (Tavakoli et al., 2018).
 * 3 independent advantage heads: exposure (9), order (3), urgency (3).
 * Factored action index = exposure * 9 + order * 3 + urgency (0-80). */
#define BRANCH_EXPOSURE_ACTIONS 9
#define BRANCH_ORDER_ACTIONS    3
#define BRANCH_URGENCY_ACTIONS  3
#define BRANCH_MAX_ACTIONS      BRANCH_EXPOSURE_ACTIONS  /* largest branch */

/* Curiosity forward model */
#define CUR_INPUT          (MARKET_DIM + 3) /* MARKET_DIM features + 3-class one-hot */
#define CUR_HIDDEN         128
#define CUR_OUTPUT         MARKET_DIM       /* predict all market features */

/* Auxiliary state */
#define DIVERSITY_WINDOW  100
#define PORTFOLIO_STATE_SIZE 8
#define BARRIER_STATE_SIZE   5

/* ------------------------------------------------------------------ */
/*  Device Helper Functions                                           */
/* ------------------------------------------------------------------ */

/** LCG random — returns float in [0, 1). */
__device__ __forceinline__ float gpu_random(unsigned int* state) {
    *state = *state * 1664525u + 1013904223u;
    return (float)(*state & 0x00FFFFFFu) / 16777216.0f;
}

/**
 * Box-Muller pair — returns two independent Normal(0, 1) samples.
 * Uses sincosf() for a single MUFU instruction (vs separate cosf + wasted sinf).
 * Callers generating multiple Gaussians should use this to halve RNG overhead.
 */
__device__ __forceinline__ void gpu_random_gaussian_pair(unsigned int* rng, float* g1, float* g2) {
    float u1 = gpu_random(rng);
    float u2 = gpu_random(rng);
    if (u1 < 1e-10f) u1 = 1e-10f;  /* avoid log(0) */
    float r = sqrtf(-2.0f * logf(u1));
    float sin_val, cos_val;
    sincosf(2.0f * 3.14159265f * u2, &sin_val, &cos_val);
    *g1 = r * cos_val;
    *g2 = r * sin_val;
}

/**
 * Box-Muller transform — returns Normal(0, 1) sample.
 * Wraps gpu_random_gaussian_pair; second sample is discarded.
 * Prefer gpu_random_gaussian_pair() in loops generating multiple Gaussians.
 */
__device__ __forceinline__ float gpu_random_gaussian(unsigned int* rng) {
    float g1, g2;
    gpu_random_gaussian_pair(rng, &g1, &g2);
    return g1;
}

/**
 * Factorized noise transform: f(x) = sign(x) * sqrt(|x|).
 * Used in NoisyNet (Fortunato et al., 2018) to reduce noise parameter count.
 */
__device__ __forceinline__ float factorized_noise_fn(float x) {
    return (x >= 0.0f ? 1.0f : -1.0f) * sqrtf(fabsf(x));
}

/** Leaky ReLU with alpha = 0.01. */
__device__ __forceinline__ float leaky_relu(float x) {
    return (x >= 0.0f) ? x : 0.01f * x;
}

/**
 * Matrix-vector multiply: output = W * input + b, with optional LeakyReLU.
 *
 * Weight layout (Candle convention): W[out_dim, in_dim] row-major,
 * so element (j, i) is at W[j * in_dim + i].
 */
__device__ void matvec_leaky_relu(
    const float* __restrict__ W,
    const float* __restrict__ b,
    const float* input,
    float* output,
    int in_dim,
    int out_dim,
    int activate
) {
    for (int j = 0; j < out_dim; j++) {
        float acc = b[j];
        const float* row = W + j * in_dim;
        for (int i = 0; i < in_dim; i++) {
            acc += row[i] * input[i];
        }
        output[j] = activate ? leaky_relu(acc) : acc;
    }
}

/* Maximum dimension for noise scratch arrays.
 * Overridable via NVRTC #define injection to match actual network widths. */
#ifndef NOISY_MAX_DIM
#define NOISY_MAX_DIM 256
#endif

/**
 * Noisy matrix-vector multiply with factorized Gaussian noise + LeakyReLU.
 *
 * Implements NoisyNet (Fortunato et al., 2018):
 *   y[j] = (W[j,:] + sigma_w * f(eps_j) * f(eps_i)) · x + (b[j] + sigma_b * f(eps_j))
 * where sigma_w = sigma_b = sigma_init / sqrt(in_dim).
 *
 * Generates noise on-the-fly using the thread's RNG state.
 * Only applied to the ONLINE network; target uses clean matvec.
 */
__device__ void noisy_matvec_leaky_relu(
    const float* __restrict__ W,
    const float* __restrict__ b,
    const float* input,
    float* output,
    int in_dim,
    int out_dim,
    int activate,
    float sigma_init,
    unsigned int* rng
) {
    float sigma_scale = sigma_init / sqrtf((float)in_dim);

    /* Generate factorized noise vectors (paired to halve Box-Muller calls) */
    float eps_in[NOISY_MAX_DIM];
    float eps_out[NOISY_MAX_DIM];

    for (int i = 0; i < in_dim && i < NOISY_MAX_DIM; i += 2) {
        float g1, g2;
        gpu_random_gaussian_pair(rng, &g1, &g2);
        eps_in[i] = factorized_noise_fn(g1);
        if (i + 1 < in_dim && i + 1 < NOISY_MAX_DIM)
            eps_in[i + 1] = factorized_noise_fn(g2);
    }
    for (int j = 0; j < out_dim && j < NOISY_MAX_DIM; j += 2) {
        float g1, g2;
        gpu_random_gaussian_pair(rng, &g1, &g2);
        eps_out[j] = factorized_noise_fn(g1);
        if (j + 1 < out_dim && j + 1 < NOISY_MAX_DIM)
            eps_out[j + 1] = factorized_noise_fn(g2);
    }

    for (int j = 0; j < out_dim; j++) {
        float ej = (j < NOISY_MAX_DIM) ? eps_out[j] : 0.0f;
        float acc = b[j] + sigma_scale * ej;
        const float* row = W + j * in_dim;
        for (int i = 0; i < in_dim; i++) {
            float ei = (i < NOISY_MAX_DIM) ? eps_in[i] : 0.0f;
            acc += (row[i] + sigma_scale * ej * ei) * input[i];
        }
        output[j] = activate ? leaky_relu(acc) : acc;
    }
}

/* ------------------------------------------------------------------ */
/*  Shared Memory Weight Caching                                      */
/* ------------------------------------------------------------------ */

/* Shared memory tile size: rows per tile × max input dim.
 * 64 rows × 256 cols = 16384 floats = 64 KB. Fits H100's 228 KB shmem.
 * Overridable via NVRTC #define injection. */
#ifndef SHMEM_TILE_ROWS
#define SHMEM_TILE_ROWS 64
#endif
#ifndef SHMEM_MAX_IN_DIM
#define SHMEM_MAX_IN_DIM 256  /* max(STATE_DIM, SHARED_H1, SHARED_H2) */
#endif

/**
 * Cooperatively load a weight tile from global to shared memory (float4 path).
 * All threads in the block participate. Must call __syncthreads() after.
 * Uses float4 vectorized loads where possible for 4x bandwidth.
 *
 * This is the pre-Hopper fallback used on sm_70/sm_80/sm_86/sm_89.
 */
__device__ __forceinline__ void cooperative_load_tile_float4(
    float* __restrict__ shmem_dst,
    const float* __restrict__ global_src,
    int count  /* total floats to load */
) {
    int vec4_count = count / 4;
    int remainder = count % 4;
    float4* dst4 = (float4*)shmem_dst;
    const float4* src4 = (const float4*)global_src;
    for (int i = threadIdx.x; i < vec4_count; i += blockDim.x) {
        dst4[i] = src4[i];
    }
    int base = vec4_count * 4;
    for (int i = threadIdx.x; i < remainder; i += blockDim.x) {
        shmem_dst[base + i] = global_src[base + i];
    }
}

#if __CUDA_ARCH__ >= 900
/**
 * TMA-accelerated tile load from global to shared memory (Hopper sm_90+).
 *
 * Uses cp.async.bulk.shared::cta.global with mbarrier synchronization to
 * perform bulk asynchronous copies from global (HBM) to shared memory via
 * the Tensor Memory Accelerator.  Only thread 0 issues the copy + mbarrier
 * instructions; all other threads are freed for compute overlap.
 *
 * Hopper TMA uses mbarrier (NOT commit_group/wait_group — those are for the
 * older cp.async non-bulk instruction family from Ampere).  The protocol:
 *   1. mbarrier.init — set expected arrival count (1 = thread 0)
 *   2. mbarrier.arrive.expect_tx — arrive AND declare expected TX bytes
 *   3. cp.async.bulk ... [mbar] — each copy auto-adds bytes to TX counter
 *   4. mbarrier.try_wait.parity — spin until arrivals + TX bytes complete
 *
 * Falls back to float4 for:
 *   - Non-16-byte-aligned sizes (cp.async.bulk requires size % 16 == 0)
 *   - Very large tiles (>256KB, should not happen in practice)
 *
 * Constraints:
 *   - src must be 16-byte aligned (guaranteed by cudarc device alloc)
 *   - shmem_dst must be 16-byte aligned (guaranteed by __shared__ layout)
 *   - count must be > 0
 *   - Caller MUST issue __syncthreads() after return
 */
__device__ __forceinline__ void cooperative_load_tile_tma(
    float* __restrict__ shmem_dst,
    const float* __restrict__ global_src,
    int count  /* total floats to load */
) {
    const int byte_count = count * (int)sizeof(float);

    /* cp.async.bulk requires byte count to be a multiple of 16.
     * Bias tiles (e.g., NUM_ACTIONS=5 → 20 bytes) violate this.
     * Also guard against very large tiles (>256KB). */
    if ((byte_count & 15) != 0 || byte_count > 256 * 1024) {
        cooperative_load_tile_float4(shmem_dst, global_src, count);
        return;
    }

    /* Shared mbarrier for TMA completion tracking.
     * Re-initialized per call; callers separate calls with __syncwarp().
     * Static __shared__ — one per block, all inline sites share it. */
    __shared__ __align__(8) unsigned long long __tma_mbar;

    /* Compute mbar shared address for all threads (same address). */
    const unsigned int mbar_addr =
        (unsigned int)__cvta_generic_to_shared(&__tma_mbar);

    if (threadIdx.x == 0) {
        /* Initialize mbarrier: 1 expected arrival (thread 0). */
        asm volatile(
            "mbarrier.init.shared::cta.b64 [%0], 1;\n"
            : : "r"(mbar_addr) : "memory"
        );

        /* Arrive with expected TX byte count.
         * This counts as the 1 arrival AND sets the expected TX bytes.
         * The mbarrier will complete once all cp.async.bulk copies deliver
         * exactly this many bytes.
         * Use a temp .b64 register for the required output operand. */
        asm volatile(
            "{\n"
            "  .reg .b64 _st;\n"
            "  mbarrier.arrive.expect_tx.shared::cta.b64 _st, [%0], %1;\n"
            "}\n"
            : : "r"(mbar_addr), "r"((unsigned int)byte_count) : "memory"
        );

        /* Issue bulk copies in 64KB chunks (H100 TMA optimal).
         * cp.async.bulk.shared::cta.global [dst], [src], size, [mbar];
         * Each copy auto-increments the mbarrier TX counter on completion. */
        const int CHUNK_FLOATS = 16384;         /* 64 KB per chunk */
        int remaining = count;
        int offset = 0;
        while (remaining > 0) {
            int chunk = remaining < CHUNK_FLOATS ? remaining : CHUNK_FLOATS;
            int chunk_bytes = chunk * (int)sizeof(float);
            asm volatile(
                "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];\n"
                :
                : "r"((unsigned int)__cvta_generic_to_shared(shmem_dst + offset)),
                  "l"((unsigned long long)(global_src + offset)),
                  "r"((unsigned int)chunk_bytes),
                  "r"(mbar_addr)
                : "memory"
            );
            offset += chunk;
            remaining -= chunk;
        }
    }

    /* Ensure all threads see the mbarrier init before waiting.
     * Critical on Hopper with independent thread scheduling — without
     * this fence, other threads could reach try_wait before thread 0
     * writes the mbarrier init.
     * Uses __syncthreads() to support multi-warp blocks (EPISODES_PER_BLOCK > 1)
     * where multiple warps share the same shmem tile. Safe for single-warp too. */
    __syncthreads();

    /* ALL threads wait for the bulk copies to complete.
     * mbarrier completes when: arrivals == init_count (1)
     * AND tx_bytes_delivered == expected_tx (byte_count).
     * Phase parity 0 matches the freshly-initialized mbarrier.
     *
     * Uses a C loop instead of a PTX label-based loop to avoid
     * duplicate-label errors when this __forceinline__ function is
     * inlined at multiple call sites in the same kernel. */
    {
        int _tma_done = 0;
        while (!_tma_done) {
            asm volatile(
                "{\n"
                "  .reg .pred p;\n"
                "  mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 p, [%1], 0;\n"
                "  selp.b32 %0, 1, 0, p;\n"
                "}\n"
                : "=r"(_tma_done) : "r"(mbar_addr) : "memory"
            );
        }
    }

    /* Ensure all threads completed the wait before caller can reinitialize
     * the mbarrier for the next tile in the same forward pass. */
    __syncthreads();
}
#endif /* __CUDA_ARCH__ >= 900 */

/**
 * Cooperatively load a weight tile from global to shared memory.
 * All threads in the block participate. Must call __syncthreads() after.
 *
 * Dispatch: on Hopper (sm_90+) uses TMA bulk async copy via thread 0.
 * On older architectures uses float4 vectorized cooperative loads.
 */
__device__ __forceinline__ void cooperative_load_tile(
    float* __restrict__ shmem_dst,
    const float* __restrict__ global_src,
    int count  /* total floats to load */
) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 && !defined(DISABLE_TMA)
    cooperative_load_tile_tma(shmem_dst, global_src, count);
#else
    cooperative_load_tile_float4(shmem_dst, global_src, count);
#endif
}

/* ------------------------------------------------------------------ */
/*  F32→BF16 Weight Conversion Kernel                                  */
/* ------------------------------------------------------------------ */

/**
 * Convert F32 weights to BF16 in-place on GPU.
 * Called once per weight sync (not per training step).
 * Grid: ceil(n/256), Block: 256.
 */
extern "C" __global__ void f32_to_bf16_kernel(
    const float* __restrict__ src,
    __nv_bfloat16* __restrict__ dst,
    int n
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n)
        dst[idx] = __float2bfloat16(src[idx]);
}

/**
 * Convert BF16 weights back to F32 on GPU.
 * Used for reverse sync (GPU weights → VarMap for checkpointing).
 */
extern "C" __global__ void bf16_to_f32_kernel(
    const __nv_bfloat16* __restrict__ src,
    float* __restrict__ dst,
    int n
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n)
        dst[idx] = __bfloat162float(src[idx]);
}

/* ------------------------------------------------------------------ */

/**
 * Matrix-vector multiply from shared memory: output = shmem_W * input + shmem_b.
 *
 * Reads weights/bias from shared memory instead of global (HBM).
 * Only computes `tile_rows` output rows starting at `tile_offset`.
 * Weight layout: shmem_W[tile_rows, in_dim] row-major (tile-local indexing).
 */
__device__ void matvec_leaky_relu_shmem(
    const float* __restrict__ shmem_W,
    const float* __restrict__ shmem_b,
    const float* input,
    float* output,
    int in_dim,
    int out_dim,
    int tile_offset,
    int tile_rows,
    int activate
) {
    (void)out_dim;  /* full output dim — unused, tile_rows is the active count */
    for (int j = 0; j < tile_rows; j++) {
        float acc = shmem_b[j];
        const float* row = shmem_W + j * in_dim;
        for (int i = 0; i < in_dim; i++) {
            acc += row[i] * input[i];
        }
        output[tile_offset + j] = activate ? leaky_relu(acc) : acc;
    }
}

/**
 * Noisy matrix-vector multiply from shared memory with factorized Gaussian noise.
 *
 * Same as noisy_matvec_leaky_relu but reads W/b from shared memory.
 * Only computes `tile_rows` output rows starting at `tile_offset`.
 */
__device__ void noisy_matvec_leaky_relu_shmem(
    const float* __restrict__ shmem_W,
    const float* __restrict__ shmem_b,
    const float* input,
    float* output,
    int in_dim,
    int out_dim,
    int tile_offset,
    int tile_rows,
    int activate,
    float sigma_init,
    unsigned int* rng
) {
    (void)out_dim;
    float sigma_scale = sigma_init / sqrtf((float)in_dim);

    /* Generate factorized noise vectors for this tile (paired to halve Box-Muller calls).
     * eps_out only needs tile_rows elements (≤ SHMEM_TILE_ROWS=64). */
    float eps_in[NOISY_MAX_DIM];
    float eps_out[SHMEM_TILE_ROWS];

    for (int i = 0; i < in_dim && i < NOISY_MAX_DIM; i += 2) {
        float g1, g2;
        gpu_random_gaussian_pair(rng, &g1, &g2);
        eps_in[i] = factorized_noise_fn(g1);
        if (i + 1 < in_dim && i + 1 < NOISY_MAX_DIM)
            eps_in[i + 1] = factorized_noise_fn(g2);
    }
    for (int j = 0; j < tile_rows && j < SHMEM_TILE_ROWS; j += 2) {
        float g1, g2;
        gpu_random_gaussian_pair(rng, &g1, &g2);
        eps_out[j] = factorized_noise_fn(g1);
        if (j + 1 < tile_rows && j + 1 < SHMEM_TILE_ROWS)
            eps_out[j + 1] = factorized_noise_fn(g2);
    }

    for (int j = 0; j < tile_rows; j++) {
        float ej = (j < SHMEM_TILE_ROWS) ? eps_out[j] : 0.0f;
        float acc = shmem_b[j] + sigma_scale * ej;
        const float* row = shmem_W + j * in_dim;
        for (int i = 0; i < in_dim; i++) {
            float ei = (i < NOISY_MAX_DIM) ? eps_in[i] : 0.0f;
            acc += (row[i] + sigma_scale * ej * ei) * input[i];
        }
        output[tile_offset + j] = activate ? leaky_relu(acc) : acc;
    }
}

/* ------------------------------------------------------------------ */
/*  Warp-Cooperative Matrix-Vector Multiply                           */
/* ------------------------------------------------------------------ */

/**
 * Warp-cooperative reduce: sum across 32 lanes using shuffle-down.
 * Result is valid in ALL lanes (full butterfly reduction).
 */
__device__ float warp_reduce_sum_all(float val) {
    for (int offset = 16; offset > 0; offset >>= 1)
        val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
    return val;
}

/**
 * Warp-cooperative matrix-vector multiply with LeakyReLU, shared-memory tiled.
 *
 * Unlike the per-thread version, a full warp (32 lanes) cooperates on each
 * dot product. Input and output vectors are distributed across lanes in
 * strided layout: lane k holds elements k, k+32, k+64, ...
 *
 * This reduces per-thread local memory from in_dim floats (~2KB for dim=512)
 * to ceil(in_dim/32) floats (~64B), eliminating register spills on H100.
 *
 * @param shmem_W    Shared memory weight tile [tile_rows × in_dim]
 * @param shmem_b    Shared memory bias tile [tile_rows]
 * @param input_dist Per-lane distributed input [ceil(in_dim/32)]
 * @param output_dist Per-lane distributed output [ceil(out_dim/32)] (written)
 * @param in_dim     Full input dimension
 * @param out_dim    Full output dimension (only used for bounds context)
 * @param tile_offset Global output offset for this tile
 * @param tile_rows  Number of output rows in this tile (≤ SHMEM_TILE_ROWS)
 * @param activate   1 = apply LeakyReLU, 0 = linear
 * @param lane_id    threadIdx.x % 32
 */
__device__ void warp_matvec_leaky_relu_shmem(
    const float* __restrict__ shmem_W,
    const float* __restrict__ shmem_b,
    const float* input_dist,
    float* output_dist,
    int in_dim,
    int out_dim,
    int tile_offset,
    int tile_rows,
    int activate,
    int lane_id
) {
    (void)out_dim;
    for (int j = 0; j < tile_rows; j++) {
        const float* row = shmem_W + j * in_dim;
        /* Each lane computes partial dot product over its elements */
        float partial = 0.0f;
        for (int i = lane_id; i < in_dim; i += 32)
            partial += row[i] * input_dist[i / 32];
        /* Add bias (only once, via lane 0) */
        if (lane_id == 0)
            partial += shmem_b[j];
        /* Full warp reduction — result in all lanes */
        float sum = warp_reduce_sum_all(partial);
        /* Activation */
        if (activate) sum = leaky_relu(sum);
        /* Scatter to owning lane in strided layout */
        int gj = tile_offset + j;
        if (lane_id == (gj & 31))
            output_dist[gj >> 5] = sum;
    }
}

/**
 * Warp-cooperative noisy matrix-vector multiply with factorized Gaussian noise.
 *
 * Same as warp_matvec_leaky_relu_shmem but applies NoisyNet perturbation.
 * Noise vectors are generated on-the-fly per lane (no large eps_in array needed).
 *
 * @param sigma_init  NoisyNet sigma_0 parameter
 * @param rng         Per-thread RNG state pointer
 */
__device__ void warp_noisy_matvec_leaky_relu_shmem(
    const float* __restrict__ shmem_W,
    const float* __restrict__ shmem_b,
    const float* input_dist,
    float* output_dist,
    int in_dim,
    int out_dim,
    int tile_offset,
    int tile_rows,
    int activate,
    float sigma_init,
    unsigned int* rng,
    int lane_id
) {
    (void)out_dim;
    float sigma_scale = sigma_init / sqrtf((float)in_dim);

    /* Each lane generates its portion of input noise (strided) */
    /* With in_dim=512 and 32 lanes, each lane generates 16 noise values */
    int local_count = (in_dim + 31) / 32;

    for (int j = 0; j < tile_rows; j++) {
        /* Generate output noise for this row — same value needed by all lanes */
        /* Use lane 0's RNG and broadcast */
        float eps_out_j;
        if (lane_id == 0)
            eps_out_j = factorized_noise_fn(gpu_random_gaussian(rng));
        eps_out_j = __shfl_sync(0xFFFFFFFF, eps_out_j, 0);

        const float* row = shmem_W + j * in_dim;
        float partial = 0.0f;
        for (int k = 0; k < local_count; k++) {
            int i = lane_id + k * 32;
            if (i < in_dim) {
                float eps_in_i = factorized_noise_fn(gpu_random_gaussian(rng));
                float w_noisy = row[i] + sigma_scale * eps_out_j * eps_in_i;
                partial += w_noisy * input_dist[k];
            }
        }
        /* Add noisy bias (lane 0 only) */
        if (lane_id == 0)
            partial += shmem_b[j] + sigma_scale * eps_out_j;
        /* Full warp reduction */
        float sum = warp_reduce_sum_all(partial);
        if (activate) sum = leaky_relu(sum);
        int gj = tile_offset + j;
        if (lane_id == (gj & 31))
            output_dist[gj >> 5] = sum;
    }
}

/**
 * Warp-cooperative matvec with broadcast output (non-distributed).
 *
 * Same as warp_matvec_leaky_relu_shmem but ALL lanes store every output
 * element (no strided scatter). Used for small output layers where all
 * lanes need the full result (e.g., C51 atom logits).
 *
 * @param output  Non-distributed output array — all lanes write all elements
 */
__device__ void warp_matvec_broadcast_shmem(
    const float* __restrict__ shmem_W,
    const float* __restrict__ shmem_b,
    const float* input_dist,
    float* output,          /* NOT distributed — all lanes get full copy */
    int in_dim,
    int out_dim,
    int tile_offset,
    int tile_rows,
    int activate,
    int lane_id
) {
    (void)out_dim;
    for (int j = 0; j < tile_rows; j++) {
        const float* row = shmem_W + j * in_dim;
        float partial = 0.0f;
        for (int i = lane_id; i < in_dim; i += 32)
            partial += row[i] * input_dist[i / 32];
        if (lane_id == 0)
            partial += shmem_b[j];
        float sum = warp_reduce_sum_all(partial);
        if (activate) sum = leaky_relu(sum);
        /* Broadcast: ALL lanes store (redundant but correct for small arrays) */
        output[tile_offset + j] = sum;
    }
}

/**
 * Warp-cooperative noisy matvec with broadcast output.
 * Same as warp_noisy_matvec_leaky_relu_shmem but all lanes get full output.
 */
__device__ void warp_noisy_matvec_broadcast_shmem(
    const float* __restrict__ shmem_W,
    const float* __restrict__ shmem_b,
    const float* input_dist,
    float* output,          /* NOT distributed — all lanes get full copy */
    int in_dim,
    int out_dim,
    int tile_offset,
    int tile_rows,
    int activate,
    float sigma_init,
    unsigned int* rng,
    int lane_id
) {
    (void)out_dim;
    float sigma_scale = sigma_init / sqrtf((float)in_dim);
    int local_count = (in_dim + 31) / 32;

    for (int j = 0; j < tile_rows; j++) {
        float eps_out_j;
        if (lane_id == 0)
            eps_out_j = factorized_noise_fn(gpu_random_gaussian(rng));
        eps_out_j = __shfl_sync(0xFFFFFFFF, eps_out_j, 0);

        const float* row = shmem_W + j * in_dim;
        float partial = 0.0f;
        for (int k = 0; k < local_count; k++) {
            int i = lane_id + k * 32;
            if (i < in_dim) {
                float eps_in_i = factorized_noise_fn(gpu_random_gaussian(rng));
                float w_noisy = row[i] + sigma_scale * eps_out_j * eps_in_i;
                partial += w_noisy * input_dist[k];
            }
        }
        if (lane_id == 0)
            partial += shmem_b[j] + sigma_scale * eps_out_j;
        float sum = warp_reduce_sum_all(partial);
        if (activate) sum = leaky_relu(sum);
        output[tile_offset + j] = sum;
    }
}

/**
 * Warp-cooperative in-place RMSNorm on distributed data.
 *
 * Each lane holds elements in strided layout. Sum-of-squares is computed
 * via warp reduction, then each lane normalizes its own elements.
 *
 * @param data_dist  Distributed data [DIST_SIZE(dim)] per lane (modified in place)
 * @param gamma      Non-distributed gamma weights [dim] in global memory
 * @param dim        Full dimension
 * @param lane_id    Lane index (0..31)
 */
__device__ void warp_rmsnorm_inplace(float* data_dist, const float* __restrict__ gamma, int dim, int lane_id) {
    /* Each lane computes partial sum of squares over its elements */
    float partial_sq = 0.0f;
    for (int i = lane_id; i < dim; i += 32) {
        float val = data_dist[i / 32];
        partial_sq += val * val;
    }
    /* Warp reduction for total sum of squares */
    float total_sq = warp_reduce_sum_all(partial_sq);
    float rms = sqrtf(total_sq / (float)dim + 1e-6f);
    float inv_rms = 1.0f / rms;

    /* Each lane normalizes its own elements */
    for (int i = lane_id; i < dim; i += 32) {
        data_dist[i / 32] = data_dist[i / 32] * inv_rms * gamma[i];
    }
}

/**
 * In-place RMSNorm: data[i] = (data[i] / RMS(data)) * gamma[i].
 *
 * RMS = sqrt(mean(data^2) + eps), eps = 1e-6.
 * Used by DistributionalDuelingQNetwork for activation normalization.
 */
__device__ void rmsnorm_inplace(float* data, const float* __restrict__ gamma, int dim) {
    /* Compute mean(x^2) */
    float sum_sq = 0.0f;
    for (int i = 0; i < dim; i++) {
        sum_sq += data[i] * data[i];
    }
    float rms = sqrtf(sum_sq / (float)dim + 1e-6f);
    float inv_rms = 1.0f / rms;

    for (int i = 0; i < dim; i++) {
        data[i] = data[i] * inv_rms * gamma[i];
    }
}

/* ------------------------------------------------------------------ */
/*  Domain Device Functions                                           */
/* ------------------------------------------------------------------ */

/**
 * Map DQN action index (0-4) to target exposure fraction.
 *
 * DQN outputs 5 exposure actions directly:
 *   0=Short100=-1.0, 1=Short50=-0.5, 2=Flat=0.0, 3=Long50=0.5, 4=Long100=1.0
 */
__device__ __forceinline__ float action_to_exposure(int action_idx) {
    switch (action_idx) {
        case 0: return -1.0f;   /* Short100 */
        case 1: return -0.5f;   /* Short50  */
        case 2: return  0.0f;   /* Flat     */
        case 3: return  0.5f;   /* Long50   */
        case 4: return  1.0f;   /* Long100  */
        default: return 0.0f;
    }
}

/**
 * Map DQN action index (0-4) to transaction cost rate.
 *
 * With 5-action DQN, order type is determined by the CPU-side OrderRouter,
 * not by the DQN action. GPU kernel uses a fixed Market cost for all actions.
 */
__device__ __forceinline__ float action_to_tx_cost(int action_idx) {
    (void)action_idx;
    return 0.0001f;  /* Market rate for all DQN exposure actions */
}

/**
 * Map factored action index (0-44) to target exposure fraction.
 *
 * Both PPO and Branching DQN use the 45-action factored space (5 exposures x 3 order x 3 urgency).
 * Exposure = action_index / 9.
 */
__device__ __forceinline__ float factored_action_to_exposure(int action_idx) {
    int exposure_level = action_idx / 9;  /* 0-4 */
    switch (exposure_level) {
        case 0: return -1.0f;   /* Short100 */
        case 1: return -0.5f;   /* Short50  */
        case 2: return  0.0f;   /* Flat     */
        case 3: return  0.5f;   /* Long50   */
        case 4: return  1.0f;   /* Long100  */
        default: return 0.0f;
    }
}

/**
 * PPO: Map factored action index (0-44) to transaction cost rate.
 *
 * OrderType encoded as (action_index / 3) % 3:
 *   0 = Market (0.0001), 1 = LimitMaker (0.00005), 2 = IoC (0.00015)
 */
__device__ __forceinline__ float ppo_action_to_tx_cost(int action_idx) {
    int order_type = (action_idx / 3) % 3;
    switch (order_type) {
        case 0: return 0.0001f;    /* Market    */
        case 1: return 0.00005f;   /* LimitMaker */
        case 2: return 0.00015f;   /* IoC       */
        default: return 0.0001f;
    }
}

/**
 * Initialize triple-barrier on position open.
 *
 * barrier_state[5]: entry_price, upper, lower, expiry_step, result
 * barrier_config[3]: profit_mult, loss_mult, max_holding_bars
 */
__device__ __forceinline__ void barrier_init(
    float* barrier_state,
    const float* __restrict__ barrier_config,
    float entry_price,
    int current_step
) {
    float profit_mult = barrier_config[0];   /* e.g. 1.02 */
    float loss_mult   = barrier_config[1];   /* e.g. 0.98 */
    int   max_hold    = (int)barrier_config[2];

    barrier_state[0] = entry_price;
    barrier_state[1] = entry_price * profit_mult;          /* upper */
    barrier_state[2] = entry_price * loss_mult;            /* lower */
    barrier_state[3] = (float)(current_step + max_hold);   /* expiry */
    barrier_state[4] = 0.0f;                               /* result: 0=pending */
}

/**
 * Check triple-barrier conditions.
 *
 * Returns label: +1 = profit, -1 = loss, 0 = pending/expired.
 * Sets barrier_state[4] = label if triggered.
 */
__device__ __forceinline__ int barrier_check(
    float* barrier_state,
    float price,
    int current_step,
    float position
) {
    /* No active barrier */
    if (barrier_state[0] <= 0.0f) return 0;

    float upper  = barrier_state[1];
    float lower  = barrier_state[2];
    int   expiry = (int)barrier_state[3];
    float sign   = (position >= 0.0f) ? 1.0f : -1.0f;

    int label = 0;
    if (sign > 0.0f) {
        /* Long: upper = profit, lower = loss */
        if (price >= upper) label = 1;
        else if (price <= lower) label = -1;
    } else {
        /* Short: lower = profit, upper = loss */
        if (price <= lower) label = 1;
        else if (price >= upper) label = -1;
    }

    /* Expiry */
    if (label == 0 && current_step >= expiry) {
        float entry = barrier_state[0];
        float pnl   = (price - entry) * sign;
        label = (pnl > 0.0f) ? 1 : -1;
    }

    if (label != 0) {
        barrier_state[4] = (float)label;
    }
    return label;
}

/** Reset barrier state to inactive. */
__device__ __forceinline__ void barrier_reset(float* barrier_state) {
    barrier_state[0] = 0.0f;
    barrier_state[1] = 0.0f;
    barrier_state[2] = 0.0f;
    barrier_state[3] = 0.0f;
    barrier_state[4] = 0.0f;
}

/**
 * Diversity penalty based on Shannon entropy over a sliding window
 * of recent actions, tracked at the exposure level (9 categories).
 *
 * With DQN_NUM_ACTIONS=9, each action IS an exposure level directly.
 * Max entropy for 9 categories = log2(9) = 3.17.
 * Caller applies diversity_scale as a multiplier.
 *
 * diversity_window[DIVERSITY_WINDOW]: ring buffer of exposure indices (0-8).
 * diversity_meta[2]: [write_pos, count].
 */
__device__ float diversity_entropy(
    int* diversity_window,
    int* diversity_meta,
    int action_idx
) {
    /* With 9-action DQN, action_idx IS the exposure level (0-8) */
    int category = action_idx;
    if (category < 0) category = 0;
    if (category >= DQN_NUM_ACTIONS) category = DQN_NUM_ACTIONS - 1;

    int pos   = diversity_meta[0];
    int count = diversity_meta[1];

    /* Write into ring buffer */
    diversity_window[pos] = category;
    diversity_meta[0] = (pos + 1) % DIVERSITY_WINDOW;
    if (count < DIVERSITY_WINDOW) {
        count++;
        diversity_meta[1] = count;
    }

    if (count < 2) return 0.0f;

    /* Count categories (5 exposure levels) */
    int counts[DQN_NUM_ACTIONS];
    for (int i = 0; i < DQN_NUM_ACTIONS; i++) counts[i] = 0;
    for (int i = 0; i < count; i++) {
        int c = diversity_window[i];
        if (c >= 0 && c < DQN_NUM_ACTIONS) counts[c]++;
    }

    /* Shannon entropy: -sum(p * log2(p)) */
    float entropy = 0.0f;
    float inv_n = 1.0f / (float)count;
    for (int c = 0; c < DQN_NUM_ACTIONS; c++) {
        if (counts[c] > 0) {
            float p = (float)counts[c] * inv_n;
            entropy -= p * log2f(p);
        }
    }

    /* Proportional diversity penalty: continuous gradient instead of cliff.
     * Max entropy for 5 categories = log2(5) ≈ 2.32.
     * Penalty scales linearly with how far entropy is below maximum.
     * Caller applies diversity_scale as multiplier. */
    float max_entropy = log2f((float)DQN_NUM_ACTIONS);  /* 3.17 for 9 actions */
    if (entropy < max_entropy) {
        return -((max_entropy - entropy) / max_entropy) * 0.1f;
    }
    return 0.0f;
}

/**
 * Curiosity forward model inference.
 *
 * Predicts next_state features from (state, action_category).
 * Returns MSE prediction error clamped to [0, max_reward].
 *
 * Architecture: [CUR_INPUT=MARKET_DIM+3] -> [CUR_HIDDEN=128] LReLU -> [CUR_OUTPUT=MARKET_DIM]
 * Input: MARKET_DIM state features (first MARKET_DIM of state[STATE_DIM]) + 3-class one-hot.
 * Action mapping: exposure/9: <=1 -> Short(0), ==2 -> Flat(1), >=3 -> Long(2).
 */
__device__ float curiosity_inference(
    const float* state,
    const float* next_state,
    int action_idx,
    const float* __restrict__ w_c1,    /* [CUR_HIDDEN, CUR_INPUT] */
    const float* __restrict__ b_c1,    /* [CUR_HIDDEN]            */
    const float* __restrict__ w_c2,    /* [CUR_OUTPUT, CUR_HIDDEN] */
    const float* __restrict__ b_c2,    /* [CUR_OUTPUT]            */
    float* scratch,                    /* [CUR_HIDDEN]            */
    float max_reward
) {
    /* Build input: MARKET_DIM state features + 3 one-hot */
    float input[CUR_INPUT];
    for (int i = 0; i < MARKET_DIM; i++) {
        input[i] = state[i];
    }

    /* Action to category one-hot (5-action DQN: action IS exposure level) */
    int category;
    if (action_idx <= 1)      category = 0;  /* Short100/Short50 */
    else if (action_idx == 2) category = 1;  /* Flat */
    else                      category = 2;  /* Long50/Long100 */

    input[MARKET_DIM + 0] = (category == 0) ? 1.0f : 0.0f;
    input[MARKET_DIM + 1] = (category == 1) ? 1.0f : 0.0f;
    input[MARKET_DIM + 2] = (category == 2) ? 1.0f : 0.0f;

    /* Hidden layer */
    matvec_leaky_relu(w_c1, b_c1, input, scratch, CUR_INPUT, CUR_HIDDEN, 1);

    /* Output layer (no activation) */
    float pred[CUR_OUTPUT];
    matvec_leaky_relu(w_c2, b_c2, scratch, pred, CUR_HIDDEN, CUR_OUTPUT, 0);

    /* MSE against actual next_state features (first MARKET_DIM) */
    float mse = 0.0f;
    for (int i = 0; i < CUR_OUTPUT; i++) {
        float diff = pred[i] - next_state[i];
        mse += diff * diff;
    }
    mse /= (float)CUR_OUTPUT;

    return fminf(mse, max_reward);
}

/* ------------------------------------------------------------------ */
/*  Differential Sharpe Ratio (Moody & Saffell 2001)                  */
/* ------------------------------------------------------------------ */

/**
 * Compute one DSR step and update EMA accumulators in-place.
 *
 * D_t = (B_prev * delta_A - 0.5 * A_prev * delta_B) / (B_prev - A_prev^2)^{3/2}
 *
 * where delta_A = r_t - A_prev, delta_B = r_t^2 - B_prev.
 * Returns clamped DSR in [-5, 5].
 */
__device__ __forceinline__ float dsr_step(
    float r_t,
    float* ema_A,     /* EMA of returns (in/out) */
    float* ema_B,     /* EMA of squared returns (in/out) */
    int*   dsr_init,  /* 0 = first call, set to 1 after init */
    float  eta        /* EMA decay rate (e.g. 0.01) */
) {
    if (!(*dsr_init)) {
        *ema_A = r_t;
        *ema_B = r_t * r_t + 1e-8f;  /* avoid zero variance */
        *dsr_init = 1;
        return 0.0f;  /* No DSR on first step */
    }

    float a_prev = *ema_A;
    float b_prev = *ema_B;

    /* Update EMAs */
    *ema_A = a_prev + eta * (r_t - a_prev);
    *ema_B = b_prev + eta * (r_t * r_t - b_prev);

    /* DSR formula */
    float delta_a = r_t - a_prev;
    float delta_b = r_t * r_t - b_prev;
    float variance = b_prev - a_prev * a_prev;

    if (variance < 1e-12f) {
        /* Fallback: scale raw PnL to approximate DSR magnitude range */
        return fmaxf(-5.0f, fminf(5.0f, r_t * 100.0f));
    }

    float denom = powf(variance, 1.5f);
    float dsr = (b_prev * delta_a - 0.5f * a_prev * delta_b) / denom;

    return fmaxf(-5.0f, fminf(5.0f, dsr));
}

/* ------------------------------------------------------------------ */
/*  N-Step Return Accumulator                                         */
/* ------------------------------------------------------------------ */

/* Maximum n-step lookahead (compile-time cap for ring buffer sizing) */
#ifndef N_STEPS_MAX
#define N_STEPS_MAX 8
#endif

/**
 * Accumulate n-step discounted return using a ring buffer.
 *
 * On each call, pushes new reward into the ring. When the ring has n entries,
 * computes G = sum_{k=0}^{n-1} gamma^k * ring[k] and returns it.
 * Before the ring is full, returns the partially accumulated sum.
 *
 * @param ring      Ring buffer of rewards [N_STEPS_MAX] (per-episode, persistent)
 * @param ring_idx  Current write index (per-episode, persistent)
 * @param ring_len  Current number of valid entries (per-episode, persistent)
 * @param reward    New reward to push
 * @param gamma     Discount factor
 * @param n_steps   Number of lookahead steps (1 = standard TD, >1 = n-step)
 * @return          Accumulated n-step return (may be partial if ring_len < n_steps)
 */
__device__ __forceinline__ float nstep_push_and_sum(
    float* ring,
    int*   ring_idx,
    int*   ring_len,
    float  reward,
    float  gamma,
    int    n_steps
) {
    /* Push into ring buffer */
    int idx = *ring_idx;
    ring[idx] = reward;
    *ring_idx = (idx + 1) % n_steps;
    if (*ring_len < n_steps) (*ring_len)++;

    /* Compute discounted sum over ring contents */
    float G = 0.0f;
    float discount = 1.0f;
    int len = *ring_len;
    int start = (*ring_idx - len + n_steps) % n_steps;
    for (int k = 0; k < len; k++) {
        int i = (start + k) % n_steps;
        G += discount * ring[i];
        discount *= gamma;
    }
    return G;
}

/**
 * Reset n-step ring buffer (call on episode boundary).
 */
__device__ __forceinline__ void nstep_reset(
    float* ring,
    int*   ring_idx,
    int*   ring_len,
    int    n_steps
) {
    *ring_idx = 0;
    *ring_len = 0;
    for (int i = 0; i < n_steps && i < N_STEPS_MAX; i++)
        ring[i] = 0.0f;
}

/* ------------------------------------------------------------------ */
/*  Order Routing (GPU-side OrderRouter::route)                       */
/* ------------------------------------------------------------------ */

/**
 * Deterministic order routing based on market microstructure.
 *
 * Maps (exposure_idx, spread, volatility) → (order_type, urgency).
 * Matches CPU `OrderRouter::route()` in ml-core/src/order_router.rs exactly.
 *
 * Order types: 0=LimitMaker, 1=IoC, 2=Market
 * Urgency:     0=Patient, 1=Normal, 2=Aggressive
 */
__device__ __forceinline__ void route_order(
    float spread,
    float median_spread,
    float volatility,
    float median_vol,
    int* out_order_type,
    int* out_urgency
) {
    /* Order type selection based on spread vs median */
    int order;
    if (median_spread <= 0.0f) {
        order = 2;  /* Market — no spread reference */
    } else if (spread < median_spread) {
        order = 0;  /* LimitMaker — tight spread, save costs */
    } else if (spread > 2.0f * median_spread) {
        order = 2;  /* Market — wide spread, ensure fill */
    } else {
        order = 1;  /* IoC — moderate spread */
    }

    /* Urgency selection based on volatility vs median */
    int urgency;
    if (median_vol <= 0.0f) {
        urgency = 1;  /* Normal — no vol reference */
    } else if (volatility > 1.5f * median_vol) {
        urgency = 2;  /* Aggressive — high vol */
    } else if (volatility < 0.5f * median_vol) {
        urgency = 0;  /* Patient — low vol */
    } else {
        urgency = 1;  /* Normal */
    }

    *out_order_type = order;
    *out_urgency = urgency;
}

/**
 * Order-type-dependent transaction cost rate.
 *
 * Matches CPU FillSimulator::spread_cost() cost structure:
 *   Market    = spread_cost_frac * spread_bps / 10000
 *   IoC       = spread_cost_frac * 0.5 * spread_bps / 10000
 *   LimitMaker = -(spread_capture_frac * spread_bps / 10000) [rebate]
 *
 * Returns cost as a fraction of trade value (positive = cost, negative = rebate).
 */
__device__ __forceinline__ float order_type_tx_cost(
    int order_type,
    float spread_bps,
    float spread_cost_frac,
    float spread_capture_frac
) {
    switch (order_type) {
        case 0:  /* LimitMaker — maker rebate */
            return -(spread_bps * spread_capture_frac / 10000.0f);
        case 1:  /* IoC — half spread cost */
            return spread_bps * spread_cost_frac * 0.5f / 10000.0f;
        case 2:  /* Market — full spread cost */
            return spread_bps * spread_cost_frac / 10000.0f;
        default:
            return spread_bps * spread_cost_frac / 10000.0f;
    }
}

/* ------------------------------------------------------------------ */
/*  Fill Simulation (GPU-side FillSimulator::simulate_fill)            */
/* ------------------------------------------------------------------ */

/**
 * Splitmix64-style deterministic hash: (step, action_idx) → [0.0, 1.0).
 *
 * Bit-exact match with CPU splitmix64_unit() in fill_simulator.rs.
 * Uses the same Weyl-sequence seed combining and two-round finalizer.
 */
__device__ __forceinline__ float splitmix64_fill(int step, int action_idx) {
    /* Combine step and action_idx with Weyl sequence constants */
    unsigned long long z = (unsigned long long)step * 0x9e3779b97f4a7c15ULL
        + (unsigned long long)action_idx * 0x6a09e667f3bcc908ULL;

    /* Splitmix64 two-round finalizer */
    z = (z ^ (z >> 30)) * 0xbf58476d1ce4e5b9ULL;
    z = (z ^ (z >> 27)) * 0x94d049bb133111ebULL;
    z ^= (z >> 31);

    /* Convert to [0.0, 1.0) — upper 53 bits as double mantissa, then narrow to float */
    double d = (double)(z >> 11) / (double)(1ULL << 53);
    return (float)d;
}

/* ------------------------------------------------------------------ */
/*  Kernel Weight Pack — reduces ~48 weight pointer params to 1 struct */
/* ------------------------------------------------------------------ */

/**
 * Packed weight pointer structs for the DQN experience collection kernel.
 *
 * Passed by value in CUDA constant memory (384 bytes total, well within
 * the 4 KB per-kernel limit). Each field is a device pointer to a
 * pre-uploaded weight buffer on GPU.
 *
 * Layout must match the Rust #[repr(C)] structs in gpu_weights.rs EXACTLY.
 */

struct DuelingWeightPtrs {
    const float* w_s1; const float* b_s1;   /* shared layer 0 */
    const float* w_s2; const float* b_s2;   /* shared layer 1 */
    const float* w_v1; const float* b_v1;   /* value FC       */
    const float* w_v2; const float* b_v2;   /* value out      */
    const float* w_a1; const float* b_a1;   /* advantage FC   */
    const float* w_a2; const float* b_a2;   /* advantage out  */
};

struct CuriosityWeightPtrs {
    const float* w1; const float* b1;       /* FC layer 1     */
    const float* w2; const float* b2;       /* FC layer 2     */
};

struct RMSNormWeightPtrs {
    const float* s0;                        /* shared_rmsnorm_0 gamma */
    const float* s1;                        /* shared_rmsnorm_1 gamma */
    const float* v;                         /* value_rmsnorm gamma    */
    const float* a;                         /* advantage_rmsnorm gamma */
};

struct BranchWeightPtrs {
    const float* w_o1; const float* b_o1;   /* order FC       */
    const float* w_o2; const float* b_o2;   /* order out      */
    const float* w_u1; const float* b_u1;   /* urgency FC     */
    const float* w_u2; const float* b_u2;   /* urgency out    */
};

struct KernelWeightPack {
    struct DuelingWeightPtrs   online;      /* 12 ptrs = 96 B  */
    struct DuelingWeightPtrs   target;      /* 12 ptrs = 96 B  */
    struct CuriosityWeightPtrs curiosity;   /*  4 ptrs = 32 B  */
    struct RMSNormWeightPtrs   rmsnorm;     /*  4 ptrs = 32 B  */
    struct BranchWeightPtrs    online_br;   /*  8 ptrs = 64 B  */
    struct BranchWeightPtrs    target_br;   /*  8 ptrs = 64 B  */
};                                          /* total  = 384 B  */

/**
 * Unpack KernelWeightPack into local __restrict__ aliases matching the
 * variable names used by q_forward_* helper functions.
 *
 * Call at the top of each kernel body after receiving `struct KernelWeightPack wp`.
 * The __restrict__ qualifier enables the compiler to assume no aliasing.
 */
#define UNPACK_WEIGHT_PTRS(wp) \
    /* Online dueling (12) */ \
    const float* __restrict__ on_w_s1 = wp.online.w_s1; \
    const float* __restrict__ on_b_s1 = wp.online.b_s1; \
    const float* __restrict__ on_w_s2 = wp.online.w_s2; \
    const float* __restrict__ on_b_s2 = wp.online.b_s2; \
    const float* __restrict__ on_w_v1 = wp.online.w_v1; \
    const float* __restrict__ on_b_v1 = wp.online.b_v1; \
    const float* __restrict__ on_w_v2 = wp.online.w_v2; \
    const float* __restrict__ on_b_v2 = wp.online.b_v2; \
    const float* __restrict__ on_w_a1 = wp.online.w_a1; \
    const float* __restrict__ on_b_a1 = wp.online.b_a1; \
    const float* __restrict__ on_w_a2 = wp.online.w_a2; \
    const float* __restrict__ on_b_a2 = wp.online.b_a2; \
    /* Target dueling (12) */ \
    const float* __restrict__ tg_w_s1 = wp.target.w_s1; \
    const float* __restrict__ tg_b_s1 = wp.target.b_s1; \
    const float* __restrict__ tg_w_s2 = wp.target.w_s2; \
    const float* __restrict__ tg_b_s2 = wp.target.b_s2; \
    const float* __restrict__ tg_w_v1 = wp.target.w_v1; \
    const float* __restrict__ tg_b_v1 = wp.target.b_v1; \
    const float* __restrict__ tg_w_v2 = wp.target.w_v2; \
    const float* __restrict__ tg_b_v2 = wp.target.b_v2; \
    const float* __restrict__ tg_w_a1 = wp.target.w_a1; \
    const float* __restrict__ tg_b_a1 = wp.target.b_a1; \
    const float* __restrict__ tg_w_a2 = wp.target.w_a2; \
    const float* __restrict__ tg_b_a2 = wp.target.b_a2; \
    /* Curiosity (4) */ \
    const float* __restrict__ cur_w1 = wp.curiosity.w1; \
    const float* __restrict__ cur_b1 = wp.curiosity.b1; \
    const float* __restrict__ cur_w2 = wp.curiosity.w2; \
    const float* __restrict__ cur_b2 = wp.curiosity.b2; \
    /* RMSNorm (4) */ \
    const float* __restrict__ rms_s0_gamma = wp.rmsnorm.s0; \
    const float* __restrict__ rms_s1_gamma = wp.rmsnorm.s1; \
    const float* __restrict__ rms_v_gamma  = wp.rmsnorm.v;  \
    const float* __restrict__ rms_a_gamma  = wp.rmsnorm.a;  \
    /* Online branching (8) */ \
    const float* __restrict__ on_w_bo1 = wp.online_br.w_o1; \
    const float* __restrict__ on_b_bo1 = wp.online_br.b_o1; \
    const float* __restrict__ on_w_bo2 = wp.online_br.w_o2; \
    const float* __restrict__ on_b_bo2 = wp.online_br.b_o2; \
    const float* __restrict__ on_w_bu1 = wp.online_br.w_u1; \
    const float* __restrict__ on_b_bu1 = wp.online_br.b_u1; \
    const float* __restrict__ on_w_bu2 = wp.online_br.w_u2; \
    const float* __restrict__ on_b_bu2 = wp.online_br.b_u2; \
    /* Target branching (8) */ \
    const float* __restrict__ tg_w_bo1 = wp.target_br.w_o1; \
    const float* __restrict__ tg_b_bo1 = wp.target_br.b_o1; \
    const float* __restrict__ tg_w_bo2 = wp.target_br.w_o2; \
    const float* __restrict__ tg_b_bo2 = wp.target_br.b_o2; \
    const float* __restrict__ tg_w_bu1 = wp.target_br.w_u1; \
    const float* __restrict__ tg_b_bu1 = wp.target_br.b_u1; \
    const float* __restrict__ tg_w_bu2 = wp.target_br.w_u2; \
    const float* __restrict__ tg_b_bu2 = wp.target_br.b_u2;

/**
 * Compute fill probability for a given order type and market conditions.
 *
 * Matches CPU FillSimulator::fill_probability() exactly.
 *
 * Parameters:
 *   order_type: 0=LimitMaker, 1=IoC, 2=Market
 *   urgency:    0=Patient, 1=Normal, 2=Aggressive
 *   normalized_vol: volatility / median_volatility (1.0 = average)
 *   ioc_fill_prob: base IoC fill probability (default 0.85)
 *   limit_fill_min: minimum limit fill probability (default 0.30)
 *   limit_fill_max: maximum limit fill probability (default 0.80)
 */
__device__ __forceinline__ float fill_probability(
    int order_type,
    int urgency,
    float normalized_vol,
    float ioc_fill_prob,
    float limit_fill_min,
    float limit_fill_max
) {
    float base_prob;
    switch (order_type) {
        case 2:  /* Market — always fills */
            base_prob = 1.0f;
            break;
        case 1:  /* IoC */
            base_prob = ioc_fill_prob;
            break;
        case 0:  /* LimitMaker — vol-dependent */
        default: {
            float vol_clamped = fminf(fmaxf(normalized_vol, 0.0f), 3.0f);
            float vol_factor = vol_clamped / 3.0f;
            base_prob = limit_fill_min + vol_factor * (limit_fill_max - limit_fill_min);
            break;
        }
    }

    /* Urgency modifier */
    float urgency_mod;
    switch (urgency) {
        case 0:  urgency_mod = 0.85f;  break;  /* Patient */
        case 2:  urgency_mod = 1.10f;  break;  /* Aggressive */
        default: urgency_mod = 1.0f;   break;  /* Normal */
    }

    float prob = base_prob * urgency_mod;
    return fminf(fmaxf(prob, 0.0f), 1.0f);
}

/**
 * Simulate order fill using deterministic hash.
 *
 * Matches CPU FillSimulator::simulate_fill() exactly.
 *
 * Returns: 1 if filled, 0 if not filled.
 * Writes cost adjustment to *out_cost_adj (positive=cost, negative=rebate, 0=no fill).
 */
__device__ __forceinline__ int simulate_fill_check(
    int order_type,
    int urgency,
    float normalized_vol,
    float spread_bps,
    int step,
    int action_idx,
    float ioc_fill_prob,
    float limit_fill_min,
    float limit_fill_max,
    float spread_cost_frac,
    float spread_capture_frac,
    float* out_cost_adj
) {
    float prob = fill_probability(order_type, urgency, normalized_vol,
                                  ioc_fill_prob, limit_fill_min, limit_fill_max);
    float pseudo_random = splitmix64_fill(step, action_idx);
    int filled = (pseudo_random < prob) ? 1 : 0;

    if (filled) {
        *out_cost_adj = order_type_tx_cost(order_type, spread_bps,
                                           spread_cost_frac, spread_capture_frac);
    } else {
        *out_cost_adj = 0.0f;
    }
    return filled;
}

/* ── Warp-cooperative forward pass utilities ────────────────────────────── */

#ifndef SHMEM_MIN
#define SHMEM_MIN(a, b) (((a) < (b)) ? (a) : (b))
#endif

#ifndef DIST_SIZE
#define DIST_SIZE(dim) (((dim) + 31) / 32)
#endif

/* Warp-cooperative tiled layer: cooperative load + warp_matvec */
#ifndef TILE_LAYER_WARP_CLEAN
#define TILE_LAYER_WARP_CLEAN(W_global, b_global, input_dist, output_dist, in_d, out_d, act, shmem_w, shmem_b, lane) \
    do { \
        for (int _tile = 0; _tile < ((out_d) + SHMEM_TILE_ROWS - 1) / SHMEM_TILE_ROWS; _tile++) { \
            int _ts = _tile * SHMEM_TILE_ROWS; \
            int _tr = SHMEM_MIN(SHMEM_TILE_ROWS, (out_d) - _ts); \
            cooperative_load_tile(shmem_w, (W_global) + _ts * (in_d), _tr * (in_d)); \
            cooperative_load_tile(shmem_b, (b_global) + _ts, _tr); \
            __syncwarp(0xFFFFFFFF); \
            warp_matvec_leaky_relu_shmem(shmem_w, shmem_b, input_dist, output_dist, \
                                          in_d, out_d, _ts, _tr, act, lane); \
            __syncwarp(0xFFFFFFFF); \
        } \
    } while (0)
#endif

/* Guard: only compile when DQN architecture constants are defined */
#if defined(SHARED_H1) && defined(SHARED_H2) && defined(VALUE_H) && defined(ADV_H)

/**
 * Warp-cooperative dueling Q-network forward pass (clean, no noise).
 *
 * Runs the full DQN forward pass: shared layers → value head → advantage head
 * → Q(s,a) = V(s) + A(s,a) - mean(A).  All 32 lanes cooperate on matrix-vector
 * products; all lanes get all NUM_ACTIONS Q-values for argmax.
 */
__device__ void q_forward_dueling_warp_shmem(
    const float* state_dist,
    const float* __restrict__ w_s1, const float* __restrict__ b_s1,
    const float* __restrict__ w_s2, const float* __restrict__ b_s2,
    const float* __restrict__ w_v1, const float* __restrict__ b_v1,
    const float* __restrict__ w_v2, const float* __restrict__ b_v2,
    const float* __restrict__ w_a1, const float* __restrict__ b_a1,
    const float* __restrict__ w_a2, const float* __restrict__ b_a2,
    float* scratch1_dist,
    float* scratch2_dist,
    float* q_values,
    int lane_id,
    float* shmem_weights,
    float* shmem_bias
) {
    /* Shared layers: state -> scratch1 -> scratch2 */
    TILE_LAYER_WARP_CLEAN(w_s1, b_s1, state_dist, scratch1_dist,
                          STATE_DIM, SHARED_H1, 1, shmem_weights, shmem_bias, lane_id);
    TILE_LAYER_WARP_CLEAN(w_s2, b_s2, scratch1_dist, scratch2_dist,
                          SHARED_H1, SHARED_H2, 1, shmem_weights, shmem_bias, lane_id);

    /* Value head hidden layer: scratch2 -> scratch_v (reuse scratch1_dist) */
    float* scratch_v_dist = scratch1_dist;
    TILE_LAYER_WARP_CLEAN(w_v1, b_v1, scratch2_dist, scratch_v_dist,
                          SHARED_H2, VALUE_H, 1, shmem_weights, shmem_bias, lane_id);

    /* Value output: [1, VALUE_H] — single neuron */
    float value;
    {
        cooperative_load_tile(shmem_weights, w_v2, VALUE_H);
        cooperative_load_tile(shmem_bias, b_v2, 1);
        __syncthreads();
        float partial = 0.0f;
        for (int i = lane_id; i < VALUE_H; i += 32)
            partial += shmem_weights[i] * scratch_v_dist[i / 32];
        if (lane_id == 0)
            partial += shmem_bias[0];
        value = warp_reduce_sum_all(partial);
        __syncthreads();
    }

    /* Advantage head hidden layer: scratch2 -> scratch_a */
    float* scratch_a_dist = scratch1_dist + DIST_SIZE(VALUE_H);
    TILE_LAYER_WARP_CLEAN(w_a1, b_a1, scratch2_dist, scratch_a_dist,
                          SHARED_H2, ADV_H, 1, shmem_weights, shmem_bias, lane_id);

    /* Advantage output: [NUM_ACTIONS, ADV_H] — all lanes get full copy */
    float adv[NUM_ACTIONS];
    {
        cooperative_load_tile(shmem_weights, w_a2, NUM_ACTIONS * ADV_H);
        cooperative_load_tile(shmem_bias, b_a2, NUM_ACTIONS);
        __syncthreads();
        for (int a = 0; a < NUM_ACTIONS; a++) {
            const float* row = shmem_weights + a * ADV_H;
            float partial = 0.0f;
            for (int i = lane_id; i < ADV_H; i += 32)
                partial += row[i] * scratch_a_dist[i / 32];
            if (lane_id == 0)
                partial += shmem_bias[a];
            adv[a] = warp_reduce_sum_all(partial);
        }
        __syncthreads();
    }

    /* Q(s,a) = V(s) + A(s,a) - mean(A) */
    float adv_mean = 0.0f;
    for (int i = 0; i < NUM_ACTIONS; i++) adv_mean += adv[i];
    adv_mean /= (float)NUM_ACTIONS;
    for (int i = 0; i < NUM_ACTIONS; i++) {
        q_values[i] = value + adv[i] - adv_mean;
    }
}

#endif /* SHARED_H1 && SHARED_H2 && VALUE_H && ADV_H */
