/**
 * experience_kernels.cu
 *
 * Three focused CUDA kernels that handle the non-Q-forward parts of DQN
 * experience collection.  The Q-network forward pass is performed by cuBLAS
 * SGEMM from the Rust host between kernel launches.
 *
 * Kernels:
 *   1. experience_state_gather  — assemble batch state tensor for cuBLAS
 *   2. experience_action_select — epsilon-greedy on cuBLAS Q-values
 *   3. experience_env_step      — portfolio simulation, reward, done detection
 *
 * Design:
 *   - Grid: ceil(N/256), Block: 256.  One thread per episode — trivially parallel.
 *   - No shared memory required.
 *   - Standalone file: does NOT include common_device_functions.cuh.
 *     Constants normally injected via NVRTC are guarded with #ifndef
 *     fallback defaults; production use injects the correct values at
 *     NVRTC compile time.
 *
 * State layout (matches common_device_functions.cuh):
 *   [0 .. market_dim)                      : market features
 *   [market_dim .. market_dim+3)           : portfolio features (value_norm, position, cash_norm)
 *   [market_dim+3 .. state_dim)            : zero-pad for tensor-core alignment
 *
 * Portfolio state layout used by experience kernels ([N, PORTFOLIO_STRIDE=20]):
 *   [0]  position        — current contract position (signed)
 *   [1]  cash            — cash balance
 *   [2]  portfolio_value — mark-to-market total value (cash + position * price)
 *   [3]  (reserved)      — was dsr_A, unused by reward v5
 *   [4]  (reserved)      — was dsr_B, unused by reward v5
 *   [5]  (reserved)      — was pnl_ema, unused by reward v5
 *   [6]  (reserved)      — was pnl_var, unused by reward v5
 *   [7]  peak_equity     — high-water mark (init to initial_capital)
 *   [8]  flat_counter    — consecutive flat steps (float for GPU simplicity)
 *   [9]  prev_equity     — equity at previous step (init to initial_capital)
 *   [10] hold_time       — consecutive steps with position (total, not just losing)
 *   [11] realized_pnl    — cumulative realized PnL
 *
 * Branching DQN (Tavakoli et al., 2018):
 *   3 independent advantage heads: exposure (b0_size=5), order (b1_size=3),
 *   urgency (b2_size=3).  cuBLAS produces q_values [N, 11] = concat([b0],[b1],[b2]).
 *   Factored action index = a0 * b1_size * b2_size + a1 * b2_size + a2.
 *   Exposure mapping: 0→-1.0, 1→-0.5, 2→0.0, 3→0.5, 4→1.0.
 *
 * Advanced features deferred to follow-up tasks (Phase 4+):
 *   curiosity, barrier tracking, diversity entropy, fill simulation,
 *   NoisyNet, C51, RMSNorm, N-step returns, action masking.
 */

/* ------------------------------------------------------------------ */
/* Portfolio stride for experience kernels (12 floats per episode).    */
/* portfolio_sim_kernel uses its own stride of 8 — do NOT change it.  */
/* ------------------------------------------------------------------ */
#define PORTFOLIO_STRIDE 20

/* ------------------------------------------------------------------ */
/* Compile-time constants — overridable via NVRTC #define injection.  */
/* ------------------------------------------------------------------ */

#ifndef MARKET_DIM
#define MARKET_DIM 42
#endif

#ifndef PORTFOLIO_DIM
#define PORTFOLIO_DIM 3
#endif

#ifndef STATE_DIM
#define STATE_DIM 48
#endif

/* ------------------------------------------------------------------ */
/* Inline helpers — self-contained, no external dependencies.         */
/* ------------------------------------------------------------------ */

/**
 * LCG random — returns float in [0, 1).
 *
 * Identical implementation to gpu_random() in common_device_functions.cuh
 * so episodes share the same statistical properties when mixing old and
 * new kernels during the transition period.
 */
__device__ __forceinline__ float lcg_random(unsigned int* state) {
    *state = *state * 1664525u + 1013904223u;
    return (float)(*state & 0x00FFFFFFu) / 16777216.0f;
}

/**
 * Argmax over a float array of length n.
 * Returns the index of the maximum element; ties broken in favour of the
 * lowest index (first maximum found).
 */
__device__ __forceinline__ int argmax_n(const float* arr, int n) {
    int best_idx = 0;
    float best_val = arr[0];
    for (int j = 1; j < n; j++) {
        if (arr[j] > best_val) {
            best_val = arr[j];
            best_idx = j;
        }
    }
    return best_idx;
}

/**
 * Map exposure branch index (0–8) to target exposure fraction.
 *   0 → -1.00  (S100)
 *   1 → -0.75  (S75)
 *   2 → -0.50  (S50)
 *   3 → -0.25  (S25)
 *   4 →  0.00  (Flat)
 *   5 → +0.25  (L25)
 *   6 → +0.50  (L50)
 *   7 → +0.75  (L75)
 *   8 → +1.00  (L100)
 *
 * Formula: -1.0 + idx * (2.0 / (b0_size - 1)).
 * With b0_size=9: step = 0.25. Flat = index 4.
 */
__device__ __forceinline__ float exposure_idx_to_fraction(int idx) {
    return -1.0f + (float)idx * 0.25f;
}

/* ================================================================== */
/* Kernel 1: experience_state_gather                                  */
/* ================================================================== */

/**
 * Gather market features and portfolio state into a flat batch tensor
 * suitable for cuBLAS SGEMM (Q-network forward pass).
 *
 * Grid: ceil(N / 256), Block: 256.  One thread per episode.
 *
 * @param market_features   [total_bars, market_dim]  raw market feature matrix
 * @param episode_starts    [N]                        global bar index of episode start
 * @param current_timesteps [N]                        timestep offset within episode (read-only)
 * @param portfolio_states  [N, 3]                     {position, cash, portfolio_value}
 * @param batch_states      [N, state_dim]             output: assembled state batch
 * @param N                                            number of episodes
 * @param total_bars                                   number of bars in market_features
 * @param state_dim                                    full state dimension (tensor-core aligned)
 * @param market_dim                                   number of market features per bar
 */
