/**
 * Zero-Roundtrip PPO Experience Collection Kernel
 *
 * Requires common_device_functions.cuh prepended via NVRTC source concatenation.
 * Launch config: grid=(ceil(N/32),1,1), block=(32,1,1).
 * Each thread runs one independent episode.
 *
 * Phase A: Forward rollout (L timesteps) — actor, critic, portfolio, rewards
 * Phase B: Backward GAE scan — advantages and returns
 */

/* Override NUM_ACTIONS for PPO: use full 45-action factored space */
#undef NUM_ACTIONS
#define NUM_ACTIONS PPO_NUM_ACTIONS

/* PPO Actor layer sizes: 54 -> 128 (ReLU) -> 64 (ReLU) -> 45 */
#define ACTOR_H1  128
#define ACTOR_H2   64

/* PPO Critic layer sizes: 54 -> 512 -> 384 -> 256 -> 128 -> 64 -> 1 (all ReLU) */
#define CRITIC_H1 512
#define CRITIC_H2 384
#define CRITIC_H3 256
#define CRITIC_H4 128
#define CRITIC_H5  64

/* Maximum GAE rollout length (matches timesteps_per_episode default) */
#define MAX_GAE_LEN 500

/* ------------------------------------------------------------------ */
/*  PPO-Specific Device Functions                                      */
/* ------------------------------------------------------------------ */

/**
 * PPO Actor MLP forward pass.
 *
 *   state[54] -> h1[128] (LeakyReLU) -> h2[64] (LeakyReLU) -> logits[45]
 *
 * @param state   Input state vector [STATE_DIM]
 * @param pw1     Layer 1 weights [ACTOR_H1, STATE_DIM]
 * @param pb1     Layer 1 biases  [ACTOR_H1]
 * @param pw2     Layer 2 weights [ACTOR_H2, ACTOR_H1]
 * @param pb2     Layer 2 biases  [ACTOR_H2]
 * @param pw3     Output weights  [NUM_ACTIONS, ACTOR_H2]
 * @param pb3     Output biases   [NUM_ACTIONS]
 * @param h1      Scratch buffer  [ACTOR_H1]
 * @param h2      Scratch buffer  [ACTOR_H2]
 * @param logits  Output logits   [NUM_ACTIONS]
 */
__device__ void ppo_actor_forward(
    const float* state,
    const float* __restrict__ pw1,   /* [ACTOR_H1, STATE_DIM] */
    const float* __restrict__ pb1,   /* [ACTOR_H1]            */
    const float* __restrict__ pw2,   /* [ACTOR_H2, ACTOR_H1]  */
    const float* __restrict__ pb2,   /* [ACTOR_H2]            */
    const float* __restrict__ pw3,   /* [NUM_ACTIONS, ACTOR_H2] */
    const float* __restrict__ pb3,   /* [NUM_ACTIONS]         */
    float* h1,                       /* [ACTOR_H1] scratch    */
    float* h2,                       /* [ACTOR_H2] scratch    */
    float* logits                    /* [NUM_ACTIONS] output  */
) {
    /* Hidden layer 1: state -> h1 with LeakyReLU */
    matvec_leaky_relu(pw1, pb1, state, h1, STATE_DIM, ACTOR_H1, 1);

    /* Hidden layer 2: h1 -> h2 with LeakyReLU */
    matvec_leaky_relu(pw2, pb2, h1, h2, ACTOR_H1, ACTOR_H2, 1);

    /* Output layer: h2 -> logits (no activation) */
    matvec_leaky_relu(pw3, pb3, h2, logits, ACTOR_H2, NUM_ACTIONS, 0);
}

/**
 * Stable softmax + categorical sampling.
 *
 * 1. Find max logit for numerical stability
 * 2. exp(logit - max) and accumulate sum
 * 3. Normalize to probabilities
 * 4. CDF scan with LCG random draw -> action index
 * 5. Compute log(p[action]) for PPO loss
 *
 * @param logits       Input logits [NUM_ACTIONS]
 * @param probs        Scratch + output probabilities [NUM_ACTIONS]
 * @param rng          Pointer to LCG RNG state
 * @param out_action   Output: selected action index
 * @param out_log_prob Output: log probability of selected action
 */
