// Vectorized backtest environment step kernel.
// One thread per walk-forward window. Each thread steps sequentially.
//
// Portfolio state layout per window [8 floats]:
//   [0] value        - current portfolio value
//   [1] position     - current position size (-1.0 to +1.0)
//   [2] cash         - cash balance
//   [3] entry_price  - entry price of current position (0 if flat)
//   [4] max_equity   - peak equity for drawdown tracking
//   [5] step_pnl     - PnL this step (for reward)
//   [6] cum_return   - cumulative log return
//   [7] step_count   - number of completed steps

#define PORTFOLIO_STATE_SIZE 8

extern "C" __global__ void backtest_env_step(
    // Market data (read-only, uploaded once)
    const float* __restrict__ prices,       // [n_windows * max_len * 4] (OHLC)
    const int*   __restrict__ window_lens,  // [n_windows]

    // Actions from model for current step
    const int*   __restrict__ actions,      // [n_windows] factored action index

    // Portfolio state (read-write, persistent across steps)
    float*       portfolio_state,           // [n_windows * PORTFOLIO_STATE_SIZE]

    // Step outputs
    float*       step_rewards,              // [n_windows]
    float*       step_returns,              // [n_windows * max_len] (accumulated)
    int*         done_flags,                // [n_windows]
    int*         actions_history,           // [n_windows * max_len] (accumulated, for metrics)

    // Config
    int          n_windows,
    int          max_len,
    float        max_position,
    float        tx_cost_bps,
    float        spread_cost,
    int          current_step,
    // Branching DQN action space sizes
    int          b0_size,                   // exposure branch size (default 9)
    int          b1_size,                   // order type branch size (default 3)
    int          b2_size                    // urgency branch size (default 3)
) {
    // Shared memory tile for coalesced portfolio reads/writes
    __shared__ float shmem_pf[256 * PORTFOLIO_STATE_SIZE];

    int w = blockIdx.x * blockDim.x + threadIdx.x;
    int local_tid = threadIdx.x;

    // Cooperative tile load: all active threads load their portfolio slice
    // in one coalesced warp transaction before divergent early-exit checks.
    int ps = w * PORTFOLIO_STATE_SIZE;
    if (w < n_windows) {
        for (int i = 0; i < PORTFOLIO_STATE_SIZE; i++) {
            shmem_pf[local_tid * PORTFOLIO_STATE_SIZE + i] = portfolio_state[ps + i];
        }
    }
    __syncthreads();

    if (w >= n_windows) return;
    if (done_flags[w]) return;

    int wlen = window_lens[w];
    if (current_step >= wlen) {
        done_flags[w] = 1;
        return;
    }

    // Read current prices
    int price_base = (w * max_len + current_step) * 4;
    float open  = prices[price_base + 0];
    float high  = prices[price_base + 1];
    float low   = prices[price_base + 2];
    float close = prices[price_base + 3];

    // Read portfolio state from shared memory tile
    float value      = shmem_pf[local_tid * PORTFOLIO_STATE_SIZE + 0];
    float position   = shmem_pf[local_tid * PORTFOLIO_STATE_SIZE + 1];
    float cash       = shmem_pf[local_tid * PORTFOLIO_STATE_SIZE + 2];
    float entry_price = shmem_pf[local_tid * PORTFOLIO_STATE_SIZE + 3];
    float max_equity = shmem_pf[local_tid * PORTFOLIO_STATE_SIZE + 4];
    float cum_return = shmem_pf[local_tid * PORTFOLIO_STATE_SIZE + 6];

    // Decode factored action to exposure index (matches experience_env_step).
    // Factored action = exposure_idx * b1_size * b2_size + order_idx * b2_size + urgency_idx
    int action_val = actions[w];
    int exposure_idx;
    if (b1_size > 0 && b2_size > 0) {
        int denom = b1_size * b2_size;
        exposure_idx = action_val / denom;
    } else {
        exposure_idx = action_val;
    }
    // Clamp to valid range
    if (exposure_idx < 0) exposure_idx = 0;
    if (exposure_idx >= b0_size) exposure_idx = b0_size - 1;

    // Map exposure index to fraction: -1.0 + idx * (2.0 / (b0_size - 1))
    // With b0_size=9: step=0.25, range [-1.0, +1.0]
    float step_size = (b0_size > 1) ? 2.0f / (float)(b0_size - 1) : 0.0f;
    float target_exposure = -1.0f + (float)exposure_idx * step_size;
    target_exposure *= max_position;

    // Suppress unused variable warnings
    (void)open; (void)high; (void)low;

    // Decode order_type and urgency from factored action (for tx cost differentiation)
    int order_type_idx = (b1_size > 0 && b2_size > 0)
        ? (action_val / b2_size) % b1_size : 0;

    // Execute trade if position changes
    float delta = target_exposure - position;
    float trade_cost = 0.0f;
    if (fabsf(delta) > 0.001f && close > 0.0f) {
        /* EXACT match of training kernel (experience_env_step) tx cost formula:
         *   tx_cost = |delta| * close * (multiplier * 0.0001 * spread_scale * impact + premium)
         * - spread_scale: sqrt(|delta|/max_pos) (Almgren-Chriss square-root impact)
         * - order_premium: Market=0, IoC=+2bps, LimitMaker=-5bps rebate */
        float spread_scale = sqrtf(fabsf(delta) / fmaxf(max_position, 0.01f));
        if (spread_scale < 1.0f) spread_scale = 1.0f;
        float impact_scale = 1.0f; /* simplified: no dynamic spread widening in backtest */
        float order_premium = (order_type_idx == 0) ? 0.0f
                            : (order_type_idx == 1) ? 0.0002f   /* IoC: +2 bps */
                            : -0.0005f;                          /* LimitMaker: -5 bps rebate */
        trade_cost = fabsf(delta) * close * (tx_cost_bps * 0.0001f * spread_scale * impact_scale + order_premium)
                   + fabsf(delta) * spread_cost * 0.5f;
        cash -= trade_cost;

        // Mark-to-market old position
        if (fabsf(position) > 0.001f && entry_price > 0.0f) {
            float pnl = position * (close - entry_price);
            cash += pnl;
        }

        position = target_exposure;
        entry_price = close;
    }

    // Mark-to-market current position
    float unrealized = 0.0f;
    if (fabsf(position) > 0.001f && entry_price > 0.0f) {
        unrealized = position * (close - entry_price);
    }
    float new_value = cash + unrealized;

    // Step return
    float step_ret = (value > 0.0f) ? (new_value - value) / value : 0.0f;
    float new_cum_return = cum_return + step_ret;

    // Update max equity for drawdown
    float new_max = fmaxf(max_equity, new_value);

    // Write portfolio state directly to global memory (no shmem write-back needed —
    // threads with done_flags[w]==1 exited early and cannot participate in __syncthreads)
    portfolio_state[ps + 0] = new_value;
    portfolio_state[ps + 1] = position;
    portfolio_state[ps + 2] = cash;
    portfolio_state[ps + 3] = entry_price;
    portfolio_state[ps + 4] = new_max;
    portfolio_state[ps + 5] = step_ret;       // step PnL (for reward)
    portfolio_state[ps + 6] = new_cum_return;
    portfolio_state[ps + 7] += 1.0f;          // step count

    // Outputs
    step_rewards[w] = step_ret;
    step_returns[w * max_len + current_step] = step_ret;
    actions_history[w * max_len + current_step] = actions[w];
}