extern "C" __global__ void experience_state_gather(
    const float* __restrict__ market_features,
    const int*   __restrict__ episode_starts,
    const int*   __restrict__ current_timesteps,
    const float* __restrict__ portfolio_states,
    float*                    batch_states,
    int N,
    int total_bars,
    int state_dim,
    int market_dim
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= N) return;

    int bar_idx = episode_starts[i] + current_timesteps[i];

    float* out = batch_states + (long long)i * state_dim;

    /* Out-of-data: write zero state; cuBLAS will produce Q-values of ~0. */
    if (bar_idx >= total_bars) {
        for (int k = 0; k < state_dim; k++)
            out[k] = 0.0f;
        return;
    }

    /* -- Market features: [0 .. market_dim) -- */
    const float* mf_row = market_features + (long long)bar_idx * market_dim;
    for (int k = 0; k < market_dim; k++)
        out[k] = mf_row[k];

    /* -- Portfolio features: [market_dim .. market_dim+8) --
     *
     * 8 informative features that let the Q-network SEE its trading context.
     * Without these, the model is blind to risk, P&L, and position state.
     *
     * slot +0: position / max_position        — normalized position [-1, +1]
     * slot +1: unrealized_pnl / equity         — how is THIS trade doing?
     * slot +2: drawdown                        — (peak - equity) / peak [0, 1]
     * slot +3: hold_time / 100.0               — normalized holding duration
     * slot +4: realized_pnl / equity           — overall session P&L
     * slot +5: distance_to_floor               — (equity - floor) / equity [0, 1]
     * slot +6: trade_return                    — trade P&L / equity
     * slot +7: cash_ratio                      — cash / equity
     *
     * max_position is read from kernel arg (line ~370). We use a constant
     * approximation here (max_pos not available in gather kernel).
     * The network learns the actual scale via the position feature.
     */
    const float* ps          = portfolio_states + (long long)i * PORTFOLIO_STRIDE;
    float position           = ps[0];
    float cash               = ps[1];
    float portfolio_value    = ps[2];
    float peak_equity        = ps[7];
    float prev_equity        = ps[9];
    float hold_time          = ps[10];
    float realized_pnl       = ps[11];
    float entry_price        = ps[12];
    float trade_start_pnl    = ps[13];

    float equity = (portfolio_value > 1.0f) ? portfolio_value : 1.0f;
    float drawdown = (peak_equity > 1.0f)
                   ? (peak_equity - portfolio_value) / peak_equity
                   : 0.0f;
    drawdown = (drawdown > 0.0f) ? drawdown : 0.0f;

    /* Unrealized P&L: current position mark-to-market minus entry cost */
    float unrealized_pnl = (entry_price > 0.0f && position != 0.0f)
                         ? position * (market_features[(long long)bar_idx * market_dim] - entry_price)
                         : 0.0f;

    /* Trade return since entry */
    float trade_return = (trade_start_pnl != 0.0f || realized_pnl != 0.0f)
                       ? (realized_pnl - trade_start_pnl + unrealized_pnl) / equity
                       : 0.0f;

    /* Capital floor distance (75% of peak = capital floor) */
    float floor = peak_equity * 0.75f;
    float floor_dist = (equity > floor && equity > 1.0f)
                     ? (equity - floor) / equity
                     : 0.0f;

    int portfolio_base = market_dim;
    if (portfolio_base + 7 < state_dim) {
        out[portfolio_base + 0] = position;                        /* raw position */
        out[portfolio_base + 1] = unrealized_pnl / equity;         /* trade P&L signal */
        out[portfolio_base + 2] = drawdown;                        /* risk: how deep are we? */
        out[portfolio_base + 3] = hold_time / 100.0f;              /* how long in trade? */
        out[portfolio_base + 4] = realized_pnl / equity;           /* session P&L */
        out[portfolio_base + 5] = floor_dist;                      /* distance to game over */
        out[portfolio_base + 6] = trade_return;                    /* this trade's return */
        out[portfolio_base + 7] = cash / equity;                   /* available capital */
    }

    /* -- Multi-timeframe features: [market_dim+8 .. market_dim+8+16) --
     *
     * 4 lookback windows × 4 features = 16 GPU-native features.
     * Computed directly from the market data already on GPU — no CPU.
     * The model sees price action at multiple timescales simultaneously.
     *
     * Lookback windows: 5, 15, 60, 240 bars (≈ 5min, 15min, 1hr, 4hr)
     * Features per window:
     *   +0: return = (close_now - close_N) / close_N
     *   +1: volatility = max_high - min_low over N bars / close_now
     *   +2: volume_trend = volume_now / avg_volume_N (>1 = increasing)
     *   +3: momentum = close_now vs midpoint of N-bar range (0=bottom, 1=top)
     */
    const int lookbacks[4] = {5, 15, 60, 240};
    int mtf_base = market_dim + 8;
    for (int lb = 0; lb < 4; lb++) {
        int N = lookbacks[lb];
        int past_idx = bar_idx - N;
        int slot = mtf_base + lb * 4;

        if (past_idx >= 0 && slot + 3 < state_dim) {
            /* Close prices: feature[0] of each bar is typically close or a return.
             * Use raw market_features — index 0 is the first feature per bar. */
            const float* now_row = market_features + (long long)bar_idx * market_dim;
            const float* past_row = market_features + (long long)past_idx * market_dim;
            float close_now = now_row[0];
            float close_past = past_row[0];

            /* Return over N bars */
            float ret = (close_past > 0.0f) ? (close_now - close_past) / close_past : 0.0f;
            float scaled_ret = ret * 100.0f;
            out[slot + 0] = fmaxf(-10.0f, fminf(10.0f, scaled_ret)); /* clamp ±10% */

            /* Volatility: scan high/low over window (approx from close changes) */
            float max_val = close_now;
            float min_val = close_now;
            float vol_sum = 0.0f;
            int vol_count = 0;
            for (int j = past_idx; j <= bar_idx && j < total_bars; j++) {
                float v = market_features[(long long)j * market_dim];
                if (v > max_val) max_val = v;
                if (v < min_val) min_val = v;
                /* Volume proxy: use feature index 4 if it exists (commonly volume) */
                if (market_dim > 4) {
                    vol_sum += market_features[(long long)j * market_dim + 4];
                    vol_count++;
                }
            }
            float range = (close_now > 0.0f) ? (max_val - min_val) / close_now : 0.0f;
            float scaled_range = range * 100.0f;
            out[slot + 1] = fmaxf(0.0f, fminf(10.0f, scaled_range)); /* clamp 0-10% */

            /* Volume trend: current vs average */
            float avg_vol = (vol_count > 0) ? vol_sum / (float)vol_count : 1.0f;
            float cur_vol = (market_dim > 4) ? now_row[4] : 1.0f;
            float vol_ratio = (avg_vol > 0.0f) ? cur_vol / avg_vol : 1.0f;
            out[slot + 2] = fmaxf(0.0f, fminf(5.0f, vol_ratio)); /* clamp 0-5x */

            /* Momentum: position within range [0=bottom, 1=top] */
            float range_size = max_val - min_val;
            out[slot + 3] = (range_size > 0.0f)
                          ? (close_now - min_val) / range_size
                          : 0.5f;
        } else {
            /* Not enough history — zero pad */
            if (slot + 3 < state_dim) {
                out[slot + 0] = 0.0f;
                out[slot + 1] = 0.0f;
                out[slot + 2] = 0.0f;
                out[slot + 3] = 0.0f;
            }
        }
    }

    /* -- OFI features (if present) are AFTER multi-timeframe -- */
    /* The OFI features are uploaded to a separate buffer and appended
     * by the experience collector AFTER this kernel runs. The zero-pad
     * below covers the OFI + alignment slots. */

    /* -- Zero-pad remaining dimensions (tensor-core alignment) -- */
    int filled = market_dim + 8 + 16; /* 42 market + 8 portfolio + 16 multi-tf */
    for (int k = filled; k < state_dim; k++)
        out[k] = 0.0f;
}

/* ================================================================== */
/* Kernel 2: experience_action_select                                 */
/* ================================================================== */

/**
 * Epsilon-greedy action selection on Q-values produced by cuBLAS.
 *
 * Branching Dueling DQN — q_values layout: [N, b0_size + b1_size + b2_size]
 *   Branch 0 (exposure): q_values[i*(b0+b1+b2) + 0              .. +b0)
 *   Branch 1 (order):    q_values[i*(b0+b1+b2) + b0             .. +b0+b1)
 *   Branch 2 (urgency):  q_values[i*(b0+b1+b2) + b0+b1          .. +b0+b1+b2)
 *   Factored action = a0 * b1_size * b2_size + a1 * b2_size + a2.
 *
 * Q-gap conviction filter: when greedy, if Q(best_exposure) - Q(flat) <
 * q_gap_threshold, force exposure to flat (index 2). This reduces trade
 * frequency by requiring high conviction for position entry.
 *
 * Grid: ceil(N / 256), Block: 256.  One thread per episode.
 *
 * @param q_values          [N, q_stride]  Q-values from cuBLAS
 * @param out_actions       [N]            selected factored action index (output)
 * @param rng_states        [N]            per-episode LCG RNG counter (updated in place)
 * @param epsilon                          exploration probability in [0, 1]
 * @param N                                number of episodes
 * @param b0_size                          exposure branch size (default 5)
 * @param b1_size                          order branch size (default 3)
 * @param b2_size                          urgency branch size (default 3)
 * @param q_gap_threshold                  min Q-gap for trade entry (0.0 = disabled)
 */