__device__ void softmax_sample(
    const float* logits,
    float* probs,
    unsigned int* rng,
    int* out_action,
    float* out_log_prob
) {
    /* Step 1: Find max logit for numerical stability */
    float max_logit = logits[0];
    for (int i = 1; i < NUM_ACTIONS; i++) {
        if (logits[i] > max_logit) max_logit = logits[i];
    }

    /* Step 2: exp(logit - max) and accumulate sum */
    float sum_exp = 0.0f;
    for (int i = 0; i < NUM_ACTIONS; i++) {
        probs[i] = expf(logits[i] - max_logit);
        sum_exp += probs[i];
    }

    /* Step 3: Normalize to probabilities */
    float inv_sum = 1.0f / fmaxf(sum_exp, 1e-8f);
    for (int i = 0; i < NUM_ACTIONS; i++) {
        probs[i] *= inv_sum;
    }

    /* Step 4: CDF scan + random draw -> action index */
    float u = gpu_random(rng);
    float cdf = 0.0f;
    int action = NUM_ACTIONS - 1;  /* default to last action */
    for (int i = 0; i < NUM_ACTIONS; i++) {
        cdf += probs[i];
        if (u < cdf) {
            action = i;
            break;
        }
    }

    /* Step 5: log probability for PPO loss */
    float p = fmaxf(probs[action], 1e-8f);  /* clamp for log safety */
    *out_action   = action;
    *out_log_prob = logf(p);
}

/**
 * PPO Critic (value network) forward pass — 5-layer deep with ping-pong buffers.
 *
 *   state[54] -> 512(LReLU) -> 384(LReLU) -> 256(LReLU) -> 128(LReLU) -> 64(LReLU) -> 1
 *
 * Uses two scratch buffers (scratch_a[512], scratch_b[512]) that alternate
 * between layers to avoid extra memory.
 *
 * @param state      Input state vector [STATE_DIM]
 * @param vw1-vw6    Weight matrices for each layer
 * @param vb1-vb6    Bias vectors for each layer
 * @param scratch_a  Ping-pong scratch buffer A [CRITIC_H1] (512 wide)
 * @param scratch_b  Ping-pong scratch buffer B [CRITIC_H1] (512 wide)
 * @return           Scalar value estimate V(s)
 */
__device__ float ppo_critic_forward(
    const float* state,
    const float* __restrict__ vw1,   /* [CRITIC_H1, STATE_DIM]  = [512, 54]  */
    const float* __restrict__ vb1,   /* [CRITIC_H1]             = [512]      */
    const float* __restrict__ vw2,   /* [CRITIC_H2, CRITIC_H1]  = [384, 512] */
    const float* __restrict__ vb2,   /* [CRITIC_H2]             = [384]      */
    const float* __restrict__ vw3,   /* [CRITIC_H3, CRITIC_H2]  = [256, 384] */
    const float* __restrict__ vb3,   /* [CRITIC_H3]             = [256]      */
    const float* __restrict__ vw4,   /* [CRITIC_H4, CRITIC_H3]  = [128, 256] */
    const float* __restrict__ vb4,   /* [CRITIC_H4]             = [128]      */
    const float* __restrict__ vw5,   /* [CRITIC_H5, CRITIC_H4]  = [64, 128]  */
    const float* __restrict__ vb5,   /* [CRITIC_H5]             = [64]       */
    const float* __restrict__ vw6,   /* [1, CRITIC_H5]          = [1, 64]    */
    const float* __restrict__ vb6,   /* [1]                                  */
    float* scratch_a,                /* [CRITIC_H1] = [512] ping-pong A      */
    float* scratch_b                 /* [CRITIC_H1] = [512] ping-pong B      */
) {
    /* Layer 1: state[54] -> scratch_a[512] with LeakyReLU */
    matvec_leaky_relu(vw1, vb1, state, scratch_a, STATE_DIM, CRITIC_H1, 1);

    /* Layer 2: scratch_a[512] -> scratch_b[384] with LeakyReLU */
    matvec_leaky_relu(vw2, vb2, scratch_a, scratch_b, CRITIC_H1, CRITIC_H2, 1);

    /* Layer 3: scratch_b[384] -> scratch_a[256] with LeakyReLU */
    matvec_leaky_relu(vw3, vb3, scratch_b, scratch_a, CRITIC_H2, CRITIC_H3, 1);

    /* Layer 4: scratch_a[256] -> scratch_b[128] with LeakyReLU */
    matvec_leaky_relu(vw4, vb4, scratch_a, scratch_b, CRITIC_H3, CRITIC_H4, 1);

    /* Layer 5: scratch_b[128] -> scratch_a[64] with LeakyReLU */
    matvec_leaky_relu(vw5, vb5, scratch_b, scratch_a, CRITIC_H4, CRITIC_H5, 1);

    /* Output layer: scratch_a[64] -> scalar (no activation) */
    float value = vb6[0];
    for (int i = 0; i < CRITIC_H5; i++) {
        value += vw6[i] * scratch_a[i];
    }

    return value;
}