extern "C" __global__ void experience_action_select(
    const float*  __restrict__ q_values,
    int*                       out_actions,
    unsigned int*              rng_states,
    float*                     out_q_gaps,    /* [N] output: Q-gap per episode for conviction sizing */
    float epsilon,
    int N,
    int b0_size,
    int b1_size,
    int b2_size,
    float q_gap_threshold
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= N) return;

    unsigned int rng = rng_states[i];

    int q_stride      = b0_size + b1_size + b2_size;
    const float* q_b0 = q_values + (long long)i * q_stride;
    const float* q_b1 = q_b0 + b0_size;
    const float* q_b2 = q_b1 + b1_size;

    int a0, a1, a2;

    /* Branch 0: exposure — with Q-gap conviction filter.
     * When greedy (not exploring), only trade if Q(best) - Q(flat)
     * exceeds q_gap_threshold. Otherwise default to flat (index 2).
     * This teaches the model to trade only with high conviction. */
    if (lcg_random(&rng) < epsilon) {
        int r = (int)(lcg_random(&rng) * (float)b0_size);
        a0 = (r >= b0_size) ? b0_size - 1 : r;
    } else {
        a0 = argmax_n(q_b0, b0_size);
        /* Q-gap filter: force flat when conviction is low */
        if (q_gap_threshold > 0.0f && b0_size > 2) {
            int flat_idx = b0_size / 2; /* center = Flat (index 4 for 9-action) */
            float q_best = q_b0[a0];
            float q_flat = q_b0[flat_idx];
            if (q_best - q_flat < q_gap_threshold) {
                a0 = flat_idx;
            }
        }
    }

    /* Branch 1: order type */
    if (lcg_random(&rng) < epsilon) {
        int r = (int)(lcg_random(&rng) * (float)b1_size);
        a1 = (r >= b1_size) ? b1_size - 1 : r;
    } else {
        a1 = argmax_n(q_b1, b1_size);
    }

    /* Branch 2: urgency */
    if (lcg_random(&rng) < epsilon) {
        int r = (int)(lcg_random(&rng) * (float)b2_size);
        a2 = (r >= b2_size) ? b2_size - 1 : r;
    } else {
        a2 = argmax_n(q_b2, b2_size);
    }

    /* Compose factored action (Tavakoli et al., 2018, §3.1) */
    int action_idx = a0 * b1_size * b2_size + a1 * b2_size + a2;

    out_actions[i] = action_idx;
    rng_states[i]  = rng;

    /* Output Q-gap for conviction-based position sizing.
     * Q-gap = Q(best_exposure) - Q(flat). Higher = more conviction.
     * Used by env_step to scale target_position. */
    if (out_q_gaps != NULL) {
        int flat_idx = b0_size / 2;
        float q_best = q_b0[argmax_n(q_b0, b0_size)];
        float q_flat = q_b0[flat_idx];
        float gap = q_best - q_flat;
        out_q_gaps[i] = (gap > 0.0f) ? gap : 0.0f;
    }
}

/* ================================================================== */
/* Kernel 3: experience_env_step                                      */
/* ================================================================== */

/**
 * Portfolio simulation, reward v5 (trade-aware hybrid), and done detection.
 *
 * For each episode i this kernel:
 *   1. Reads the current-bar raw prices from targets[bar_idx].
 *   2. Decodes the factored action to an exposure branch index.
 *   3. Computes target position = exposure_fraction * max_position.
 *   4. Applies position adjustment delta with volatility-scaled tx cost.
 *   5. Runs dynamic trailing stop (regime-adaptive).
 *   6. Computes reward v5:
 *        dense  = raw_pnl / equity  (only when in trade, keeps gradients flowing)
 *        sparse = trade_return * patience_mult  (at trade exit, primary signal)
 *        reward = 0.1 * dense + 2.0 * sparse
 *        if reward < 0: reward *= loss_aversion  (prospect theory)
 *   7. Writes (batch_states, action, reward, done) to output replay buffer.
 *   8. Updates portfolio_states[0..19] in place.
 *   9. Increments current_timesteps[i].
 *
 * Grid: ceil(N / 256), Block: 256.  One thread per episode.
 *
 * targets layout: [total_bars, 4]
 *   [0] preproc_close  — log-return-normalised close (network only)
 *   [1] preproc_next   — log-return-normalised next close (unused here)
 *   [2] raw_close      — raw close price (for position cost + tx)
 *   [3] raw_next       — raw next-bar close (for mark-to-market PnL)
 *
 * portfolio_states layout: [N, PORTFOLIO_STRIDE=20]  (read-write)
 *   See file header for field definitions.
 *
 * @param targets            [total_bars, 4]   price data
 * @param episode_starts     [N]               global bar offset per episode
 * @param current_timesteps  [N]               current episode step (incremented)
 * @param actions            [N]               factored action from action_select
 * @param portfolio_states   [N, 20]           full per-episode state (read-write)
 * @param out_states         [N, L, state_dim] replay output: states
 * @param out_actions        [N, L]            replay output: actions
 * @param out_rewards        [N, L]            replay output: rewards
 * @param out_dones          [N, L]            replay output: done flags
 * @param batch_states       [N, state_dim]    current state batch from gather kernel
 * @param max_position                         max allowed contract position (absolute)
 * @param tx_cost_multiplier                   proportional transaction cost rate
 * @param loss_aversion                        asymmetric loss scaling factor (~1.5)
 * @param features           [total_bars, market_dim] market features for regime detection
 * @param market_dim                           MARKET_DIM (42)
 * @param L                                    episode length (replay buffer columns)
 * @param N                                    number of episodes
 * @param total_bars                           length of targets/market_features
 * @param state_dim                            state vector length
 * @param b0_size                              exposure branch size (default 9)
 * @param b1_size                              order branch size (default 3)
 * @param b2_size                              urgency branch size (default 3)
 * @param current_t                            timestep index (0-based) in this batch
 * @param cvar_scales        [N] or NULL       CVaR position scaling
 * @param q_gaps             [N] or NULL       Q-gap conviction scaling
 * @param raw_returns_out    [N, L] or NULL    true per-bar portfolio return (unshaped)
 */