/**
 * Generalized Advantage Estimation (GAE) backward scan.
 *
 * Computes advantages and returns by scanning backwards through
 * the collected rollout data:
 *
 *   for t = L-1 down to 0:
 *       delta = rewards[t] + gamma * values[t+1] * (1-dones[t]) - values[t]
 *       gae   = delta + gamma * lambda * (1-dones[t]) * gae
 *       advantages[t] = gae
 *       returns[t]    = gae + values[t]
 *
 * @param rewards     Per-timestep rewards [L]
 * @param values      Per-timestep value estimates [L+1] (values[L] is bootstrap)
 * @param dones       Per-timestep done flags [L] (1.0 = done, 0.0 = not done)
 * @param advantages  Output advantage estimates [L]
 * @param returns     Output return targets [L]
 * @param L           Number of timesteps
 * @param gamma       Discount factor
 * @param lambda      GAE lambda parameter
 */
__device__ void compute_gae_backward(
    const float* rewards,
    const float* values,
    const float* dones,
    float* advantages,
    float* returns,
    int L,
    float gamma,
    float lambda
) {
    float gae = 0.0f;

    for (int t = L - 1; t >= 0; t--) {
        float not_done = 1.0f - dones[t];
        float delta = rewards[t] + gamma * values[t + 1] * not_done - values[t];
        gae = delta + gamma * lambda * not_done * gae;
        advantages[t] = gae;
        returns[t]    = gae + values[t];
    }
}

/* ------------------------------------------------------------------ */
/*  Main Kernel                                                        */
/* ------------------------------------------------------------------ */

/**
 * Full PPO experience collection kernel.
 *
 * Each thread runs one independent episode of L timesteps (Phase A),
 * then performs a backward GAE scan (Phase B).
 *
 * Grid: (ceil(N/32), 1, 1),  Block: (32, 1, 1).
 */
extern "C" __global__ void ppo_full_experience_kernel(
    /* ---- Market data ---- */
    const float* __restrict__ market_features,   /* [total_bars, MARKET_DIM] */
    const float* __restrict__ targets,           /* [total_bars, 4]         */
    const int*   __restrict__ episode_starts,    /* [N]                     */

    /* ---- Actor weights (6 pointers) ---- */
    const float* __restrict__ pw1,    /* [ACTOR_H1, STATE_DIM]    = [128, 54] */
    const float* __restrict__ pb1,    /* [ACTOR_H1]               = [128]     */
    const float* __restrict__ pw2,    /* [ACTOR_H2, ACTOR_H1]     = [64, 128] */
    const float* __restrict__ pb2,    /* [ACTOR_H2]               = [64]      */
    const float* __restrict__ pw3,    /* [NUM_ACTIONS, ACTOR_H2]  = [45, 64]  */
    const float* __restrict__ pb3,    /* [NUM_ACTIONS]            = [45]      */

    /* ---- Critic weights (12 pointers) ---- */
    const float* __restrict__ vw1,    /* [CRITIC_H1, STATE_DIM]   = [512, 54]  */
    const float* __restrict__ vb1,    /* [CRITIC_H1]              = [512]      */
    const float* __restrict__ vw2,    /* [CRITIC_H2, CRITIC_H1]   = [384, 512] */
    const float* __restrict__ vb2,    /* [CRITIC_H2]              = [384]      */
    const float* __restrict__ vw3,    /* [CRITIC_H3, CRITIC_H2]   = [256, 384] */
    const float* __restrict__ vb3,    /* [CRITIC_H3]              = [256]      */
    const float* __restrict__ vw4,    /* [CRITIC_H4, CRITIC_H3]   = [128, 256] */
    const float* __restrict__ vb4,    /* [CRITIC_H4]              = [128]      */
    const float* __restrict__ vw5,    /* [CRITIC_H5, CRITIC_H4]   = [64, 128]  */
    const float* __restrict__ vb5,    /* [CRITIC_H5]              = [64]       */
    const float* __restrict__ vw6,    /* [1, CRITIC_H5]           = [1, 64]    */
    const float* __restrict__ vb6,    /* [1]                                   */

    /* ---- Curiosity model weights (4 pointers) ---- */
    const float* __restrict__ cur_w1,  /* [CUR_HIDDEN, CUR_INPUT]  */
    const float* __restrict__ cur_b1,  /* [CUR_HIDDEN]             */
    const float* __restrict__ cur_w2,  /* [CUR_OUTPUT, CUR_HIDDEN] */
    const float* __restrict__ cur_b2,  /* [CUR_OUTPUT]             */

    /* ---- Per-episode mutable state arrays ---- */
    float* portfolio_states,        /* [N, PORTFOLIO_STATE_SIZE] */
    float* barrier_states,          /* [N, BARRIER_STATE_SIZE]   */
    int*   diversity_windows,       /* [N, DIVERSITY_WINDOW]     */
    int*   diversity_metas,         /* [N, 2]                    */

    /* ---- Barrier config (shared) ---- */
    const float* __restrict__ barrier_config,  /* [3]: profit_mult, loss_mult, max_bars */

    /* ---- Scalar configs ---- */
    float max_position,
    int   episode_length,
    int   total_bars,
    int   L,                  /* timesteps per episode (rollout length) */
    float gamma,
    float gae_lambda,
    float curiosity_max_reward,
    int   N,                  /* total number of episodes */
    float barrier_scale,
    float diversity_scale,
    float curiosity_scale,
    float risk_weight,

    /* ---- RNG states [N] ---- */
    unsigned int* rng_states,

    /* ---- Output arrays ---- */
    float* out_states,        /* [N, L, STATE_DIM]  */
    int*   out_actions,       /* [N, L]             */
    float* out_log_probs,     /* [N, L]             */
    float* out_advantages,    /* [N, L]             */
    float* out_returns,       /* [N, L]             */
    int*   out_dones          /* [N, L]             */
) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    if (tid >= N) return;

    /* ---- Load per-thread portfolio state ---- */
    int ps_off = tid * PORTFOLIO_STATE_SIZE;
    float cash        = portfolio_states[ps_off + 0];
    float position    = portfolio_states[ps_off + 1];
    float entry_price = portfolio_states[ps_off + 2];
    float initial_cap = portfolio_states[ps_off + 3];
    float spread      = portfolio_states[ps_off + 4];
    float last_price  = portfolio_states[ps_off + 5];
    float reserve_pct = portfolio_states[ps_off + 6];
    float cum_costs   = portfolio_states[ps_off + 7];

    /* ---- Load per-thread barrier state ---- */
    int bs_off = tid * BARRIER_STATE_SIZE;
    float barrier_st[BARRIER_STATE_SIZE];
    for (int i = 0; i < BARRIER_STATE_SIZE; i++)
        barrier_st[i] = barrier_states[bs_off + i];

    /* ---- Load per-thread diversity window ---- */
    int dw_off = tid * DIVERSITY_WINDOW;
    int div_window[DIVERSITY_WINDOW];
    for (int i = 0; i < DIVERSITY_WINDOW; i++)
        div_window[i] = diversity_windows[dw_off + i];

    int dm_off = tid * 2;
    int div_meta[2];
    div_meta[0] = diversity_metas[dm_off + 0];
    div_meta[1] = diversity_metas[dm_off + 1];

    unsigned int rng = rng_states[tid];
    int ep_start = episode_starts[tid];

    /* ---- Per-thread scratch buffers ---- */
    float state[STATE_DIM];
    float actor_h1[ACTOR_H1];
    float actor_h2[ACTOR_H2];
    float logits[NUM_ACTIONS];
    float probs[NUM_ACTIONS];
    float critic_a[CRITIC_H1];   /* ping-pong buffer A (512 wide) */
    float critic_b[CRITIC_H1];   /* ping-pong buffer B (512 wide) */
    float next_state[STATE_DIM];
    float cur_scratch[CUR_HIDDEN];

    /* GAE accumulation arrays */
    float gae_values[MAX_GAE_LEN + 1];   /* values[L] is bootstrap */
    float gae_rewards[MAX_GAE_LEN];
    float gae_dones[MAX_GAE_LEN];
    float gae_advantages[MAX_GAE_LEN];
    float gae_returns[MAX_GAE_LEN];

    int step_in_episode = 0;
    int actual_L = (L <= MAX_GAE_LEN) ? L : MAX_GAE_LEN;

    /* ================================================================ */
    /*  Phase A: Forward Rollout (L timesteps)                          */
    /* ================================================================ */

    for (int t = 0; t < actual_L; t++) {
        int global_bar = ep_start + t;
        int out_off = tid * actual_L + t;

        /* Handle out-of-data */
        if (global_bar >= total_bars - 1) {
            for (int i = 0; i < STATE_DIM; i++)
                out_states[out_off * STATE_DIM + i] = 0.0f;
            out_actions[out_off]   = 0;
            out_log_probs[out_off] = 0.0f;
            out_dones[out_off]     = 1;
            gae_values[t]  = 0.0f;
            gae_rewards[t] = 0.0f;
            gae_dones[t]   = 1.0f;
            continue;
        }

        /* ---- Step 1: Read 51 market features from global memory ---- */
        int mf_off = global_bar * MARKET_DIM;
        for (int i = 0; i < MARKET_DIM; i++)
            state[i] = market_features[mf_off + i];

        /* ---- Step 2: Compute 3 portfolio features ---- */
        int t_off = global_bar * 4;
        float current_close     = targets[t_off + 0];
        float next_close        = targets[t_off + 1];
        float current_close_raw = targets[t_off + 2];
        float next_close_raw    = targets[t_off + 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 max_pos_norm  = (price > 0.0f) ? initial_cap / price : 1.0f;
        float pos_norm      = position / max_pos_norm;

        state[MARKET_DIM + 0] = current_norm;   /* normalized value    */
        state[MARKET_DIM + 1] = pos_norm;        /* normalized position */
        state[MARKET_DIM + 2] = spread;          /* spread              */

        /* ---- Step 3: Actor forward -> logits -> softmax_sample ---- */
        ppo_actor_forward(
            state,
            pw1, pb1, pw2, pb2, pw3, pb3,
            actor_h1, actor_h2, logits
        );

        int action_idx;
        float log_prob;
        softmax_sample(logits, probs, &rng, &action_idx, &log_prob);

        /* ---- Step 4: Critic forward -> value estimate ---- */
        float value = ppo_critic_forward(
            state,
            vw1, vb1, vw2, vb2, vw3, vb3,
            vw4, vb4, vw5, vb5, vw6, vb6,
            critic_a, critic_b
        );

        /* Store value for GAE */
        gae_values[t] = value;

        /* ---- Step 5: Portfolio simulation ---- */
        float target_exposure  = factored_action_to_exposure(action_idx);
        float target_position  = target_exposure * max_position;
        float tx_rate          = ppo_action_to_tx_cost(action_idx);

        /* Detect reversal (sign change) */
        int is_reversal = (position > 0.0f && target_position < 0.0f) ||
                          (position < 0.0f && target_position > 0.0f);

        if (is_reversal) {
            /* Phase 1: Close current position */
            float close_cash = position * price;
            float close_cost = fabsf(position) * price * tx_rate;
            cash += close_cash - close_cost;
            cum_costs += close_cost;

            /* Phase 2: Open opposite position */
            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 {
            /* Non-reversal: adjust position directly */
            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;

                /* Cash reserve check for buys */
                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, (price > 0.0f) ? floorf(affordable / price) : 0.0f);
                    }
                }

                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;

        /* ---- Step 6: Barrier tracking ---- */
        float old_barrier_entry = barrier_st[0];
        if (entry_price > 0.0f && old_barrier_entry <= 0.0f) {
            barrier_init(barrier_st, barrier_config, entry_price, global_bar);
        }

        int barrier_label = barrier_check(barrier_st, price, global_bar, position);
        if (barrier_label != 0) {
            barrier_reset(barrier_st);
        }

        /* ---- Step 7: Diversity entropy penalty ---- */
        float div_penalty = diversity_entropy(div_window, div_meta, action_idx);

        /* ---- Step 8: Mark-to-market + build next_state for curiosity ---- */
        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;

        int next_bar = global_bar + 1;
        if (next_bar < total_bars) {
            int nmf_off = next_bar * MARKET_DIM;
            for (int i = 0; i < MARKET_DIM; i++)
                next_state[i] = market_features[nmf_off + i];
        } else {
            for (int i = 0; i < MARKET_DIM; i++)
                next_state[i] = state[i];
        }
        float next_max_pos_norm = (next_price > 0.0f) ? initial_cap / next_price : 1.0f;
        next_state[MARKET_DIM + 0] = next_norm;
        next_state[MARKET_DIM + 1] = position / next_max_pos_norm;
        next_state[MARKET_DIM + 2] = spread;

        /* ---- Step 9: Curiosity inference ---- */
        float curiosity_reward = 0.0f;
        if (cur_w1 != 0) {
            curiosity_reward = curiosity_inference(
                state, next_state, action_idx,
                cur_w1, cur_b1, cur_w2, cur_b2,
                cur_scratch, curiosity_max_reward
            );
        }

        /* ---- Step 10: Risk penalty (drawdown) ---- */
        float pnl_reward = 0.0f;
        if (current_norm > 0.0f) {
            pnl_reward = (next_norm - current_norm) / current_norm;
        }

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

        /* ---- Step 11: Reward combination ---- */
        float barrier_mult = 1.0f;
        if (barrier_label != 0) {
            barrier_mult = 1.0f + barrier_scale * (float)barrier_label;
        }

        float combined_reward = pnl_reward * barrier_mult
                              + diversity_scale * div_penalty
                              + curiosity_scale * curiosity_reward;

        /* ---- Step 12: Episode done check ---- */
        step_in_episode++;
        int time_done    = (step_in_episode >= episode_length) ? 1 : 0;
        int barrier_done = (barrier_label != 0) ? 1 : 0;
        int data_done    = (next_bar >= total_bars) ? 1 : 0;
        int done = (time_done || barrier_done || data_done) ? 1 : 0;

        /* ---- Step 13: Store per-timestep outputs ---- */
        for (int i = 0; i < STATE_DIM; i++)
            out_states[out_off * STATE_DIM + i] = state[i];
        out_actions[out_off]   = action_idx;
        out_log_probs[out_off] = log_prob;
        out_dones[out_off]     = done;

        /* Store for GAE backward scan */
        gae_rewards[t] = combined_reward;
        gae_dones[t]   = (float)done;

        /* ---- Step 14: Episode reset on done ---- */
        if (done) {
            cash = initial_cap;
            position = 0.0f;
            entry_price = 0.0f;
            cum_costs = 0.0f;
            last_price = 0.0f;
            step_in_episode = 0;
            barrier_reset(barrier_st);
            /* Reset diversity window */
            for (int i = 0; i < DIVERSITY_WINDOW; i++)
                div_window[i] = 0;
            div_meta[0] = 0;
            div_meta[1] = 0;
        }
    } /* end Phase A timestep loop */

    /* ================================================================ */
    /*  Phase B: Backward GAE Scan                                      */
    /* ================================================================ */

    /* Step 1: Compute bootstrap value for GAE */
    int last_t = actual_L - 1;
    int last_out_off = tid * actual_L + last_t;
    if (out_dones[last_out_off] == 1) {
        /* Last step was terminal — bootstrap value is 0 */
        gae_values[actual_L] = 0.0f;
    } else {
        /* Last step was not terminal — run critic on next_state for bootstrap */
        int last_global_bar = ep_start + last_t;
        int next_bar_boot = last_global_bar + 1;

        /* Build bootstrap next_state */
        float boot_state[STATE_DIM];
        if (next_bar_boot < total_bars) {
            int nmf_off = next_bar_boot * MARKET_DIM;
            for (int i = 0; i < MARKET_DIM; i++)
                boot_state[i] = market_features[nmf_off + i];
        } else {
            /* Reuse last state market features from output */
            for (int i = 0; i < MARKET_DIM; i++)
                boot_state[i] = out_states[last_out_off * STATE_DIM + i];
        }

        /* Approximate portfolio features from current thread state */
        float boot_price_raw = 0.0f;
        if (next_bar_boot < total_bars) {
            int t_off_boot = next_bar_boot * 4;
            boot_price_raw = targets[t_off_boot + 2];
            if (boot_price_raw <= 0.0f) boot_price_raw = targets[t_off_boot + 0];
        }
        if (boot_price_raw <= 0.0f) boot_price_raw = 1.0f;

        float boot_value = cash + position * boot_price_raw;
        float boot_norm  = boot_value / initial_cap;
        float boot_max_pos = (boot_price_raw > 0.0f) ? initial_cap / boot_price_raw : 1.0f;

        boot_state[MARKET_DIM + 0] = boot_norm;
        boot_state[MARKET_DIM + 1] = position / boot_max_pos;
        boot_state[MARKET_DIM + 2] = spread;

        gae_values[actual_L] = ppo_critic_forward(
            boot_state,
            vw1, vb1, vw2, vb2, vw3, vb3,
            vw4, vb4, vw5, vb5, vw6, vb6,
            critic_a, critic_b
        );
    }

    /* Step 2: Run GAE backward scan */
    compute_gae_backward(
        gae_rewards, gae_values, gae_dones,
        gae_advantages, gae_returns,
        actual_L, gamma, gae_lambda
    );

    /* Step 3: Write advantages and returns to output buffers */
    for (int t = 0; t < actual_L; t++) {
        int out_off = tid * actual_L + t;
        out_advantages[out_off] = gae_advantages[t];
        out_returns[out_off]    = gae_returns[t];
    }

    /* ---- Write back per-thread state ---- */
    portfolio_states[ps_off + 0] = cash;
    portfolio_states[ps_off + 1] = position;
    portfolio_states[ps_off + 2] = entry_price;
    portfolio_states[ps_off + 3] = initial_cap;
    portfolio_states[ps_off + 4] = spread;
    portfolio_states[ps_off + 5] = last_price;
    portfolio_states[ps_off + 6] = reserve_pct;
    portfolio_states[ps_off + 7] = cum_costs;

    for (int i = 0; i < BARRIER_STATE_SIZE; i++)
        barrier_states[bs_off + i] = barrier_st[i];

    for (int i = 0; i < DIVERSITY_WINDOW; i++)
        diversity_windows[dw_off + i] = div_window[i];

    diversity_metas[dm_off + 0] = div_meta[0];
    diversity_metas[dm_off + 1] = div_meta[1];

    rng_states[tid] = rng;
}