extern "C" __global__ void experience_env_step(
    const float* __restrict__ targets,
    const int*   __restrict__ episode_starts,
    int*                      current_timesteps,
    const int*   __restrict__ actions,
    float*                    portfolio_states,
    float*                    out_states,
    int*                      out_actions,
    float*                    out_rewards,
    float*                    out_dones,
    const float* __restrict__ batch_states,
    float max_position,
    float tx_cost_multiplier,
    float loss_aversion,
    const float* __restrict__ features,
    int market_dim,
    int L,
    int N,
    int total_bars,
    int state_dim,
    int b0_size,
    int b1_size,
    int b2_size,
    int current_t,
    const float* __restrict__ cvar_scales,  /* [N] or NULL — CVaR position scaling */
    const float* __restrict__ q_gaps,       /* [N] or NULL — Q-gap conviction scaling */
    float*                    raw_returns_out /* [N, L] output: true per-bar portfolio return (unshapen) */
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= N) return;

    int t       = current_timesteps[i];
    int bar_idx = episode_starts[i] + t;

    /* Flat index into [N, L] output arrays. */
    long long out_off = (long long)i * L + current_t;

    /* ---- Write current state to replay buffer ---- */
    const float* src = batch_states + (long long)i * state_dim;
    float*       dst = out_states   + out_off * state_dim;
    for (int k = 0; k < state_dim; k++)
        dst[k] = src[k];

    /* ---- Write action ---- */
    int action_idx = actions[i];
    out_actions[out_off] = action_idx;

    /* ---- Out-of-data check ---- */
    if (bar_idx >= total_bars - 1) {
        out_rewards[out_off] = 0.0f;
        out_dones[out_off]   = 1.0f;
        /* Do not advance timestep: episode is at the data boundary. */
        return;
    }

    /* ---- Read raw prices from targets ---- */
    /* Target layout: [preproc_close, preproc_next, raw_close, raw_next]
     *   [0:1] = preprocessed (log-return normalized) — for Q-network
     *   [2:3] = raw dollar prices — for portfolio simulation P&L + tx costs */
    const float* tgt = targets + (long long)bar_idx * 4;
    float raw_close  = tgt[2];
    float raw_next   = tgt[3];

    /* Guard against degenerate prices from data gaps. */
    if (raw_close <= 0.0f) raw_close = 1.0f;
    if (raw_next  <= 0.0f) raw_next  = raw_close;

    /* ---- Read full portfolio state (PORTFOLIO_STRIDE=20) ---- */
    float* ps            = portfolio_states + (long long)i * PORTFOLIO_STRIDE;
    float position       = ps[0];
    float cash           = ps[1];
    /* ps[2] = portfolio_value (updated at end) */
    /* ps[3:6] reserved (unused by reward v5) */
    float peak_equity    = ps[7];
    float flat_counter   = ps[8];
    float prev_equity    = ps[9];
    float hold_time      = ps[10];
    /* ps[11] = realized_pnl (cumulative, updated at end) */
    float entry_price    = ps[12];   /* price when trade was entered */
    float trade_start_pnl = ps[13]; /* realized_pnl snapshot at trade entry */
    float win_count      = ps[14];  /* Kelly: number of profitable trade exits */
    float loss_count     = ps[15];  /* Kelly: number of losing trade exits */
    float sum_wins       = ps[16];  /* Kelly: cumulative profit from winners */
    float sum_losses     = ps[17];  /* Kelly: cumulative |loss| from losers */
    float sum_returns    = ps[18];  /* Kelly: cumulative net returns (for μ) */
    float sum_sq_returns = ps[19];  /* Kelly: cumulative squared returns (for σ²) */

    /* ---- Decode exposure index from factored action ---- */
    int exposure_idx;
    if (b1_size > 0 && b2_size > 0) {
        /* Branching: exposure branch = action / (b1_size * b2_size) */
        int denom_act = b1_size * b2_size;
        exposure_idx = action_idx / denom_act;
    } else {
        /* Flat: action IS the exposure index */
        exposure_idx = action_idx;
    }
    /* Clamp to valid range — defence against caller bugs. */
    if (exposure_idx < 0)        exposure_idx = 0;
    if (exposure_idx >= b0_size) exposure_idx = b0_size - 1;

    float target_exposure = exposure_idx_to_fraction(exposure_idx);
    float target_position = target_exposure * max_position;

    /* Risk management as ENVIRONMENT PHYSICS (not action overrides).
     *
     * CVaR, conviction, and Kelly scale the target position, but the agent
     * CAN OBSERVE these scales in its portfolio state features (position,
     * floor_distance, etc). The agent learns to select exposure levels that
     * account for the environment's risk scaling.
     *
     * This is like friction in physics: the agent selects force (exposure),
     * the environment applies friction (risk scaling), and the outcome
     * (position + reward) reflects both. The agent learns the physics. */

    /* CVaR position scaling: scale down in high-risk quantiles */
    if (cvar_scales != NULL) {
        float cvar_scale = cvar_scales[i];
        if (cvar_scale > 0.0f && cvar_scale < 1.0f) {
            target_position *= cvar_scale;
        }
    }

    /* Q-gap conviction: scale position by trading conviction */
    if (q_gaps != NULL) {
        float q_gap = q_gaps[i];
        float conviction = fminf(fmaxf(q_gap * 0.5f, 0.25f), 1.0f);
        target_position *= conviction;
    }

    /* Kelly criterion: scale by optimal fraction after sufficient samples.
     * win_count/loss_count already declared earlier from ps[14]/ps[15]. */
    {
        float total_trades = win_count + loss_count;
        if (total_trades >= 20.0f) {
            float sum_wins_val  = ps[16];
            float sum_losses_val = ps[17];
        float avg_win  = sum_wins_val / fmaxf(win_count, 1.0f);
        float avg_loss = sum_losses_val / fmaxf(loss_count, 1.0f);
        float win_rate = win_count / total_trades;
        /* Kelly fraction: f* = (bp - q) / b where b=avg_win/avg_loss */
        float b = avg_win / fmaxf(avg_loss, 0.001f);
        float kelly_f = (b * win_rate - (1.0f - win_rate)) / fmaxf(b, 0.001f);
        kelly_f = fminf(fmaxf(kelly_f, 0.0f), 1.0f); /* clamp [0, 1] */
        /* Half-Kelly for safety */
        kelly_f *= 0.5f;
        if (kelly_f > 0.01f) {
            target_position *= kelly_f;
        }
    } /* end Kelly block */
    } /* end Kelly scope */

    /* Final NaN/Inf guard on target_position.
     * If exposure_idx_to_fraction produced NaN, zero the position. */
    if (isnan(target_position) || isinf(target_position)) {
        target_position = 0.0f;
    }

    /* ---- Position adjustment with volatility-scaled transaction cost ---- */
    /* Real spread widens in volatile markets. CUSUM direction (feature[41])
     * is used here as a market-stress proxy for spread scaling — this is
     * CORRECT (high |CUSUM| = trending = wider spreads). NOT used as vol
     * proxy for reward normalization — ATR(14) at feature[9] handles that. */
    float cusum_raw = 0.0f;
    if (features != NULL && bar_idx < total_bars && market_dim > 41) {
        cusum_raw = features[(long long)bar_idx * market_dim + 41];
    }
    float spread_scale = cusum_raw / 0.5f;
    spread_scale = (spread_scale < 0.5f) ? 0.5f : ((spread_scale > 2.0f) ? 2.0f : spread_scale);

    float delta   = target_position - position;
    float tx_cost = 0.0f;

    if (delta != 0.0f) {
        /* Square-root market impact (Almgren & Chriss, 2000).
         * impact_scale = 1 + sqrt(|delta|/max_position) — standard academic model.
         * √x grows slower than x² → less penalty for moderate sizes,
         * more realistic than quadratic for futures markets. */
        float size_ratio = fabsf(delta) / (max_position > 0.0f ? max_position : 1.0f);
        float impact_scale = 1.0f + sqrtf(size_ratio);

        /* Order-type cost differentiation from branching action:
         * Factored action = exposure * 9 + order_type * 3 + urgency
         * order_type: 0=Market (+0 bps), 1=IoC (+2 bps), 2=LimitMaker (-5 bps rebate)
         * LimitMaker gets a REBATE — teaches the model to prefer limit orders. */
        int order_type_idx = (action_idx / b2_size) % b1_size;
        float order_premium = (order_type_idx == 0) ? 0.0f
                            : (order_type_idx == 1) ? 0.0002f   /* IoC: +2 bps */
                            : -0.0005f;                          /* LimitMaker: -5 bps rebate */

        tx_cost   = fabsf(delta) * raw_close * (tx_cost_multiplier * 0.0001f * spread_scale * impact_scale + order_premium);
        cash     -= delta * raw_close;
        cash     -= tx_cost;
        position  = target_position;
    }

    /* ---- Mark-to-market PnL for this timestep ---- */
    float raw_pnl = position * (raw_next - raw_close);
    /* Old position's PnL — needed for reversal trade tracking.
     * raw_pnl uses the NEW position (after action). On reversals,
     * the old position's PnL determines trade segment return. */
    float old_pos_pnl = ps[0] * (raw_next - raw_close);

    /* ---- Mark-to-market equity (needed for peak_equity and reward) ---- */
    float equity = cash + position * raw_next;
    peak_equity = (equity > peak_equity) ? equity : peak_equity;

    /* ==== Trade lifecycle tracking (v4: segment-based, reversal-aware) ==== */
    /*
     * A "trade segment" = contiguous period with the same position sign.
     * Segment ends when: (1) exit to flat, (2) reversal (sign change), (3) episode end.
     *
     * A reversal (e.g. S100->L50) is TWO half-trades:
     *   - Close old direction: book P&L, update Kelly, fire sparse reward
     *   - Open new direction: fresh entry_price, fresh trade_start_pnl
     *
     * This matches institutional P&L accounting (round-trip segments).
     */

    float is_flat = (fabsf(position) < 0.001f) ? 1.0f : 0.0f;

    /* Sign of position: -1 (short), 0 (flat), +1 (long) */
    int prev_sign = (ps[0] > 0.001f) ? 1 : ((ps[0] < -0.001f) ? -1 : 0);
    int curr_sign = (position > 0.001f) ? 1 : ((position < -0.001f) ? -1 : 0);

    int entering_trade  = (prev_sign == 0 && curr_sign != 0);
    int exiting_trade   = (prev_sign != 0 && curr_sign == 0);
    int reversing_trade = (prev_sign != 0 && curr_sign != 0 && prev_sign != curr_sign);

    /* A segment completes on exit OR reversal */
    int segment_complete = exiting_trade || reversing_trade;

    /* On reversal: close old segment, open new one.
     * reversal_return is saved for the sparse reward block below
     * (trade_start_pnl is overwritten here, so recomputing would yield 0). */
    float reversal_return = 0.0f;
    if (reversing_trade) {
        /* Closing P&L for the completed segment.
         * CRITICAL: raw_pnl uses the NEW position (already updated at line 677).
         * old_pos_pnl (saved above) has the OLD position's final-bar PnL. */
        float closing_pnl = ps[11] + old_pos_pnl - trade_start_pnl;
        reversal_return = closing_pnl / (prev_equity > 1.0f ? prev_equity : 1.0f);

        /* Kelly stats for completed segment */
        if (reversal_return > 0.0f) {
            win_count += 1.0f;
            sum_wins += reversal_return;
        } else {
            loss_count += 1.0f;
            sum_losses += fabsf(reversal_return);
        }
        sum_returns += reversal_return;
        sum_sq_returns += reversal_return * reversal_return;

        /* Reset for new segment — realized_pnl snapshot after booking the close */
        entry_price = raw_close;
        trade_start_pnl = ps[11] + old_pos_pnl;
    }

    /* On entry from flat (unchanged) */
    if (entering_trade) {
        entry_price = raw_close;
        trade_start_pnl = ps[11]; /* cumulative realized_pnl at entry */
    }

    /* ════════════════════════════════════════════════════════════════════
     * DYNAMIC TRAILING STOP — regime-adaptive, locks in profits.
     *
     * Only activates when in a trade, profitable, and held > 2 bars.
     * Trail distance adapts to regime via ADX (trend) and CUSUM (vol).
     * ════════════════════════════════════════════════════════════════════ */
    float unrealized_trade_pnl = 0.0f;
    if (prev_sign != 0 && hold_time > 0.0f) {
        unrealized_trade_pnl = (ps[11] + raw_pnl - trade_start_pnl) / (prev_equity > 1.0f ? prev_equity : 1.0f);
    }

    /* Trail distance adapts to regime */
    float trail_adx = 0.0f, trail_cusum = 0.0f;
    if (bar_idx < total_bars && market_dim >= 42) {
        trail_adx   = features[(long long)bar_idx * market_dim + 40];
        trail_cusum = features[(long long)bar_idx * market_dim + 41];
    }
    float vol_scale = 1.0f + (trail_cusum > 0.5f ? trail_cusum : 0.0f);
    vol_scale = fminf(vol_scale, 2.5f);
    float trend_scale = 1.0f + (trail_adx > 25.0f ? (trail_adx - 25.0f) / 50.0f : 0.0f);
    trend_scale = fminf(trend_scale, 2.0f);

    float trail_distance = 0.005f * vol_scale * trend_scale; /* 0.5% base, widens in vol/trend */

    /* Trailing stop: only when profitable and held > 2 bars */
    if (prev_sign != 0 && hold_time > 2.0f && peak_equity > 1.0f) {
        float peak_return = (peak_equity - prev_equity) / (prev_equity > 1.0f ? prev_equity : 1.0f);
        if (peak_return > trail_distance) {
            float trail_floor = peak_return - trail_distance;
            if (unrealized_trade_pnl < trail_floor) {
                /* Trail triggered — force exit */
                position = 0.0f;
                float exit_cost = fabsf(ps[0]) * raw_close * tx_cost_multiplier * 0.0001f;
                cash = ps[1] + ps[0] * raw_close - exit_cost;
                is_flat = 1.0f;
                exiting_trade = 1;
            }
        }
    }

    /* Early-exit override: model tries to exit to FLAT before 5 bars — force hold.
     * Only blocks flat exits, NOT reversals — reversals always allowed.
     * CRITICAL: fix the stored action to match what actually executed.
     * Without this, the replay buffer stores (state, action=Flat, reward_from_Hold)
     * which corrupts Q-values for Flat — classic action aliasing bug.
     * TODO(hyperopt): min_hold_bars=5 is hardcoded — candidate for hyperopt parameter. */
    if (prev_sign != 0 && curr_sign == 0 && hold_time < 5.0f && hold_time > 0.0f && !exiting_trade) {
        position = ps[0];
        cash = ps[1];
        is_flat = 0.0f;
        exiting_trade = 0;
        float prev_exposure_frac = ps[0] / (max_position > 0.0f ? max_position : 1.0f);
        int held_exposure = (int)((prev_exposure_frac + 1.0f) / 0.25f + 0.5f);
        if (held_exposure < 0) held_exposure = 0;
        if (held_exposure >= b0_size) held_exposure = b0_size - 1;
        int original_order = (action_idx / b2_size) % b1_size;
        int original_urgency = action_idx % b2_size;
        action_idx = held_exposure * b1_size * b2_size + original_order * b2_size + original_urgency;
        out_actions[out_off] = action_idx;
    }

    /* Recompute segment_complete after trailing stop / early-exit may have
     * modified exiting_trade. Trailing stop sets exiting_trade=1; early-exit
     * override sets exiting_trade=0. reversing_trade is never modified. */
    segment_complete = exiting_trade || reversing_trade;

    /* Save hold_time BEFORE reset — sparse reward needs the pre-reset value
     * for the patience multiplier. Without this, hold_time is 0 at exit
     * and the sparse reward condition (hold_time > 0) never fires. */
    float segment_hold_time = hold_time;

    /* Hold time tracking: counts TOTAL bars in position (not just losing).
     * Reset to 0 when flat OR on reversal (new segment starts fresh).
     * This feeds the patience multiplier in reward v5. */
    if (is_flat < 0.5f && !reversing_trade) {
        hold_time += 1.0f;
    } else if (reversing_trade) {
        /* New segment starts: count this bar as bar 1 of the new direction */
        hold_time = 1.0f;
    } else {
        hold_time = 0.0f;
    }

    /* Idle counter (flat bars) — used for state features, no penalty */
    if (is_flat > 0.5f) {
        flat_counter += 1.0f;
    } else {
        flat_counter = 0.0f;
    }

    /* ════════════════════════════════════════════════════════════════════
     * REWARD v6: Sparse Trade-Completion Only
     *
     * Validated by ETDQN (Takara et al., 2023) — outperformed standard
     * DQN by 1.46-7.13x on cumulative returns with this exact approach.
     *
     * Design:
     *   DURING TRADE: reward = 0.0  (ZERO — no noise, no signal)
     *   WHEN FLAT:    reward = 0.0  (no trade, no signal)
     *   AT EXIT:      reward = SCALE × vol_normalized(trade_return)
     *   TURNOVER:     reward -= PENALTY × |delta_position|
     *
     * Why zero during trade: per-bar ES returns (~0.001) have SNR of 0.001
     * after weighting — mathematically unlearnable. The model trains on
     * noise 90% of the time. Removing it lets the Q-network focus on
     * what matters: which entries lead to profitable exits.
     *
     * Credit assignment: n-step returns (n=3) propagate exit reward
     * backward 3 bars. Earlier bars learn via Q(s') bootstrap over
     * training iterations. HER provides counterfactual learning.
     *
     * Vol normalization (Zhang 2020): makes rewards comparable across
     * trending vs ranging regimes. The model targets risk-adjusted
     * returns (Sharpe per trade), not raw P&L.
     * ════════════════════════════════════════════════════════════════════ */

    float reward = 0.0f;

    /* ---- Sparse: trade completion reward (vol-normalized) ---- */
    if (segment_complete && segment_hold_time > 0.0f) {
        float segment_pnl = ps[11] + raw_pnl - trade_start_pnl;
        float segment_return = segment_pnl / (prev_equity > 1.0f ? prev_equity : 1.0f);

        /* Kelly statistics for EXIT (reversal Kelly already done above) */
        if (exiting_trade) {
            if (segment_return > 0.0f) {
                win_count += 1.0f;
                sum_wins += segment_return;
            } else {
                loss_count += 1.0f;
                sum_losses += fabsf(segment_return);
            }
            sum_returns += segment_return;
            sum_sq_returns += segment_return * segment_return;
        }

        /* For reversals, use saved reversal_return (trade_start_pnl was overwritten) */
        if (reversing_trade) {
            segment_return = reversal_return;
        }

        /* Vol normalization: ATR(14) proxy for realized volatility.
         * Feature[9] stores safe_normalize(ln(atr), -7.0, 9.0) = (ln(atr)+7)/16.
         * Reverse: ln(atr) = feature[9]*16 - 7, atr = exp(ln(atr)).
         * atr_pct = atr / price gives per-bar percentage volatility.
         * Makes reward comparable across trending and ranging regimes. */
        float atr_norm = 0.0f;
        if (features != NULL && bar_idx < total_bars && market_dim > 9) {
            atr_norm = features[(long long)bar_idx * market_dim + 9];
        }
        float log_atr = atr_norm * 16.0f - 7.0f;
        float atr_pct = expf(log_atr) / fmaxf(raw_close, 1.0f);
        float vol_proxy = fmaxf(atr_pct, 0.0001f);
        float vol_norm = vol_proxy * sqrtf(fmaxf(segment_hold_time, 1.0f));
        float vol_normalized_return = segment_return / vol_norm;

        /* Scale to learnable magnitude: target reward range ~[-10, +10].
         * With Xavier init weight noise ~0.001, a reward of 1.0 has SNR=1000.
         * ETDQN showed this range produces stable C51 learning. */
        reward = 10.0f * vol_normalized_return;

        /* Loss aversion BEFORE squash (prospect theory): negative rewards
         * weighted 1.5x. Applied pre-squash so the asymmetry is smooth —
         * post-squash would create [-15, +10] hard asymmetry. */
        if (reward < 0.0f) {
            reward *= loss_aversion;
        }

        /* Soft squash: tanh maps [-inf,+inf] -> [-10,+10] smoothly.
         * Preserves tail information (5% win > 1% win) unlike hard clamp
         * which maps both to 10.0. Gradient never exactly zero. */
        reward = 10.0f * tanhf(reward / 10.0f);
    }

    /* Turnover penalty REMOVED (reward v6): tx costs already deducted from
     * cash at line ~679 (Almgren-Chriss impact model). The old 0.05*|delta|
     * penalty double-counted costs and over-penalized necessary rebalancing. */

    /* ---- Portfolio value for floor check ---- */
    float new_portfolio_value_pre_floor = cash + position * raw_next;

    /* ---- Capital floor protection (PDT $25K rule) ---- */
    /* If equity drops to 25% drawdown from peak, terminate the episode.
     * With $35K initial, this triggers at ~$26.25K — above the $25K PDT minimum.
     * Uses peak_equity (not initial_capital) so it adapts as account grows.
     * The model learns: approaching the floor = game over = zero future reward. */
    float new_portfolio_value = new_portfolio_value_pre_floor;
    float capital_floor = peak_equity * 0.75f;
    if (new_portfolio_value < capital_floor) {
        /* Force flat — emergency exit all positions */
        position = 0.0f;
        cash = new_portfolio_value;
        new_portfolio_value = cash;
        /* Massive penalty — approaching the floor is catastrophic */
        reward = -1.0f;
    }

    /* ---- Done detection ---- */
    int next_bar = bar_idx + 1;
    int done     = (next_bar >= total_bars || new_portfolio_value < capital_floor) ? 1 : 0;

    /* ---- Update full portfolio state (PORTFOLIO_STRIDE=20) ---- */
    ps[0]  = position;
    ps[1]  = cash;
    ps[2]  = new_portfolio_value;
    /* ps[3:6] reserved — reward v5 does not use DSR/PnL EMA */
    ps[7]  = peak_equity;
    ps[8]  = flat_counter;
    ps[9]  = new_portfolio_value;  /* prev_equity = current equity for next step */
    ps[10] = hold_time;
    /* realized_pnl accumulates. On reversal bars, use old_pos_pnl
     * (saved before position was overwritten at ps[0] = position).
     * For non-reversal bars, raw_pnl (= current position's PnL) is correct. */
    if (reversing_trade) {
        ps[11] = ps[11] + old_pos_pnl; /* old position's PnL (saved at line ~700) */
    } else {
        ps[11] = ps[11] + raw_pnl; /* normal: current position's PnL */
    }
    ps[12] = entry_price;          /* preserved across bars of a trade */
    ps[13] = trade_start_pnl;      /* realized_pnl snapshot at trade entry */
    ps[14] = win_count;            /* Kelly: accumulated across episode */
    ps[15] = loss_count;
    ps[16] = sum_wins;
    ps[17] = sum_losses;
    ps[18] = sum_returns;          /* Kelly continuous: Σ returns */
    ps[19] = sum_sq_returns;       /* Kelly continuous: Σ returns² */

    /* ---- NO global reward clamp ---- */
    /* Reward v5: dense in [-small, +small], sparse in [-moderate, +moderate].
     * The combined reward is unbounded in principle but practically small.
     * C51 v_range should cover Q = reward / (1-gamma). */

    /* ---- Write reward and done flag ---- */
    /* Final NaN guard — if reward is NaN/Inf, write 0.0 instead of poisoning
     * the replay buffer. This prevents gradient explosion from propagating. */
    if (isnan(reward) || isinf(reward)) reward = 0.0f;
    out_rewards[out_off] = reward;
    out_dones[out_off]   = (float)done;

    /* ---- Write RAW portfolio return (unshaped) for accurate Sharpe/MaxDD ---- */
    /* True fractional return: (equity_t - equity_{t-1}) / equity_{t-1}.
     * No dense/sparse weighting, no loss_aversion, no patience multiplier.
     * This is what risk management (Sharpe, Sortino, MaxDD, VaR) must use. */
    if (raw_returns_out != NULL) {
        float portfolio_return = (prev_equity > 1.0f)
            ? (new_portfolio_value - prev_equity) / prev_equity
            : 0.0f;
        if (isnan(portfolio_return) || isinf(portfolio_return)) portfolio_return = 0.0f;
        raw_returns_out[out_off] = portfolio_return;
    }

    /* ---- Advance timestep counter ---- */
    current_timesteps[i] = t + 1;
}

/* ══════════════════════════════════════════════════════════════════════════
 * LEGACY: Portfolio simulation kernel (used by GpuPortfolioSimulator)
 *
 * Single-thread sequential processing for hyperopt/backtest paths.
 * NOT part of the timestep-loop experience collection — kept for
 * GpuPortfolioSimulator compatibility in gpu_portfolio.rs.
 * ══════════════════════════════════════════════════════════════════════════ */

__device__ __forceinline__ float action_to_exposure(int action_idx) {
    /* 9 levels: same formula as exposure_idx_to_fraction */
    return -1.0f + (float)action_idx * 0.25f;
}

__device__ __forceinline__ float action_to_tx_cost(int action_idx) {
    (void)action_idx;
    return 0.0001f;
}

extern "C" __global__ void portfolio_sim_kernel(
    const float* __restrict__ targets,
    const int*   __restrict__ actions,
    float*       portfolio_state,
    float*       portfolio_out,
    float*       rewards_out,
    int*         done_out,
    int          batch_start,
    int          batch_size,
    float        max_position,
    int          episode_length,
    int          total_bars
) {
    if (threadIdx.x != 0 || blockIdx.x != 0) return;

    float cash          = portfolio_state[0];
    float position      = portfolio_state[1];
    float entry_price   = portfolio_state[2];
    float initial_cap   = portfolio_state[3];
    float spread        = portfolio_state[4];
    float last_price    = portfolio_state[5];
    float reserve_pct   = portfolio_state[6];
    float cum_costs     = portfolio_state[7];

    for (int b = 0; b < batch_size; b++) {
        int global_idx = batch_start + b;
        int action_idx = actions[b];

        int t_offset = global_idx * 4;
        float current_close     = targets[t_offset + 0];
        float next_close        = targets[t_offset + 1];
        float current_close_raw = targets[t_offset + 2];
        float next_close_raw    = targets[t_offset + 3];

        float price = (current_close_raw != 0.0f) ? current_close_raw : current_close;
        if (price <= 0.0f) price = 1.0f;

        float current_value = cash + position * price;
        float current_norm = current_value / initial_cap;

        float target_exposure = action_to_exposure(action_idx);
        float target_position = target_exposure * max_position;
        float tx_rate = action_to_tx_cost(action_idx);

        bool is_reversal = (position > 0.0f && target_position < 0.0f) ||
                           (position < 0.0f && target_position > 0.0f);

        if (is_reversal) {
            float close_cash = position * price;
            float close_cost = fabsf(position) * price * tx_rate;
            cash += close_cash - close_cost;
            cum_costs += close_cost;

            float reserve = (reserve_pct > 0.0f) ? current_value * (reserve_pct / 100.0f) : 0.0f;
            float affordable = fmaxf(cash - reserve, 0.0f);
            float max_contracts = (price > 0.0f) ? affordable / (price * (1.0f + tx_rate)) : 0.0f;
            max_contracts = floorf(max_contracts);
            float actual = fminf(max_contracts, fabsf(target_position));

            if (actual > 0.0f) {
                float new_pos = (target_position > 0.0f) ? actual : -actual;
                float open_cost = actual * price * tx_rate;
                cash -= new_pos * price + open_cost;
                cum_costs += open_cost;
                position = new_pos;
                entry_price = price;
            } else {
                position = 0.0f;
                entry_price = 0.0f;
            }
        } else {
            float delta = target_position - position;
            if (fabsf(delta) > 0.0f) {
                float trade_cost = fabsf(delta) * price * tx_rate;
                cum_costs += trade_cost;
                cash -= trade_cost;

                if (delta > 0.0f && reserve_pct > 0.0f) {
                    float pv = cash + position * price;
                    float reserve = pv * (reserve_pct / 100.0f);
                    float buy_cost = delta * price;
                    if (cash - buy_cost < reserve) {
                        float affordable = fmaxf(cash - reserve, 0.0f);
                        delta = fminf(delta, floorf(affordable / price));
                    }
                }

                if (delta > 0.0f) {
                    entry_price = price;
                } else if (target_position == 0.0f) {
                    entry_price = 0.0f;
                }
                cash -= delta * price;
                position = position + delta;
            }
        }

        last_price = price;

        float next_price = (next_close_raw != 0.0f) ? next_close_raw : next_close;
        if (next_price <= 0.0f) next_price = price;
        float next_value = cash + position * next_price;
        float next_norm = next_value / initial_cap;

        float max_pos_norm = (price > 0.0f) ? initial_cap / price : 1.0f;
        float pos_norm = position / max_pos_norm;

        portfolio_out[b * 3 + 0] = next_norm;
        portfolio_out[b * 3 + 1] = pos_norm;
        portfolio_out[b * 3 + 2] = spread;

        float reward = 0.0f;
        if (current_norm > 0.0f) {
            reward = (next_norm - current_norm) / current_norm;
        }

        float abs_pos = fabsf(pos_norm);
        if (abs_pos > 0.8f) {
            reward -= (abs_pos - 0.8f) * 5.0f * 0.1f;
        }

        rewards_out[b] = reward;

        int step = global_idx + 1;
        int time_done = (step % episode_length == 0) ? 1 : 0;
        int data_done = (step >= total_bars) ? 1 : 0;
        done_out[b] = (time_done || data_done) ? 1 : 0;

        if (done_out[b]) {
            cash = initial_cap;
            position = 0.0f;
            entry_price = 0.0f;
            cum_costs = 0.0f;
        }
    }

    portfolio_state[0] = cash;
    portfolio_state[1] = position;
    portfolio_state[2] = entry_price;
    portfolio_state[3] = initial_cap;
    portfolio_state[4] = spread;
    portfolio_state[5] = last_price;
    portfolio_state[6] = reserve_pct;
    portfolio_state[7] = cum_costs;
}

/* ================================================================== */
/* Kernel 4: compute_expected_q                                        */
/* ================================================================== */

/**
 * Compute expected Q-values from C51 distributional logits.
 *
 * For each sample and each branch action, combines value + advantage
 * logits via dueling aggregation, applies softmax over atoms, and
 * computes expected Q = sum(softmax(logits) * z_support).
 *
 * When num_atoms == 1, v_logits and b_logits are already plain
 * Q-values and this kernel is not needed (pass them directly
 * to action_select).
 *
 * Grid: ceil(N / 256), Block: 256.  One thread per episode.
 *
 * @param v_logits    [N, NA]                 value head logits
 * @param b_logits    [N, (B0+B1+B2)*NA]      branch advantage logits
 * @param q_values    [N, B0+B1+B2]           output: expected Q per action
 * @param N                                    number of samples
 * @param num_atoms                            C51 atom count (NA)
 * @param b0_size                              exposure branch size
 * @param b1_size                              order branch size
 * @param b2_size                              urgency branch size
 * @param v_min                                C51 minimum support
 * @param v_max                                C51 maximum support
 */
extern "C" __global__ void compute_expected_q(
    const float* __restrict__ v_logits,
    const float* __restrict__ b_logits,
    float*                    q_values,
    int N,
    int num_atoms,
    int b0_size,
    int b1_size,
    int b2_size,
    float v_min,
    float v_max
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= N) return;

    int total_actions = b0_size + b1_size + b2_size;
    float dz = (num_atoms > 1) ? (v_max - v_min) / (float)(num_atoms - 1) : 0.0f;

    /* Per-sample value logits */
    const float* v_row = v_logits + (long long)i * num_atoms;
    const float* b_row = b_logits + (long long)i * total_actions * num_atoms;

    /* Compute advantage mean across all actions for dueling subtraction */
    /* For each action a: dueling_logits[a][z] = v[z] + adv[a][z] - mean_adv[z]
     * Expected Q = sum_z( softmax(dueling_logits[a])[z] * z_support[z] )
     *
     * For efficiency, compute mean advantage per atom first. */

    for (int a = 0; a < total_actions; a++) {
        const float* adv_a = b_row + (long long)a * num_atoms;

        /* Combine value + advantage (no mean subtraction for action selection —
         * the argmax is invariant to the shared value baseline). */
        /* Softmax + expectation over atoms */
        float max_logit = -1e30f;
        for (int z = 0; z < num_atoms; z++) {
            float logit = v_row[z] + adv_a[z];
            if (logit > max_logit) max_logit = logit;
        }

        float sum_exp = 0.0f;
        for (int z = 0; z < num_atoms; z++) {
            float logit = v_row[z] + adv_a[z];
            sum_exp += expf(logit - max_logit);
        }

        float expected_q = 0.0f;
        for (int z = 0; z < num_atoms; z++) {
            float logit = v_row[z] + adv_a[z];
            float prob = expf(logit - max_logit) / sum_exp;
            float z_val = v_min + (float)z * dz;
            expected_q += prob * z_val;
        }

        q_values[(long long)i * total_actions + a] = expected_q;
    }
}

/* ================================================================== */
/* Kernel 5: expert_action_override                                    */
/* ================================================================== */

/**
 * GPU-native expert action override using inline MA crossover + ADX filter.
 *
 * For each episode, reads prices from `targets` and ADX from `features` at the
 * current bar position. Computes a simple fast/slow price crossover signal:
 *   - If close > open (bullish bar) AND ADX > 25 → Long (a0 = b0-1)
 *   - If close < open (bearish bar) AND ADX > 25 → Short (a0 = 0)
 *   - Otherwise → no override
 *
 * With probability `expert_ratio`, overrides the Q-network's exposure action.
 * Order and urgency branches are preserved from the Q-network.
 *
 * Zero CPU involvement — all data read from GPU-resident targets/features buffers.
 *
 * Grid: ceil(N / 256), Block: 256.  One thread per episode.
 */
extern "C" __global__ void expert_action_override(
    int*                       out_actions,       /* [N] Q-network actions (overwritten) */
    const float* __restrict__  targets,           /* [total_bars, 4] OHLCV price data */
    const float* __restrict__  features,          /* [total_bars, MARKET_DIM] market features */
    const float* __restrict__  portfolio_states,  /* [N, PORTFOLIO_STRIDE] per-episode state */
    const int*   __restrict__  episode_starts,    /* [N] bar offset per episode */
    const int*   __restrict__  current_timesteps, /* [N] current step in episode */
    unsigned int*              rng_states,        /* [N] per-episode RNG (updated) */
    float expert_ratio,                           /* override probability [0,1] */
    int N,
    int total_bars,
    int market_dim,
    int b0_size,                                  /* exposure branch size (9) */
    int b1_size,                                  /* order branch size (3) */
    int b2_size                                   /* urgency branch size (3) */
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= N) return;
    if (expert_ratio <= 0.0f) return;

    unsigned int rng = rng_states[i];
    float r = lcg_random(&rng);
    rng_states[i] = rng;

    if (r >= expert_ratio) return;  /* Keep Q-network action */

    /* Look up current bar position */
    int bar = episode_starts[i] + current_timesteps[i];
    if (bar < 1 || bar >= total_bars) return;

    /* Read prices: compute multi-bar momentum as EMA proxy.
     * True EMA requires sequential scan (O(n) per thread). Instead, use
     * 5-bar price momentum + 20-bar momentum divergence as a crossover proxy.
     * When short-term momentum > long-term momentum AND ADX is strong → trend.
     * This approximates fast/slow EMA crossover without sequential scan. */
    if (bar < 20) return;  /* Need 20 bars of history */

    float adx = features[bar * market_dim + 40];
    float cusum = features[bar * market_dim + 41];

    /* 5-bar return (fast proxy) */
    float close_now = targets[bar * 4 + 3];
    float close_5   = targets[(bar - 5) * 4 + 3];
    float close_20  = targets[(bar - 20) * 4 + 3];
    float ret_5  = (close_now - close_5)  / (close_5  + 1e-8f);
    float ret_20 = (close_now - close_20) / (close_20 + 1e-8f);

    /* Expert signal: momentum crossover + ADX strength + CUSUM direction */
    int expert_a0 = -1;  /* -1 = no opinion */
    if (adx > 25.0f) {
        /* Bullish: fast momentum > slow momentum AND CUSUM positive */
        if (ret_5 > ret_20 && ret_5 > 0.001f && cusum > 0.0f) {
            expert_a0 = b0_size - 1;  /* Long100 */
        }
        /* Bearish: fast momentum < slow momentum AND CUSUM negative */
        else if (ret_5 < ret_20 && ret_5 < -0.001f && cusum < 0.0f) {
            expert_a0 = 0;            /* Short100 */
        }
    }

    if (expert_a0 < 0) return;  /* No clear signal */

    /* Check if current position already matches expert signal — skip if redundant.
     * portfolio_states[i * PORTFOLIO_STRIDE + 0] = current position (-1 to +1).
     * If expert says Long and we're already Long (position > 0.5), don't waste
     * the override. Same for Short. */
    float current_pos = portfolio_states[i * 20 + 0];  /* PORTFOLIO_STRIDE=20, pos at idx 0 */
    if (expert_a0 == b0_size - 1 && current_pos > 0.5f) return;  /* Already Long */
    if (expert_a0 == 0 && current_pos < -0.5f) return;           /* Already Short */

    /* Override exposure branch, keep Q-network's order and urgency */
    int current = out_actions[i];
    int a1 = (current / b2_size) % b1_size;
    int a2 = current % b2_size;
    out_actions[i] = expert_a0 * b1_size * b2_size + a1 * b2_size + a2;
}
